Compare commits
58 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| e5d6194947 | |||
| a263d20e58 | |||
| a18fc3493a | |||
| 29968897c5 | |||
| 5cf6382e42 | |||
| bf7b972edb | |||
| 771f968401 | |||
| 792b85c2be | |||
| b801c0da5f | |||
| 704b6d6c29 | |||
| 5c5440d776 | |||
| f57e0243dd | |||
| 31c675bdec | |||
| 95b7876057 | |||
| 63aade0a60 | |||
| 2a35ade607 | |||
| 41defa998a | |||
| 1f4f1fa557 | |||
| a15086688a | |||
| b66702606c | |||
| d854020322 | |||
| 131dfed3b4 | |||
| 96e471882a | |||
| c100bbfdfb | |||
| b9c0d8304f | |||
| c68a6cc32d | |||
| 0c4d2bebea | |||
| 4111d1ee00 | |||
| 0003d228e5 | |||
| 656bd6469c | |||
| 87f51a9c32 | |||
| 6f60e2fb6d | |||
| b0314fee57 | |||
| f91dbb6b6d | |||
| ccae51c466 | |||
| 4acc7916dd | |||
| 4981277eef | |||
| d5e04593f1 | |||
| a0bec521b6 | |||
| 89bf80e3d8 | |||
| 78243a6fb8 | |||
| cda3200987 | |||
| b81ac7c29f | |||
| 9c13e4ca07 | |||
| 7eee2a3b24 | |||
| 0293610b23 | |||
| 7d32135bb9 | |||
| 3bf8e228b8 | |||
| 2eabbad738 | |||
| f07f6f17ca | |||
| 8db2516001 | |||
| 20d272b6c3 | |||
| 85010a3ea5 | |||
| 1299cfe0b7 | |||
| bb5fc086ae | |||
| bc05af2fd7 | |||
| 8211e57b61 | |||
| 9c4006b9aa |
@@ -0,0 +1,30 @@
|
||||
# PostgreSQL Configuration
|
||||
POSTGRES_HOST=localhost
|
||||
POSTGRES_PORT=5432
|
||||
POSTGRES_DB=polymarket_tracker
|
||||
POSTGRES_USER=tracker
|
||||
POSTGRES_PASSWORD=dev_password
|
||||
|
||||
# Constructed database URL (for application use)
|
||||
DATABASE_URL=postgresql://tracker:dev_password@localhost:5432/polymarket_tracker
|
||||
|
||||
# Redis Configuration
|
||||
REDIS_HOST=localhost
|
||||
REDIS_PORT=6379
|
||||
REDIS_URL=redis://localhost:6379
|
||||
|
||||
# Optional: Development tool ports
|
||||
ADMINER_PORT=8080
|
||||
REDIS_INSIGHT_PORT=5540
|
||||
|
||||
# Polygon RPC Configuration
|
||||
POLYGON_RPC_URL=https://polygon-rpc.com
|
||||
POLYGON_FALLBACK_RPC_URL=https://polygon-bor.publicnode.com
|
||||
|
||||
# Polymarket API Configuration
|
||||
POLYMARKET_WS_URL=wss://ws-subscriptions-clob.polymarket.com/ws/market
|
||||
|
||||
# Alert Webhooks (optional)
|
||||
DISCORD_WEBHOOK_URL=
|
||||
TELEGRAM_BOT_TOKEN=
|
||||
TELEGRAM_CHAT_ID=
|
||||
@@ -0,0 +1,98 @@
|
||||
name: CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
|
||||
jobs:
|
||||
lint:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
cache: "pip"
|
||||
|
||||
- name: Install dependencies
|
||||
run: pip install -e ".[dev]"
|
||||
|
||||
- name: Lint with ruff
|
||||
run: ruff check src/ tests/
|
||||
|
||||
- name: Check formatting with ruff
|
||||
run: ruff format --check src/ tests/
|
||||
|
||||
type-check:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
cache: "pip"
|
||||
|
||||
- name: Install dependencies
|
||||
run: pip install -e ".[dev]"
|
||||
|
||||
- name: Type check with mypy
|
||||
# TODO: Remove continue-on-error after fixing #48
|
||||
continue-on-error: true
|
||||
run: mypy src/
|
||||
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:15
|
||||
env:
|
||||
POSTGRES_DB: test_db
|
||||
POSTGRES_USER: test
|
||||
POSTGRES_PASSWORD: test
|
||||
ports:
|
||||
- "5432:5432"
|
||||
options: >-
|
||||
--health-cmd pg_isready
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 5
|
||||
|
||||
redis:
|
||||
image: redis:7
|
||||
ports:
|
||||
- "6379:6379"
|
||||
options: >-
|
||||
--health-cmd "redis-cli ping"
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 5
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
cache: "pip"
|
||||
|
||||
- name: Install dependencies
|
||||
run: pip install -e ".[dev]"
|
||||
|
||||
- name: Run tests with coverage
|
||||
run: pytest --cov=src --cov-report=xml --cov-report=term-missing
|
||||
env:
|
||||
DATABASE_URL: postgresql://test:test@localhost:5432/test_db
|
||||
REDIS_URL: redis://localhost:6379
|
||||
|
||||
- name: Upload coverage to Codecov
|
||||
uses: codecov/codecov-action@v4
|
||||
with:
|
||||
files: ./coverage.xml
|
||||
fail_ci_if_error: false
|
||||
env:
|
||||
CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
|
||||
+118
@@ -0,0 +1,118 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# PyInstaller
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
.python-version
|
||||
|
||||
# PEP 582
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# Environment
|
||||
.env
|
||||
.env.local
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# Ruff
|
||||
.ruff_cache/
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Local development
|
||||
*.local
|
||||
.direnv/
|
||||
@@ -0,0 +1,23 @@
|
||||
repos:
|
||||
- repo: https://github.com/pre-commit/pre-commit-hooks
|
||||
rev: v4.5.0
|
||||
hooks:
|
||||
- id: trailing-whitespace
|
||||
- id: end-of-file-fixer
|
||||
- id: check-yaml
|
||||
- id: check-added-large-files
|
||||
- id: check-merge-conflict
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.1.9
|
||||
hooks:
|
||||
- id: ruff
|
||||
args: [--fix]
|
||||
- id: ruff-format
|
||||
|
||||
- repo: https://github.com/pre-commit/mirrors-mypy
|
||||
rev: v1.8.0
|
||||
hooks:
|
||||
- id: mypy
|
||||
additional_dependencies:
|
||||
- pydantic>=2.0.0
|
||||
@@ -120,15 +120,49 @@ cp .env.example .env
|
||||
# Edit .env with your API keys
|
||||
|
||||
# Start infrastructure (PostgreSQL, Redis)
|
||||
docker-compose up -d
|
||||
docker compose up -d
|
||||
|
||||
# Wait for services to be healthy
|
||||
docker compose ps
|
||||
|
||||
# Install Python dependencies
|
||||
pip install -e .
|
||||
|
||||
# Run database migrations
|
||||
alembic upgrade head
|
||||
|
||||
# Run the tracker
|
||||
python -m src.main
|
||||
```
|
||||
|
||||
### Docker Services
|
||||
|
||||
The development stack includes:
|
||||
|
||||
| Service | Port | Description |
|
||||
|---------|------|-------------|
|
||||
| PostgreSQL 15 | 5432 | Primary database |
|
||||
| Redis 7 | 6379 | Caching and pub/sub |
|
||||
| Adminer | 8080 | Database admin UI (optional) |
|
||||
| RedisInsight | 5540 | Redis admin UI (optional) |
|
||||
|
||||
```bash
|
||||
# Start core services only
|
||||
docker compose up -d
|
||||
|
||||
# Start with development tools (Adminer, RedisInsight)
|
||||
docker compose --profile tools up -d
|
||||
|
||||
# View logs
|
||||
docker compose logs -f
|
||||
|
||||
# Stop all services
|
||||
docker compose down
|
||||
|
||||
# Stop and remove volumes (reset data)
|
||||
docker compose down -v
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
```bash
|
||||
|
||||
+79
@@ -0,0 +1,79 @@
|
||||
# Alembic Configuration File
|
||||
|
||||
[alembic]
|
||||
# Path to migration scripts
|
||||
script_location = alembic
|
||||
|
||||
# Template used to generate migration files
|
||||
file_template = %%(year)d%%(month).2d%%(day).2d_%%(hour).2d%%(minute).2d_%%(slug)s
|
||||
|
||||
# Prepend timestamp to migration file names
|
||||
prepend_date = True
|
||||
|
||||
# Timezone to use when rendering the date within the migration file
|
||||
# as well as the filename. Use UTC for consistency.
|
||||
timezone = UTC
|
||||
|
||||
# Max length of characters to apply to the "slug" field
|
||||
truncate_slug_length = 40
|
||||
|
||||
# Set to 'true' to run the environment during the 'revision' command
|
||||
revision_environment = false
|
||||
|
||||
# Set to 'true' to allow .pyc and .pyo files without a source .py file
|
||||
sourceless = false
|
||||
|
||||
# Version location specification
|
||||
version_locations = %(here)s/alembic/versions
|
||||
|
||||
# Version path separator
|
||||
version_path_separator = os
|
||||
|
||||
# Database URL - override with SQLALCHEMY_DATABASE_URL environment variable
|
||||
sqlalchemy.url = postgresql://localhost/polymarket_tracker
|
||||
|
||||
|
||||
[post_write_hooks]
|
||||
# Post write hooks define scripts to run after generating new revision files
|
||||
|
||||
# Format using "black" - only if available
|
||||
# hooks = black
|
||||
# black.type = console_scripts
|
||||
# black.entrypoint = black
|
||||
# black.options = -q
|
||||
|
||||
|
||||
# Logging configuration
|
||||
[loggers]
|
||||
keys = root,sqlalchemy,alembic
|
||||
|
||||
[handlers]
|
||||
keys = console
|
||||
|
||||
[formatters]
|
||||
keys = generic
|
||||
|
||||
[logger_root]
|
||||
level = WARN
|
||||
handlers = console
|
||||
qualname =
|
||||
|
||||
[logger_sqlalchemy]
|
||||
level = WARN
|
||||
handlers =
|
||||
qualname = sqlalchemy.engine
|
||||
|
||||
[logger_alembic]
|
||||
level = INFO
|
||||
handlers =
|
||||
qualname = alembic
|
||||
|
||||
[handler_console]
|
||||
class = StreamHandler
|
||||
args = (sys.stderr,)
|
||||
level = NOTSET
|
||||
formatter = generic
|
||||
|
||||
[formatter_generic]
|
||||
format = %(levelname)-5.5s [%(name)s] %(message)s
|
||||
datefmt = %H:%M:%S
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Alembic migration environment configuration."""
|
||||
|
||||
import os
|
||||
from logging.config import fileConfig
|
||||
|
||||
from alembic import context
|
||||
from sqlalchemy import engine_from_config, pool
|
||||
|
||||
from polymarket_insider_tracker.storage.models import Base
|
||||
|
||||
# this is the Alembic Config object
|
||||
config = context.config
|
||||
|
||||
# Interpret the config file for Python logging
|
||||
if config.config_file_name is not None:
|
||||
fileConfig(config.config_file_name)
|
||||
|
||||
# Target metadata for 'autogenerate' support
|
||||
target_metadata = Base.metadata
|
||||
|
||||
# Get database URL from environment variable or config
|
||||
database_url = os.environ.get("DATABASE_URL")
|
||||
if database_url:
|
||||
config.set_main_option("sqlalchemy.url", database_url)
|
||||
|
||||
|
||||
def run_migrations_offline() -> None:
|
||||
"""Run migrations in 'offline' mode.
|
||||
|
||||
This configures the context with just a URL and not an Engine,
|
||||
though an Engine is acceptable here as well. By skipping the Engine
|
||||
creation we don't even need a DBAPI to be available.
|
||||
|
||||
Calls to context.execute() here emit the given string to the
|
||||
script output.
|
||||
"""
|
||||
url = config.get_main_option("sqlalchemy.url")
|
||||
context.configure(
|
||||
url=url,
|
||||
target_metadata=target_metadata,
|
||||
literal_binds=True,
|
||||
dialect_opts={"paramstyle": "named"},
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
def run_migrations_online() -> None:
|
||||
"""Run migrations in 'online' mode.
|
||||
|
||||
In this scenario we need to create an Engine and associate a
|
||||
connection with the context.
|
||||
"""
|
||||
connectable = engine_from_config(
|
||||
config.get_section(config.config_ini_section, {}),
|
||||
prefix="sqlalchemy.",
|
||||
poolclass=pool.NullPool,
|
||||
)
|
||||
|
||||
with connectable.connect() as connection:
|
||||
context.configure(connection=connection, target_metadata=target_metadata)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
if context.is_offline_mode():
|
||||
run_migrations_offline()
|
||||
else:
|
||||
run_migrations_online()
|
||||
@@ -0,0 +1,26 @@
|
||||
"""${message}
|
||||
|
||||
Revision ID: ${up_revision}
|
||||
Revises: ${down_revision | comma,n}
|
||||
Create Date: ${create_date}
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
${imports if imports else ""}
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = ${repr(up_revision)}
|
||||
down_revision: Union[str, None] = ${repr(down_revision)}
|
||||
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
${downgrades if downgrades else "pass"}
|
||||
@@ -0,0 +1,87 @@
|
||||
"""Initial schema for wallet profiles and funding transfers.
|
||||
|
||||
Revision ID: 001_initial
|
||||
Revises:
|
||||
Create Date: 2026-01-04 00:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "001_initial"
|
||||
down_revision: Union[str, None] = None
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Wallet profiles table
|
||||
op.create_table(
|
||||
"wallet_profiles",
|
||||
sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
|
||||
sa.Column("address", sa.String(42), nullable=False),
|
||||
sa.Column("nonce", sa.Integer(), nullable=False),
|
||||
sa.Column("first_seen_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("is_fresh", sa.Boolean(), nullable=False),
|
||||
sa.Column("matic_balance", sa.Numeric(30, 0), nullable=True),
|
||||
sa.Column("usdc_balance", sa.Numeric(20, 6), nullable=True),
|
||||
sa.Column("analyzed_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint("address"),
|
||||
)
|
||||
op.create_index("idx_wallet_profiles_address", "wallet_profiles", ["address"])
|
||||
|
||||
# Funding transfers table
|
||||
op.create_table(
|
||||
"funding_transfers",
|
||||
sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
|
||||
sa.Column("from_address", sa.String(42), nullable=False),
|
||||
sa.Column("to_address", sa.String(42), nullable=False),
|
||||
sa.Column("amount", sa.Numeric(30, 6), nullable=False),
|
||||
sa.Column("token", sa.String(10), nullable=False),
|
||||
sa.Column("tx_hash", sa.String(66), nullable=False),
|
||||
sa.Column("block_number", sa.Integer(), nullable=False),
|
||||
sa.Column("timestamp", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint("tx_hash"),
|
||||
)
|
||||
op.create_index("idx_funding_transfers_to", "funding_transfers", ["to_address"])
|
||||
op.create_index("idx_funding_transfers_from", "funding_transfers", ["from_address"])
|
||||
op.create_index("idx_funding_transfers_block", "funding_transfers", ["block_number"])
|
||||
|
||||
# Wallet relationships table
|
||||
op.create_table(
|
||||
"wallet_relationships",
|
||||
sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
|
||||
sa.Column("wallet_a", sa.String(42), nullable=False),
|
||||
sa.Column("wallet_b", sa.String(42), nullable=False),
|
||||
sa.Column("relationship_type", sa.String(20), nullable=False),
|
||||
sa.Column("confidence", sa.Numeric(3, 2), nullable=False),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint(
|
||||
"wallet_a", "wallet_b", "relationship_type", name="uq_wallet_relationship"
|
||||
),
|
||||
)
|
||||
op.create_index("idx_wallet_relationships_a", "wallet_relationships", ["wallet_a"])
|
||||
op.create_index("idx_wallet_relationships_b", "wallet_relationships", ["wallet_b"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("idx_wallet_relationships_b", table_name="wallet_relationships")
|
||||
op.drop_index("idx_wallet_relationships_a", table_name="wallet_relationships")
|
||||
op.drop_table("wallet_relationships")
|
||||
|
||||
op.drop_index("idx_funding_transfers_block", table_name="funding_transfers")
|
||||
op.drop_index("idx_funding_transfers_from", table_name="funding_transfers")
|
||||
op.drop_index("idx_funding_transfers_to", table_name="funding_transfers")
|
||||
op.drop_table("funding_transfers")
|
||||
|
||||
op.drop_index("idx_wallet_profiles_address", table_name="wallet_profiles")
|
||||
op.drop_table("wallet_profiles")
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
"""Pytest configuration for the polymarket-insider-tracker tests."""
|
||||
|
||||
import pytest
|
||||
|
||||
# Configure pytest-asyncio
|
||||
pytest_plugins = ["pytest_asyncio"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def event_loop_policy():
|
||||
"""Use default event loop policy."""
|
||||
import asyncio
|
||||
|
||||
return asyncio.DefaultEventLoopPolicy()
|
||||
@@ -0,0 +1,72 @@
|
||||
version: "3.8"
|
||||
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:15
|
||||
container_name: polymarket-postgres
|
||||
ports:
|
||||
- "${POSTGRES_PORT:-5432}:5432"
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
environment:
|
||||
POSTGRES_DB: ${POSTGRES_DB:-polymarket_tracker}
|
||||
POSTGRES_USER: ${POSTGRES_USER:-tracker}
|
||||
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-dev_password}
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER:-tracker} -d ${POSTGRES_DB:-polymarket_tracker}"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
start_period: 10s
|
||||
restart: unless-stopped
|
||||
|
||||
redis:
|
||||
image: redis:7
|
||||
container_name: polymarket-redis
|
||||
ports:
|
||||
- "${REDIS_PORT:-6379}:6379"
|
||||
volumes:
|
||||
- redis_data:/data
|
||||
command: redis-server --appendonly yes
|
||||
healthcheck:
|
||||
test: ["CMD", "redis-cli", "ping"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
start_period: 5s
|
||||
restart: unless-stopped
|
||||
|
||||
# Optional development tools - use with: docker compose --profile tools up
|
||||
adminer:
|
||||
image: adminer:latest
|
||||
container_name: polymarket-adminer
|
||||
ports:
|
||||
- "${ADMINER_PORT:-8080}:8080"
|
||||
depends_on:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
profiles:
|
||||
- tools
|
||||
restart: unless-stopped
|
||||
|
||||
redis-insight:
|
||||
image: redis/redisinsight:latest
|
||||
container_name: polymarket-redis-insight
|
||||
ports:
|
||||
- "${REDIS_INSIGHT_PORT:-5540}:5540"
|
||||
volumes:
|
||||
- redis_insight_data:/data
|
||||
depends_on:
|
||||
redis:
|
||||
condition: service_healthy
|
||||
profiles:
|
||||
- tools
|
||||
restart: unless-stopped
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
driver: local
|
||||
redis_data:
|
||||
driver: local
|
||||
redis_insight_data:
|
||||
driver: local
|
||||
+103
@@ -0,0 +1,103 @@
|
||||
[project]
|
||||
name = "polymarket-insider-tracker"
|
||||
version = "0.1.0"
|
||||
description = "Detect insider activity on Polymarket"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"py-clob-client>=0.1.0",
|
||||
"web3>=6.0.0",
|
||||
"httpx>=0.25.0",
|
||||
"redis>=5.0.0",
|
||||
"sqlalchemy>=2.0.0",
|
||||
"alembic>=1.13.0",
|
||||
"pydantic>=2.0.0",
|
||||
"pydantic-settings>=2.0.0",
|
||||
"python-dotenv>=1.0.0",
|
||||
"websockets>=12.0",
|
||||
"prometheus-client>=0.19.0",
|
||||
"aiohttp>=3.9.0",
|
||||
"scikit-learn>=1.3.0",
|
||||
"numpy>=1.24.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=7.0.0",
|
||||
"pytest-asyncio>=0.21.0",
|
||||
"pytest-cov>=4.0.0",
|
||||
"ruff>=0.1.0",
|
||||
"mypy>=1.7.0",
|
||||
"pre-commit>=3.0.0",
|
||||
"aiosqlite>=0.19.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 100
|
||||
target-version = "py311"
|
||||
src = ["src", "tests"]
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"E", # pycodestyle errors
|
||||
"W", # pycodestyle warnings
|
||||
"F", # Pyflakes
|
||||
"I", # isort
|
||||
"B", # flake8-bugbear
|
||||
"C4", # flake8-comprehensions
|
||||
"UP", # pyupgrade
|
||||
"ARG", # flake8-unused-arguments
|
||||
"SIM", # flake8-simplify
|
||||
]
|
||||
ignore = [
|
||||
"E501", # line too long (handled by formatter)
|
||||
]
|
||||
|
||||
[tool.ruff.lint.isort]
|
||||
known-first-party = ["polymarket_insider_tracker"]
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.11"
|
||||
strict = true
|
||||
warn_return_any = true
|
||||
warn_unused_ignores = true
|
||||
disallow_untyped_defs = true
|
||||
disallow_incomplete_defs = true
|
||||
check_untyped_defs = true
|
||||
disallow_untyped_decorators = true
|
||||
no_implicit_optional = true
|
||||
warn_redundant_casts = true
|
||||
warn_unused_configs = true
|
||||
plugins = ["pydantic.mypy"]
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = ["py_clob_client.*", "web3.*", "redis.*", "sklearn.*", "prometheus_client.*"]
|
||||
ignore_missing_imports = true
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests"]
|
||||
addopts = "-v --tb=short"
|
||||
markers = [
|
||||
"slow: marks tests as slow (deselect with '-m \"not slow\"')",
|
||||
"integration: marks tests as integration tests",
|
||||
]
|
||||
|
||||
[tool.coverage.run]
|
||||
source = ["src"]
|
||||
branch = true
|
||||
|
||||
[tool.coverage.report]
|
||||
exclude_lines = [
|
||||
"pragma: no cover",
|
||||
"def __repr__",
|
||||
"raise NotImplementedError",
|
||||
"if TYPE_CHECKING:",
|
||||
]
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Polymarket Insider Tracker - Detect insider activity on Polymarket."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,302 @@
|
||||
"""CLI entry point for Polymarket Insider Tracker.
|
||||
|
||||
This module provides the main entry point for running the tracker
|
||||
from the command line.
|
||||
|
||||
Usage:
|
||||
python -m polymarket_insider_tracker [options]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
import logging.config
|
||||
import sys
|
||||
from typing import NoReturn
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from polymarket_insider_tracker import __version__
|
||||
from polymarket_insider_tracker.config import Settings, clear_settings_cache, get_settings
|
||||
from polymarket_insider_tracker.pipeline import Pipeline
|
||||
from polymarket_insider_tracker.shutdown import GracefulShutdown
|
||||
|
||||
# Application info
|
||||
APP_NAME = "Polymarket Insider Tracker"
|
||||
APP_VERSION = __version__
|
||||
|
||||
# Exit codes
|
||||
EXIT_SUCCESS = 0
|
||||
EXIT_ERROR = 1
|
||||
EXIT_CONFIG_ERROR = 2
|
||||
EXIT_INTERRUPTED = 130
|
||||
|
||||
|
||||
def create_parser() -> argparse.ArgumentParser:
|
||||
"""Create the argument parser for the CLI.
|
||||
|
||||
Returns:
|
||||
Configured ArgumentParser instance.
|
||||
"""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="polymarket-insider-tracker",
|
||||
description="Detect insider trading activity on Polymarket prediction markets.",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
python -m polymarket_insider_tracker Run full pipeline
|
||||
python -m polymarket_insider_tracker --config-check Validate config and exit
|
||||
python -m polymarket_insider_tracker --dry-run Run without sending alerts
|
||||
python -m polymarket_insider_tracker --log-level DEBUG Enable debug logging
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--version",
|
||||
action="version",
|
||||
version=f"%(prog)s {APP_VERSION}",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--config-check",
|
||||
action="store_true",
|
||||
help="Validate configuration and exit without running pipeline",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--log-level",
|
||||
choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
|
||||
default=None,
|
||||
help="Override logging level (default: from settings)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Run pipeline but don't send alerts",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--health-port",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Override health check port (default: from settings)",
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def configure_logging(level: str) -> None:
|
||||
"""Configure structured logging for the application.
|
||||
|
||||
Args:
|
||||
level: Logging level string (DEBUG, INFO, etc.)
|
||||
"""
|
||||
config = {
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
"formatters": {
|
||||
"standard": {
|
||||
"format": "%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||||
"datefmt": "%Y-%m-%d %H:%M:%S",
|
||||
},
|
||||
"detailed": {
|
||||
"format": (
|
||||
"%(asctime)s [%(levelname)s] %(name)s (%(filename)s:%(lineno)d): %(message)s"
|
||||
),
|
||||
"datefmt": "%Y-%m-%d %H:%M:%S",
|
||||
},
|
||||
},
|
||||
"handlers": {
|
||||
"console": {
|
||||
"class": "logging.StreamHandler",
|
||||
"level": level,
|
||||
"formatter": "detailed" if level == "DEBUG" else "standard",
|
||||
"stream": "ext://sys.stdout",
|
||||
},
|
||||
},
|
||||
"root": {
|
||||
"level": level,
|
||||
"handlers": ["console"],
|
||||
},
|
||||
# Quieter logging for noisy libraries
|
||||
"loggers": {
|
||||
"httpx": {"level": "WARNING"},
|
||||
"httpcore": {"level": "WARNING"},
|
||||
"websockets": {"level": "WARNING"},
|
||||
"web3": {"level": "WARNING"},
|
||||
"urllib3": {"level": "WARNING"},
|
||||
},
|
||||
}
|
||||
logging.config.dictConfig(config)
|
||||
|
||||
|
||||
def print_banner() -> None:
|
||||
"""Print the application startup banner."""
|
||||
banner = f"""
|
||||
╔══════════════════════════════════════════════════════════════╗
|
||||
║ ║
|
||||
║ {APP_NAME:^56} ║
|
||||
║ {"v" + APP_VERSION:^56} ║
|
||||
║ ║
|
||||
╚══════════════════════════════════════════════════════════════╝
|
||||
"""
|
||||
print(banner)
|
||||
|
||||
|
||||
def print_config_summary(settings: Settings, dry_run: bool) -> None:
|
||||
"""Print a summary of the configuration.
|
||||
|
||||
Args:
|
||||
settings: Application settings.
|
||||
dry_run: Whether dry-run mode is enabled.
|
||||
"""
|
||||
summary = settings.redacted_summary()
|
||||
print("Configuration:")
|
||||
print(f" Database: {summary['database_url']}")
|
||||
print(f" Redis: {summary['redis_url']}")
|
||||
print(f" Log Level: {summary['log_level']}")
|
||||
print(f" Health Port: {summary['health_port']}")
|
||||
print(f" Dry Run: {dry_run}")
|
||||
print(f" Discord: {'enabled' if summary['discord_enabled'] == 'True' else 'disabled'}")
|
||||
print(f" Telegram: {'enabled' if summary['telegram_enabled'] == 'True' else 'disabled'}")
|
||||
print()
|
||||
|
||||
|
||||
def validate_config() -> Settings | None:
|
||||
"""Validate and load configuration.
|
||||
|
||||
Returns:
|
||||
Settings instance if valid, None if invalid.
|
||||
"""
|
||||
try:
|
||||
# Clear cache to force reload
|
||||
clear_settings_cache()
|
||||
return get_settings()
|
||||
except ValidationError as e:
|
||||
print("Configuration validation failed:", file=sys.stderr)
|
||||
for error in e.errors():
|
||||
field = ".".join(str(loc) for loc in error["loc"])
|
||||
msg = error["msg"]
|
||||
print(f" {field}: {msg}", file=sys.stderr)
|
||||
return None
|
||||
|
||||
|
||||
def run_config_check(settings: Settings) -> int:
|
||||
"""Run configuration check and exit.
|
||||
|
||||
Args:
|
||||
settings: Validated settings.
|
||||
|
||||
Returns:
|
||||
Exit code (0 for success).
|
||||
"""
|
||||
print("Configuration is valid!")
|
||||
print()
|
||||
print_config_summary(settings, dry_run=False)
|
||||
|
||||
# Test component availability
|
||||
print("Checking component availability...")
|
||||
|
||||
# Check Discord
|
||||
if settings.discord.enabled:
|
||||
print(" Discord: configured")
|
||||
else:
|
||||
print(" Discord: not configured")
|
||||
|
||||
# Check Telegram
|
||||
if settings.telegram.enabled:
|
||||
print(" Telegram: configured")
|
||||
else:
|
||||
print(" Telegram: not configured")
|
||||
|
||||
print()
|
||||
print("All checks passed. Ready to run.")
|
||||
return EXIT_SUCCESS
|
||||
|
||||
|
||||
async def run_pipeline(
|
||||
settings: Settings,
|
||||
dry_run: bool,
|
||||
shutdown_timeout: float = 30.0,
|
||||
) -> int:
|
||||
"""Run the main pipeline with graceful shutdown handling.
|
||||
|
||||
Args:
|
||||
settings: Application settings.
|
||||
dry_run: Whether to skip sending alerts.
|
||||
shutdown_timeout: Maximum time to wait for graceful shutdown.
|
||||
|
||||
Returns:
|
||||
Exit code.
|
||||
"""
|
||||
logger = logging.getLogger(__name__)
|
||||
shutdown = GracefulShutdown(timeout=shutdown_timeout)
|
||||
|
||||
try:
|
||||
async with shutdown:
|
||||
pipeline = Pipeline(settings, dry_run=dry_run)
|
||||
|
||||
# Register pipeline cleanup
|
||||
shutdown.register_cleanup(pipeline.stop)
|
||||
|
||||
logger.info("Starting pipeline...")
|
||||
await pipeline.start()
|
||||
|
||||
logger.info("Pipeline running. Press Ctrl+C to stop.")
|
||||
|
||||
# Wait for shutdown signal
|
||||
await shutdown.wait()
|
||||
|
||||
logger.info("Shutdown signal received, stopping pipeline...")
|
||||
await pipeline.stop()
|
||||
|
||||
return EXIT_SUCCESS
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Interrupted by user")
|
||||
return EXIT_INTERRUPTED
|
||||
except Exception as e:
|
||||
logger.exception("Pipeline failed: %s", e)
|
||||
return EXIT_ERROR
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> NoReturn:
|
||||
"""Main entry point for the CLI.
|
||||
|
||||
Args:
|
||||
argv: Command line arguments (defaults to sys.argv[1:]).
|
||||
"""
|
||||
parser = create_parser()
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
# Validate configuration first
|
||||
settings = validate_config()
|
||||
if settings is None:
|
||||
sys.exit(EXIT_CONFIG_ERROR)
|
||||
|
||||
# Determine effective log level
|
||||
log_level = args.log_level or settings.log_level
|
||||
configure_logging(log_level)
|
||||
|
||||
# Print banner
|
||||
print_banner()
|
||||
|
||||
# Config check mode
|
||||
if args.config_check:
|
||||
sys.exit(run_config_check(settings))
|
||||
|
||||
# Determine dry-run mode
|
||||
dry_run = args.dry_run or settings.dry_run
|
||||
|
||||
# Print config summary
|
||||
print_config_summary(settings, dry_run)
|
||||
|
||||
# Run pipeline
|
||||
exit_code = asyncio.run(run_pipeline(settings, dry_run))
|
||||
sys.exit(exit_code)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Alerting layer - Real-time notification delivery."""
|
||||
|
||||
from polymarket_insider_tracker.alerter.channels.discord import DiscordChannel
|
||||
from polymarket_insider_tracker.alerter.channels.telegram import TelegramChannel
|
||||
from polymarket_insider_tracker.alerter.dispatcher import (
|
||||
AlertChannel,
|
||||
AlertDispatcher,
|
||||
CircuitBreakerState,
|
||||
DispatchResult,
|
||||
)
|
||||
from polymarket_insider_tracker.alerter.formatter import AlertFormatter
|
||||
from polymarket_insider_tracker.alerter.history import AlertHistory, AlertRecord
|
||||
from polymarket_insider_tracker.alerter.models import FormattedAlert
|
||||
|
||||
__all__ = [
|
||||
"AlertChannel",
|
||||
"AlertDispatcher",
|
||||
"AlertFormatter",
|
||||
"AlertHistory",
|
||||
"AlertRecord",
|
||||
"CircuitBreakerState",
|
||||
"DiscordChannel",
|
||||
"DispatchResult",
|
||||
"FormattedAlert",
|
||||
"TelegramChannel",
|
||||
]
|
||||
@@ -0,0 +1,9 @@
|
||||
"""Alert channel implementations for various platforms."""
|
||||
|
||||
from polymarket_insider_tracker.alerter.channels.discord import DiscordChannel
|
||||
from polymarket_insider_tracker.alerter.channels.telegram import TelegramChannel
|
||||
|
||||
__all__ = [
|
||||
"DiscordChannel",
|
||||
"TelegramChannel",
|
||||
]
|
||||
@@ -0,0 +1,116 @@
|
||||
"""Discord webhook channel implementation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from polymarket_insider_tracker.alerter.models import FormattedAlert
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DiscordChannel:
|
||||
"""Discord webhook channel for sending alerts.
|
||||
|
||||
Sends formatted alerts to Discord via webhook URL with rate limiting
|
||||
and retry support.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
webhook_url: str,
|
||||
*,
|
||||
rate_limit_per_minute: int = 30,
|
||||
max_retries: int = 3,
|
||||
retry_delay: float = 1.0,
|
||||
timeout: float = 10.0,
|
||||
) -> None:
|
||||
"""Initialize Discord channel.
|
||||
|
||||
Args:
|
||||
webhook_url: Discord webhook URL.
|
||||
rate_limit_per_minute: Maximum messages per minute (Discord limit is 30).
|
||||
max_retries: Maximum retry attempts on failure.
|
||||
retry_delay: Base delay between retries (exponential backoff).
|
||||
timeout: HTTP request timeout in seconds.
|
||||
"""
|
||||
self.webhook_url = webhook_url
|
||||
self.rate_limit_per_minute = rate_limit_per_minute
|
||||
self.max_retries = max_retries
|
||||
self.retry_delay = retry_delay
|
||||
self.timeout = timeout
|
||||
self.name = "discord"
|
||||
|
||||
# Rate limiting state
|
||||
self._request_times: list[float] = []
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def _wait_for_rate_limit(self) -> None:
|
||||
"""Wait if rate limit is exceeded."""
|
||||
async with self._lock:
|
||||
now = asyncio.get_event_loop().time()
|
||||
# Remove requests older than 1 minute
|
||||
self._request_times = [t for t in self._request_times if now - t < 60]
|
||||
|
||||
if len(self._request_times) >= self.rate_limit_per_minute:
|
||||
# Wait until the oldest request expires
|
||||
wait_time = 60 - (now - self._request_times[0])
|
||||
if wait_time > 0:
|
||||
logger.debug(f"Discord rate limit hit, waiting {wait_time:.2f}s")
|
||||
await asyncio.sleep(wait_time)
|
||||
|
||||
self._request_times.append(now)
|
||||
|
||||
async def send(self, alert: FormattedAlert) -> bool:
|
||||
"""Send alert to Discord webhook.
|
||||
|
||||
Args:
|
||||
alert: Formatted alert with discord_embed.
|
||||
|
||||
Returns:
|
||||
True if delivery succeeded, False otherwise.
|
||||
"""
|
||||
await self._wait_for_rate_limit()
|
||||
|
||||
payload = {
|
||||
"embeds": [alert.discord_embed],
|
||||
}
|
||||
|
||||
for attempt in range(self.max_retries):
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
self.webhook_url,
|
||||
json=payload,
|
||||
)
|
||||
|
||||
if response.status_code == 204:
|
||||
logger.info("Discord alert delivered successfully")
|
||||
return True
|
||||
|
||||
if response.status_code == 429:
|
||||
# Rate limited by Discord
|
||||
retry_after = response.json().get("retry_after", 1.0)
|
||||
logger.warning(f"Discord rate limited, retry after {retry_after}s")
|
||||
await asyncio.sleep(retry_after)
|
||||
continue
|
||||
|
||||
logger.error(f"Discord webhook failed: {response.status_code} {response.text}")
|
||||
|
||||
except httpx.TimeoutException:
|
||||
logger.warning(f"Discord webhook timeout (attempt {attempt + 1})")
|
||||
except httpx.HTTPError as e:
|
||||
logger.error(f"Discord webhook error: {e}")
|
||||
|
||||
# Exponential backoff
|
||||
if attempt < self.max_retries - 1:
|
||||
delay = self.retry_delay * (2**attempt)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
logger.error("Discord delivery failed after all retries")
|
||||
return False
|
||||
@@ -0,0 +1,131 @@
|
||||
"""Telegram Bot API channel implementation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from polymarket_insider_tracker.alerter.models import FormattedAlert
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TELEGRAM_API_BASE = "https://api.telegram.org/bot{token}/sendMessage"
|
||||
|
||||
|
||||
class TelegramChannel:
|
||||
"""Telegram Bot API channel for sending alerts.
|
||||
|
||||
Sends formatted alerts to Telegram via Bot API with rate limiting
|
||||
and retry support.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bot_token: str,
|
||||
chat_id: str,
|
||||
*,
|
||||
rate_limit_per_minute: int = 20,
|
||||
max_retries: int = 3,
|
||||
retry_delay: float = 1.0,
|
||||
timeout: float = 10.0,
|
||||
) -> None:
|
||||
"""Initialize Telegram channel.
|
||||
|
||||
Args:
|
||||
bot_token: Telegram bot token.
|
||||
chat_id: Target chat/channel ID.
|
||||
rate_limit_per_minute: Maximum messages per minute.
|
||||
max_retries: Maximum retry attempts on failure.
|
||||
retry_delay: Base delay between retries (exponential backoff).
|
||||
timeout: HTTP request timeout in seconds.
|
||||
"""
|
||||
self.bot_token = bot_token
|
||||
self.chat_id = chat_id
|
||||
self.rate_limit_per_minute = rate_limit_per_minute
|
||||
self.max_retries = max_retries
|
||||
self.retry_delay = retry_delay
|
||||
self.timeout = timeout
|
||||
self.name = "telegram"
|
||||
|
||||
self._api_url = TELEGRAM_API_BASE.format(token=bot_token)
|
||||
|
||||
# Rate limiting state
|
||||
self._request_times: list[float] = []
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def _wait_for_rate_limit(self) -> None:
|
||||
"""Wait if rate limit is exceeded."""
|
||||
async with self._lock:
|
||||
now = asyncio.get_event_loop().time()
|
||||
# Remove requests older than 1 minute
|
||||
self._request_times = [t for t in self._request_times if now - t < 60]
|
||||
|
||||
if len(self._request_times) >= self.rate_limit_per_minute:
|
||||
# Wait until the oldest request expires
|
||||
wait_time = 60 - (now - self._request_times[0])
|
||||
if wait_time > 0:
|
||||
logger.debug(f"Telegram rate limit hit, waiting {wait_time:.2f}s")
|
||||
await asyncio.sleep(wait_time)
|
||||
|
||||
self._request_times.append(now)
|
||||
|
||||
async def send(self, alert: FormattedAlert) -> bool:
|
||||
"""Send alert to Telegram channel.
|
||||
|
||||
Args:
|
||||
alert: Formatted alert with telegram_markdown.
|
||||
|
||||
Returns:
|
||||
True if delivery succeeded, False otherwise.
|
||||
"""
|
||||
await self._wait_for_rate_limit()
|
||||
|
||||
payload = {
|
||||
"chat_id": self.chat_id,
|
||||
"text": alert.telegram_markdown,
|
||||
"parse_mode": "MarkdownV2",
|
||||
"disable_web_page_preview": False,
|
||||
}
|
||||
|
||||
for attempt in range(self.max_retries):
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
self._api_url,
|
||||
json=payload,
|
||||
)
|
||||
|
||||
result = response.json()
|
||||
|
||||
if result.get("ok"):
|
||||
logger.info("Telegram alert delivered successfully")
|
||||
return True
|
||||
|
||||
error_code = result.get("error_code", 0)
|
||||
description = result.get("description", "Unknown error")
|
||||
|
||||
if error_code == 429:
|
||||
# Rate limited
|
||||
retry_after = result.get("parameters", {}).get("retry_after", 1)
|
||||
logger.warning(f"Telegram rate limited, retry after {retry_after}s")
|
||||
await asyncio.sleep(retry_after)
|
||||
continue
|
||||
|
||||
logger.error(f"Telegram API error: {error_code} - {description}")
|
||||
|
||||
except httpx.TimeoutException:
|
||||
logger.warning(f"Telegram API timeout (attempt {attempt + 1})")
|
||||
except httpx.HTTPError as e:
|
||||
logger.error(f"Telegram API error: {e}")
|
||||
|
||||
# Exponential backoff
|
||||
if attempt < self.max_retries - 1:
|
||||
delay = self.retry_delay * (2**attempt)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
logger.error("Telegram delivery failed after all retries")
|
||||
return False
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Alert dispatcher for multi-channel delivery."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Protocol
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from polymarket_insider_tracker.alerter.models import FormattedAlert
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AlertChannel(Protocol):
|
||||
"""Protocol for alert delivery channels."""
|
||||
|
||||
name: str
|
||||
|
||||
async def send(self, alert: FormattedAlert) -> bool:
|
||||
"""Send alert to channel. Returns True on success."""
|
||||
...
|
||||
|
||||
|
||||
@dataclass
|
||||
class CircuitBreakerState:
|
||||
"""State for circuit breaker pattern.
|
||||
|
||||
Tracks failures and manages open/closed state for a channel.
|
||||
"""
|
||||
|
||||
failure_count: int = 0
|
||||
last_failure_time: datetime | None = None
|
||||
is_open: bool = False
|
||||
half_open_attempts: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class DispatchResult:
|
||||
"""Result of dispatching an alert to all channels."""
|
||||
|
||||
success_count: int
|
||||
failure_count: int
|
||||
channel_results: dict[str, bool] = field(default_factory=dict)
|
||||
timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
@property
|
||||
def all_succeeded(self) -> bool:
|
||||
"""Return True if all channels succeeded."""
|
||||
return self.failure_count == 0 and self.success_count > 0
|
||||
|
||||
|
||||
class AlertDispatcher:
|
||||
"""Dispatcher for sending alerts to multiple channels.
|
||||
|
||||
Manages concurrent delivery to all configured channels with
|
||||
circuit breaker protection for failing channels.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: list[AlertChannel],
|
||||
*,
|
||||
failure_threshold: int = 5,
|
||||
recovery_timeout_seconds: int = 60,
|
||||
half_open_max_attempts: int = 3,
|
||||
) -> None:
|
||||
"""Initialize the dispatcher.
|
||||
|
||||
Args:
|
||||
channels: List of alert channels to dispatch to.
|
||||
failure_threshold: Number of consecutive failures before opening circuit.
|
||||
recovery_timeout_seconds: Time to wait before half-opening circuit.
|
||||
half_open_max_attempts: Number of test attempts in half-open state.
|
||||
"""
|
||||
self.channels = channels
|
||||
self.failure_threshold = failure_threshold
|
||||
self.recovery_timeout_seconds = recovery_timeout_seconds
|
||||
self.half_open_max_attempts = half_open_max_attempts
|
||||
|
||||
# Circuit breaker state per channel
|
||||
self._circuit_state: dict[str, CircuitBreakerState] = {
|
||||
ch.name: CircuitBreakerState() for ch in channels
|
||||
}
|
||||
|
||||
def _should_attempt(self, channel_name: str) -> bool:
|
||||
"""Check if we should attempt delivery to this channel."""
|
||||
state = self._circuit_state[channel_name]
|
||||
|
||||
if not state.is_open:
|
||||
return True
|
||||
|
||||
# Check if we should try half-open
|
||||
if state.last_failure_time:
|
||||
elapsed = (datetime.now(UTC) - state.last_failure_time).total_seconds()
|
||||
if (
|
||||
elapsed >= self.recovery_timeout_seconds
|
||||
and state.half_open_attempts < self.half_open_max_attempts
|
||||
):
|
||||
# Allow half-open attempt
|
||||
logger.info(
|
||||
f"Circuit half-open for {channel_name}, attempt {state.half_open_attempts + 1}"
|
||||
)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _record_success(self, channel_name: str) -> None:
|
||||
"""Record a successful delivery."""
|
||||
state = self._circuit_state[channel_name]
|
||||
state.failure_count = 0
|
||||
state.is_open = False
|
||||
state.half_open_attempts = 0
|
||||
state.last_failure_time = None
|
||||
logger.debug(f"Circuit closed for {channel_name}")
|
||||
|
||||
def _record_failure(self, channel_name: str) -> None:
|
||||
"""Record a failed delivery."""
|
||||
state = self._circuit_state[channel_name]
|
||||
state.failure_count += 1
|
||||
state.last_failure_time = datetime.now(UTC)
|
||||
|
||||
if state.is_open:
|
||||
# Failed during half-open, increment attempts
|
||||
state.half_open_attempts += 1
|
||||
elif state.failure_count >= self.failure_threshold:
|
||||
# Open the circuit
|
||||
state.is_open = True
|
||||
logger.warning(
|
||||
f"Circuit opened for {channel_name} after {state.failure_count} failures"
|
||||
)
|
||||
|
||||
async def _send_to_channel(
|
||||
self, channel: AlertChannel, alert: FormattedAlert
|
||||
) -> tuple[str, bool]:
|
||||
"""Send alert to a single channel with circuit breaker."""
|
||||
channel_name = channel.name
|
||||
|
||||
if not self._should_attempt(channel_name):
|
||||
logger.debug(f"Skipping {channel_name} - circuit open")
|
||||
return (channel_name, False)
|
||||
|
||||
try:
|
||||
success = await channel.send(alert)
|
||||
if success:
|
||||
self._record_success(channel_name)
|
||||
else:
|
||||
self._record_failure(channel_name)
|
||||
return (channel_name, success)
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending to {channel_name}: {e}")
|
||||
self._record_failure(channel_name)
|
||||
return (channel_name, False)
|
||||
|
||||
async def dispatch(self, alert: FormattedAlert) -> DispatchResult:
|
||||
"""Dispatch alert to all channels concurrently.
|
||||
|
||||
Args:
|
||||
alert: Formatted alert to send.
|
||||
|
||||
Returns:
|
||||
DispatchResult with per-channel status.
|
||||
"""
|
||||
if not self.channels:
|
||||
logger.warning("No channels configured for dispatch")
|
||||
return DispatchResult(success_count=0, failure_count=0)
|
||||
|
||||
# Send to all channels concurrently
|
||||
tasks = [self._send_to_channel(ch, alert) for ch in self.channels]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# Aggregate results
|
||||
channel_results = dict(results)
|
||||
success_count = sum(1 for success in channel_results.values() if success)
|
||||
failure_count = len(channel_results) - success_count
|
||||
|
||||
result = DispatchResult(
|
||||
success_count=success_count,
|
||||
failure_count=failure_count,
|
||||
channel_results=channel_results,
|
||||
)
|
||||
|
||||
logger.info(f"Dispatch complete: {success_count}/{len(channel_results)} succeeded")
|
||||
|
||||
return result
|
||||
|
||||
async def dispatch_batch(self, alerts: list[FormattedAlert]) -> list[DispatchResult]:
|
||||
"""Dispatch multiple alerts sequentially.
|
||||
|
||||
Args:
|
||||
alerts: List of formatted alerts to send.
|
||||
|
||||
Returns:
|
||||
List of DispatchResult for each alert.
|
||||
"""
|
||||
results = []
|
||||
for alert in alerts:
|
||||
result = await self.dispatch(alert)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def get_circuit_status(self) -> dict[str, dict[str, object]]:
|
||||
"""Get current circuit breaker status for all channels."""
|
||||
return {
|
||||
name: {
|
||||
"is_open": state.is_open,
|
||||
"failure_count": state.failure_count,
|
||||
"half_open_attempts": state.half_open_attempts,
|
||||
"last_failure": (
|
||||
state.last_failure_time.isoformat() if state.last_failure_time else None
|
||||
),
|
||||
}
|
||||
for name, state in self._circuit_state.items()
|
||||
}
|
||||
|
||||
def reset_circuit(self, channel_name: str) -> bool:
|
||||
"""Manually reset circuit breaker for a channel.
|
||||
|
||||
Args:
|
||||
channel_name: Name of channel to reset.
|
||||
|
||||
Returns:
|
||||
True if channel was found and reset.
|
||||
"""
|
||||
if channel_name in self._circuit_state:
|
||||
self._circuit_state[channel_name] = CircuitBreakerState()
|
||||
logger.info(f"Circuit reset for {channel_name}")
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,402 @@
|
||||
"""Alert message formatter for multi-channel delivery.
|
||||
|
||||
This module transforms RiskAssessment objects into human-readable,
|
||||
actionable alert messages optimized for Discord, Telegram, and plain text.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from decimal import Decimal
|
||||
from typing import Literal
|
||||
|
||||
from polymarket_insider_tracker.alerter.models import FormattedAlert
|
||||
from polymarket_insider_tracker.detector.models import RiskAssessment
|
||||
|
||||
# Polymarket URLs
|
||||
POLYMARKET_MARKET_URL = "https://polymarket.com/event/{slug}"
|
||||
POLYGONSCAN_ADDRESS_URL = "https://polygonscan.com/address/{address}"
|
||||
|
||||
# Discord embed colors (decimal values)
|
||||
COLOR_HIGH_RISK = 15158332 # Red (#E74C3C)
|
||||
COLOR_MEDIUM_RISK = 15105570 # Orange (#E67E22)
|
||||
COLOR_LOW_RISK = 16776960 # Yellow (#FFFF00)
|
||||
|
||||
# Risk level thresholds
|
||||
HIGH_RISK_THRESHOLD = 0.7
|
||||
MEDIUM_RISK_THRESHOLD = 0.5
|
||||
|
||||
|
||||
def truncate_address(address: str, chars: int = 4) -> str:
|
||||
"""Truncate an Ethereum address to 0x1234...5678 format."""
|
||||
if len(address) < chars * 2 + 4:
|
||||
return address
|
||||
return f"{address[: chars + 2]}...{address[-chars:]}"
|
||||
|
||||
|
||||
def format_usdc(amount: Decimal) -> str:
|
||||
"""Format a USDC amount with commas and 2 decimal places."""
|
||||
return f"${amount:,.2f}"
|
||||
|
||||
|
||||
def get_risk_level(score: float) -> str:
|
||||
"""Get human-readable risk level from score."""
|
||||
if score >= HIGH_RISK_THRESHOLD:
|
||||
return "HIGH"
|
||||
if score >= MEDIUM_RISK_THRESHOLD:
|
||||
return "MEDIUM"
|
||||
return "LOW"
|
||||
|
||||
|
||||
def get_risk_color(score: float) -> int:
|
||||
"""Get Discord embed color based on risk score."""
|
||||
if score >= HIGH_RISK_THRESHOLD:
|
||||
return COLOR_HIGH_RISK
|
||||
if score >= MEDIUM_RISK_THRESHOLD:
|
||||
return COLOR_MEDIUM_RISK
|
||||
return COLOR_LOW_RISK
|
||||
|
||||
|
||||
def get_triggered_signals(assessment: RiskAssessment) -> list[str]:
|
||||
"""Get list of triggered signal names."""
|
||||
signals = []
|
||||
if assessment.fresh_wallet_signal:
|
||||
signals.append("Fresh Wallet")
|
||||
if assessment.size_anomaly_signal:
|
||||
signals.append("Large Position")
|
||||
if assessment.size_anomaly_signal.is_niche_market:
|
||||
signals.append("Niche Market")
|
||||
return signals
|
||||
|
||||
|
||||
class AlertFormatter:
|
||||
"""Formats RiskAssessments into multi-channel alert messages.
|
||||
|
||||
Supports two verbosity levels:
|
||||
- compact: Essential info only (wallet, score, market)
|
||||
- detailed: Full context (all signals, links, trade details)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
verbosity: Literal["compact", "detailed"] = "detailed",
|
||||
) -> None:
|
||||
"""Initialize the formatter.
|
||||
|
||||
Args:
|
||||
verbosity: Level of detail in formatted messages.
|
||||
"""
|
||||
self.verbosity = verbosity
|
||||
|
||||
def format(self, assessment: RiskAssessment) -> FormattedAlert:
|
||||
"""Format a risk assessment into a multi-channel alert.
|
||||
|
||||
Args:
|
||||
assessment: The risk assessment to format.
|
||||
|
||||
Returns:
|
||||
FormattedAlert with all channel formats.
|
||||
"""
|
||||
# Build common data
|
||||
wallet_short = truncate_address(assessment.wallet_address)
|
||||
risk_level = get_risk_level(assessment.weighted_score)
|
||||
signals = get_triggered_signals(assessment)
|
||||
|
||||
# Build links
|
||||
links = self._build_links(assessment)
|
||||
|
||||
# Build title
|
||||
title = f"🚨 Suspicious Activity Detected - {risk_level} Risk"
|
||||
|
||||
# Build body based on verbosity
|
||||
body = self._build_body(assessment, wallet_short, risk_level, signals)
|
||||
|
||||
# Build channel-specific formats
|
||||
discord_embed = self._build_discord_embed(
|
||||
assessment, wallet_short, risk_level, signals, links
|
||||
)
|
||||
telegram_md = self._build_telegram_markdown(
|
||||
assessment, wallet_short, risk_level, signals, links
|
||||
)
|
||||
plain_text = self._build_plain_text(assessment, wallet_short, risk_level, signals, links)
|
||||
|
||||
return FormattedAlert(
|
||||
title=title,
|
||||
body=body,
|
||||
discord_embed=discord_embed,
|
||||
telegram_markdown=telegram_md,
|
||||
plain_text=plain_text,
|
||||
links=links,
|
||||
)
|
||||
|
||||
def _build_links(self, assessment: RiskAssessment) -> dict[str, str]:
|
||||
"""Build dictionary of relevant links."""
|
||||
trade = assessment.trade_event
|
||||
links = {
|
||||
"wallet": POLYGONSCAN_ADDRESS_URL.format(address=assessment.wallet_address),
|
||||
}
|
||||
|
||||
# Add market link if we have the slug
|
||||
if trade.market_slug:
|
||||
links["market"] = POLYMARKET_MARKET_URL.format(slug=trade.market_slug)
|
||||
|
||||
return links
|
||||
|
||||
def _build_body(
|
||||
self,
|
||||
assessment: RiskAssessment,
|
||||
wallet_short: str,
|
||||
risk_level: str,
|
||||
signals: list[str],
|
||||
) -> str:
|
||||
"""Build the main body text."""
|
||||
trade = assessment.trade_event
|
||||
|
||||
if self.verbosity == "compact":
|
||||
return (
|
||||
f"Wallet {wallet_short} made a {trade.side} trade "
|
||||
f"({format_usdc(trade.notional_value)}) with risk score "
|
||||
f"{assessment.weighted_score:.2f} ({risk_level})"
|
||||
)
|
||||
|
||||
# Detailed body
|
||||
lines = [
|
||||
f"Wallet: {wallet_short}",
|
||||
f"Risk Score: {assessment.weighted_score:.2f} ({risk_level})",
|
||||
f"Trade: {trade.side} {trade.outcome} @ ${trade.price:.3f}",
|
||||
f"Size: {format_usdc(trade.notional_value)}",
|
||||
]
|
||||
|
||||
if signals:
|
||||
lines.append(f"Signals: {', '.join(signals)}")
|
||||
|
||||
if trade.event_title:
|
||||
lines.append(f"Market: {trade.event_title}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def _build_discord_embed(
|
||||
self,
|
||||
assessment: RiskAssessment,
|
||||
wallet_short: str,
|
||||
risk_level: str,
|
||||
signals: list[str],
|
||||
links: dict[str, str],
|
||||
) -> dict[str, object]:
|
||||
"""Build Discord-optimized embed format."""
|
||||
trade = assessment.trade_event
|
||||
color = get_risk_color(assessment.weighted_score)
|
||||
|
||||
# Get wallet age if available
|
||||
wallet_age_str = ""
|
||||
if assessment.fresh_wallet_signal:
|
||||
age_hours = assessment.fresh_wallet_signal.wallet_profile.age_hours
|
||||
if age_hours is not None:
|
||||
if age_hours < 1:
|
||||
wallet_age_str = f" (Age: {int(age_hours * 60)}m)"
|
||||
else:
|
||||
wallet_age_str = f" (Age: {age_hours:.0f}h)"
|
||||
|
||||
fields: list[dict[str, object]] = [
|
||||
{
|
||||
"name": "Wallet",
|
||||
"value": f"`{wallet_short}`{wallet_age_str}",
|
||||
"inline": True,
|
||||
},
|
||||
{
|
||||
"name": "Risk Score",
|
||||
"value": f"{assessment.weighted_score:.2f} ({risk_level})",
|
||||
"inline": True,
|
||||
},
|
||||
]
|
||||
|
||||
# Market field
|
||||
market_title = trade.event_title or trade.market_slug or "Unknown Market"
|
||||
market_value = market_title
|
||||
if "market" in links:
|
||||
market_value = f"[{market_title}]({links['market']})"
|
||||
fields.append({"name": "Market", "value": market_value, "inline": False})
|
||||
|
||||
# Trade details
|
||||
trade_detail = (
|
||||
f"{trade.side} {trade.outcome} @ ${trade.price:.3f} | "
|
||||
f"{format_usdc(trade.notional_value)}"
|
||||
)
|
||||
fields.append({"name": "Trade", "value": trade_detail, "inline": False})
|
||||
|
||||
# Signals (if any)
|
||||
if signals:
|
||||
fields.append(
|
||||
{
|
||||
"name": "Signals",
|
||||
"value": ", ".join(signals),
|
||||
"inline": False,
|
||||
}
|
||||
)
|
||||
|
||||
# Add detailed info for detailed verbosity
|
||||
if self.verbosity == "detailed":
|
||||
# Add confidence breakdown
|
||||
confidences = []
|
||||
if assessment.fresh_wallet_signal:
|
||||
conf = assessment.fresh_wallet_signal.confidence
|
||||
confidences.append(f"Fresh Wallet: {conf:.0%}")
|
||||
if assessment.size_anomaly_signal:
|
||||
conf = assessment.size_anomaly_signal.confidence
|
||||
confidences.append(f"Size Anomaly: {conf:.0%}")
|
||||
|
||||
if confidences:
|
||||
fields.append(
|
||||
{
|
||||
"name": "Confidence",
|
||||
"value": " | ".join(confidences),
|
||||
"inline": False,
|
||||
}
|
||||
)
|
||||
|
||||
embed: dict[str, object] = {
|
||||
"title": "🚨 Suspicious Activity Detected",
|
||||
"color": color,
|
||||
"fields": fields,
|
||||
"footer": {"text": "Polymarket Insider Tracker"},
|
||||
}
|
||||
|
||||
# Add wallet link as URL if available
|
||||
if "wallet" in links:
|
||||
embed["url"] = links["wallet"]
|
||||
|
||||
return embed
|
||||
|
||||
def _build_telegram_markdown(
|
||||
self,
|
||||
assessment: RiskAssessment,
|
||||
wallet_short: str,
|
||||
risk_level: str,
|
||||
signals: list[str],
|
||||
links: dict[str, str],
|
||||
) -> str:
|
||||
"""Build Telegram-optimized markdown format."""
|
||||
trade = assessment.trade_event
|
||||
|
||||
lines = ["🚨 *Suspicious Activity Detected*", ""]
|
||||
|
||||
# Wallet with link
|
||||
wallet_line = f"*Wallet:* `{wallet_short}`"
|
||||
if assessment.fresh_wallet_signal:
|
||||
age_hours = assessment.fresh_wallet_signal.wallet_profile.age_hours
|
||||
if age_hours is not None:
|
||||
if age_hours < 1:
|
||||
wallet_line += f" \\(Age: {int(age_hours * 60)}m\\)"
|
||||
else:
|
||||
wallet_line += f" \\(Age: {age_hours:.0f}h\\)"
|
||||
lines.append(wallet_line)
|
||||
|
||||
# Risk score
|
||||
lines.append(f"*Risk Score:* {assessment.weighted_score:.2f} \\({risk_level}\\)")
|
||||
|
||||
# Market
|
||||
market_title = trade.event_title or trade.market_slug or "Unknown Market"
|
||||
# Escape special Telegram markdown characters
|
||||
market_title_escaped = self._escape_telegram_markdown(market_title)
|
||||
if "market" in links:
|
||||
lines.append(f"*Market:* [{market_title_escaped}]({links['market']})")
|
||||
else:
|
||||
lines.append(f"*Market:* {market_title_escaped}")
|
||||
|
||||
# Trade details
|
||||
usdc_value = format_usdc(trade.notional_value).replace("$", "\\$")
|
||||
lines.append(
|
||||
f"*Trade:* {trade.side} {trade.outcome} @ \\${trade.price:.3f} \\| {usdc_value}"
|
||||
)
|
||||
|
||||
# Signals
|
||||
if signals:
|
||||
lines.append(f"*Signals:* {', '.join(signals)}")
|
||||
|
||||
# Links
|
||||
lines.append("")
|
||||
if "wallet" in links:
|
||||
lines.append(f"[View Wallet]({links['wallet']})")
|
||||
if "market" in links:
|
||||
lines.append(f"[View Market]({links['market']})")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def _escape_telegram_markdown(self, text: str) -> str:
|
||||
"""Escape special Telegram MarkdownV2 characters."""
|
||||
special_chars = [
|
||||
"_",
|
||||
"*",
|
||||
"[",
|
||||
"]",
|
||||
"(",
|
||||
")",
|
||||
"~",
|
||||
"`",
|
||||
">",
|
||||
"#",
|
||||
"+",
|
||||
"-",
|
||||
"=",
|
||||
"|",
|
||||
"{",
|
||||
"}",
|
||||
".",
|
||||
"!",
|
||||
]
|
||||
for char in special_chars:
|
||||
text = text.replace(char, f"\\{char}")
|
||||
return text
|
||||
|
||||
def _build_plain_text(
|
||||
self,
|
||||
assessment: RiskAssessment,
|
||||
wallet_short: str,
|
||||
risk_level: str,
|
||||
signals: list[str],
|
||||
links: dict[str, str],
|
||||
) -> str:
|
||||
"""Build plain text format for generic channels."""
|
||||
trade = assessment.trade_event
|
||||
|
||||
lines = [
|
||||
"SUSPICIOUS ACTIVITY DETECTED",
|
||||
"=" * 30,
|
||||
"",
|
||||
]
|
||||
|
||||
# Wallet info
|
||||
wallet_line = f"Wallet: {wallet_short}"
|
||||
if assessment.fresh_wallet_signal:
|
||||
age_hours = assessment.fresh_wallet_signal.wallet_profile.age_hours
|
||||
if age_hours is not None:
|
||||
if age_hours < 1:
|
||||
wallet_line += f" (Age: {int(age_hours * 60)}m)"
|
||||
else:
|
||||
wallet_line += f" (Age: {age_hours:.0f}h)"
|
||||
lines.append(wallet_line)
|
||||
|
||||
# Risk
|
||||
lines.append(f"Risk Score: {assessment.weighted_score:.2f} ({risk_level})")
|
||||
|
||||
# Market
|
||||
market_title = trade.event_title or trade.market_slug or "Unknown Market"
|
||||
lines.append(f"Market: {market_title}")
|
||||
|
||||
# Trade
|
||||
lines.append(
|
||||
f"Trade: {trade.side} {trade.outcome} @ ${trade.price:.3f} | "
|
||||
f"{format_usdc(trade.notional_value)}"
|
||||
)
|
||||
|
||||
# Signals
|
||||
if signals:
|
||||
lines.append(f"Signals: {', '.join(signals)}")
|
||||
|
||||
# Links
|
||||
lines.append("")
|
||||
if "wallet" in links:
|
||||
lines.append(f"Wallet: {links['wallet']}")
|
||||
if "market" in links:
|
||||
lines.append(f"Market: {links['market']}")
|
||||
|
||||
return "\n".join(lines)
|
||||
@@ -0,0 +1,396 @@
|
||||
"""Alert history tracking and deduplication.
|
||||
|
||||
This module provides alert history management with deduplication
|
||||
to prevent spam and enable analytics on alert patterns.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from polymarket_insider_tracker.detector.models import RiskAssessment
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AlertRecord:
|
||||
"""Record of a sent alert.
|
||||
|
||||
Attributes:
|
||||
alert_id: Unique identifier for this alert.
|
||||
wallet_address: Trader's wallet address.
|
||||
market_id: Market condition ID.
|
||||
risk_score: Final weighted risk score.
|
||||
signals_triggered: List of signal names that triggered.
|
||||
channels_attempted: List of channels we tried to send to.
|
||||
channels_succeeded: List of channels that succeeded.
|
||||
dedup_key: Key used for deduplication.
|
||||
feedback_useful: User feedback on alert usefulness.
|
||||
created_at: When the alert was sent.
|
||||
"""
|
||||
|
||||
alert_id: str
|
||||
wallet_address: str
|
||||
market_id: str
|
||||
risk_score: float
|
||||
signals_triggered: list[str]
|
||||
channels_attempted: list[str]
|
||||
channels_succeeded: list[str]
|
||||
dedup_key: str
|
||||
feedback_useful: bool | None = None
|
||||
created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialize to dictionary for storage."""
|
||||
return {
|
||||
"alert_id": self.alert_id,
|
||||
"wallet_address": self.wallet_address,
|
||||
"market_id": self.market_id,
|
||||
"risk_score": self.risk_score,
|
||||
"signals_triggered": self.signals_triggered,
|
||||
"channels_attempted": self.channels_attempted,
|
||||
"channels_succeeded": self.channels_succeeded,
|
||||
"dedup_key": self.dedup_key,
|
||||
"feedback_useful": self.feedback_useful,
|
||||
"created_at": self.created_at.isoformat(),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> AlertRecord:
|
||||
"""Deserialize from dictionary."""
|
||||
created_at = data.get("created_at")
|
||||
if isinstance(created_at, str):
|
||||
created_at = datetime.fromisoformat(created_at)
|
||||
elif created_at is None:
|
||||
created_at = datetime.now(UTC)
|
||||
|
||||
return cls(
|
||||
alert_id=data["alert_id"],
|
||||
wallet_address=data["wallet_address"],
|
||||
market_id=data["market_id"],
|
||||
risk_score=float(data["risk_score"]),
|
||||
signals_triggered=data.get("signals_triggered", []),
|
||||
channels_attempted=data.get("channels_attempted", []),
|
||||
channels_succeeded=data.get("channels_succeeded", []),
|
||||
dedup_key=data["dedup_key"],
|
||||
feedback_useful=data.get("feedback_useful"),
|
||||
created_at=created_at,
|
||||
)
|
||||
|
||||
|
||||
def _generate_dedup_key(wallet_address: str, market_id: str, hour: datetime) -> str:
|
||||
"""Generate deduplication key for wallet/market/hour combination."""
|
||||
hour_str = hour.strftime("%Y%m%d%H")
|
||||
return f"{wallet_address}:{market_id}:{hour_str}"
|
||||
|
||||
|
||||
def _get_signals_from_assessment(assessment: RiskAssessment) -> list[str]:
|
||||
"""Extract triggered signal names from assessment."""
|
||||
signals = []
|
||||
if assessment.fresh_wallet_signal:
|
||||
signals.append("fresh_wallet")
|
||||
if assessment.size_anomaly_signal:
|
||||
signals.append("size_anomaly")
|
||||
if assessment.size_anomaly_signal.is_niche_market:
|
||||
signals.append("niche_market")
|
||||
return signals
|
||||
|
||||
|
||||
class AlertHistory:
|
||||
"""Tracks alert history and provides deduplication.
|
||||
|
||||
Uses Redis for storage with configurable dedup window.
|
||||
"""
|
||||
|
||||
# Redis key prefixes
|
||||
KEY_PREFIX_DEDUP = "alert:dedup:"
|
||||
KEY_PREFIX_ALERT = "alert:record:"
|
||||
KEY_PREFIX_FEEDBACK = "alert:feedback:"
|
||||
KEY_INDEX_TIME = "alert:index:time"
|
||||
KEY_INDEX_WALLET = "alert:index:wallet:"
|
||||
KEY_INDEX_MARKET = "alert:index:market:"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis: Any,
|
||||
*,
|
||||
dedup_window_hours: int = 1,
|
||||
retention_days: int = 30,
|
||||
) -> None:
|
||||
"""Initialize alert history.
|
||||
|
||||
Args:
|
||||
redis: Redis client (async).
|
||||
dedup_window_hours: Hours to deduplicate alerts for same wallet/market.
|
||||
retention_days: Days to retain alert history.
|
||||
"""
|
||||
self.redis = redis
|
||||
self.dedup_window_hours = dedup_window_hours
|
||||
self.retention_days = retention_days
|
||||
self._dedup_ttl = dedup_window_hours * 3600
|
||||
self._retention_ttl = retention_days * 86400
|
||||
|
||||
def _get_dedup_key(self, assessment: RiskAssessment) -> str:
|
||||
"""Get deduplication key for an assessment."""
|
||||
now = datetime.now(UTC)
|
||||
return _generate_dedup_key(
|
||||
assessment.wallet_address,
|
||||
assessment.market_id,
|
||||
now,
|
||||
)
|
||||
|
||||
async def should_send(self, assessment: RiskAssessment) -> bool:
|
||||
"""Check if alert should be sent (not a duplicate).
|
||||
|
||||
Args:
|
||||
assessment: Risk assessment to check.
|
||||
|
||||
Returns:
|
||||
True if alert should be sent, False if duplicate.
|
||||
"""
|
||||
dedup_key = self._get_dedup_key(assessment)
|
||||
redis_key = f"{self.KEY_PREFIX_DEDUP}{dedup_key}"
|
||||
|
||||
# Check if key exists
|
||||
exists = await self.redis.exists(redis_key)
|
||||
if exists:
|
||||
logger.debug(f"Duplicate alert for {dedup_key}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
async def record_sent(
|
||||
self,
|
||||
assessment: RiskAssessment,
|
||||
channels_attempted: list[str],
|
||||
channels_succeeded: dict[str, bool],
|
||||
) -> str:
|
||||
"""Record that an alert was sent.
|
||||
|
||||
Args:
|
||||
assessment: The risk assessment that was alerted.
|
||||
channels_attempted: List of channels we tried to send to.
|
||||
channels_succeeded: Dict of channel name -> success status.
|
||||
|
||||
Returns:
|
||||
The alert_id for this record.
|
||||
"""
|
||||
alert_id = str(uuid.uuid4())
|
||||
dedup_key = self._get_dedup_key(assessment)
|
||||
now = datetime.now(UTC)
|
||||
|
||||
# Create record
|
||||
record = AlertRecord(
|
||||
alert_id=alert_id,
|
||||
wallet_address=assessment.wallet_address,
|
||||
market_id=assessment.market_id,
|
||||
risk_score=assessment.weighted_score,
|
||||
signals_triggered=_get_signals_from_assessment(assessment),
|
||||
channels_attempted=channels_attempted,
|
||||
channels_succeeded=[ch for ch, success in channels_succeeded.items() if success],
|
||||
dedup_key=dedup_key,
|
||||
created_at=now,
|
||||
)
|
||||
|
||||
# Store in Redis with pipeline
|
||||
async with self.redis.pipeline() as pipe:
|
||||
# Store dedup key with TTL
|
||||
dedup_redis_key = f"{self.KEY_PREFIX_DEDUP}{dedup_key}"
|
||||
pipe.set(dedup_redis_key, "1", ex=self._dedup_ttl)
|
||||
|
||||
# Store alert record
|
||||
alert_redis_key = f"{self.KEY_PREFIX_ALERT}{alert_id}"
|
||||
pipe.set(
|
||||
alert_redis_key,
|
||||
json.dumps(record.to_dict()),
|
||||
ex=self._retention_ttl,
|
||||
)
|
||||
|
||||
# Add to time index (sorted set with timestamp as score)
|
||||
timestamp_score = now.timestamp()
|
||||
pipe.zadd(self.KEY_INDEX_TIME, {alert_id: timestamp_score})
|
||||
|
||||
# Add to wallet index
|
||||
wallet_index_key = f"{self.KEY_INDEX_WALLET}{assessment.wallet_address}"
|
||||
pipe.zadd(wallet_index_key, {alert_id: timestamp_score})
|
||||
pipe.expire(wallet_index_key, self._retention_ttl)
|
||||
|
||||
# Add to market index
|
||||
market_index_key = f"{self.KEY_INDEX_MARKET}{assessment.market_id}"
|
||||
pipe.zadd(market_index_key, {alert_id: timestamp_score})
|
||||
pipe.expire(market_index_key, self._retention_ttl)
|
||||
|
||||
await pipe.execute()
|
||||
|
||||
logger.info(f"Recorded alert {alert_id} for {assessment.wallet_address}")
|
||||
return alert_id
|
||||
|
||||
async def record_feedback(self, alert_id: str, useful: bool) -> bool:
|
||||
"""Record user feedback on alert usefulness.
|
||||
|
||||
Args:
|
||||
alert_id: The alert to provide feedback on.
|
||||
useful: Whether the alert was useful.
|
||||
|
||||
Returns:
|
||||
True if feedback was recorded, False if alert not found.
|
||||
"""
|
||||
alert_redis_key = f"{self.KEY_PREFIX_ALERT}{alert_id}"
|
||||
|
||||
# Get existing record
|
||||
data = await self.redis.get(alert_redis_key)
|
||||
if not data:
|
||||
logger.warning(f"Alert {alert_id} not found for feedback")
|
||||
return False
|
||||
|
||||
# Update record
|
||||
record_dict = json.loads(data)
|
||||
record_dict["feedback_useful"] = useful
|
||||
|
||||
# Get remaining TTL
|
||||
ttl = await self.redis.ttl(alert_redis_key)
|
||||
if ttl < 0:
|
||||
ttl = self._retention_ttl
|
||||
|
||||
# Store updated record
|
||||
await self.redis.set(alert_redis_key, json.dumps(record_dict), ex=ttl)
|
||||
logger.info(f"Recorded feedback for alert {alert_id}: useful={useful}")
|
||||
return True
|
||||
|
||||
async def get_alert(self, alert_id: str) -> AlertRecord | None:
|
||||
"""Get a specific alert record.
|
||||
|
||||
Args:
|
||||
alert_id: The alert ID to retrieve.
|
||||
|
||||
Returns:
|
||||
AlertRecord if found, None otherwise.
|
||||
"""
|
||||
alert_redis_key = f"{self.KEY_PREFIX_ALERT}{alert_id}"
|
||||
data = await self.redis.get(alert_redis_key)
|
||||
if not data:
|
||||
return None
|
||||
return AlertRecord.from_dict(json.loads(data))
|
||||
|
||||
async def get_alerts(
|
||||
self,
|
||||
start: datetime,
|
||||
end: datetime,
|
||||
wallet: str | None = None,
|
||||
market: str | None = None,
|
||||
limit: int = 100,
|
||||
) -> list[AlertRecord]:
|
||||
"""Query alert history.
|
||||
|
||||
Args:
|
||||
start: Start of time range.
|
||||
end: End of time range.
|
||||
wallet: Optional wallet address filter.
|
||||
market: Optional market ID filter.
|
||||
limit: Maximum number of results.
|
||||
|
||||
Returns:
|
||||
List of matching AlertRecord objects.
|
||||
"""
|
||||
start_score = start.timestamp()
|
||||
end_score = end.timestamp()
|
||||
|
||||
# Determine which index to use
|
||||
if wallet:
|
||||
index_key = f"{self.KEY_INDEX_WALLET}{wallet}"
|
||||
elif market:
|
||||
index_key = f"{self.KEY_INDEX_MARKET}{market}"
|
||||
else:
|
||||
index_key = self.KEY_INDEX_TIME
|
||||
|
||||
# Get alert IDs from index
|
||||
alert_ids = await self.redis.zrangebyscore(
|
||||
index_key,
|
||||
start_score,
|
||||
end_score,
|
||||
start=0,
|
||||
num=limit,
|
||||
)
|
||||
|
||||
if not alert_ids:
|
||||
return []
|
||||
|
||||
# Fetch all records
|
||||
records = []
|
||||
for alert_id in alert_ids:
|
||||
if isinstance(alert_id, bytes):
|
||||
alert_id = alert_id.decode()
|
||||
record = await self.get_alert(alert_id)
|
||||
if record:
|
||||
# Apply additional filters if needed
|
||||
if wallet and record.wallet_address != wallet:
|
||||
continue
|
||||
if market and record.market_id != market:
|
||||
continue
|
||||
records.append(record)
|
||||
|
||||
return records
|
||||
|
||||
async def get_recent_count(
|
||||
self,
|
||||
hours: int = 24,
|
||||
wallet: str | None = None,
|
||||
) -> int:
|
||||
"""Get count of alerts in recent hours.
|
||||
|
||||
Args:
|
||||
hours: Number of hours to look back.
|
||||
wallet: Optional wallet address filter.
|
||||
|
||||
Returns:
|
||||
Number of alerts in time period.
|
||||
"""
|
||||
end = datetime.now(UTC)
|
||||
start = end - timedelta(hours=hours)
|
||||
|
||||
index_key = f"{self.KEY_INDEX_WALLET}{wallet}" if wallet else self.KEY_INDEX_TIME
|
||||
|
||||
count = await self.redis.zcount(
|
||||
index_key,
|
||||
start.timestamp(),
|
||||
end.timestamp(),
|
||||
)
|
||||
return int(count)
|
||||
|
||||
async def cleanup_old_alerts(self) -> int:
|
||||
"""Remove alerts older than retention period.
|
||||
|
||||
Returns:
|
||||
Number of alerts removed.
|
||||
"""
|
||||
cutoff = datetime.now(UTC) - timedelta(days=self.retention_days)
|
||||
cutoff_score = cutoff.timestamp()
|
||||
|
||||
# Get old alert IDs
|
||||
old_ids = await self.redis.zrangebyscore(
|
||||
self.KEY_INDEX_TIME,
|
||||
"-inf",
|
||||
cutoff_score,
|
||||
)
|
||||
|
||||
if not old_ids:
|
||||
return 0
|
||||
|
||||
# Remove from time index
|
||||
removed = await self.redis.zremrangebyscore(
|
||||
self.KEY_INDEX_TIME,
|
||||
"-inf",
|
||||
cutoff_score,
|
||||
)
|
||||
|
||||
# Note: Individual alert records will expire via TTL
|
||||
# Wallet/market indexes will also expire via TTL
|
||||
logger.info(f"Cleaned up {removed} old alert references")
|
||||
return int(removed)
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Data models for the alerter module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FormattedAlert:
|
||||
"""A formatted alert message ready for delivery across multiple channels.
|
||||
|
||||
Attributes:
|
||||
title: Short alert title/headline.
|
||||
body: Main alert body text.
|
||||
discord_embed: Discord-optimized embed dictionary.
|
||||
telegram_markdown: Telegram-formatted markdown string.
|
||||
plain_text: Plain text fallback for other channels.
|
||||
links: Dictionary of relevant links (e.g., market, wallet explorer).
|
||||
"""
|
||||
|
||||
title: str
|
||||
body: str
|
||||
discord_embed: dict[str, object]
|
||||
telegram_markdown: str
|
||||
plain_text: str
|
||||
links: dict[str, str] = field(default_factory=dict)
|
||||
@@ -0,0 +1,263 @@
|
||||
"""Configuration management service with Pydantic Settings.
|
||||
|
||||
This module provides centralized configuration management for the
|
||||
Polymarket Insider Tracker application, loading and validating
|
||||
environment variables at startup.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field, SecretStr, field_validator
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class DatabaseSettings(BaseSettings):
|
||||
"""Database connection settings."""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="")
|
||||
|
||||
url: str = Field(
|
||||
alias="DATABASE_URL",
|
||||
description="PostgreSQL connection string",
|
||||
)
|
||||
|
||||
@field_validator("url")
|
||||
@classmethod
|
||||
def validate_url(cls, v: str) -> str:
|
||||
"""Validate database URL format."""
|
||||
if not v.startswith(("postgresql://", "postgresql+asyncpg://")):
|
||||
raise ValueError("DATABASE_URL must be a PostgreSQL connection string")
|
||||
return v
|
||||
|
||||
|
||||
class RedisSettings(BaseSettings):
|
||||
"""Redis connection settings."""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="")
|
||||
|
||||
url: str = Field(
|
||||
default="redis://localhost:6379",
|
||||
alias="REDIS_URL",
|
||||
description="Redis connection string",
|
||||
)
|
||||
|
||||
@field_validator("url")
|
||||
@classmethod
|
||||
def validate_url(cls, v: str) -> str:
|
||||
"""Validate Redis URL format."""
|
||||
if not v.startswith("redis://"):
|
||||
raise ValueError("REDIS_URL must start with redis://")
|
||||
return v
|
||||
|
||||
|
||||
class PolygonSettings(BaseSettings):
|
||||
"""Polygon blockchain RPC settings."""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="POLYGON_")
|
||||
|
||||
rpc_url: str = Field(
|
||||
default="https://polygon-rpc.com",
|
||||
alias="POLYGON_RPC_URL",
|
||||
description="Primary Polygon RPC endpoint",
|
||||
)
|
||||
fallback_rpc_url: str | None = Field(
|
||||
default=None,
|
||||
alias="POLYGON_FALLBACK_RPC_URL",
|
||||
description="Fallback Polygon RPC endpoint",
|
||||
)
|
||||
|
||||
@field_validator("rpc_url", "fallback_rpc_url")
|
||||
@classmethod
|
||||
def validate_url(cls, v: str | None) -> str | None:
|
||||
"""Validate RPC URL format."""
|
||||
if v is None:
|
||||
return v
|
||||
if not v.startswith(("http://", "https://")):
|
||||
raise ValueError("RPC URL must be an HTTP(S) endpoint")
|
||||
return v
|
||||
|
||||
|
||||
class PolymarketSettings(BaseSettings):
|
||||
"""Polymarket API settings."""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="POLYMARKET_")
|
||||
|
||||
ws_url: str = Field(
|
||||
default="wss://ws-subscriptions-clob.polymarket.com/ws/market",
|
||||
alias="POLYMARKET_WS_URL",
|
||||
description="Polymarket WebSocket URL for live data",
|
||||
)
|
||||
api_key: SecretStr | None = Field(
|
||||
default=None,
|
||||
alias="POLYMARKET_API_KEY",
|
||||
description="Optional Polymarket API key",
|
||||
)
|
||||
|
||||
@field_validator("ws_url")
|
||||
@classmethod
|
||||
def validate_ws_url(cls, v: str) -> str:
|
||||
"""Validate WebSocket URL format."""
|
||||
if not v.startswith(("ws://", "wss://")):
|
||||
raise ValueError("WebSocket URL must start with ws:// or wss://")
|
||||
return v
|
||||
|
||||
|
||||
class DiscordSettings(BaseSettings):
|
||||
"""Discord notification settings."""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="DISCORD_")
|
||||
|
||||
webhook_url: SecretStr | None = Field(
|
||||
default=None,
|
||||
alias="DISCORD_WEBHOOK_URL",
|
||||
description="Discord webhook URL for alerts",
|
||||
)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
"""Check if Discord notifications are enabled."""
|
||||
return self.webhook_url is not None
|
||||
|
||||
|
||||
class TelegramSettings(BaseSettings):
|
||||
"""Telegram notification settings."""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="TELEGRAM_")
|
||||
|
||||
bot_token: SecretStr | None = Field(
|
||||
default=None,
|
||||
alias="TELEGRAM_BOT_TOKEN",
|
||||
description="Telegram bot token",
|
||||
)
|
||||
chat_id: str | None = Field(
|
||||
default=None,
|
||||
alias="TELEGRAM_CHAT_ID",
|
||||
description="Telegram chat ID for alerts",
|
||||
)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
"""Check if Telegram notifications are enabled."""
|
||||
return self.bot_token is not None and self.chat_id is not None
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
"""Main application settings.
|
||||
|
||||
Loads configuration from environment variables with support for
|
||||
.env files via python-dotenv.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from polymarket_insider_tracker.config import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
print(settings.database.url)
|
||||
print(settings.log_level)
|
||||
```
|
||||
"""
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=".env",
|
||||
env_file_encoding="utf-8",
|
||||
extra="ignore",
|
||||
)
|
||||
|
||||
# Nested configuration groups
|
||||
database: DatabaseSettings = Field(default_factory=DatabaseSettings)
|
||||
redis: RedisSettings = Field(default_factory=RedisSettings)
|
||||
polygon: PolygonSettings = Field(default_factory=PolygonSettings)
|
||||
polymarket: PolymarketSettings = Field(default_factory=PolymarketSettings)
|
||||
discord: DiscordSettings = Field(default_factory=DiscordSettings)
|
||||
telegram: TelegramSettings = Field(default_factory=TelegramSettings)
|
||||
|
||||
# Application settings
|
||||
log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = Field(
|
||||
default="INFO",
|
||||
alias="LOG_LEVEL",
|
||||
description="Logging level",
|
||||
)
|
||||
health_port: int = Field(
|
||||
default=8080,
|
||||
alias="HEALTH_PORT",
|
||||
description="HTTP port for health check endpoints",
|
||||
ge=1,
|
||||
le=65535,
|
||||
)
|
||||
dry_run: bool = Field(
|
||||
default=False,
|
||||
alias="DRY_RUN",
|
||||
description="Run without sending actual alerts",
|
||||
)
|
||||
|
||||
def get_logging_level(self) -> int:
|
||||
"""Get the numeric logging level."""
|
||||
level: int = getattr(logging, self.log_level)
|
||||
return level
|
||||
|
||||
def redacted_summary(self) -> dict[str, str | dict[str, str]]:
|
||||
"""Get a summary of settings with secrets redacted.
|
||||
|
||||
Returns:
|
||||
Dictionary of settings with sensitive values masked.
|
||||
"""
|
||||
return {
|
||||
"database_url": self._redact_url(self.database.url),
|
||||
"redis_url": self._redact_url(self.redis.url),
|
||||
"polygon": {
|
||||
"rpc_url": self.polygon.rpc_url,
|
||||
"fallback_rpc_url": self.polygon.fallback_rpc_url or "(not set)",
|
||||
},
|
||||
"polymarket": {
|
||||
"ws_url": self.polymarket.ws_url,
|
||||
"api_key": "(set)" if self.polymarket.api_key else "(not set)",
|
||||
},
|
||||
"discord_enabled": str(self.discord.enabled),
|
||||
"telegram_enabled": str(self.telegram.enabled),
|
||||
"log_level": self.log_level,
|
||||
"health_port": str(self.health_port),
|
||||
"dry_run": str(self.dry_run),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _redact_url(url: str) -> str:
|
||||
"""Redact password from URL if present."""
|
||||
if "@" in url and "://" in url:
|
||||
# URL has credentials - redact the password
|
||||
protocol_end = url.index("://") + 3
|
||||
at_pos = url.index("@")
|
||||
creds_part = url[protocol_end:at_pos]
|
||||
if ":" in creds_part:
|
||||
username = creds_part.split(":")[0]
|
||||
return f"{url[:protocol_end]}{username}:***@{url[at_pos + 1 :]}"
|
||||
return url
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_settings() -> Settings:
|
||||
"""Get the application settings singleton.
|
||||
|
||||
Uses LRU cache to ensure settings are loaded only once and
|
||||
reused across the application.
|
||||
|
||||
Returns:
|
||||
The Settings instance.
|
||||
|
||||
Raises:
|
||||
ValidationError: If required environment variables are missing
|
||||
or have invalid values.
|
||||
"""
|
||||
return Settings()
|
||||
|
||||
|
||||
def clear_settings_cache() -> None:
|
||||
"""Clear the settings cache.
|
||||
|
||||
Useful for testing when you need to reload settings with
|
||||
different environment variables.
|
||||
"""
|
||||
get_settings.cache_clear()
|
||||
@@ -0,0 +1,24 @@
|
||||
"""Anomaly detection layer - Suspicious activity identification."""
|
||||
|
||||
from polymarket_insider_tracker.detector.fresh_wallet import FreshWalletDetector
|
||||
from polymarket_insider_tracker.detector.models import (
|
||||
FreshWalletSignal,
|
||||
RiskAssessment,
|
||||
SizeAnomalySignal,
|
||||
SniperClusterSignal,
|
||||
)
|
||||
from polymarket_insider_tracker.detector.scorer import RiskScorer, SignalBundle
|
||||
from polymarket_insider_tracker.detector.size_anomaly import SizeAnomalyDetector
|
||||
from polymarket_insider_tracker.detector.sniper import SniperDetector
|
||||
|
||||
__all__ = [
|
||||
"FreshWalletDetector",
|
||||
"FreshWalletSignal",
|
||||
"RiskAssessment",
|
||||
"RiskScorer",
|
||||
"SignalBundle",
|
||||
"SizeAnomalyDetector",
|
||||
"SizeAnomalySignal",
|
||||
"SniperClusterSignal",
|
||||
"SniperDetector",
|
||||
]
|
||||
@@ -0,0 +1,238 @@
|
||||
"""Fresh wallet detection algorithm.
|
||||
|
||||
This module provides the FreshWalletDetector class that identifies trades
|
||||
from fresh wallets and generates alert signals with confidence scores.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from decimal import Decimal
|
||||
|
||||
from polymarket_insider_tracker.detector.models import FreshWalletSignal
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
from polymarket_insider_tracker.profiler.analyzer import WalletAnalyzer
|
||||
from polymarket_insider_tracker.profiler.models import WalletProfile
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_MIN_TRADE_SIZE = Decimal("1000") # $1,000 minimum trade size
|
||||
DEFAULT_MAX_NONCE = 5 # Max nonce to be considered fresh
|
||||
DEFAULT_MAX_AGE_HOURS = 48.0 # Max age in hours to be considered fresh
|
||||
|
||||
# Confidence scoring constants
|
||||
BASE_CONFIDENCE = 0.5
|
||||
BRAND_NEW_BONUS = 0.2 # nonce == 0
|
||||
VERY_YOUNG_BONUS = 0.1 # age < 2 hours
|
||||
LARGE_TRADE_BONUS = 0.1 # trade size > $10,000
|
||||
LARGE_TRADE_THRESHOLD = Decimal("10000")
|
||||
|
||||
|
||||
class FreshWalletDetector:
|
||||
"""Detector for fresh wallet trading patterns.
|
||||
|
||||
This detector analyzes trade events for fresh wallet signals. A trade
|
||||
is flagged as suspicious if:
|
||||
1. The wallet meets freshness criteria (low nonce, recent activity)
|
||||
2. The trade size meets minimum threshold
|
||||
|
||||
The detector produces confidence scores based on multiple factors:
|
||||
- Wallet nonce (brand new = higher confidence)
|
||||
- Wallet age (very young = higher confidence)
|
||||
- Trade size (larger = higher confidence)
|
||||
|
||||
Example:
|
||||
```python
|
||||
analyzer = WalletAnalyzer(polygon_client, redis=redis)
|
||||
detector = FreshWalletDetector(analyzer)
|
||||
|
||||
# Analyze a single trade
|
||||
signal = await detector.analyze(trade_event)
|
||||
if signal is not None:
|
||||
print(f"Fresh wallet detected! Confidence: {signal.confidence}")
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
wallet_analyzer: WalletAnalyzer,
|
||||
*,
|
||||
min_trade_size: Decimal = DEFAULT_MIN_TRADE_SIZE,
|
||||
max_nonce: int = DEFAULT_MAX_NONCE,
|
||||
max_age_hours: float = DEFAULT_MAX_AGE_HOURS,
|
||||
) -> None:
|
||||
"""Initialize the fresh wallet detector.
|
||||
|
||||
Args:
|
||||
wallet_analyzer: WalletAnalyzer instance for wallet profiling.
|
||||
min_trade_size: Minimum trade size to analyze (default $1,000).
|
||||
max_nonce: Maximum nonce to consider wallet fresh (default 5).
|
||||
max_age_hours: Maximum age in hours to consider fresh (default 48).
|
||||
"""
|
||||
self._analyzer = wallet_analyzer
|
||||
self._min_trade_size = min_trade_size
|
||||
self._max_nonce = max_nonce
|
||||
self._max_age_hours = max_age_hours
|
||||
|
||||
async def analyze(self, trade: TradeEvent) -> FreshWalletSignal | None:
|
||||
"""Analyze a trade event for fresh wallet signals.
|
||||
|
||||
This method:
|
||||
1. Filters out trades below the minimum size threshold
|
||||
2. Analyzes the trader's wallet profile
|
||||
3. Determines if the wallet is fresh
|
||||
4. Calculates confidence score based on multiple factors
|
||||
|
||||
Args:
|
||||
trade: TradeEvent to analyze.
|
||||
|
||||
Returns:
|
||||
FreshWalletSignal if the trade is from a fresh wallet,
|
||||
None otherwise.
|
||||
"""
|
||||
# Filter by minimum trade size
|
||||
if trade.notional_value < self._min_trade_size:
|
||||
logger.debug(
|
||||
"Trade %s below minimum size: %s < %s",
|
||||
trade.trade_id,
|
||||
trade.notional_value,
|
||||
self._min_trade_size,
|
||||
)
|
||||
return None
|
||||
|
||||
# Get wallet profile
|
||||
try:
|
||||
profile = await self._analyzer.analyze(trade.wallet_address)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to analyze wallet %s for trade %s: %s",
|
||||
trade.wallet_address,
|
||||
trade.trade_id,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
|
||||
# Check if wallet is fresh
|
||||
if not self._is_wallet_fresh(profile):
|
||||
logger.debug(
|
||||
"Wallet %s is not fresh (nonce=%d, age=%s)",
|
||||
trade.wallet_address,
|
||||
profile.nonce,
|
||||
profile.age_hours,
|
||||
)
|
||||
return None
|
||||
|
||||
# Calculate confidence score
|
||||
confidence, factors = self.calculate_confidence(profile, trade)
|
||||
|
||||
logger.info(
|
||||
"Fresh wallet signal: wallet=%s, market=%s, size=%s, confidence=%.2f",
|
||||
trade.wallet_address[:10] + "...",
|
||||
trade.market_id[:10] + "...",
|
||||
trade.notional_value,
|
||||
confidence,
|
||||
)
|
||||
|
||||
return FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=profile,
|
||||
confidence=confidence,
|
||||
factors=factors,
|
||||
)
|
||||
|
||||
def _is_wallet_fresh(self, profile: WalletProfile) -> bool:
|
||||
"""Check if wallet meets freshness criteria.
|
||||
|
||||
A wallet is considered fresh if:
|
||||
1. Nonce is at or below max_nonce threshold
|
||||
2. Age is unknown OR within max_age_hours
|
||||
|
||||
Args:
|
||||
profile: Wallet profile to check.
|
||||
|
||||
Returns:
|
||||
True if wallet is fresh.
|
||||
"""
|
||||
# Must have few transactions
|
||||
if profile.nonce > self._max_nonce:
|
||||
return False
|
||||
|
||||
# If age is known, must be recent
|
||||
return not (profile.age_hours is not None and profile.age_hours > self._max_age_hours)
|
||||
|
||||
def calculate_confidence(
|
||||
self,
|
||||
profile: WalletProfile,
|
||||
trade: TradeEvent,
|
||||
) -> tuple[float, dict[str, float]]:
|
||||
"""Calculate confidence score based on multiple factors.
|
||||
|
||||
Confidence scoring:
|
||||
- Base: 0.5 (fresh wallet detected)
|
||||
- +0.2 if nonce == 0 (brand new wallet)
|
||||
- +0.1 if age < 2 hours (very young)
|
||||
- +0.1 if trade size > $10,000 (large trade)
|
||||
|
||||
Final confidence is clamped to [0.0, 1.0].
|
||||
|
||||
Args:
|
||||
profile: Wallet profile with nonce and age data.
|
||||
trade: Trade event with size data.
|
||||
|
||||
Returns:
|
||||
Tuple of (confidence_score, factors_dict).
|
||||
"""
|
||||
factors: dict[str, float] = {"base": BASE_CONFIDENCE}
|
||||
confidence = BASE_CONFIDENCE
|
||||
|
||||
# Brand new wallet bonus
|
||||
if profile.nonce == 0:
|
||||
factors["brand_new"] = BRAND_NEW_BONUS
|
||||
confidence += BRAND_NEW_BONUS
|
||||
|
||||
# Very young wallet bonus
|
||||
if profile.age_hours is not None and profile.age_hours < 2.0:
|
||||
factors["very_young"] = VERY_YOUNG_BONUS
|
||||
confidence += VERY_YOUNG_BONUS
|
||||
|
||||
# Large trade bonus
|
||||
if trade.notional_value > LARGE_TRADE_THRESHOLD:
|
||||
factors["large_trade"] = LARGE_TRADE_BONUS
|
||||
confidence += LARGE_TRADE_BONUS
|
||||
|
||||
# Clamp to valid range
|
||||
confidence = max(0.0, min(1.0, confidence))
|
||||
|
||||
return confidence, factors
|
||||
|
||||
async def analyze_batch(
|
||||
self,
|
||||
trades: list[TradeEvent],
|
||||
) -> list[FreshWalletSignal]:
|
||||
"""Analyze multiple trades for fresh wallet signals.
|
||||
|
||||
Processes trades in parallel for efficiency.
|
||||
|
||||
Args:
|
||||
trades: List of trades to analyze.
|
||||
|
||||
Returns:
|
||||
List of FreshWalletSignal for trades from fresh wallets.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
tasks = [self.analyze(trade) for trade in trades]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
signals: list[FreshWalletSignal] = []
|
||||
for trade, result in zip(trades, results, strict=True):
|
||||
if isinstance(result, BaseException):
|
||||
logger.warning(
|
||||
"Failed to analyze trade %s: %s",
|
||||
trade.trade_id,
|
||||
result,
|
||||
)
|
||||
continue
|
||||
if result is not None:
|
||||
signals.append(result)
|
||||
|
||||
return signals
|
||||
@@ -0,0 +1,274 @@
|
||||
"""Data models for the detector module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
|
||||
from polymarket_insider_tracker.ingestor.models import MarketMetadata, TradeEvent
|
||||
from polymarket_insider_tracker.profiler.models import WalletProfile
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FreshWalletSignal:
|
||||
"""Signal emitted when a fresh wallet makes a suspicious trade.
|
||||
|
||||
This signal combines trade event data with wallet profile analysis
|
||||
to produce a confidence score indicating the likelihood of suspicious
|
||||
activity.
|
||||
|
||||
Attributes:
|
||||
trade_event: The original trade event that triggered this signal.
|
||||
wallet_profile: Analyzed profile of the trader's wallet.
|
||||
confidence: Overall confidence score (0.0 to 1.0).
|
||||
factors: Individual factor scores contributing to confidence.
|
||||
timestamp: When this signal was generated.
|
||||
"""
|
||||
|
||||
trade_event: TradeEvent
|
||||
wallet_profile: WalletProfile
|
||||
confidence: float
|
||||
factors: dict[str, float]
|
||||
timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
@property
|
||||
def wallet_address(self) -> str:
|
||||
"""Return the wallet address from the trade event."""
|
||||
return self.trade_event.wallet_address
|
||||
|
||||
@property
|
||||
def market_id(self) -> str:
|
||||
"""Return the market ID from the trade event."""
|
||||
return self.trade_event.market_id
|
||||
|
||||
@property
|
||||
def trade_size_usdc(self) -> Decimal:
|
||||
"""Return the trade size in USDC (notional value)."""
|
||||
return self.trade_event.notional_value
|
||||
|
||||
@property
|
||||
def is_high_confidence(self) -> bool:
|
||||
"""Return True if confidence exceeds 0.7."""
|
||||
return self.confidence >= 0.7
|
||||
|
||||
@property
|
||||
def is_very_high_confidence(self) -> bool:
|
||||
"""Return True if confidence exceeds 0.85."""
|
||||
return self.confidence >= 0.85
|
||||
|
||||
def to_dict(self) -> dict[str, object]:
|
||||
"""Serialize to dictionary for Redis stream publishing."""
|
||||
return {
|
||||
"wallet_address": self.wallet_address,
|
||||
"market_id": self.market_id,
|
||||
"trade_id": self.trade_event.trade_id,
|
||||
"trade_size": str(self.trade_size_usdc),
|
||||
"trade_side": self.trade_event.side,
|
||||
"trade_price": str(self.trade_event.price),
|
||||
"wallet_nonce": self.wallet_profile.nonce,
|
||||
"wallet_age_hours": self.wallet_profile.age_hours,
|
||||
"wallet_is_fresh": self.wallet_profile.is_fresh,
|
||||
"confidence": self.confidence,
|
||||
"factors": self.factors,
|
||||
"timestamp": self.timestamp.isoformat(),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SizeAnomalySignal:
|
||||
"""Signal emitted when a trade has unusually large position size.
|
||||
|
||||
This signal is generated when a trade's size significantly impacts
|
||||
the market volume or order book depth, indicating potential informed
|
||||
trading activity.
|
||||
|
||||
Attributes:
|
||||
trade_event: The original trade event that triggered this signal.
|
||||
market_metadata: Metadata about the market being traded.
|
||||
volume_impact: Trade size as fraction of 24h volume (0.0 if unknown).
|
||||
book_impact: Trade size as fraction of order book depth (0.0 if unknown).
|
||||
is_niche_market: Whether the market is considered niche/low-volume.
|
||||
confidence: Overall confidence score (0.0 to 1.0).
|
||||
factors: Individual factor scores contributing to confidence.
|
||||
timestamp: When this signal was generated.
|
||||
"""
|
||||
|
||||
trade_event: TradeEvent
|
||||
market_metadata: MarketMetadata
|
||||
volume_impact: float
|
||||
book_impact: float
|
||||
is_niche_market: bool
|
||||
confidence: float
|
||||
factors: dict[str, float]
|
||||
timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
@property
|
||||
def wallet_address(self) -> str:
|
||||
"""Return the wallet address from the trade event."""
|
||||
return self.trade_event.wallet_address
|
||||
|
||||
@property
|
||||
def market_id(self) -> str:
|
||||
"""Return the market ID from the trade event."""
|
||||
return self.trade_event.market_id
|
||||
|
||||
@property
|
||||
def trade_size_usdc(self) -> Decimal:
|
||||
"""Return the trade size in USDC (notional value)."""
|
||||
return self.trade_event.notional_value
|
||||
|
||||
@property
|
||||
def is_high_confidence(self) -> bool:
|
||||
"""Return True if confidence exceeds 0.7."""
|
||||
return self.confidence >= 0.7
|
||||
|
||||
@property
|
||||
def is_very_high_confidence(self) -> bool:
|
||||
"""Return True if confidence exceeds 0.85."""
|
||||
return self.confidence >= 0.85
|
||||
|
||||
def to_dict(self) -> dict[str, object]:
|
||||
"""Serialize to dictionary for Redis stream publishing."""
|
||||
return {
|
||||
"wallet_address": self.wallet_address,
|
||||
"market_id": self.market_id,
|
||||
"trade_id": self.trade_event.trade_id,
|
||||
"trade_size": str(self.trade_size_usdc),
|
||||
"trade_side": self.trade_event.side,
|
||||
"trade_price": str(self.trade_event.price),
|
||||
"market_category": self.market_metadata.category,
|
||||
"volume_impact": self.volume_impact,
|
||||
"book_impact": self.book_impact,
|
||||
"is_niche_market": self.is_niche_market,
|
||||
"confidence": self.confidence,
|
||||
"factors": self.factors,
|
||||
"timestamp": self.timestamp.isoformat(),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SniperClusterSignal:
|
||||
"""Signal emitted when a wallet is identified as part of a sniper cluster.
|
||||
|
||||
Sniper clusters are groups of wallets that consistently enter markets
|
||||
within minutes of their creation, suggesting coordinated insider activity.
|
||||
|
||||
Attributes:
|
||||
wallet_address: The wallet identified as a sniper.
|
||||
cluster_id: Unique identifier for this cluster.
|
||||
cluster_size: Number of wallets in the cluster.
|
||||
avg_entry_delta_seconds: Average time (seconds) from market creation to entry.
|
||||
markets_in_common: Number of markets where cluster members overlap.
|
||||
confidence: Confidence score (0.0 to 1.0) based on clustering strength.
|
||||
timestamp: When this signal was generated.
|
||||
"""
|
||||
|
||||
wallet_address: str
|
||||
cluster_id: str
|
||||
cluster_size: int
|
||||
avg_entry_delta_seconds: float
|
||||
markets_in_common: int
|
||||
confidence: float
|
||||
timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
@property
|
||||
def is_high_confidence(self) -> bool:
|
||||
"""Return True if confidence exceeds 0.7."""
|
||||
return self.confidence >= 0.7
|
||||
|
||||
@property
|
||||
def is_very_high_confidence(self) -> bool:
|
||||
"""Return True if confidence exceeds 0.85."""
|
||||
return self.confidence >= 0.85
|
||||
|
||||
def to_dict(self) -> dict[str, object]:
|
||||
"""Serialize to dictionary for Redis stream publishing."""
|
||||
return {
|
||||
"wallet_address": self.wallet_address,
|
||||
"cluster_id": self.cluster_id,
|
||||
"cluster_size": self.cluster_size,
|
||||
"avg_entry_delta_seconds": self.avg_entry_delta_seconds,
|
||||
"markets_in_common": self.markets_in_common,
|
||||
"confidence": self.confidence,
|
||||
"timestamp": self.timestamp.isoformat(),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RiskAssessment:
|
||||
"""Combined risk assessment aggregating all signal types.
|
||||
|
||||
This represents the final scoring output that determines whether
|
||||
a trade should trigger an alert, combining signals from multiple
|
||||
detectors with configurable weights.
|
||||
|
||||
Attributes:
|
||||
trade_event: The original trade event being assessed.
|
||||
wallet_address: The trader's wallet address.
|
||||
market_id: The market condition ID.
|
||||
fresh_wallet_signal: Signal from fresh wallet detector, if triggered.
|
||||
size_anomaly_signal: Signal from size anomaly detector, if triggered.
|
||||
signals_triggered: Count of how many signal types fired.
|
||||
weighted_score: Final weighted combination of all signals (0.0 to 1.0).
|
||||
should_alert: Whether this assessment meets alert threshold.
|
||||
assessment_id: Unique identifier for this assessment.
|
||||
timestamp: When this assessment was generated.
|
||||
"""
|
||||
|
||||
trade_event: TradeEvent
|
||||
wallet_address: str
|
||||
market_id: str
|
||||
|
||||
# Individual signals (None if not triggered)
|
||||
fresh_wallet_signal: FreshWalletSignal | None
|
||||
size_anomaly_signal: SizeAnomalySignal | None
|
||||
|
||||
# Combined scoring
|
||||
signals_triggered: int
|
||||
weighted_score: float
|
||||
should_alert: bool
|
||||
|
||||
# Metadata
|
||||
assessment_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
@property
|
||||
def is_high_risk(self) -> bool:
|
||||
"""Return True if weighted score exceeds 0.7."""
|
||||
return self.weighted_score >= 0.7
|
||||
|
||||
@property
|
||||
def is_very_high_risk(self) -> bool:
|
||||
"""Return True if weighted score exceeds 0.85."""
|
||||
return self.weighted_score >= 0.85
|
||||
|
||||
@property
|
||||
def trade_size_usdc(self) -> Decimal:
|
||||
"""Return the trade size in USDC (notional value)."""
|
||||
return self.trade_event.notional_value
|
||||
|
||||
def to_dict(self) -> dict[str, object]:
|
||||
"""Serialize to dictionary for Redis stream publishing."""
|
||||
return {
|
||||
"assessment_id": self.assessment_id,
|
||||
"wallet_address": self.wallet_address,
|
||||
"market_id": self.market_id,
|
||||
"trade_id": self.trade_event.trade_id,
|
||||
"trade_size": str(self.trade_size_usdc),
|
||||
"trade_side": self.trade_event.side,
|
||||
"trade_price": str(self.trade_event.price),
|
||||
"signals_triggered": self.signals_triggered,
|
||||
"weighted_score": self.weighted_score,
|
||||
"should_alert": self.should_alert,
|
||||
"has_fresh_wallet_signal": self.fresh_wallet_signal is not None,
|
||||
"has_size_anomaly_signal": self.size_anomaly_signal is not None,
|
||||
"fresh_wallet_confidence": (
|
||||
self.fresh_wallet_signal.confidence if self.fresh_wallet_signal else None
|
||||
),
|
||||
"size_anomaly_confidence": (
|
||||
self.size_anomaly_signal.confidence if self.size_anomaly_signal else None
|
||||
),
|
||||
"timestamp": self.timestamp.isoformat(),
|
||||
}
|
||||
@@ -0,0 +1,304 @@
|
||||
"""Composite risk scorer combining all detector signals.
|
||||
|
||||
This module provides the RiskScorer class that aggregates signals from
|
||||
multiple detectors into a unified risk assessment with weighted scoring.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from polymarket_insider_tracker.detector.models import (
|
||||
FreshWalletSignal,
|
||||
RiskAssessment,
|
||||
SizeAnomalySignal,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_ALERT_THRESHOLD = 0.6
|
||||
DEFAULT_DEDUP_WINDOW_SECONDS = 3600 # 1 hour
|
||||
DEFAULT_REDIS_KEY_PREFIX = "polymarket:dedup:"
|
||||
|
||||
# Default weights for each signal type
|
||||
DEFAULT_WEIGHTS = {
|
||||
"fresh_wallet": 0.40,
|
||||
"size_anomaly": 0.35,
|
||||
"niche_market": 0.25,
|
||||
}
|
||||
|
||||
# Multi-signal bonuses
|
||||
MULTI_SIGNAL_BONUS_2 = 1.2 # 20% bonus for 2 signals
|
||||
MULTI_SIGNAL_BONUS_3 = 1.3 # 30% bonus for 3+ signals
|
||||
|
||||
|
||||
@dataclass
|
||||
class SignalBundle:
|
||||
"""Bundle of signals for a single trade.
|
||||
|
||||
Collects all available signals for a trade event to pass to the scorer.
|
||||
"""
|
||||
|
||||
trade_event: TradeEvent
|
||||
fresh_wallet_signal: FreshWalletSignal | None = None
|
||||
size_anomaly_signal: SizeAnomalySignal | None = None
|
||||
|
||||
@property
|
||||
def wallet_address(self) -> str:
|
||||
"""Return the wallet address from the trade event."""
|
||||
return self.trade_event.wallet_address
|
||||
|
||||
@property
|
||||
def market_id(self) -> str:
|
||||
"""Return the market ID from the trade event."""
|
||||
return self.trade_event.market_id
|
||||
|
||||
|
||||
class RiskScorer:
|
||||
"""Composite risk scorer combining signals into unified assessments.
|
||||
|
||||
This scorer:
|
||||
- Aggregates signals from multiple detectors for the same trade
|
||||
- Applies configurable weights based on signal type
|
||||
- Calculates multi-signal bonuses for correlated signals
|
||||
- Enforces deduplication to prevent alert spam
|
||||
- Produces RiskAssessment objects for downstream alerting
|
||||
|
||||
Scoring Formula:
|
||||
weighted_score = sum(signal.confidence * weight[type] for signal in signals)
|
||||
|
||||
# Multi-signal bonus
|
||||
if signals >= 2: weighted_score *= 1.2
|
||||
if signals >= 3: weighted_score *= 1.3
|
||||
|
||||
# Cap at 1.0
|
||||
final_score = min(weighted_score, 1.0)
|
||||
|
||||
should_alert = final_score >= alert_threshold AND not deduplicated
|
||||
|
||||
Example:
|
||||
```python
|
||||
redis = Redis.from_url("redis://localhost:6379")
|
||||
scorer = RiskScorer(redis)
|
||||
|
||||
bundle = SignalBundle(
|
||||
trade_event=trade,
|
||||
fresh_wallet_signal=fresh_signal,
|
||||
size_anomaly_signal=size_signal,
|
||||
)
|
||||
|
||||
assessment = await scorer.assess(bundle)
|
||||
if assessment.should_alert:
|
||||
await send_alert(assessment)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis: Redis,
|
||||
*,
|
||||
weights: dict[str, float] | None = None,
|
||||
alert_threshold: float = DEFAULT_ALERT_THRESHOLD,
|
||||
dedup_window_seconds: int = DEFAULT_DEDUP_WINDOW_SECONDS,
|
||||
key_prefix: str = DEFAULT_REDIS_KEY_PREFIX,
|
||||
) -> None:
|
||||
"""Initialize the risk scorer.
|
||||
|
||||
Args:
|
||||
redis: Redis async client for deduplication.
|
||||
weights: Custom weights for signal types. Defaults to DEFAULT_WEIGHTS.
|
||||
alert_threshold: Minimum score to trigger alert (default 0.6).
|
||||
dedup_window_seconds: Window for deduplication (default 3600 = 1 hour).
|
||||
key_prefix: Redis key prefix for dedup keys.
|
||||
"""
|
||||
self._redis = redis
|
||||
self._weights = weights or DEFAULT_WEIGHTS.copy()
|
||||
self._alert_threshold = alert_threshold
|
||||
self._dedup_window = dedup_window_seconds
|
||||
self._key_prefix = key_prefix
|
||||
|
||||
async def assess(self, bundle: SignalBundle) -> RiskAssessment:
|
||||
"""Assess a trade's risk based on all available signals.
|
||||
|
||||
This method:
|
||||
1. Counts triggered signals
|
||||
2. Calculates weighted score with bonuses
|
||||
3. Checks deduplication
|
||||
4. Creates RiskAssessment
|
||||
|
||||
Args:
|
||||
bundle: SignalBundle containing trade and all signals.
|
||||
|
||||
Returns:
|
||||
RiskAssessment with final scoring and alert decision.
|
||||
"""
|
||||
# Calculate weighted score
|
||||
weighted_score, signals_triggered = self.calculate_weighted_score(bundle)
|
||||
|
||||
# Determine if should alert (before dedup check)
|
||||
meets_threshold = weighted_score >= self._alert_threshold
|
||||
|
||||
# Check deduplication
|
||||
is_duplicate = False
|
||||
if meets_threshold:
|
||||
is_duplicate = await self._check_and_set_dedup(
|
||||
bundle.wallet_address,
|
||||
bundle.market_id,
|
||||
)
|
||||
|
||||
should_alert = meets_threshold and not is_duplicate
|
||||
|
||||
# Log assessment
|
||||
if should_alert:
|
||||
logger.info(
|
||||
"Risk assessment triggered alert: wallet=%s, market=%s, score=%.2f, signals=%d",
|
||||
bundle.wallet_address[:10] + "...",
|
||||
bundle.market_id[:10] + "...",
|
||||
weighted_score,
|
||||
signals_triggered,
|
||||
)
|
||||
elif is_duplicate:
|
||||
logger.debug(
|
||||
"Risk assessment deduplicated: wallet=%s, market=%s",
|
||||
bundle.wallet_address[:10] + "...",
|
||||
bundle.market_id[:10] + "...",
|
||||
)
|
||||
|
||||
return RiskAssessment(
|
||||
trade_event=bundle.trade_event,
|
||||
wallet_address=bundle.wallet_address,
|
||||
market_id=bundle.market_id,
|
||||
fresh_wallet_signal=bundle.fresh_wallet_signal,
|
||||
size_anomaly_signal=bundle.size_anomaly_signal,
|
||||
signals_triggered=signals_triggered,
|
||||
weighted_score=weighted_score,
|
||||
should_alert=should_alert,
|
||||
)
|
||||
|
||||
def calculate_weighted_score(self, bundle: SignalBundle) -> tuple[float, int]:
|
||||
"""Calculate weighted score from all signals.
|
||||
|
||||
Applies per-signal weights and multi-signal bonuses.
|
||||
|
||||
Args:
|
||||
bundle: SignalBundle with all available signals.
|
||||
|
||||
Returns:
|
||||
Tuple of (weighted_score, signals_triggered_count).
|
||||
"""
|
||||
score = 0.0
|
||||
signals_triggered = 0
|
||||
|
||||
# Fresh wallet signal
|
||||
if bundle.fresh_wallet_signal is not None:
|
||||
weight = self._weights.get("fresh_wallet", 0.0)
|
||||
score += bundle.fresh_wallet_signal.confidence * weight
|
||||
signals_triggered += 1
|
||||
|
||||
# Size anomaly signal
|
||||
if bundle.size_anomaly_signal is not None:
|
||||
weight = self._weights.get("size_anomaly", 0.0)
|
||||
score += bundle.size_anomaly_signal.confidence * weight
|
||||
signals_triggered += 1
|
||||
|
||||
# Additional niche market weight
|
||||
if bundle.size_anomaly_signal.is_niche_market:
|
||||
niche_weight = self._weights.get("niche_market", 0.0)
|
||||
score += bundle.size_anomaly_signal.confidence * niche_weight
|
||||
|
||||
# Apply multi-signal bonus
|
||||
if signals_triggered >= 3:
|
||||
score *= MULTI_SIGNAL_BONUS_3
|
||||
elif signals_triggered >= 2:
|
||||
score *= MULTI_SIGNAL_BONUS_2
|
||||
|
||||
# Cap at 1.0
|
||||
score = min(score, 1.0)
|
||||
|
||||
return score, signals_triggered
|
||||
|
||||
async def _check_and_set_dedup(
|
||||
self,
|
||||
wallet_address: str,
|
||||
market_id: str,
|
||||
) -> bool:
|
||||
"""Check if this wallet/market combo was recently alerted.
|
||||
|
||||
If not a duplicate, sets the dedup key with TTL.
|
||||
|
||||
Args:
|
||||
wallet_address: The trader's wallet address.
|
||||
market_id: The market condition ID.
|
||||
|
||||
Returns:
|
||||
True if this is a duplicate (already alerted), False otherwise.
|
||||
"""
|
||||
key = f"{self._key_prefix}{wallet_address}:{market_id}"
|
||||
|
||||
# Try to set with NX (only if not exists)
|
||||
was_set = await self._redis.set(
|
||||
key,
|
||||
datetime.now(UTC).isoformat(),
|
||||
nx=True,
|
||||
ex=self._dedup_window,
|
||||
)
|
||||
|
||||
# If was_set is None/False, key already existed = duplicate
|
||||
return not was_set
|
||||
|
||||
async def clear_dedup(
|
||||
self,
|
||||
wallet_address: str,
|
||||
market_id: str,
|
||||
) -> bool:
|
||||
"""Clear dedup key for a wallet/market combo.
|
||||
|
||||
Useful for testing or manual override.
|
||||
|
||||
Args:
|
||||
wallet_address: The trader's wallet address.
|
||||
market_id: The market condition ID.
|
||||
|
||||
Returns:
|
||||
True if key was deleted, False if it didn't exist.
|
||||
"""
|
||||
key = f"{self._key_prefix}{wallet_address}:{market_id}"
|
||||
deleted = await self._redis.delete(key)
|
||||
return int(deleted) > 0
|
||||
|
||||
async def assess_batch(self, bundles: list[SignalBundle]) -> list[RiskAssessment]:
|
||||
"""Assess multiple trade bundles.
|
||||
|
||||
Args:
|
||||
bundles: List of SignalBundles to assess.
|
||||
|
||||
Returns:
|
||||
List of RiskAssessments.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
tasks = [self.assess(bundle) for bundle in bundles]
|
||||
return await asyncio.gather(*tasks)
|
||||
|
||||
def get_weights(self) -> dict[str, float]:
|
||||
"""Get current signal weights.
|
||||
|
||||
Returns:
|
||||
Copy of the weights dictionary.
|
||||
"""
|
||||
return self._weights.copy()
|
||||
|
||||
def set_weights(self, weights: dict[str, float]) -> None:
|
||||
"""Update signal weights.
|
||||
|
||||
Useful for A/B testing different weight configurations.
|
||||
|
||||
Args:
|
||||
weights: New weights dictionary.
|
||||
"""
|
||||
self._weights = weights.copy()
|
||||
logger.info("Updated risk scorer weights: %s", self._weights)
|
||||
@@ -0,0 +1,354 @@
|
||||
"""Position size anomaly detection algorithm.
|
||||
|
||||
This module provides the SizeAnomalyDetector class that identifies trades
|
||||
with unusually large position sizes relative to market liquidity.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from decimal import Decimal
|
||||
|
||||
from polymarket_insider_tracker.detector.models import SizeAnomalySignal
|
||||
from polymarket_insider_tracker.ingestor.metadata_sync import MarketMetadataSync
|
||||
from polymarket_insider_tracker.ingestor.models import MarketMetadata, TradeEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_VOLUME_THRESHOLD = 0.02 # 2% of daily volume
|
||||
DEFAULT_BOOK_THRESHOLD = 0.05 # 5% of order book depth
|
||||
DEFAULT_NICHE_VOLUME_THRESHOLD = Decimal("50000") # $50k daily volume
|
||||
|
||||
# Niche market categories - markets in these categories with low specificity
|
||||
# are more likely to have insider information value
|
||||
NICHE_PRONE_CATEGORIES = frozenset({"science", "tech", "finance", "other"})
|
||||
|
||||
|
||||
class SizeAnomalyDetector:
|
||||
"""Detector for unusually large trade sizes.
|
||||
|
||||
This detector analyzes trade events for size anomalies by comparing
|
||||
the trade size against market liquidity metrics:
|
||||
- Volume impact: trade size / 24h volume
|
||||
- Book impact: trade size / order book depth
|
||||
|
||||
When volume data is unavailable, the detector uses category-based
|
||||
heuristics to identify niche markets where large trades are more
|
||||
significant.
|
||||
|
||||
Confidence scoring:
|
||||
- Volume impact > threshold: base score from impact ratio
|
||||
- Book impact > threshold: additional score from impact ratio
|
||||
- Niche market multiplier: 1.5x for low-volume markets
|
||||
|
||||
Example:
|
||||
```python
|
||||
sync = MarketMetadataSync(redis, clob_client)
|
||||
detector = SizeAnomalyDetector(sync)
|
||||
|
||||
# Analyze a trade
|
||||
signal = await detector.analyze(trade_event)
|
||||
if signal is not None:
|
||||
print(f"Size anomaly detected! Confidence: {signal.confidence}")
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
metadata_sync: MarketMetadataSync,
|
||||
*,
|
||||
volume_threshold: float = DEFAULT_VOLUME_THRESHOLD,
|
||||
book_threshold: float = DEFAULT_BOOK_THRESHOLD,
|
||||
niche_volume_threshold: Decimal = DEFAULT_NICHE_VOLUME_THRESHOLD,
|
||||
) -> None:
|
||||
"""Initialize the size anomaly detector.
|
||||
|
||||
Args:
|
||||
metadata_sync: MarketMetadataSync for fetching market metadata.
|
||||
volume_threshold: Threshold for volume impact (default 0.02 = 2%).
|
||||
book_threshold: Threshold for book impact (default 0.05 = 5%).
|
||||
niche_volume_threshold: Volume below which market is niche ($50k).
|
||||
"""
|
||||
self._metadata_sync = metadata_sync
|
||||
self._volume_threshold = volume_threshold
|
||||
self._book_threshold = book_threshold
|
||||
self._niche_volume_threshold = niche_volume_threshold
|
||||
|
||||
async def analyze(
|
||||
self,
|
||||
trade: TradeEvent,
|
||||
*,
|
||||
daily_volume: Decimal | None = None,
|
||||
book_depth: Decimal | None = None,
|
||||
) -> SizeAnomalySignal | None:
|
||||
"""Analyze a trade event for size anomalies.
|
||||
|
||||
This method:
|
||||
1. Fetches market metadata
|
||||
2. Calculates volume and book impact (if data available)
|
||||
3. Determines if market is niche
|
||||
4. Calculates confidence score
|
||||
|
||||
Args:
|
||||
trade: TradeEvent to analyze.
|
||||
daily_volume: Optional 24h volume in USDC. If provided, enables
|
||||
volume impact calculation.
|
||||
book_depth: Optional order book depth in USDC. If provided,
|
||||
enables book impact calculation.
|
||||
|
||||
Returns:
|
||||
SizeAnomalySignal if the trade triggers anomaly detection,
|
||||
None otherwise.
|
||||
"""
|
||||
# Get market metadata
|
||||
try:
|
||||
metadata = await self._metadata_sync.get_market(trade.market_id)
|
||||
if metadata is None:
|
||||
logger.warning(
|
||||
"No metadata found for market %s, creating minimal metadata",
|
||||
trade.market_id,
|
||||
)
|
||||
metadata = self._create_minimal_metadata(trade)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to get metadata for market %s: %s",
|
||||
trade.market_id,
|
||||
e,
|
||||
)
|
||||
metadata = self._create_minimal_metadata(trade)
|
||||
|
||||
trade_size = trade.notional_value
|
||||
|
||||
# Calculate impacts
|
||||
volume_impact = self._calculate_volume_impact(trade_size, daily_volume)
|
||||
book_impact = self._calculate_book_impact(trade_size, book_depth)
|
||||
|
||||
# Determine if niche market
|
||||
is_niche = self._is_niche_market(metadata, daily_volume)
|
||||
|
||||
# Check if any threshold exceeded
|
||||
exceeds_volume = volume_impact > self._volume_threshold
|
||||
exceeds_book = book_impact > self._book_threshold
|
||||
|
||||
if not exceeds_volume and not exceeds_book and not is_niche:
|
||||
logger.debug(
|
||||
"Trade %s does not exceed thresholds: volume=%.4f, book=%.4f",
|
||||
trade.trade_id,
|
||||
volume_impact,
|
||||
book_impact,
|
||||
)
|
||||
return None
|
||||
|
||||
# Calculate confidence score
|
||||
confidence, factors = self.calculate_confidence(
|
||||
volume_impact=volume_impact,
|
||||
book_impact=book_impact,
|
||||
is_niche=is_niche,
|
||||
)
|
||||
|
||||
# Only emit signal if confidence is meaningful
|
||||
if confidence < 0.1:
|
||||
return None
|
||||
|
||||
logger.info(
|
||||
"Size anomaly signal: market=%s, size=%s, volume_impact=%.4f, "
|
||||
"book_impact=%.4f, niche=%s, confidence=%.2f",
|
||||
trade.market_id[:10] + "...",
|
||||
trade_size,
|
||||
volume_impact,
|
||||
book_impact,
|
||||
is_niche,
|
||||
confidence,
|
||||
)
|
||||
|
||||
return SizeAnomalySignal(
|
||||
trade_event=trade,
|
||||
market_metadata=metadata,
|
||||
volume_impact=volume_impact,
|
||||
book_impact=book_impact,
|
||||
is_niche_market=is_niche,
|
||||
confidence=confidence,
|
||||
factors=factors,
|
||||
)
|
||||
|
||||
def _create_minimal_metadata(self, trade: TradeEvent) -> MarketMetadata:
|
||||
"""Create minimal metadata from trade event."""
|
||||
from polymarket_insider_tracker.ingestor.models import Token
|
||||
|
||||
return MarketMetadata(
|
||||
condition_id=trade.market_id,
|
||||
question=trade.event_title or "Unknown Market",
|
||||
description="",
|
||||
tokens=(
|
||||
Token(
|
||||
token_id=trade.asset_id,
|
||||
outcome=trade.outcome,
|
||||
price=trade.price,
|
||||
),
|
||||
),
|
||||
category="other",
|
||||
)
|
||||
|
||||
def _calculate_volume_impact(
|
||||
self,
|
||||
trade_size: Decimal,
|
||||
daily_volume: Decimal | None,
|
||||
) -> float:
|
||||
"""Calculate trade size as fraction of daily volume.
|
||||
|
||||
Args:
|
||||
trade_size: Trade notional value in USDC.
|
||||
daily_volume: 24h trading volume in USDC.
|
||||
|
||||
Returns:
|
||||
Volume impact ratio, or 0.0 if volume unknown.
|
||||
"""
|
||||
if daily_volume is None or daily_volume <= 0:
|
||||
return 0.0
|
||||
return float(trade_size / daily_volume)
|
||||
|
||||
def _calculate_book_impact(
|
||||
self,
|
||||
trade_size: Decimal,
|
||||
book_depth: Decimal | None,
|
||||
) -> float:
|
||||
"""Calculate trade size as fraction of order book depth.
|
||||
|
||||
Args:
|
||||
trade_size: Trade notional value in USDC.
|
||||
book_depth: Visible order book depth in USDC.
|
||||
|
||||
Returns:
|
||||
Book impact ratio, or 0.0 if depth unknown.
|
||||
"""
|
||||
if book_depth is None or book_depth <= 0:
|
||||
return 0.0
|
||||
return float(trade_size / book_depth)
|
||||
|
||||
def _is_niche_market(
|
||||
self,
|
||||
metadata: MarketMetadata,
|
||||
daily_volume: Decimal | None,
|
||||
) -> bool:
|
||||
"""Determine if market is considered niche.
|
||||
|
||||
A market is niche if:
|
||||
- Volume is below threshold ($50k), OR
|
||||
- Category is prone to insider info AND volume is unknown
|
||||
|
||||
Args:
|
||||
metadata: Market metadata with category.
|
||||
daily_volume: Optional 24h volume.
|
||||
|
||||
Returns:
|
||||
True if market is considered niche.
|
||||
"""
|
||||
# If volume known and below threshold, it's niche
|
||||
if daily_volume is not None and daily_volume < self._niche_volume_threshold:
|
||||
return True
|
||||
|
||||
# If volume unknown, use category heuristics
|
||||
return daily_volume is None and metadata.category in NICHE_PRONE_CATEGORIES
|
||||
|
||||
def calculate_confidence(
|
||||
self,
|
||||
*,
|
||||
volume_impact: float,
|
||||
book_impact: float,
|
||||
is_niche: bool,
|
||||
) -> tuple[float, dict[str, float]]:
|
||||
"""Calculate confidence score based on impact metrics.
|
||||
|
||||
Confidence scoring:
|
||||
- Volume impact: min(impact/threshold, 3) / 3 * 0.5
|
||||
- Book impact: min(impact/threshold, 3) / 3 * 0.3
|
||||
- Niche multiplier: 1.5x final score
|
||||
|
||||
Final confidence clamped to [0.0, 1.0].
|
||||
|
||||
Args:
|
||||
volume_impact: Trade size / daily volume ratio.
|
||||
book_impact: Trade size / book depth ratio.
|
||||
is_niche: Whether market is niche.
|
||||
|
||||
Returns:
|
||||
Tuple of (confidence_score, factors_dict).
|
||||
"""
|
||||
factors: dict[str, float] = {}
|
||||
confidence = 0.0
|
||||
|
||||
# Volume impact component
|
||||
if volume_impact > self._volume_threshold:
|
||||
ratio = min(volume_impact / self._volume_threshold, 3.0)
|
||||
volume_score = ratio / 3.0 * 0.5
|
||||
factors["volume_impact"] = volume_score
|
||||
confidence += volume_score
|
||||
|
||||
# Book impact component
|
||||
if book_impact > self._book_threshold:
|
||||
ratio = min(book_impact / self._book_threshold, 3.0)
|
||||
book_score = ratio / 3.0 * 0.3
|
||||
factors["book_impact"] = book_score
|
||||
confidence += book_score
|
||||
|
||||
# Niche market multiplier
|
||||
if is_niche and confidence > 0:
|
||||
factors["niche_multiplier"] = 1.5
|
||||
confidence *= 1.5
|
||||
|
||||
# If niche but no other signals, give small base confidence
|
||||
if is_niche and confidence == 0:
|
||||
factors["niche_base"] = 0.2
|
||||
confidence = 0.2
|
||||
|
||||
# Clamp to valid range
|
||||
confidence = max(0.0, min(1.0, confidence))
|
||||
|
||||
return confidence, factors
|
||||
|
||||
async def analyze_batch(
|
||||
self,
|
||||
trades: list[TradeEvent],
|
||||
*,
|
||||
volume_data: dict[str, Decimal] | None = None,
|
||||
book_data: dict[str, Decimal] | None = None,
|
||||
) -> list[SizeAnomalySignal]:
|
||||
"""Analyze multiple trades for size anomalies.
|
||||
|
||||
Processes trades in parallel for efficiency.
|
||||
|
||||
Args:
|
||||
trades: List of trades to analyze.
|
||||
volume_data: Optional dict mapping market_id to 24h volume.
|
||||
book_data: Optional dict mapping market_id to book depth.
|
||||
|
||||
Returns:
|
||||
List of SizeAnomalySignal for trades with anomalies.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
volume_data = volume_data or {}
|
||||
book_data = book_data or {}
|
||||
|
||||
tasks = [
|
||||
self.analyze(
|
||||
trade,
|
||||
daily_volume=volume_data.get(trade.market_id),
|
||||
book_depth=book_data.get(trade.market_id),
|
||||
)
|
||||
for trade in trades
|
||||
]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
signals: list[SizeAnomalySignal] = []
|
||||
for trade, result in zip(trades, results, strict=True):
|
||||
if isinstance(result, BaseException):
|
||||
logger.warning(
|
||||
"Failed to analyze trade %s: %s",
|
||||
trade.trade_id,
|
||||
result,
|
||||
)
|
||||
continue
|
||||
if result is not None:
|
||||
signals.append(result)
|
||||
|
||||
return signals
|
||||
@@ -0,0 +1,471 @@
|
||||
"""Sniper cluster detection using DBSCAN clustering.
|
||||
|
||||
This module identifies wallets that exhibit coordinated "sniper" behavior -
|
||||
consistently entering markets within minutes of their creation, suggesting
|
||||
advance knowledge of market creation times.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
from sklearn.cluster import DBSCAN
|
||||
|
||||
from polymarket_insider_tracker.detector.models import SniperClusterSignal
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MarketEntry:
|
||||
"""Record of a wallet's entry into a market.
|
||||
|
||||
Attributes:
|
||||
wallet_address: The wallet that entered the market.
|
||||
market_id: The market condition ID.
|
||||
entry_delta_seconds: Time between market creation and wallet entry.
|
||||
position_size: Size of the initial position in USDC.
|
||||
timestamp: When the entry occurred.
|
||||
"""
|
||||
|
||||
wallet_address: str
|
||||
market_id: str
|
||||
entry_delta_seconds: float
|
||||
position_size: Decimal
|
||||
timestamp: datetime
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClusterInfo:
|
||||
"""Information about a detected sniper cluster.
|
||||
|
||||
Attributes:
|
||||
cluster_id: Unique identifier for this cluster.
|
||||
wallet_addresses: Set of wallet addresses in the cluster.
|
||||
avg_entry_delta: Average entry delay in seconds across the cluster.
|
||||
markets_in_common: Number of markets where cluster members overlap.
|
||||
created_at: When this cluster was first detected.
|
||||
"""
|
||||
|
||||
cluster_id: str
|
||||
wallet_addresses: set[str]
|
||||
avg_entry_delta: float
|
||||
markets_in_common: int
|
||||
created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
|
||||
class SniperDetector:
|
||||
"""Detects sniper clusters using DBSCAN clustering algorithm.
|
||||
|
||||
The detector tracks wallet entries across markets and periodically
|
||||
runs DBSCAN clustering to identify groups of wallets with similar
|
||||
timing patterns (consistently entering markets early after creation).
|
||||
|
||||
Attributes:
|
||||
entry_threshold_seconds: Maximum seconds after market creation to be
|
||||
considered a "sniper" entry (default 300 = 5 minutes).
|
||||
min_cluster_size: Minimum wallets to form a cluster (default 3).
|
||||
eps: DBSCAN epsilon parameter for neighborhood distance (default 0.5).
|
||||
min_samples: DBSCAN minimum samples for core point (default 2).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
entry_threshold_seconds: int = 300,
|
||||
min_cluster_size: int = 3,
|
||||
eps: float = 0.5,
|
||||
min_samples: int = 2,
|
||||
min_entries_per_wallet: int = 2,
|
||||
) -> None:
|
||||
"""Initialize the sniper detector.
|
||||
|
||||
Args:
|
||||
entry_threshold_seconds: Max seconds for sniper entry (default 300).
|
||||
min_cluster_size: Minimum cluster size (default 3).
|
||||
eps: DBSCAN epsilon (default 0.5).
|
||||
min_samples: DBSCAN min samples (default 2).
|
||||
min_entries_per_wallet: Minimum entries to include wallet (default 2).
|
||||
"""
|
||||
self.entry_threshold_seconds = entry_threshold_seconds
|
||||
self.min_cluster_size = min_cluster_size
|
||||
self.eps = eps
|
||||
self.min_samples = min_samples
|
||||
self.min_entries_per_wallet = min_entries_per_wallet
|
||||
|
||||
# Entry tracking
|
||||
self._entries: list[MarketEntry] = []
|
||||
self._wallet_entries: dict[str, list[MarketEntry]] = defaultdict(list)
|
||||
self._market_wallets: dict[str, set[str]] = defaultdict(set)
|
||||
|
||||
# Cluster tracking
|
||||
self._known_clusters: dict[str, ClusterInfo] = {}
|
||||
self._wallet_cluster_map: dict[str, str] = {}
|
||||
|
||||
# Previously signaled wallets (to avoid duplicate signals)
|
||||
self._signaled_wallets: set[str] = set()
|
||||
|
||||
def record_entry(
|
||||
self,
|
||||
trade: TradeEvent,
|
||||
market_created_at: datetime,
|
||||
) -> None:
|
||||
"""Record a market entry for clustering analysis.
|
||||
|
||||
Only records entries that occur within the threshold time after
|
||||
market creation (potential sniper behavior).
|
||||
|
||||
Args:
|
||||
trade: The trade event representing market entry.
|
||||
market_created_at: When the market was created.
|
||||
"""
|
||||
# Calculate entry delta
|
||||
entry_time = trade.timestamp
|
||||
delta = (entry_time - market_created_at).total_seconds()
|
||||
|
||||
# Only track entries within threshold (potential snipers)
|
||||
if delta < 0 or delta > self.entry_threshold_seconds:
|
||||
return
|
||||
|
||||
entry = MarketEntry(
|
||||
wallet_address=trade.wallet_address.lower(),
|
||||
market_id=trade.market_id,
|
||||
entry_delta_seconds=delta,
|
||||
position_size=trade.notional_value,
|
||||
timestamp=entry_time,
|
||||
)
|
||||
|
||||
self._entries.append(entry)
|
||||
self._wallet_entries[entry.wallet_address].append(entry)
|
||||
self._market_wallets[entry.market_id].add(entry.wallet_address)
|
||||
|
||||
logger.debug(
|
||||
"Recorded sniper entry: wallet=%s market=%s delta=%.1fs",
|
||||
entry.wallet_address[:10],
|
||||
entry.market_id[:10],
|
||||
delta,
|
||||
)
|
||||
|
||||
def run_clustering(self) -> list[SniperClusterSignal]:
|
||||
"""Run DBSCAN clustering and return new sniper signals.
|
||||
|
||||
Clusters wallets based on their entry timing patterns across markets.
|
||||
Returns signals only for newly identified cluster members.
|
||||
|
||||
Returns:
|
||||
List of SniperClusterSignal for newly detected cluster members.
|
||||
"""
|
||||
# Filter wallets with enough entries
|
||||
eligible_wallets = [
|
||||
wallet
|
||||
for wallet, entries in self._wallet_entries.items()
|
||||
if len(entries) >= self.min_entries_per_wallet
|
||||
]
|
||||
|
||||
if len(eligible_wallets) < self.min_cluster_size:
|
||||
logger.debug(
|
||||
"Not enough eligible wallets for clustering: %d < %d",
|
||||
len(eligible_wallets),
|
||||
self.min_cluster_size,
|
||||
)
|
||||
return []
|
||||
|
||||
# Build feature matrix
|
||||
feature_vectors, wallet_index = self._build_feature_matrix(eligible_wallets)
|
||||
|
||||
if len(feature_vectors) == 0:
|
||||
return []
|
||||
|
||||
# Run DBSCAN
|
||||
clustering = DBSCAN(
|
||||
eps=self.eps,
|
||||
min_samples=self.min_samples,
|
||||
metric="euclidean",
|
||||
).fit(feature_vectors)
|
||||
|
||||
# Process clusters
|
||||
signals = self._process_clustering_results(
|
||||
clustering.labels_,
|
||||
wallet_index,
|
||||
)
|
||||
|
||||
return signals
|
||||
|
||||
def _build_feature_matrix(
|
||||
self,
|
||||
wallets: list[str],
|
||||
) -> tuple[np.ndarray, dict[int, str]]:
|
||||
"""Build feature matrix for DBSCAN clustering.
|
||||
|
||||
Features per entry:
|
||||
- Normalized market hash (0-1 range)
|
||||
- Normalized entry delta (in hours, typically 0-0.083)
|
||||
- Log-normalized position size
|
||||
|
||||
Args:
|
||||
wallets: List of wallet addresses to include.
|
||||
|
||||
Returns:
|
||||
Tuple of (feature_matrix, wallet_index_map).
|
||||
"""
|
||||
features = []
|
||||
wallet_index: dict[int, str] = {}
|
||||
row_idx = 0
|
||||
|
||||
for wallet in wallets:
|
||||
entries = self._wallet_entries[wallet]
|
||||
for entry in entries:
|
||||
# Normalize market ID to 0-1 range
|
||||
market_hash = (
|
||||
(
|
||||
int(
|
||||
hashlib.md5( # noqa: S324
|
||||
entry.market_id.encode()
|
||||
).hexdigest()[:8],
|
||||
16,
|
||||
)
|
||||
% 1000
|
||||
)
|
||||
/ 1000.0
|
||||
)
|
||||
|
||||
# Normalize entry delta to hours (0-5 mins = 0-0.083 hours)
|
||||
delta_hours = entry.entry_delta_seconds / 3600.0
|
||||
|
||||
# Log-normalize position size
|
||||
log_size = float(np.log10(max(float(entry.position_size), 1.0)))
|
||||
|
||||
features.append([market_hash, delta_hours, log_size])
|
||||
wallet_index[row_idx] = wallet
|
||||
row_idx += 1
|
||||
|
||||
return np.array(features), wallet_index
|
||||
|
||||
def _process_clustering_results(
|
||||
self,
|
||||
labels: np.ndarray,
|
||||
wallet_index: dict[int, str],
|
||||
) -> list[SniperClusterSignal]:
|
||||
"""Process DBSCAN clustering results into signals.
|
||||
|
||||
Args:
|
||||
labels: Cluster labels from DBSCAN (-1 = noise).
|
||||
wallet_index: Map from row index to wallet address.
|
||||
|
||||
Returns:
|
||||
List of signals for newly detected cluster members.
|
||||
"""
|
||||
# Group rows by cluster
|
||||
cluster_rows: dict[int, list[int]] = defaultdict(list)
|
||||
for row_idx, label in enumerate(labels):
|
||||
if label != -1: # Skip noise
|
||||
cluster_rows[label].append(row_idx)
|
||||
|
||||
signals: list[SniperClusterSignal] = []
|
||||
|
||||
for _cluster_label, rows in cluster_rows.items():
|
||||
# Get unique wallets in this cluster
|
||||
cluster_wallets = {wallet_index[row] for row in rows}
|
||||
|
||||
if len(cluster_wallets) < self.min_cluster_size:
|
||||
continue
|
||||
|
||||
# Calculate cluster statistics
|
||||
cluster_stats = self._calculate_cluster_stats(cluster_wallets)
|
||||
|
||||
# Generate or reuse cluster ID
|
||||
cluster_id = self._get_or_create_cluster_id(cluster_wallets)
|
||||
|
||||
# Update cluster info
|
||||
self._known_clusters[cluster_id] = ClusterInfo(
|
||||
cluster_id=cluster_id,
|
||||
wallet_addresses=cluster_wallets,
|
||||
avg_entry_delta=cluster_stats["avg_delta"],
|
||||
markets_in_common=int(cluster_stats["markets_in_common"]),
|
||||
)
|
||||
|
||||
# Update wallet-cluster mapping
|
||||
for wallet in cluster_wallets:
|
||||
self._wallet_cluster_map[wallet] = cluster_id
|
||||
|
||||
# Generate signals for new cluster members
|
||||
for wallet in cluster_wallets:
|
||||
if wallet not in self._signaled_wallets:
|
||||
confidence = self._calculate_confidence(
|
||||
cluster_wallets,
|
||||
cluster_stats,
|
||||
)
|
||||
|
||||
signal = SniperClusterSignal(
|
||||
wallet_address=wallet,
|
||||
cluster_id=cluster_id,
|
||||
cluster_size=len(cluster_wallets),
|
||||
avg_entry_delta_seconds=cluster_stats["avg_delta"],
|
||||
markets_in_common=int(cluster_stats["markets_in_common"]),
|
||||
confidence=confidence,
|
||||
)
|
||||
|
||||
signals.append(signal)
|
||||
self._signaled_wallets.add(wallet)
|
||||
|
||||
logger.info(
|
||||
"New sniper detected: wallet=%s cluster=%s confidence=%.2f",
|
||||
wallet[:10],
|
||||
cluster_id[:8],
|
||||
confidence,
|
||||
)
|
||||
|
||||
return signals
|
||||
|
||||
def _calculate_cluster_stats(
|
||||
self,
|
||||
cluster_wallets: set[str],
|
||||
) -> dict[str, float | int]:
|
||||
"""Calculate statistics for a cluster of wallets.
|
||||
|
||||
Args:
|
||||
cluster_wallets: Set of wallet addresses in the cluster.
|
||||
|
||||
Returns:
|
||||
Dict with avg_delta, markets_in_common statistics.
|
||||
"""
|
||||
# Calculate average entry delta
|
||||
all_deltas: list[float] = []
|
||||
for wallet in cluster_wallets:
|
||||
for entry in self._wallet_entries[wallet]:
|
||||
all_deltas.append(entry.entry_delta_seconds)
|
||||
|
||||
avg_delta = sum(all_deltas) / len(all_deltas) if all_deltas else 0.0
|
||||
|
||||
# Calculate markets in common
|
||||
wallet_markets: list[set[str]] = []
|
||||
for wallet in cluster_wallets:
|
||||
markets = {e.market_id for e in self._wallet_entries[wallet]}
|
||||
wallet_markets.append(markets)
|
||||
|
||||
if len(wallet_markets) >= 2:
|
||||
common_markets = set.intersection(*wallet_markets)
|
||||
markets_in_common = len(common_markets)
|
||||
else:
|
||||
markets_in_common = 0
|
||||
|
||||
return {
|
||||
"avg_delta": avg_delta,
|
||||
"markets_in_common": markets_in_common,
|
||||
}
|
||||
|
||||
def _get_or_create_cluster_id(self, wallets: set[str]) -> str:
|
||||
"""Get existing cluster ID or create new one.
|
||||
|
||||
Checks if majority of wallets belong to an existing cluster
|
||||
and returns that ID, otherwise creates a new ID.
|
||||
|
||||
Args:
|
||||
wallets: Set of wallet addresses.
|
||||
|
||||
Returns:
|
||||
Cluster ID string.
|
||||
"""
|
||||
# Check if majority belongs to existing cluster
|
||||
existing_clusters: dict[str, int] = defaultdict(int)
|
||||
for wallet in wallets:
|
||||
if wallet in self._wallet_cluster_map:
|
||||
existing_clusters[self._wallet_cluster_map[wallet]] += 1
|
||||
|
||||
if existing_clusters:
|
||||
best_cluster = max(existing_clusters, key=lambda k: existing_clusters[k])
|
||||
if existing_clusters[best_cluster] >= len(wallets) // 2:
|
||||
return best_cluster
|
||||
|
||||
return str(uuid.uuid4())
|
||||
|
||||
def _calculate_confidence(
|
||||
self,
|
||||
cluster_wallets: set[str],
|
||||
stats: dict[str, float | int],
|
||||
) -> float:
|
||||
"""Calculate confidence score for a cluster.
|
||||
|
||||
Higher confidence when:
|
||||
- Larger cluster size
|
||||
- Lower average entry delta (faster entries)
|
||||
- More markets in common
|
||||
|
||||
Args:
|
||||
cluster_wallets: Wallets in the cluster.
|
||||
stats: Cluster statistics dict.
|
||||
|
||||
Returns:
|
||||
Confidence score from 0.0 to 1.0.
|
||||
"""
|
||||
# Size factor: more wallets = higher confidence
|
||||
size_factor = min(1.0, len(cluster_wallets) / 10.0)
|
||||
|
||||
# Speed factor: faster entries = higher confidence
|
||||
# 0 seconds = 1.0, 300 seconds = 0.0
|
||||
avg_delta = float(stats["avg_delta"])
|
||||
speed_factor = max(0.0, 1.0 - (avg_delta / self.entry_threshold_seconds))
|
||||
|
||||
# Overlap factor: more markets in common = higher confidence
|
||||
markets_common = int(stats["markets_in_common"])
|
||||
overlap_factor = min(1.0, markets_common / 5.0)
|
||||
|
||||
# Weighted combination
|
||||
confidence = 0.3 * size_factor + 0.4 * speed_factor + 0.3 * overlap_factor
|
||||
|
||||
return round(min(1.0, confidence), 3)
|
||||
|
||||
def is_sniper(self, wallet_address: str) -> bool:
|
||||
"""Check if a wallet is in any known sniper cluster.
|
||||
|
||||
Args:
|
||||
wallet_address: Wallet address to check.
|
||||
|
||||
Returns:
|
||||
True if wallet is a known sniper.
|
||||
"""
|
||||
return wallet_address.lower() in self._wallet_cluster_map
|
||||
|
||||
def get_cluster_for_wallet(self, wallet_address: str) -> ClusterInfo | None:
|
||||
"""Get cluster info for a wallet if it belongs to one.
|
||||
|
||||
Args:
|
||||
wallet_address: Wallet address to look up.
|
||||
|
||||
Returns:
|
||||
ClusterInfo if wallet is in a cluster, None otherwise.
|
||||
"""
|
||||
cluster_id = self._wallet_cluster_map.get(wallet_address.lower())
|
||||
if cluster_id:
|
||||
return self._known_clusters.get(cluster_id)
|
||||
return None
|
||||
|
||||
def get_entry_count(self) -> int:
|
||||
"""Return the total number of tracked entries."""
|
||||
return len(self._entries)
|
||||
|
||||
def get_wallet_count(self) -> int:
|
||||
"""Return the number of unique wallets tracked."""
|
||||
return len(self._wallet_entries)
|
||||
|
||||
def get_cluster_count(self) -> int:
|
||||
"""Return the number of detected clusters."""
|
||||
return len(self._known_clusters)
|
||||
|
||||
def clear_entries(self) -> None:
|
||||
"""Clear all tracked entries (for periodic cleanup)."""
|
||||
self._entries.clear()
|
||||
self._wallet_entries.clear()
|
||||
self._market_wallets.clear()
|
||||
logger.info("Cleared all sniper detector entries")
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Data ingestion layer - Real-time Polymarket trade streaming."""
|
||||
|
||||
from polymarket_insider_tracker.ingestor.clob_client import (
|
||||
ClobClient,
|
||||
ClobClientError,
|
||||
RetryError,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.health import (
|
||||
HealthMonitor,
|
||||
HealthReport,
|
||||
HealthStatus,
|
||||
StreamHealth,
|
||||
StreamStatus,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.metadata_sync import (
|
||||
MarketMetadataSync,
|
||||
MetadataSyncError,
|
||||
SyncState,
|
||||
SyncStats,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.models import (
|
||||
Market,
|
||||
MarketMetadata,
|
||||
Orderbook,
|
||||
OrderbookLevel,
|
||||
Token,
|
||||
TradeEvent,
|
||||
derive_category,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.publisher import (
|
||||
ConsumerGroupExistsError,
|
||||
EventPublisher,
|
||||
PublisherError,
|
||||
StreamEntry,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.websocket import (
|
||||
ConnectionState,
|
||||
TradeStreamError,
|
||||
TradeStreamHandler,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.websocket import (
|
||||
StreamStats as WebSocketStreamStats,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# CLOB Client
|
||||
"ClobClient",
|
||||
"ClobClientError",
|
||||
"RetryError",
|
||||
# Health Monitor
|
||||
"HealthMonitor",
|
||||
"HealthReport",
|
||||
"HealthStatus",
|
||||
"StreamHealth",
|
||||
"StreamStatus",
|
||||
# Metadata Sync
|
||||
"MarketMetadataSync",
|
||||
"MetadataSyncError",
|
||||
"SyncState",
|
||||
"SyncStats",
|
||||
# Models
|
||||
"Market",
|
||||
"MarketMetadata",
|
||||
"Orderbook",
|
||||
"OrderbookLevel",
|
||||
"Token",
|
||||
"TradeEvent",
|
||||
"derive_category",
|
||||
# Publisher
|
||||
"ConsumerGroupExistsError",
|
||||
"EventPublisher",
|
||||
"PublisherError",
|
||||
"StreamEntry",
|
||||
# WebSocket
|
||||
"ConnectionState",
|
||||
"WebSocketStreamStats",
|
||||
"TradeStreamError",
|
||||
"TradeStreamHandler",
|
||||
]
|
||||
@@ -0,0 +1,343 @@
|
||||
"""Wrapper around py-clob-client with rate limiting and retry logic."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
from typing import ParamSpec, TypeVar
|
||||
|
||||
from py_clob_client.client import ClobClient as BaseClobClient
|
||||
from py_clob_client.clob_types import BookParams
|
||||
|
||||
from polymarket_insider_tracker.ingestor.models import Market, Orderbook
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
P = ParamSpec("P")
|
||||
T = TypeVar("T")
|
||||
|
||||
# Constants
|
||||
DEFAULT_HOST = "https://clob.polymarket.com"
|
||||
MAX_REQUESTS_PER_SECOND = 10
|
||||
MIN_REQUEST_INTERVAL = 1.0 / MAX_REQUESTS_PER_SECOND # 0.1 seconds
|
||||
|
||||
DEFAULT_MAX_RETRIES = 3
|
||||
DEFAULT_RETRY_BASE_DELAY = 1.0
|
||||
RETRY_STATUS_CODES = (429, 500, 502, 503, 504)
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""Token bucket rate limiter for API requests."""
|
||||
|
||||
def __init__(self, max_requests_per_second: float = MAX_REQUESTS_PER_SECOND) -> None:
|
||||
"""Initialize the rate limiter.
|
||||
|
||||
Args:
|
||||
max_requests_per_second: Maximum requests allowed per second.
|
||||
"""
|
||||
self._min_interval = 1.0 / max_requests_per_second
|
||||
self._last_request_time: float = 0.0
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def acquire(self) -> None:
|
||||
"""Wait until a request slot is available."""
|
||||
async with self._lock:
|
||||
now = time.monotonic()
|
||||
elapsed = now - self._last_request_time
|
||||
if elapsed < self._min_interval:
|
||||
wait_time = self._min_interval - elapsed
|
||||
await asyncio.sleep(wait_time)
|
||||
self._last_request_time = time.monotonic()
|
||||
|
||||
def acquire_sync(self) -> None:
|
||||
"""Synchronous version of acquire for sync operations."""
|
||||
now = time.monotonic()
|
||||
elapsed = now - self._last_request_time
|
||||
if elapsed < self._min_interval:
|
||||
wait_time = self._min_interval - elapsed
|
||||
time.sleep(wait_time)
|
||||
self._last_request_time = time.monotonic()
|
||||
|
||||
|
||||
class RetryError(Exception):
|
||||
"""Raised when all retry attempts are exhausted."""
|
||||
|
||||
def __init__(self, message: str, last_exception: Exception | None = None) -> None:
|
||||
super().__init__(message)
|
||||
self.last_exception = last_exception
|
||||
|
||||
|
||||
def with_retry(
|
||||
max_retries: int = DEFAULT_MAX_RETRIES,
|
||||
base_delay: float = DEFAULT_RETRY_BASE_DELAY,
|
||||
retry_on: tuple[type[Exception], ...] = (Exception,),
|
||||
) -> Callable[[Callable[P, T]], Callable[P, T]]:
|
||||
"""Decorator for adding retry logic with exponential backoff.
|
||||
|
||||
Args:
|
||||
max_retries: Maximum number of retry attempts.
|
||||
base_delay: Base delay in seconds (doubles with each retry).
|
||||
retry_on: Tuple of exception types to retry on.
|
||||
|
||||
Returns:
|
||||
Decorated function with retry logic.
|
||||
"""
|
||||
|
||||
def decorator(func: Callable[P, T]) -> Callable[P, T]:
|
||||
@wraps(func)
|
||||
def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
|
||||
last_exception: Exception | None = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
except retry_on as e:
|
||||
last_exception = e
|
||||
if attempt == max_retries:
|
||||
break
|
||||
|
||||
delay = base_delay * (2**attempt)
|
||||
logger.warning(
|
||||
"Attempt %d/%d failed: %s. Retrying in %.1f seconds...",
|
||||
attempt + 1,
|
||||
max_retries + 1,
|
||||
str(e),
|
||||
delay,
|
||||
)
|
||||
time.sleep(delay)
|
||||
|
||||
raise RetryError(
|
||||
f"All {max_retries + 1} attempts failed for {func.__name__}",
|
||||
last_exception=last_exception,
|
||||
)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
class ClobClientError(Exception):
|
||||
"""Base exception for ClobClient errors."""
|
||||
|
||||
|
||||
class ClobClient:
|
||||
"""Wrapper around py-clob-client with rate limiting and retry logic.
|
||||
|
||||
This client provides a clean interface for querying Polymarket CLOB data
|
||||
with built-in rate limiting (10 requests/second) and automatic retry
|
||||
with exponential backoff on transient errors.
|
||||
|
||||
Example:
|
||||
>>> client = ClobClient() # Uses POLYMARKET_API_KEY env var
|
||||
>>> markets = client.get_markets()
|
||||
>>> orderbook = client.get_orderbook("token_id_here")
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
host: str = DEFAULT_HOST,
|
||||
max_retries: int = DEFAULT_MAX_RETRIES,
|
||||
requests_per_second: float = MAX_REQUESTS_PER_SECOND,
|
||||
) -> None:
|
||||
"""Initialize the CLOB client.
|
||||
|
||||
Args:
|
||||
api_key: Polymarket API key. If not provided, reads from
|
||||
POLYMARKET_API_KEY environment variable.
|
||||
host: CLOB API endpoint URL.
|
||||
max_retries: Maximum retry attempts for failed requests.
|
||||
requests_per_second: Rate limit for API requests.
|
||||
"""
|
||||
self._api_key = api_key or os.environ.get("POLYMARKET_API_KEY")
|
||||
self._host = host
|
||||
self._max_retries = max_retries
|
||||
self._rate_limiter = RateLimiter(requests_per_second)
|
||||
|
||||
# Initialize the underlying client (read-only, no auth needed for queries)
|
||||
self._client = BaseClobClient(host)
|
||||
|
||||
logger.info(
|
||||
"Initialized ClobClient with host=%s, rate_limit=%.1f req/s",
|
||||
host,
|
||||
requests_per_second,
|
||||
)
|
||||
|
||||
def _with_rate_limit(self, func: Callable[P, T]) -> Callable[P, T]:
|
||||
"""Wrap a function with rate limiting."""
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
|
||||
self._rate_limiter.acquire_sync()
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
@with_retry()
|
||||
def get_markets(self, active_only: bool = True) -> list[Market]:
|
||||
"""Fetch all markets from the CLOB.
|
||||
|
||||
Args:
|
||||
active_only: If True, only return active (non-closed) markets.
|
||||
|
||||
Returns:
|
||||
List of Market objects.
|
||||
"""
|
||||
self._rate_limiter.acquire_sync()
|
||||
|
||||
all_markets: list[Market] = []
|
||||
cursor: str | None = None
|
||||
|
||||
while True:
|
||||
if cursor:
|
||||
response = self._client.get_simplified_markets(cursor)
|
||||
else:
|
||||
response = self._client.get_simplified_markets()
|
||||
|
||||
data = response.get("data", [])
|
||||
for market_data in data:
|
||||
market = Market.from_dict(market_data)
|
||||
if active_only and market.closed:
|
||||
continue
|
||||
all_markets.append(market)
|
||||
|
||||
next_cursor = response.get("next_cursor")
|
||||
if not next_cursor or next_cursor == "LTE=":
|
||||
break
|
||||
cursor = next_cursor
|
||||
|
||||
# Rate limit between pagination requests
|
||||
self._rate_limiter.acquire_sync()
|
||||
|
||||
logger.debug("Fetched %d markets", len(all_markets))
|
||||
return all_markets
|
||||
|
||||
@with_retry()
|
||||
def get_market(self, condition_id: str) -> Market:
|
||||
"""Fetch a specific market by its condition ID.
|
||||
|
||||
Args:
|
||||
condition_id: The market's condition ID.
|
||||
|
||||
Returns:
|
||||
Market object.
|
||||
|
||||
Raises:
|
||||
ClobClientError: If the market is not found.
|
||||
"""
|
||||
self._rate_limiter.acquire_sync()
|
||||
|
||||
try:
|
||||
response = self._client.get_market(condition_id)
|
||||
return Market.from_dict(response)
|
||||
except Exception as e:
|
||||
raise ClobClientError(f"Failed to fetch market {condition_id}: {e}") from e
|
||||
|
||||
@with_retry()
|
||||
def get_orderbook(self, token_id: str) -> Orderbook:
|
||||
"""Fetch the orderbook for a specific token.
|
||||
|
||||
Args:
|
||||
token_id: The token ID to fetch the orderbook for.
|
||||
|
||||
Returns:
|
||||
Orderbook object with bids, asks, and spread information.
|
||||
"""
|
||||
self._rate_limiter.acquire_sync()
|
||||
|
||||
try:
|
||||
orderbook = self._client.get_order_book(token_id)
|
||||
return Orderbook.from_clob_orderbook(orderbook)
|
||||
except Exception as e:
|
||||
raise ClobClientError(f"Failed to fetch orderbook for {token_id}: {e}") from e
|
||||
|
||||
@with_retry()
|
||||
def get_orderbooks(self, token_ids: list[str]) -> list[Orderbook]:
|
||||
"""Fetch orderbooks for multiple tokens in a single request.
|
||||
|
||||
Args:
|
||||
token_ids: List of token IDs to fetch orderbooks for.
|
||||
|
||||
Returns:
|
||||
List of Orderbook objects.
|
||||
"""
|
||||
self._rate_limiter.acquire_sync()
|
||||
|
||||
params = [BookParams(token_id=tid) for tid in token_ids]
|
||||
|
||||
try:
|
||||
orderbooks = self._client.get_order_books(params)
|
||||
return [Orderbook.from_clob_orderbook(ob) for ob in orderbooks]
|
||||
except Exception as e:
|
||||
raise ClobClientError(f"Failed to fetch orderbooks: {e}") from e
|
||||
|
||||
@with_retry()
|
||||
def get_midpoint(self, token_id: str) -> str | None:
|
||||
"""Fetch the midpoint price for a token.
|
||||
|
||||
Args:
|
||||
token_id: The token ID.
|
||||
|
||||
Returns:
|
||||
Midpoint price as a string, or None if unavailable.
|
||||
"""
|
||||
self._rate_limiter.acquire_sync()
|
||||
|
||||
try:
|
||||
response = self._client.get_midpoint(token_id)
|
||||
mid = response.get("mid")
|
||||
return str(mid) if mid is not None else None
|
||||
except Exception as e:
|
||||
logger.warning("Failed to get midpoint for %s: %s", token_id, e)
|
||||
return None
|
||||
|
||||
@with_retry()
|
||||
def get_price(self, token_id: str, side: str = "BUY") -> str | None:
|
||||
"""Fetch the best price for a token on a given side.
|
||||
|
||||
Args:
|
||||
token_id: The token ID.
|
||||
side: Either "BUY" or "SELL".
|
||||
|
||||
Returns:
|
||||
Best price as a string, or None if unavailable.
|
||||
"""
|
||||
self._rate_limiter.acquire_sync()
|
||||
|
||||
try:
|
||||
response = self._client.get_price(token_id, side=side)
|
||||
price = response.get("price")
|
||||
return str(price) if price is not None else None
|
||||
except Exception as e:
|
||||
logger.warning("Failed to get %s price for %s: %s", side, token_id, e)
|
||||
return None
|
||||
|
||||
def health_check(self) -> bool:
|
||||
"""Check if the CLOB API is reachable.
|
||||
|
||||
Returns:
|
||||
True if the API responds with "OK", False otherwise.
|
||||
"""
|
||||
try:
|
||||
self._rate_limiter.acquire_sync()
|
||||
result = self._client.get_ok()
|
||||
return str(result) == "OK"
|
||||
except Exception as e:
|
||||
logger.error("Health check failed: %s", e)
|
||||
return False
|
||||
|
||||
def get_server_time(self) -> int | None:
|
||||
"""Get the server timestamp.
|
||||
|
||||
Returns:
|
||||
Server timestamp in milliseconds, or None on error.
|
||||
"""
|
||||
try:
|
||||
self._rate_limiter.acquire_sync()
|
||||
result = self._client.get_server_time()
|
||||
return int(result) if result is not None else None
|
||||
except Exception as e:
|
||||
logger.error("Failed to get server time: %s", e)
|
||||
return None
|
||||
@@ -0,0 +1,510 @@
|
||||
"""Connection health monitor with metrics and HTTP endpoints.
|
||||
|
||||
This module provides health monitoring for the data ingestion layer,
|
||||
tracking connection states, event throughput, and staleness detection.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import copy
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from aiohttp import web
|
||||
from prometheus_client import Counter, Gauge, Histogram, generate_latest
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_STALE_THRESHOLD_SECONDS = 60 # No events for 60s = stale
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL = 5 # seconds
|
||||
DEFAULT_HTTP_PORT = 8080
|
||||
|
||||
|
||||
class HealthStatus(Enum):
|
||||
"""Overall health status."""
|
||||
|
||||
HEALTHY = "healthy"
|
||||
DEGRADED = "degraded"
|
||||
UNHEALTHY = "unhealthy"
|
||||
|
||||
|
||||
class StreamStatus(Enum):
|
||||
"""Status of an individual stream."""
|
||||
|
||||
ACTIVE = "active"
|
||||
STALE = "stale"
|
||||
DISCONNECTED = "disconnected"
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamHealth:
|
||||
"""Health status for an individual stream."""
|
||||
|
||||
name: str
|
||||
status: StreamStatus = StreamStatus.DISCONNECTED
|
||||
last_event_time: float | None = None
|
||||
events_received: int = 0
|
||||
events_per_second: float = 0.0
|
||||
connected_since: float | None = None
|
||||
last_error: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class HealthReport:
|
||||
"""Comprehensive health report for all streams."""
|
||||
|
||||
status: HealthStatus
|
||||
streams: dict[str, StreamHealth] = field(default_factory=dict)
|
||||
total_events_received: int = 0
|
||||
total_events_per_second: float = 0.0
|
||||
uptime_seconds: float = 0.0
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
|
||||
|
||||
# Type aliases
|
||||
HealthCallback = Callable[[HealthReport], Awaitable[None]]
|
||||
|
||||
|
||||
# Prometheus metrics
|
||||
EVENTS_TOTAL = Counter(
|
||||
"polymarket_events_total",
|
||||
"Total number of events received",
|
||||
["stream"],
|
||||
)
|
||||
|
||||
EVENTS_PER_SECOND = Gauge(
|
||||
"polymarket_events_per_second",
|
||||
"Current events per second rate",
|
||||
["stream"],
|
||||
)
|
||||
|
||||
STREAM_STATUS = Gauge(
|
||||
"polymarket_stream_status",
|
||||
"Stream status (1=active, 0.5=stale, 0=disconnected)",
|
||||
["stream"],
|
||||
)
|
||||
|
||||
LAST_EVENT_TIMESTAMP = Gauge(
|
||||
"polymarket_last_event_timestamp",
|
||||
"Unix timestamp of last event received",
|
||||
["stream"],
|
||||
)
|
||||
|
||||
EVENT_LATENCY = Histogram(
|
||||
"polymarket_event_latency_seconds",
|
||||
"Event processing latency in seconds",
|
||||
["stream"],
|
||||
buckets=[0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0],
|
||||
)
|
||||
|
||||
HEALTH_STATUS = Gauge(
|
||||
"polymarket_health_status",
|
||||
"Overall health status (1=healthy, 0.5=degraded, 0=unhealthy)",
|
||||
)
|
||||
|
||||
|
||||
class HealthMonitor:
|
||||
"""Monitor connection health and expose metrics.
|
||||
|
||||
This class tracks the health of multiple streams, calculates throughput,
|
||||
detects stale streams, and exposes Prometheus-compatible metrics.
|
||||
|
||||
Example:
|
||||
```python
|
||||
monitor = HealthMonitor(stale_threshold_seconds=60)
|
||||
await monitor.start()
|
||||
|
||||
# Record events
|
||||
monitor.record_event("trades", processing_time=0.001)
|
||||
|
||||
# Update connection state
|
||||
monitor.set_stream_connected("trades")
|
||||
|
||||
# Get health report
|
||||
report = monitor.get_health_report()
|
||||
|
||||
# HTTP endpoints: /health and /metrics
|
||||
# Start HTTP server with monitor.start_http_server(port=8080)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
stale_threshold_seconds: float = DEFAULT_STALE_THRESHOLD_SECONDS,
|
||||
health_check_interval: float = DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
on_health_change: HealthCallback | None = None,
|
||||
) -> None:
|
||||
"""Initialize the health monitor.
|
||||
|
||||
Args:
|
||||
stale_threshold_seconds: Seconds without events before stream is stale.
|
||||
health_check_interval: Seconds between health check updates.
|
||||
on_health_change: Optional callback when health status changes.
|
||||
"""
|
||||
self._stale_threshold = stale_threshold_seconds
|
||||
self._health_check_interval = health_check_interval
|
||||
self._on_health_change = on_health_change
|
||||
|
||||
self._streams: dict[str, StreamHealth] = {}
|
||||
self._start_time: float | None = None
|
||||
self._running = False
|
||||
self._health_task: asyncio.Task[None] | None = None
|
||||
self._last_health_status: HealthStatus | None = None
|
||||
|
||||
# For throughput calculation
|
||||
self._event_windows: dict[str, list[float]] = {}
|
||||
self._window_duration = 10.0 # 10 second sliding window
|
||||
|
||||
# HTTP server
|
||||
self._app: web.Application | None = None
|
||||
self._runner: web.AppRunner | None = None
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Return True if the monitor is running."""
|
||||
return self._running
|
||||
|
||||
def register_stream(self, name: str) -> None:
|
||||
"""Register a stream for monitoring.
|
||||
|
||||
Args:
|
||||
name: Unique name for the stream.
|
||||
"""
|
||||
if name not in self._streams:
|
||||
self._streams[name] = StreamHealth(name=name)
|
||||
self._event_windows[name] = []
|
||||
logger.info("Registered stream for monitoring: %s", name)
|
||||
|
||||
def set_stream_connected(self, name: str) -> None:
|
||||
"""Mark a stream as connected.
|
||||
|
||||
Args:
|
||||
name: Stream name.
|
||||
"""
|
||||
self.register_stream(name)
|
||||
stream = self._streams[name]
|
||||
stream.status = StreamStatus.ACTIVE
|
||||
stream.connected_since = time.time()
|
||||
stream.last_error = None
|
||||
STREAM_STATUS.labels(stream=name).set(1.0)
|
||||
logger.debug("Stream connected: %s", name)
|
||||
|
||||
def set_stream_disconnected(self, name: str, error: str | None = None) -> None:
|
||||
"""Mark a stream as disconnected.
|
||||
|
||||
Args:
|
||||
name: Stream name.
|
||||
error: Optional error message.
|
||||
"""
|
||||
self.register_stream(name)
|
||||
stream = self._streams[name]
|
||||
stream.status = StreamStatus.DISCONNECTED
|
||||
stream.connected_since = None
|
||||
stream.last_error = error
|
||||
STREAM_STATUS.labels(stream=name).set(0.0)
|
||||
logger.debug("Stream disconnected: %s (error: %s)", name, error)
|
||||
|
||||
def record_event(
|
||||
self,
|
||||
stream_name: str,
|
||||
*,
|
||||
processing_time: float | None = None,
|
||||
) -> None:
|
||||
"""Record an event received from a stream.
|
||||
|
||||
Args:
|
||||
stream_name: Name of the stream.
|
||||
processing_time: Optional processing latency in seconds.
|
||||
"""
|
||||
self.register_stream(stream_name)
|
||||
now = time.time()
|
||||
|
||||
stream = self._streams[stream_name]
|
||||
stream.events_received += 1
|
||||
stream.last_event_time = now
|
||||
stream.status = StreamStatus.ACTIVE
|
||||
|
||||
# Update metrics
|
||||
EVENTS_TOTAL.labels(stream=stream_name).inc()
|
||||
LAST_EVENT_TIMESTAMP.labels(stream=stream_name).set(now)
|
||||
STREAM_STATUS.labels(stream=stream_name).set(1.0)
|
||||
|
||||
if processing_time is not None:
|
||||
EVENT_LATENCY.labels(stream=stream_name).observe(processing_time)
|
||||
|
||||
# Add to sliding window for throughput
|
||||
window = self._event_windows[stream_name]
|
||||
window.append(now)
|
||||
|
||||
# Clean old entries from window
|
||||
cutoff = now - self._window_duration
|
||||
self._event_windows[stream_name] = [t for t in window if t > cutoff]
|
||||
|
||||
def _calculate_throughput(self, stream_name: str) -> float:
|
||||
"""Calculate events per second for a stream.
|
||||
|
||||
Args:
|
||||
stream_name: Name of the stream.
|
||||
|
||||
Returns:
|
||||
Events per second rate.
|
||||
"""
|
||||
window = self._event_windows.get(stream_name, [])
|
||||
if not window:
|
||||
return 0.0
|
||||
|
||||
now = time.time()
|
||||
cutoff = now - self._window_duration
|
||||
|
||||
# Count events in window
|
||||
recent_events = [t for t in window if t > cutoff]
|
||||
if not recent_events:
|
||||
return 0.0
|
||||
|
||||
# Calculate rate based on window
|
||||
window_span = now - cutoff
|
||||
return len(recent_events) / window_span if window_span > 0 else 0.0
|
||||
|
||||
def _check_stream_staleness(self) -> None:
|
||||
"""Check all streams for staleness."""
|
||||
now = time.time()
|
||||
|
||||
for name, stream in self._streams.items():
|
||||
if stream.status == StreamStatus.DISCONNECTED:
|
||||
continue
|
||||
|
||||
if stream.last_event_time is None:
|
||||
# Connected but no events yet - check connection time
|
||||
if stream.connected_since:
|
||||
since_connect = now - stream.connected_since
|
||||
if since_connect > self._stale_threshold:
|
||||
stream.status = StreamStatus.STALE
|
||||
STREAM_STATUS.labels(stream=name).set(0.5)
|
||||
else:
|
||||
since_event = now - stream.last_event_time
|
||||
if since_event > self._stale_threshold:
|
||||
stream.status = StreamStatus.STALE
|
||||
STREAM_STATUS.labels(stream=name).set(0.5)
|
||||
else:
|
||||
stream.status = StreamStatus.ACTIVE
|
||||
STREAM_STATUS.labels(stream=name).set(1.0)
|
||||
|
||||
def _determine_overall_status(self) -> HealthStatus:
|
||||
"""Determine overall health status based on stream states.
|
||||
|
||||
Returns:
|
||||
Overall health status.
|
||||
"""
|
||||
if not self._streams:
|
||||
return HealthStatus.HEALTHY # No streams = healthy (nothing to monitor)
|
||||
|
||||
statuses = [s.status for s in self._streams.values()]
|
||||
|
||||
if all(s == StreamStatus.DISCONNECTED for s in statuses):
|
||||
return HealthStatus.UNHEALTHY
|
||||
|
||||
if any(s == StreamStatus.DISCONNECTED for s in statuses):
|
||||
return HealthStatus.DEGRADED
|
||||
|
||||
if any(s == StreamStatus.STALE for s in statuses):
|
||||
return HealthStatus.DEGRADED
|
||||
|
||||
return HealthStatus.HEALTHY
|
||||
|
||||
def get_health_report(self) -> HealthReport:
|
||||
"""Generate a comprehensive health report.
|
||||
|
||||
Returns:
|
||||
HealthReport with current status of all streams.
|
||||
"""
|
||||
self._check_stream_staleness()
|
||||
|
||||
# Update throughput metrics
|
||||
total_eps = 0.0
|
||||
for name, stream in self._streams.items():
|
||||
eps = self._calculate_throughput(name)
|
||||
stream.events_per_second = eps
|
||||
EVENTS_PER_SECOND.labels(stream=name).set(eps)
|
||||
total_eps += eps
|
||||
|
||||
overall_status = self._determine_overall_status()
|
||||
HEALTH_STATUS.set(
|
||||
1.0
|
||||
if overall_status == HealthStatus.HEALTHY
|
||||
else 0.5
|
||||
if overall_status == HealthStatus.DEGRADED
|
||||
else 0.0
|
||||
)
|
||||
|
||||
uptime = 0.0
|
||||
if self._start_time:
|
||||
uptime = time.time() - self._start_time
|
||||
|
||||
# Deep copy streams to prevent mutations affecting internal state
|
||||
streams_copy = {name: copy.copy(stream) for name, stream in self._streams.items()}
|
||||
|
||||
return HealthReport(
|
||||
status=overall_status,
|
||||
streams=streams_copy,
|
||||
total_events_received=sum(s.events_received for s in self._streams.values()),
|
||||
total_events_per_second=total_eps,
|
||||
uptime_seconds=uptime,
|
||||
)
|
||||
|
||||
async def _health_check_loop(self) -> None:
|
||||
"""Background task for periodic health checks."""
|
||||
while self._running:
|
||||
try:
|
||||
report = self.get_health_report()
|
||||
|
||||
# Notify on status change
|
||||
if self._on_health_change and report.status != self._last_health_status:
|
||||
self._last_health_status = report.status
|
||||
try:
|
||||
await self._on_health_change(report)
|
||||
except Exception as e:
|
||||
logger.error("Error in health change callback: %s", e)
|
||||
|
||||
await asyncio.sleep(self._health_check_interval)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error("Error in health check loop: %s", e)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the health monitor.
|
||||
|
||||
Begins periodic health checks and staleness detection.
|
||||
"""
|
||||
if self._running:
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._start_time = time.time()
|
||||
self._health_task = asyncio.create_task(self._health_check_loop())
|
||||
logger.info("Health monitor started")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the health monitor."""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
self._running = False
|
||||
|
||||
if self._health_task:
|
||||
self._health_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._health_task
|
||||
self._health_task = None
|
||||
|
||||
await self.stop_http_server()
|
||||
logger.info("Health monitor stopped")
|
||||
|
||||
# HTTP Server methods
|
||||
|
||||
async def _handle_health(self, _request: web.Request) -> web.Response:
|
||||
"""Handle /health endpoint."""
|
||||
report = self.get_health_report()
|
||||
|
||||
status_code = 200 if report.status == HealthStatus.HEALTHY else 503
|
||||
|
||||
body: dict[str, Any] = {
|
||||
"status": report.status.value,
|
||||
"uptime_seconds": report.uptime_seconds,
|
||||
"total_events_received": report.total_events_received,
|
||||
"total_events_per_second": round(report.total_events_per_second, 2),
|
||||
"streams": {},
|
||||
}
|
||||
|
||||
for name, stream in report.streams.items():
|
||||
body["streams"][name] = {
|
||||
"status": stream.status.value,
|
||||
"events_received": stream.events_received,
|
||||
"events_per_second": round(stream.events_per_second, 2),
|
||||
"last_event_time": stream.last_event_time,
|
||||
"last_error": stream.last_error,
|
||||
}
|
||||
|
||||
return web.json_response(body, status=status_code)
|
||||
|
||||
async def _handle_metrics(self, _request: web.Request) -> web.Response:
|
||||
"""Handle /metrics endpoint (Prometheus format)."""
|
||||
# Ensure latest values are calculated
|
||||
self.get_health_report()
|
||||
|
||||
metrics = generate_latest()
|
||||
return web.Response(
|
||||
body=metrics,
|
||||
content_type="text/plain",
|
||||
charset="utf-8",
|
||||
)
|
||||
|
||||
async def _handle_ready(self, _request: web.Request) -> web.Response:
|
||||
"""Handle /ready endpoint for k8s readiness probe."""
|
||||
report = self.get_health_report()
|
||||
|
||||
if report.status == HealthStatus.UNHEALTHY:
|
||||
return web.json_response(
|
||||
{"ready": False, "reason": "unhealthy"},
|
||||
status=503,
|
||||
)
|
||||
|
||||
return web.json_response({"ready": True}, status=200)
|
||||
|
||||
async def _handle_live(self, _request: web.Request) -> web.Response:
|
||||
"""Handle /live endpoint for k8s liveness probe."""
|
||||
# Always return 200 if the server is running
|
||||
return web.json_response({"live": True}, status=200)
|
||||
|
||||
def _create_app(self) -> web.Application:
|
||||
"""Create the aiohttp application."""
|
||||
app = web.Application()
|
||||
app.router.add_get("/health", self._handle_health)
|
||||
app.router.add_get("/metrics", self._handle_metrics)
|
||||
app.router.add_get("/ready", self._handle_ready)
|
||||
app.router.add_get("/live", self._handle_live)
|
||||
return app
|
||||
|
||||
async def start_http_server(self, port: int = DEFAULT_HTTP_PORT) -> None:
|
||||
"""Start the HTTP server for health and metrics endpoints.
|
||||
|
||||
Args:
|
||||
port: Port to listen on.
|
||||
"""
|
||||
if self._runner:
|
||||
logger.warning("HTTP server already running")
|
||||
return
|
||||
|
||||
self._app = self._create_app()
|
||||
self._runner = web.AppRunner(self._app)
|
||||
await self._runner.setup()
|
||||
|
||||
site = web.TCPSite(self._runner, "0.0.0.0", port)
|
||||
await site.start()
|
||||
|
||||
logger.info("Health HTTP server started on port %d", port)
|
||||
|
||||
async def stop_http_server(self) -> None:
|
||||
"""Stop the HTTP server."""
|
||||
if self._runner:
|
||||
await self._runner.cleanup()
|
||||
self._runner = None
|
||||
self._app = None
|
||||
logger.info("Health HTTP server stopped")
|
||||
|
||||
async def __aenter__(self) -> "HealthMonitor":
|
||||
"""Async context manager entry."""
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.stop()
|
||||
@@ -0,0 +1,362 @@
|
||||
"""Market metadata synchronizer with Redis caching.
|
||||
|
||||
This module provides a background sync service that keeps market metadata
|
||||
up-to-date in Redis, with cache-first lookups for fast access.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from enum import Enum
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from .clob_client import ClobClient
|
||||
from .models import MarketMetadata
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_SYNC_INTERVAL_SECONDS = 300 # 5 minutes
|
||||
DEFAULT_CACHE_TTL_SECONDS = 600 # 10 minutes
|
||||
DEFAULT_REDIS_KEY_PREFIX = "polymarket:market:"
|
||||
|
||||
|
||||
class SyncState(str, Enum):
|
||||
"""State of the metadata synchronizer."""
|
||||
|
||||
STOPPED = "stopped"
|
||||
STARTING = "starting"
|
||||
SYNCING = "syncing"
|
||||
IDLE = "idle"
|
||||
STOPPING = "stopping"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SyncStats:
|
||||
"""Statistics for the metadata sync process."""
|
||||
|
||||
total_syncs: int = 0
|
||||
successful_syncs: int = 0
|
||||
failed_syncs: int = 0
|
||||
markets_cached: int = 0
|
||||
last_sync_time: datetime | None = None
|
||||
last_sync_duration_seconds: float = 0.0
|
||||
last_error: str | None = None
|
||||
|
||||
|
||||
# Type aliases for callbacks
|
||||
StateCallback = Callable[[SyncState], None]
|
||||
SyncCallback = Callable[[SyncStats], None]
|
||||
|
||||
|
||||
class MetadataSyncError(Exception):
|
||||
"""Base exception for metadata sync errors."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class MarketMetadataSync:
|
||||
"""Background service that syncs market metadata to Redis.
|
||||
|
||||
This service:
|
||||
- Fetches all markets from the CLOB API on startup
|
||||
- Refreshes the cache every sync_interval_seconds (default: 5 minutes)
|
||||
- Stores market metadata in Redis with TTL-based expiration
|
||||
- Provides cache-first lookups via get_market()
|
||||
|
||||
Example:
|
||||
```python
|
||||
redis = Redis.from_url("redis://localhost:6379")
|
||||
clob = ClobClient()
|
||||
|
||||
sync = MarketMetadataSync(redis=redis, clob_client=clob)
|
||||
await sync.start()
|
||||
|
||||
# Get market metadata (cache-first)
|
||||
metadata = await sync.get_market("0x1234...")
|
||||
|
||||
await sync.stop()
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis: Redis,
|
||||
clob_client: ClobClient,
|
||||
*,
|
||||
sync_interval_seconds: int = DEFAULT_SYNC_INTERVAL_SECONDS,
|
||||
cache_ttl_seconds: int = DEFAULT_CACHE_TTL_SECONDS,
|
||||
key_prefix: str = DEFAULT_REDIS_KEY_PREFIX,
|
||||
on_state_change: StateCallback | None = None,
|
||||
on_sync_complete: SyncCallback | None = None,
|
||||
) -> None:
|
||||
"""Initialize the metadata sync service.
|
||||
|
||||
Args:
|
||||
redis: Redis async client for caching.
|
||||
clob_client: CLOB API client for fetching markets.
|
||||
sync_interval_seconds: Interval between syncs (default: 300 / 5 min).
|
||||
cache_ttl_seconds: TTL for cached entries (default: 600 / 10 min).
|
||||
key_prefix: Redis key prefix for market data.
|
||||
on_state_change: Callback for state changes.
|
||||
on_sync_complete: Callback after each sync completes.
|
||||
"""
|
||||
self._redis = redis
|
||||
self._clob = clob_client
|
||||
self._sync_interval = sync_interval_seconds
|
||||
self._cache_ttl = cache_ttl_seconds
|
||||
self._key_prefix = key_prefix
|
||||
self._on_state_change = on_state_change
|
||||
self._on_sync_complete = on_sync_complete
|
||||
|
||||
self._state = SyncState.STOPPED
|
||||
self._stats = SyncStats()
|
||||
self._sync_task: asyncio.Task[None] | None = None
|
||||
self._stop_event = asyncio.Event()
|
||||
|
||||
@property
|
||||
def state(self) -> SyncState:
|
||||
"""Current sync state."""
|
||||
return self._state
|
||||
|
||||
@property
|
||||
def stats(self) -> SyncStats:
|
||||
"""Current sync statistics."""
|
||||
return self._stats
|
||||
|
||||
def _set_state(self, new_state: SyncState) -> None:
|
||||
"""Update state and notify callback."""
|
||||
old_state = self._state
|
||||
self._state = new_state
|
||||
if self._on_state_change and old_state != new_state:
|
||||
try:
|
||||
self._on_state_change(new_state)
|
||||
except Exception as e:
|
||||
logger.warning(f"State change callback failed: {e}")
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the background sync service.
|
||||
|
||||
This will:
|
||||
1. Perform an initial sync of all markets
|
||||
2. Start a background task to periodically refresh
|
||||
"""
|
||||
if self._state != SyncState.STOPPED:
|
||||
logger.warning(f"Cannot start sync: already in state {self._state}")
|
||||
return
|
||||
|
||||
self._set_state(SyncState.STARTING)
|
||||
self._stop_event.clear()
|
||||
|
||||
# Perform initial sync
|
||||
try:
|
||||
await self._sync_all_markets()
|
||||
except Exception as e:
|
||||
logger.error(f"Initial sync failed: {e}")
|
||||
self._set_state(SyncState.ERROR)
|
||||
self._stats.last_error = str(e)
|
||||
raise MetadataSyncError(f"Failed to start: initial sync failed: {e}") from e
|
||||
|
||||
# Start background sync loop
|
||||
self._sync_task = asyncio.create_task(self._sync_loop())
|
||||
self._set_state(SyncState.IDLE)
|
||||
logger.info("Market metadata sync started")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the background sync service."""
|
||||
if self._state == SyncState.STOPPED:
|
||||
return
|
||||
|
||||
self._set_state(SyncState.STOPPING)
|
||||
self._stop_event.set()
|
||||
|
||||
if self._sync_task:
|
||||
self._sync_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._sync_task
|
||||
self._sync_task = None
|
||||
|
||||
self._set_state(SyncState.STOPPED)
|
||||
logger.info("Market metadata sync stopped")
|
||||
|
||||
async def _sync_loop(self) -> None:
|
||||
"""Background loop that periodically syncs markets."""
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
# Wait for next sync interval or stop event
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._stop_event.wait(),
|
||||
timeout=self._sync_interval,
|
||||
)
|
||||
# Stop event was set
|
||||
break
|
||||
except TimeoutError:
|
||||
# Timeout - time to sync
|
||||
pass
|
||||
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
|
||||
await self._sync_all_markets()
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Sync loop error: {e}")
|
||||
self._stats.failed_syncs += 1
|
||||
self._stats.last_error = str(e)
|
||||
self._set_state(SyncState.ERROR)
|
||||
# Continue running - will retry on next interval
|
||||
|
||||
async def _sync_all_markets(self) -> None:
|
||||
"""Fetch all markets and cache them in Redis."""
|
||||
self._set_state(SyncState.SYNCING)
|
||||
start_time = datetime.now(UTC)
|
||||
self._stats.total_syncs += 1
|
||||
|
||||
try:
|
||||
# Fetch markets from CLOB API (runs in thread pool for sync API)
|
||||
markets = await asyncio.to_thread(self._clob.get_markets, True)
|
||||
|
||||
# Cache each market in Redis
|
||||
cached_count = 0
|
||||
for market in markets:
|
||||
try:
|
||||
metadata = MarketMetadata.from_market(market)
|
||||
await self._cache_market(metadata)
|
||||
cached_count += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to cache market {market.condition_id}: {e}")
|
||||
|
||||
# Update stats
|
||||
end_time = datetime.now(UTC)
|
||||
self._stats.successful_syncs += 1
|
||||
self._stats.markets_cached = cached_count
|
||||
self._stats.last_sync_time = end_time
|
||||
self._stats.last_sync_duration_seconds = (end_time - start_time).total_seconds()
|
||||
self._stats.last_error = None
|
||||
|
||||
self._set_state(SyncState.IDLE)
|
||||
logger.info(
|
||||
f"Synced {cached_count} markets in {self._stats.last_sync_duration_seconds:.2f}s"
|
||||
)
|
||||
|
||||
# Notify callback
|
||||
if self._on_sync_complete:
|
||||
try:
|
||||
self._on_sync_complete(self._stats)
|
||||
except Exception as e:
|
||||
logger.warning(f"Sync complete callback failed: {e}")
|
||||
|
||||
except Exception as e:
|
||||
self._stats.failed_syncs += 1
|
||||
self._stats.last_error = str(e)
|
||||
self._set_state(SyncState.ERROR)
|
||||
logger.error(f"Market sync failed: {e}")
|
||||
raise
|
||||
|
||||
async def _cache_market(self, metadata: MarketMetadata) -> None:
|
||||
"""Cache a single market metadata in Redis.
|
||||
|
||||
Args:
|
||||
metadata: The market metadata to cache.
|
||||
"""
|
||||
key = f"{self._key_prefix}{metadata.condition_id}"
|
||||
value = json.dumps(metadata.to_dict())
|
||||
await self._redis.setex(key, self._cache_ttl, value)
|
||||
|
||||
async def get_market(self, condition_id: str) -> MarketMetadata | None:
|
||||
"""Get market metadata with cache-first lookup.
|
||||
|
||||
This first checks Redis cache. If not found or expired,
|
||||
it fetches from the CLOB API and caches the result.
|
||||
|
||||
Args:
|
||||
condition_id: The market condition ID.
|
||||
|
||||
Returns:
|
||||
MarketMetadata if found, None otherwise.
|
||||
"""
|
||||
# Try cache first
|
||||
key = f"{self._key_prefix}{condition_id}"
|
||||
cached = await self._redis.get(key)
|
||||
|
||||
if cached:
|
||||
try:
|
||||
data = json.loads(cached)
|
||||
return MarketMetadata.from_dict(data)
|
||||
except (json.JSONDecodeError, KeyError) as e:
|
||||
logger.warning(f"Failed to parse cached market {condition_id}: {e}")
|
||||
|
||||
# Cache miss - fetch from API
|
||||
try:
|
||||
market = await asyncio.to_thread(self._clob.get_market, condition_id)
|
||||
if market:
|
||||
metadata = MarketMetadata.from_market(market)
|
||||
await self._cache_market(metadata)
|
||||
return metadata
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to fetch market {condition_id}: {e}")
|
||||
|
||||
return None
|
||||
|
||||
async def get_markets_by_category(self, category: str) -> list[MarketMetadata]:
|
||||
"""Get all cached markets of a specific category.
|
||||
|
||||
Note: This scans all cached markets. For large datasets,
|
||||
consider using a Redis set or secondary index.
|
||||
|
||||
Args:
|
||||
category: The category to filter by.
|
||||
|
||||
Returns:
|
||||
List of matching MarketMetadata.
|
||||
"""
|
||||
results: list[MarketMetadata] = []
|
||||
pattern = f"{self._key_prefix}*"
|
||||
|
||||
cursor = 0
|
||||
while True:
|
||||
cursor, keys = await self._redis.scan(cursor, match=pattern, count=100)
|
||||
for key in keys:
|
||||
cached = await self._redis.get(key)
|
||||
if cached:
|
||||
try:
|
||||
data = json.loads(cached)
|
||||
if data.get("category") == category:
|
||||
results.append(MarketMetadata.from_dict(data))
|
||||
except (json.JSONDecodeError, KeyError):
|
||||
pass
|
||||
if cursor == 0:
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
async def invalidate_market(self, condition_id: str) -> bool:
|
||||
"""Invalidate (delete) a cached market.
|
||||
|
||||
Args:
|
||||
condition_id: The market condition ID to invalidate.
|
||||
|
||||
Returns:
|
||||
True if the key was deleted, False if it didn't exist.
|
||||
"""
|
||||
key = f"{self._key_prefix}{condition_id}"
|
||||
deleted = await self._redis.delete(key)
|
||||
return int(deleted) > 0
|
||||
|
||||
async def force_sync(self) -> None:
|
||||
"""Force an immediate sync of all markets.
|
||||
|
||||
This can be called to refresh the cache outside of the
|
||||
normal sync interval.
|
||||
"""
|
||||
await self._sync_all_markets()
|
||||
@@ -0,0 +1,490 @@
|
||||
"""Data models for the ingestor module."""
|
||||
|
||||
import contextlib
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Token:
|
||||
"""Represents a token in a Polymarket market."""
|
||||
|
||||
token_id: str
|
||||
outcome: str
|
||||
price: Decimal | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "Token":
|
||||
"""Create a Token from a dictionary."""
|
||||
price = data.get("price")
|
||||
return cls(
|
||||
token_id=str(data["token_id"]),
|
||||
outcome=str(data["outcome"]),
|
||||
price=Decimal(str(price)) if price is not None else None,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Market:
|
||||
"""Represents a Polymarket prediction market."""
|
||||
|
||||
condition_id: str
|
||||
question: str
|
||||
description: str
|
||||
tokens: tuple[Token, ...]
|
||||
end_date: datetime | None = None
|
||||
active: bool = True
|
||||
closed: bool = False
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "Market":
|
||||
"""Create a Market from a dictionary response."""
|
||||
tokens_data = data.get("tokens", [])
|
||||
tokens = tuple(Token.from_dict(t) for t in tokens_data)
|
||||
|
||||
end_date = None
|
||||
end_date_iso = data.get("end_date_iso")
|
||||
if end_date_iso:
|
||||
with contextlib.suppress(ValueError, AttributeError):
|
||||
end_date = datetime.fromisoformat(end_date_iso.replace("Z", "+00:00"))
|
||||
|
||||
return cls(
|
||||
condition_id=str(data["condition_id"]),
|
||||
question=str(data.get("question", "")),
|
||||
description=str(data.get("description", "")),
|
||||
tokens=tokens,
|
||||
end_date=end_date,
|
||||
active=bool(data.get("active", True)),
|
||||
closed=bool(data.get("closed", False)),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OrderbookLevel:
|
||||
"""Represents a single price level in an orderbook."""
|
||||
|
||||
price: Decimal
|
||||
size: Decimal
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "OrderbookLevel":
|
||||
"""Create an OrderbookLevel from a dictionary."""
|
||||
return cls(
|
||||
price=Decimal(str(data["price"])),
|
||||
size=Decimal(str(data["size"])),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Orderbook:
|
||||
"""Represents an orderbook for a Polymarket token."""
|
||||
|
||||
market: str
|
||||
asset_id: str
|
||||
bids: tuple[OrderbookLevel, ...]
|
||||
asks: tuple[OrderbookLevel, ...]
|
||||
tick_size: Decimal
|
||||
timestamp: datetime = field(default_factory=datetime.utcnow)
|
||||
|
||||
@classmethod
|
||||
def from_clob_orderbook(cls, orderbook: Any) -> "Orderbook":
|
||||
"""Create an Orderbook from a py-clob-client orderbook object."""
|
||||
bids = tuple(
|
||||
OrderbookLevel(
|
||||
price=Decimal(str(bid.price)),
|
||||
size=Decimal(str(bid.size)),
|
||||
)
|
||||
for bid in (orderbook.bids or [])
|
||||
)
|
||||
asks = tuple(
|
||||
OrderbookLevel(
|
||||
price=Decimal(str(ask.price)),
|
||||
size=Decimal(str(ask.size)),
|
||||
)
|
||||
for ask in (orderbook.asks or [])
|
||||
)
|
||||
|
||||
return cls(
|
||||
market=str(orderbook.market),
|
||||
asset_id=str(orderbook.asset_id),
|
||||
bids=bids,
|
||||
asks=asks,
|
||||
tick_size=Decimal(str(orderbook.tick_size)),
|
||||
)
|
||||
|
||||
@property
|
||||
def best_bid(self) -> Decimal | None:
|
||||
"""Return the best bid price, or None if no bids."""
|
||||
return self.bids[0].price if self.bids else None
|
||||
|
||||
@property
|
||||
def best_ask(self) -> Decimal | None:
|
||||
"""Return the best ask price, or None if no asks."""
|
||||
return self.asks[0].price if self.asks else None
|
||||
|
||||
@property
|
||||
def spread(self) -> Decimal | None:
|
||||
"""Return the bid-ask spread, or None if missing data."""
|
||||
if self.best_bid is not None and self.best_ask is not None:
|
||||
return self.best_ask - self.best_bid
|
||||
return None
|
||||
|
||||
@property
|
||||
def midpoint(self) -> Decimal | None:
|
||||
"""Return the midpoint price, or None if missing data."""
|
||||
if self.best_bid is not None and self.best_ask is not None:
|
||||
return (self.best_bid + self.best_ask) / 2
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TradeEvent:
|
||||
"""Represents a trade event from the Polymarket WebSocket feed.
|
||||
|
||||
This captures all the information about a single trade execution,
|
||||
including the market, wallet, trade details, and metadata.
|
||||
"""
|
||||
|
||||
# Core trade identifiers
|
||||
market_id: str # conditionId - the market/CTF condition ID
|
||||
trade_id: str # transactionHash - unique trade identifier
|
||||
wallet_address: str # proxyWallet - trader's wallet address
|
||||
|
||||
# Trade details
|
||||
side: Literal["BUY", "SELL"]
|
||||
outcome: str # Human-readable outcome (e.g., "Yes", "No")
|
||||
outcome_index: int # Index of the outcome (0 or 1)
|
||||
price: Decimal
|
||||
size: Decimal # Number of shares traded
|
||||
timestamp: datetime
|
||||
|
||||
# Asset information
|
||||
asset_id: str # ERC1155 token ID
|
||||
|
||||
# Market metadata
|
||||
market_slug: str = ""
|
||||
event_slug: str = ""
|
||||
event_title: str = ""
|
||||
|
||||
# Trader metadata (optional - may not be available for all trades)
|
||||
trader_name: str = ""
|
||||
trader_pseudonym: str = ""
|
||||
|
||||
@classmethod
|
||||
def from_websocket_message(cls, data: dict[str, Any]) -> "TradeEvent":
|
||||
"""Create a TradeEvent from a WebSocket activity/trade message.
|
||||
|
||||
Args:
|
||||
data: The payload from a WebSocket trade message.
|
||||
|
||||
Returns:
|
||||
TradeEvent instance.
|
||||
"""
|
||||
# Parse timestamp - it's a Unix timestamp in seconds
|
||||
raw_timestamp = data.get("timestamp", 0)
|
||||
if isinstance(raw_timestamp, int):
|
||||
timestamp = datetime.fromtimestamp(raw_timestamp, tz=UTC)
|
||||
else:
|
||||
timestamp = datetime.now(UTC)
|
||||
|
||||
# Parse side - normalize to uppercase
|
||||
side_raw = str(data.get("side", "BUY")).upper()
|
||||
side: Literal["BUY", "SELL"] = "BUY" if side_raw == "BUY" else "SELL"
|
||||
|
||||
return cls(
|
||||
market_id=str(data.get("conditionId", "")),
|
||||
trade_id=str(data.get("transactionHash", "")),
|
||||
wallet_address=str(data.get("proxyWallet", "")),
|
||||
side=side,
|
||||
outcome=str(data.get("outcome", "")),
|
||||
outcome_index=int(data.get("outcomeIndex", 0)),
|
||||
price=Decimal(str(data.get("price", 0))),
|
||||
size=Decimal(str(data.get("size", 0))),
|
||||
timestamp=timestamp,
|
||||
asset_id=str(data.get("asset", "")),
|
||||
market_slug=str(data.get("slug", "")),
|
||||
event_slug=str(data.get("eventSlug", "")),
|
||||
event_title=str(data.get("title", "")),
|
||||
trader_name=str(data.get("name", "")),
|
||||
trader_pseudonym=str(data.get("pseudonym", "")),
|
||||
)
|
||||
|
||||
@property
|
||||
def is_buy(self) -> bool:
|
||||
"""Return True if this is a buy trade."""
|
||||
return self.side == "BUY"
|
||||
|
||||
@property
|
||||
def is_sell(self) -> bool:
|
||||
"""Return True if this is a sell trade."""
|
||||
return self.side == "SELL"
|
||||
|
||||
@property
|
||||
def notional_value(self) -> Decimal:
|
||||
"""Return the notional value of the trade (price * size)."""
|
||||
return self.price * self.size
|
||||
|
||||
|
||||
# Category keywords for market classification
|
||||
_CATEGORY_KEYWORDS: dict[str, list[str]] = {
|
||||
"politics": [
|
||||
"election",
|
||||
"president",
|
||||
"congress",
|
||||
"senate",
|
||||
"house",
|
||||
"governor",
|
||||
"mayor",
|
||||
"vote",
|
||||
"ballot",
|
||||
"democrat",
|
||||
"republican",
|
||||
"trump",
|
||||
"biden",
|
||||
"political",
|
||||
"party",
|
||||
"campaign",
|
||||
"poll",
|
||||
"primary",
|
||||
"caucus",
|
||||
],
|
||||
"crypto": [
|
||||
"bitcoin",
|
||||
"ethereum",
|
||||
"crypto",
|
||||
"btc",
|
||||
"eth",
|
||||
"blockchain",
|
||||
"token",
|
||||
"defi",
|
||||
"nft",
|
||||
"altcoin",
|
||||
"solana",
|
||||
"cardano",
|
||||
"dogecoin",
|
||||
],
|
||||
"sports": [
|
||||
"nfl",
|
||||
"nba",
|
||||
"mlb",
|
||||
"nhl",
|
||||
"soccer",
|
||||
"football",
|
||||
"basketball",
|
||||
"baseball",
|
||||
"hockey",
|
||||
"tennis",
|
||||
"golf",
|
||||
"ufc",
|
||||
"boxing",
|
||||
"olympics",
|
||||
"championship",
|
||||
"super bowl",
|
||||
"world cup",
|
||||
"playoffs",
|
||||
"finals",
|
||||
],
|
||||
"entertainment": [
|
||||
"movie",
|
||||
"film",
|
||||
"oscar",
|
||||
"grammy",
|
||||
"emmy",
|
||||
"album",
|
||||
"song",
|
||||
"celebrity",
|
||||
"netflix",
|
||||
"disney",
|
||||
"streaming",
|
||||
"box office",
|
||||
"tv show",
|
||||
"series",
|
||||
"actor",
|
||||
"actress",
|
||||
"music",
|
||||
],
|
||||
"finance": [
|
||||
"stock",
|
||||
"market",
|
||||
"fed",
|
||||
"interest rate",
|
||||
"inflation",
|
||||
"gdp",
|
||||
"unemployment",
|
||||
"recession",
|
||||
"economy",
|
||||
"s&p",
|
||||
"nasdaq",
|
||||
"dow",
|
||||
"treasury",
|
||||
"bond",
|
||||
"forex",
|
||||
"gold",
|
||||
"oil",
|
||||
"commodity",
|
||||
],
|
||||
"tech": [
|
||||
"apple",
|
||||
"google",
|
||||
"microsoft",
|
||||
"amazon",
|
||||
"meta",
|
||||
"tesla",
|
||||
"ai",
|
||||
"artificial intelligence",
|
||||
"chatgpt",
|
||||
"openai",
|
||||
"semiconductor",
|
||||
"iphone",
|
||||
"android",
|
||||
"software",
|
||||
"hardware",
|
||||
"startup",
|
||||
],
|
||||
"science": [
|
||||
"nasa",
|
||||
"space",
|
||||
"climate",
|
||||
"weather",
|
||||
"vaccine",
|
||||
"covid",
|
||||
"fda",
|
||||
"drug",
|
||||
"trial",
|
||||
"research",
|
||||
"study",
|
||||
"discovery",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def derive_category(title: str) -> str:
|
||||
"""Derive a market category from the market title.
|
||||
|
||||
Args:
|
||||
title: The market question or title.
|
||||
|
||||
Returns:
|
||||
Category string, or "other" if no match found.
|
||||
"""
|
||||
title_lower = title.lower()
|
||||
|
||||
for category, keywords in _CATEGORY_KEYWORDS.items():
|
||||
for keyword in keywords:
|
||||
if keyword in title_lower:
|
||||
return category
|
||||
|
||||
return "other"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MarketMetadata:
|
||||
"""Extended market metadata with derived fields and caching support.
|
||||
|
||||
This combines the core Market data with derived metadata like category
|
||||
and is designed for efficient caching in Redis.
|
||||
"""
|
||||
|
||||
# Core market data
|
||||
condition_id: str
|
||||
question: str
|
||||
description: str
|
||||
tokens: tuple[Token, ...]
|
||||
end_date: datetime | None = None
|
||||
active: bool = True
|
||||
closed: bool = False
|
||||
|
||||
# Derived metadata
|
||||
category: str = "other"
|
||||
|
||||
# Cache metadata
|
||||
last_updated: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
@classmethod
|
||||
def from_market(cls, market: Market) -> "MarketMetadata":
|
||||
"""Create MarketMetadata from a Market object.
|
||||
|
||||
Args:
|
||||
market: The source Market object.
|
||||
|
||||
Returns:
|
||||
MarketMetadata with derived fields populated.
|
||||
"""
|
||||
return cls(
|
||||
condition_id=market.condition_id,
|
||||
question=market.question,
|
||||
description=market.description,
|
||||
tokens=market.tokens,
|
||||
end_date=market.end_date,
|
||||
active=market.active,
|
||||
closed=market.closed,
|
||||
category=derive_category(market.question),
|
||||
last_updated=datetime.now(UTC),
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialize to a dictionary for Redis storage.
|
||||
|
||||
Returns:
|
||||
Dictionary representation suitable for JSON serialization.
|
||||
"""
|
||||
return {
|
||||
"condition_id": self.condition_id,
|
||||
"question": self.question,
|
||||
"description": self.description,
|
||||
"tokens": [
|
||||
{
|
||||
"token_id": t.token_id,
|
||||
"outcome": t.outcome,
|
||||
"price": str(t.price) if t.price is not None else None,
|
||||
}
|
||||
for t in self.tokens
|
||||
],
|
||||
"end_date": self.end_date.isoformat() if self.end_date else None,
|
||||
"active": self.active,
|
||||
"closed": self.closed,
|
||||
"category": self.category,
|
||||
"last_updated": self.last_updated.isoformat(),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "MarketMetadata":
|
||||
"""Deserialize from a dictionary (from Redis storage).
|
||||
|
||||
Args:
|
||||
data: Dictionary from Redis.
|
||||
|
||||
Returns:
|
||||
MarketMetadata instance.
|
||||
"""
|
||||
tokens_data = data.get("tokens", [])
|
||||
tokens = tuple(Token.from_dict(t) for t in tokens_data)
|
||||
|
||||
end_date = None
|
||||
end_date_str = data.get("end_date")
|
||||
if end_date_str:
|
||||
with contextlib.suppress(ValueError, AttributeError):
|
||||
end_date = datetime.fromisoformat(end_date_str)
|
||||
|
||||
last_updated_str = data.get("last_updated")
|
||||
if last_updated_str:
|
||||
try:
|
||||
last_updated = datetime.fromisoformat(last_updated_str)
|
||||
except (ValueError, AttributeError):
|
||||
last_updated = datetime.now(UTC)
|
||||
else:
|
||||
last_updated = datetime.now(UTC)
|
||||
|
||||
return cls(
|
||||
condition_id=str(data["condition_id"]),
|
||||
question=str(data.get("question", "")),
|
||||
description=str(data.get("description", "")),
|
||||
tokens=tokens,
|
||||
end_date=end_date,
|
||||
active=bool(data.get("active", True)),
|
||||
closed=bool(data.get("closed", False)),
|
||||
category=str(data.get("category", "other")),
|
||||
last_updated=last_updated,
|
||||
)
|
||||
@@ -0,0 +1,424 @@
|
||||
"""Redis Streams event publisher for trade events.
|
||||
|
||||
This module provides an event publisher that writes normalized trade events
|
||||
to Redis Streams, enabling downstream consumers to process events independently.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Any, Literal
|
||||
|
||||
from redis.asyncio import Redis
|
||||
from redis.exceptions import ResponseError
|
||||
|
||||
from .models import TradeEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_STREAM_NAME = "trades"
|
||||
DEFAULT_MAX_LEN = 100_000 # 100k events
|
||||
DEFAULT_BLOCK_MS = 1000
|
||||
DEFAULT_COUNT = 10
|
||||
|
||||
|
||||
class PublisherError(Exception):
|
||||
"""Base exception for publisher errors."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ConsumerGroupExistsError(PublisherError):
|
||||
"""Raised when trying to create a consumer group that already exists."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamEntry:
|
||||
"""Represents an entry read from a Redis Stream."""
|
||||
|
||||
entry_id: str
|
||||
event: TradeEvent
|
||||
|
||||
|
||||
def _serialize_trade_event(event: TradeEvent) -> dict[str, str]:
|
||||
"""Serialize a TradeEvent to a dict suitable for Redis Streams.
|
||||
|
||||
Redis Streams require string key-value pairs, so we convert all
|
||||
values to strings.
|
||||
|
||||
Args:
|
||||
event: The TradeEvent to serialize.
|
||||
|
||||
Returns:
|
||||
Dictionary with string keys and values.
|
||||
"""
|
||||
return {
|
||||
"market_id": event.market_id,
|
||||
"trade_id": event.trade_id,
|
||||
"wallet_address": event.wallet_address,
|
||||
"side": event.side,
|
||||
"outcome": event.outcome,
|
||||
"outcome_index": str(event.outcome_index),
|
||||
"price": str(event.price),
|
||||
"size": str(event.size),
|
||||
"timestamp": event.timestamp.isoformat(),
|
||||
"asset_id": event.asset_id,
|
||||
"market_slug": event.market_slug,
|
||||
"event_slug": event.event_slug,
|
||||
"event_title": event.event_title,
|
||||
"trader_name": event.trader_name,
|
||||
"trader_pseudonym": event.trader_pseudonym,
|
||||
}
|
||||
|
||||
|
||||
def _deserialize_trade_event(data: dict[bytes | str, bytes | str]) -> TradeEvent:
|
||||
"""Deserialize a TradeEvent from Redis Stream data.
|
||||
|
||||
Args:
|
||||
data: The raw data from Redis Stream (may have bytes keys/values).
|
||||
|
||||
Returns:
|
||||
TradeEvent instance.
|
||||
"""
|
||||
# Convert bytes to strings if needed
|
||||
decoded: dict[str, str] = {}
|
||||
for k, v in data.items():
|
||||
key = k.decode() if isinstance(k, bytes) else k
|
||||
value = v.decode() if isinstance(v, bytes) else v
|
||||
decoded[key] = value
|
||||
|
||||
# Parse timestamp
|
||||
timestamp_str = decoded.get("timestamp", "")
|
||||
try:
|
||||
timestamp = datetime.fromisoformat(timestamp_str)
|
||||
except (ValueError, TypeError):
|
||||
timestamp = datetime.now(UTC)
|
||||
|
||||
# Parse side
|
||||
side_raw = decoded.get("side", "BUY").upper()
|
||||
side: Literal["BUY", "SELL"] = "BUY" if side_raw == "BUY" else "SELL"
|
||||
|
||||
return TradeEvent(
|
||||
market_id=decoded.get("market_id", ""),
|
||||
trade_id=decoded.get("trade_id", ""),
|
||||
wallet_address=decoded.get("wallet_address", ""),
|
||||
side=side,
|
||||
outcome=decoded.get("outcome", ""),
|
||||
outcome_index=int(decoded.get("outcome_index", "0")),
|
||||
price=Decimal(decoded.get("price", "0")),
|
||||
size=Decimal(decoded.get("size", "0")),
|
||||
timestamp=timestamp,
|
||||
asset_id=decoded.get("asset_id", ""),
|
||||
market_slug=decoded.get("market_slug", ""),
|
||||
event_slug=decoded.get("event_slug", ""),
|
||||
event_title=decoded.get("event_title", ""),
|
||||
trader_name=decoded.get("trader_name", ""),
|
||||
trader_pseudonym=decoded.get("trader_pseudonym", ""),
|
||||
)
|
||||
|
||||
|
||||
class EventPublisher:
|
||||
"""Event publisher using Redis Streams.
|
||||
|
||||
This class wraps Redis Streams to provide:
|
||||
- Publishing single or batch trade events
|
||||
- Consumer group management
|
||||
- Event reading for consumers
|
||||
|
||||
Example:
|
||||
```python
|
||||
redis = Redis.from_url("redis://localhost:6379")
|
||||
publisher = EventPublisher(redis)
|
||||
|
||||
# Publish events
|
||||
event_id = await publisher.publish(trade_event)
|
||||
|
||||
# Create consumer group for downstream processing
|
||||
await publisher.create_consumer_group("wallet-profiler")
|
||||
|
||||
# Read events as consumer
|
||||
entries = await publisher.read_events(
|
||||
group_name="wallet-profiler",
|
||||
consumer_name="worker-1"
|
||||
)
|
||||
for entry in entries:
|
||||
process(entry.event)
|
||||
await publisher.ack(entry.entry_id)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis: Redis,
|
||||
stream_name: str = DEFAULT_STREAM_NAME,
|
||||
*,
|
||||
max_len: int = DEFAULT_MAX_LEN,
|
||||
) -> None:
|
||||
"""Initialize the event publisher.
|
||||
|
||||
Args:
|
||||
redis: Redis async client.
|
||||
stream_name: Name of the Redis Stream.
|
||||
max_len: Maximum number of entries to keep in stream.
|
||||
"""
|
||||
self._redis = redis
|
||||
self._stream_name = stream_name
|
||||
self._max_len = max_len
|
||||
|
||||
@property
|
||||
def stream_name(self) -> str:
|
||||
"""Return the stream name."""
|
||||
return self._stream_name
|
||||
|
||||
async def publish(self, event: TradeEvent) -> str:
|
||||
"""Publish a single trade event to the stream.
|
||||
|
||||
Args:
|
||||
event: The TradeEvent to publish.
|
||||
|
||||
Returns:
|
||||
The entry ID assigned by Redis.
|
||||
"""
|
||||
data = _serialize_trade_event(event)
|
||||
# redis-py typing expects broader dict type than dict[str, str]
|
||||
entry_id = await self._redis.xadd(
|
||||
self._stream_name,
|
||||
data, # type: ignore[arg-type]
|
||||
maxlen=self._max_len,
|
||||
)
|
||||
# entry_id may be bytes or str
|
||||
if isinstance(entry_id, bytes):
|
||||
return entry_id.decode()
|
||||
return str(entry_id)
|
||||
|
||||
async def publish_batch(self, events: Sequence[TradeEvent]) -> list[str]:
|
||||
"""Publish multiple trade events atomically.
|
||||
|
||||
Uses a Redis pipeline for efficiency.
|
||||
|
||||
Args:
|
||||
events: Sequence of TradeEvents to publish.
|
||||
|
||||
Returns:
|
||||
List of entry IDs assigned by Redis.
|
||||
"""
|
||||
if not events:
|
||||
return []
|
||||
|
||||
pipe = self._redis.pipeline()
|
||||
for event in events:
|
||||
data = _serialize_trade_event(event)
|
||||
# redis-py typing expects broader dict type than dict[str, str]
|
||||
pipe.xadd(self._stream_name, data, maxlen=self._max_len) # type: ignore[arg-type]
|
||||
|
||||
results = await pipe.execute()
|
||||
|
||||
entry_ids: list[str] = []
|
||||
for entry_id in results:
|
||||
if isinstance(entry_id, bytes):
|
||||
entry_ids.append(entry_id.decode())
|
||||
else:
|
||||
entry_ids.append(str(entry_id))
|
||||
|
||||
return entry_ids
|
||||
|
||||
async def create_consumer_group(
|
||||
self,
|
||||
group_name: str,
|
||||
start_id: str = "0",
|
||||
*,
|
||||
mkstream: bool = True,
|
||||
) -> None:
|
||||
"""Create a consumer group for the stream.
|
||||
|
||||
Args:
|
||||
group_name: Name of the consumer group.
|
||||
start_id: ID to start reading from ("0" = beginning, "$" = new only).
|
||||
mkstream: Create the stream if it doesn't exist.
|
||||
|
||||
Raises:
|
||||
ConsumerGroupExistsError: If the group already exists.
|
||||
"""
|
||||
try:
|
||||
await self._redis.xgroup_create(
|
||||
self._stream_name,
|
||||
group_name,
|
||||
id=start_id,
|
||||
mkstream=mkstream,
|
||||
)
|
||||
logger.info(f"Created consumer group '{group_name}' on stream '{self._stream_name}'")
|
||||
except ResponseError as e:
|
||||
if "BUSYGROUP" in str(e):
|
||||
raise ConsumerGroupExistsError(
|
||||
f"Consumer group '{group_name}' already exists"
|
||||
) from e
|
||||
raise
|
||||
|
||||
async def ensure_consumer_group(
|
||||
self,
|
||||
group_name: str,
|
||||
start_id: str = "0",
|
||||
) -> bool:
|
||||
"""Ensure a consumer group exists, creating it if needed.
|
||||
|
||||
Args:
|
||||
group_name: Name of the consumer group.
|
||||
start_id: ID to start reading from if creating.
|
||||
|
||||
Returns:
|
||||
True if the group was created, False if it already existed.
|
||||
"""
|
||||
try:
|
||||
await self.create_consumer_group(group_name, start_id)
|
||||
return True
|
||||
except ConsumerGroupExistsError:
|
||||
return False
|
||||
|
||||
async def read_events(
|
||||
self,
|
||||
group_name: str,
|
||||
consumer_name: str,
|
||||
*,
|
||||
count: int = DEFAULT_COUNT,
|
||||
block_ms: int = DEFAULT_BLOCK_MS,
|
||||
) -> list[StreamEntry]:
|
||||
"""Read events from the stream as a consumer.
|
||||
|
||||
Args:
|
||||
group_name: Consumer group name.
|
||||
consumer_name: Name of this consumer within the group.
|
||||
count: Maximum number of entries to read.
|
||||
block_ms: Milliseconds to block waiting for new entries.
|
||||
|
||||
Returns:
|
||||
List of StreamEntry with entry ID and TradeEvent.
|
||||
"""
|
||||
# Read new entries (> means entries not delivered to this consumer)
|
||||
results = await self._redis.xreadgroup(
|
||||
group_name,
|
||||
consumer_name,
|
||||
{self._stream_name: ">"},
|
||||
count=count,
|
||||
block=block_ms,
|
||||
)
|
||||
|
||||
entries: list[StreamEntry] = []
|
||||
if not results:
|
||||
return entries
|
||||
|
||||
# Results format: [[stream_name, [(entry_id, data), ...]]]
|
||||
for _stream_name, stream_entries in results:
|
||||
for entry_id, data in stream_entries:
|
||||
# Decode entry_id
|
||||
entry_id_str = entry_id.decode() if isinstance(entry_id, bytes) else str(entry_id)
|
||||
|
||||
try:
|
||||
event = _deserialize_trade_event(data)
|
||||
entries.append(StreamEntry(entry_id=entry_id_str, event=event))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to deserialize entry {entry_id_str}: {e}")
|
||||
|
||||
return entries
|
||||
|
||||
async def read_pending(
|
||||
self,
|
||||
group_name: str,
|
||||
consumer_name: str,
|
||||
*,
|
||||
count: int = DEFAULT_COUNT,
|
||||
) -> list[StreamEntry]:
|
||||
"""Read pending (unacknowledged) entries for a consumer.
|
||||
|
||||
This is useful for recovering from crashes - entries that were
|
||||
delivered but not acknowledged will be re-read.
|
||||
|
||||
Args:
|
||||
group_name: Consumer group name.
|
||||
consumer_name: Name of this consumer.
|
||||
count: Maximum number of entries to read.
|
||||
|
||||
Returns:
|
||||
List of pending StreamEntry.
|
||||
"""
|
||||
# Read pending entries (0 means all pending entries)
|
||||
results = await self._redis.xreadgroup(
|
||||
group_name,
|
||||
consumer_name,
|
||||
{self._stream_name: "0"},
|
||||
count=count,
|
||||
)
|
||||
|
||||
entries: list[StreamEntry] = []
|
||||
if not results:
|
||||
return entries
|
||||
|
||||
for _stream_name, stream_entries in results:
|
||||
for entry_id, data in stream_entries:
|
||||
entry_id_str = entry_id.decode() if isinstance(entry_id, bytes) else str(entry_id)
|
||||
|
||||
# Skip entries with no data (already acked)
|
||||
if not data:
|
||||
continue
|
||||
|
||||
try:
|
||||
event = _deserialize_trade_event(data)
|
||||
entries.append(StreamEntry(entry_id=entry_id_str, event=event))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to deserialize pending entry {entry_id_str}: {e}")
|
||||
|
||||
return entries
|
||||
|
||||
async def ack(self, group_name: str, *entry_ids: str) -> int:
|
||||
"""Acknowledge that entries have been processed.
|
||||
|
||||
Args:
|
||||
group_name: Consumer group name.
|
||||
*entry_ids: Entry IDs to acknowledge.
|
||||
|
||||
Returns:
|
||||
Number of entries acknowledged.
|
||||
"""
|
||||
if not entry_ids:
|
||||
return 0
|
||||
result = await self._redis.xack(self._stream_name, group_name, *entry_ids)
|
||||
return int(result)
|
||||
|
||||
async def get_stream_info(self) -> dict[str, Any]:
|
||||
"""Get information about the stream.
|
||||
|
||||
Returns:
|
||||
Dictionary with stream info (length, groups, etc.).
|
||||
"""
|
||||
try:
|
||||
info = await self._redis.xinfo_stream(self._stream_name)
|
||||
return dict(info) if info else {}
|
||||
except ResponseError:
|
||||
return {}
|
||||
|
||||
async def get_stream_length(self) -> int:
|
||||
"""Get the current length of the stream.
|
||||
|
||||
Returns:
|
||||
Number of entries in the stream.
|
||||
"""
|
||||
result = await self._redis.xlen(self._stream_name)
|
||||
return int(result)
|
||||
|
||||
async def trim_stream(self, max_len: int | None = None) -> int:
|
||||
"""Trim the stream to a maximum length.
|
||||
|
||||
Args:
|
||||
max_len: Maximum entries to keep (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
Number of entries removed.
|
||||
"""
|
||||
length = max_len or self._max_len
|
||||
result = await self._redis.xtrim(self._stream_name, maxlen=length)
|
||||
return int(result)
|
||||
@@ -0,0 +1,338 @@
|
||||
"""WebSocket client for streaming Polymarket trade events."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
import websockets
|
||||
from websockets.asyncio.client import ClientConnection
|
||||
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Constants
|
||||
DEFAULT_WS_HOST = "wss://ws-live-data.polymarket.com"
|
||||
DEFAULT_PING_INTERVAL = 30 # seconds
|
||||
DEFAULT_MAX_RECONNECT_DELAY = 30 # seconds
|
||||
DEFAULT_INITIAL_RECONNECT_DELAY = 1 # seconds
|
||||
|
||||
|
||||
class ConnectionState(Enum):
|
||||
"""WebSocket connection states."""
|
||||
|
||||
DISCONNECTED = "disconnected"
|
||||
CONNECTING = "connecting"
|
||||
CONNECTED = "connected"
|
||||
RECONNECTING = "reconnecting"
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamStats:
|
||||
"""Statistics about the trade stream."""
|
||||
|
||||
trades_received: int = 0
|
||||
reconnect_count: int = 0
|
||||
last_trade_time: float | None = None
|
||||
connected_since: float | None = None
|
||||
last_error: str | None = None
|
||||
|
||||
|
||||
class TradeStreamError(Exception):
|
||||
"""Base exception for trade stream errors."""
|
||||
|
||||
|
||||
class ConnectionError(TradeStreamError):
|
||||
"""Raised when connection to WebSocket fails."""
|
||||
|
||||
|
||||
TradeCallback = Callable[[TradeEvent], Awaitable[None]]
|
||||
StateCallback = Callable[[ConnectionState], Awaitable[None]]
|
||||
|
||||
|
||||
class TradeStreamHandler:
|
||||
"""WebSocket client for streaming Polymarket trade events.
|
||||
|
||||
This handler maintains a persistent connection to Polymarket's real-time
|
||||
trade feed, automatically reconnecting on disconnection with exponential
|
||||
backoff.
|
||||
|
||||
Example:
|
||||
>>> async def on_trade(trade: TradeEvent):
|
||||
... print(f"Trade: {trade.side} {trade.size} @ {trade.price}")
|
||||
...
|
||||
>>> handler = TradeStreamHandler(on_trade=on_trade)
|
||||
>>> await handler.start() # Blocks until stop() is called
|
||||
|
||||
Attributes:
|
||||
state: Current connection state.
|
||||
stats: Statistics about the stream (trades received, reconnects, etc.).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
on_trade: TradeCallback,
|
||||
*,
|
||||
host: str = DEFAULT_WS_HOST,
|
||||
on_state_change: StateCallback | None = None,
|
||||
ping_interval: int = DEFAULT_PING_INTERVAL,
|
||||
max_reconnect_delay: int = DEFAULT_MAX_RECONNECT_DELAY,
|
||||
initial_reconnect_delay: int = DEFAULT_INITIAL_RECONNECT_DELAY,
|
||||
event_filter: str | None = None,
|
||||
market_filter: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize the trade stream handler.
|
||||
|
||||
Args:
|
||||
on_trade: Async callback invoked for each trade event.
|
||||
host: WebSocket endpoint URL.
|
||||
on_state_change: Optional callback for connection state changes.
|
||||
ping_interval: Seconds between heartbeat pings.
|
||||
max_reconnect_delay: Maximum delay between reconnection attempts.
|
||||
initial_reconnect_delay: Initial delay for reconnection backoff.
|
||||
event_filter: Optional event slug to filter trades by event.
|
||||
market_filter: Optional market slug to filter trades by market.
|
||||
"""
|
||||
self._on_trade = on_trade
|
||||
self._on_state_change = on_state_change
|
||||
self._host = host
|
||||
self._ping_interval = ping_interval
|
||||
self._max_reconnect_delay = max_reconnect_delay
|
||||
self._initial_reconnect_delay = initial_reconnect_delay
|
||||
self._event_filter = event_filter
|
||||
self._market_filter = market_filter
|
||||
|
||||
self._state = ConnectionState.DISCONNECTED
|
||||
self._stats = StreamStats()
|
||||
self._ws: ClientConnection | None = None
|
||||
self._running = False
|
||||
self._stop_event: asyncio.Event | None = None
|
||||
|
||||
@property
|
||||
def state(self) -> ConnectionState:
|
||||
"""Current connection state."""
|
||||
return self._state
|
||||
|
||||
@property
|
||||
def stats(self) -> StreamStats:
|
||||
"""Stream statistics."""
|
||||
return self._stats
|
||||
|
||||
async def _set_state(self, new_state: ConnectionState) -> None:
|
||||
"""Update state and notify callback."""
|
||||
if self._state != new_state:
|
||||
old_state = self._state
|
||||
self._state = new_state
|
||||
logger.info("Connection state: %s -> %s", old_state.value, new_state.value)
|
||||
|
||||
if self._on_state_change:
|
||||
try:
|
||||
await self._on_state_change(new_state)
|
||||
except Exception as e:
|
||||
logger.error("Error in state change callback: %s", e)
|
||||
|
||||
def _build_subscription_message(self) -> dict[str, Any]:
|
||||
"""Build the WebSocket subscription message."""
|
||||
subscription: dict[str, Any] = {
|
||||
"topic": "activity",
|
||||
"type": "trades",
|
||||
}
|
||||
|
||||
# Add filters if specified
|
||||
if self._event_filter:
|
||||
subscription["filters"] = json.dumps({"event_slug": self._event_filter})
|
||||
elif self._market_filter:
|
||||
subscription["filters"] = json.dumps({"market_slug": self._market_filter})
|
||||
|
||||
return {"subscriptions": [subscription]}
|
||||
|
||||
async def _connect(self) -> ClientConnection:
|
||||
"""Establish WebSocket connection."""
|
||||
await self._set_state(ConnectionState.CONNECTING)
|
||||
|
||||
try:
|
||||
ws = await websockets.connect(
|
||||
self._host,
|
||||
ping_interval=self._ping_interval,
|
||||
ping_timeout=self._ping_interval * 2,
|
||||
)
|
||||
|
||||
# Send subscription message
|
||||
subscribe_msg = self._build_subscription_message()
|
||||
await ws.send(json.dumps(subscribe_msg))
|
||||
|
||||
logger.info("Connected to %s and subscribed to trades", self._host)
|
||||
await self._set_state(ConnectionState.CONNECTED)
|
||||
self._stats.connected_since = time.time()
|
||||
|
||||
return ws
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to connect: %s", e)
|
||||
self._stats.last_error = str(e)
|
||||
raise ConnectionError(f"Failed to connect to {self._host}: {e}") from e
|
||||
|
||||
async def _handle_message(self, message: str) -> None:
|
||||
"""Parse and process an incoming WebSocket message."""
|
||||
try:
|
||||
data = json.loads(message)
|
||||
|
||||
# Check if this is a trade message
|
||||
topic = data.get("topic")
|
||||
msg_type = data.get("type")
|
||||
|
||||
if topic == "activity" and msg_type == "trades":
|
||||
payload = data.get("payload", {})
|
||||
trade = TradeEvent.from_websocket_message(payload)
|
||||
|
||||
self._stats.trades_received += 1
|
||||
self._stats.last_trade_time = time.time()
|
||||
|
||||
logger.debug(
|
||||
"Trade: %s %s @ %s on %s",
|
||||
trade.side,
|
||||
trade.size,
|
||||
trade.price,
|
||||
trade.market_slug,
|
||||
)
|
||||
|
||||
try:
|
||||
await self._on_trade(trade)
|
||||
except Exception as e:
|
||||
logger.error("Error in trade callback: %s", e)
|
||||
|
||||
else:
|
||||
# Log other message types for debugging
|
||||
logger.debug("Received message: topic=%s type=%s", topic, msg_type)
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning("Invalid JSON message: %s", e)
|
||||
except Exception as e:
|
||||
logger.error("Error processing message: %s", e)
|
||||
|
||||
async def _listen(self, ws: ClientConnection) -> None:
|
||||
"""Listen for messages on the WebSocket."""
|
||||
try:
|
||||
async for message in ws:
|
||||
if not self._running:
|
||||
break
|
||||
|
||||
if isinstance(message, str):
|
||||
await self._handle_message(message)
|
||||
else:
|
||||
logger.debug("Received binary message (%d bytes)", len(message))
|
||||
|
||||
except websockets.ConnectionClosed as e:
|
||||
logger.warning("Connection closed: %s", e)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error("Error in message loop: %s", e)
|
||||
raise
|
||||
|
||||
async def _reconnect_loop(self) -> None:
|
||||
"""Handle reconnection with exponential backoff."""
|
||||
delay = self._initial_reconnect_delay
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
await self._set_state(ConnectionState.RECONNECTING)
|
||||
|
||||
logger.info("Reconnecting in %.1f seconds...", delay)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
if not self._running:
|
||||
break
|
||||
|
||||
self._ws = await self._connect()
|
||||
self._stats.reconnect_count += 1
|
||||
delay = self._initial_reconnect_delay # Reset delay on success
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Reconnection failed: %s", e)
|
||||
self._stats.last_error = str(e)
|
||||
|
||||
# Exponential backoff with jitter
|
||||
delay = min(delay * 2, self._max_reconnect_delay)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Connect and begin streaming trades.
|
||||
|
||||
This method blocks until stop() is called. It automatically
|
||||
handles reconnection on disconnection.
|
||||
|
||||
Raises:
|
||||
ConnectionError: If initial connection fails.
|
||||
"""
|
||||
if self._running:
|
||||
logger.warning("Handler already running")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._stop_event = asyncio.Event()
|
||||
|
||||
try:
|
||||
# Initial connection
|
||||
self._ws = await self._connect()
|
||||
|
||||
# Main loop
|
||||
while self._running:
|
||||
try:
|
||||
await self._listen(self._ws)
|
||||
except (websockets.ConnectionClosed, Exception) as e:
|
||||
if not self._running:
|
||||
break
|
||||
|
||||
logger.warning("Connection lost: %s", e)
|
||||
await self._set_state(ConnectionState.DISCONNECTED)
|
||||
|
||||
# Attempt reconnection
|
||||
await self._reconnect_loop()
|
||||
|
||||
if not self._running or self._ws is None:
|
||||
break
|
||||
|
||||
finally:
|
||||
await self._cleanup()
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Gracefully disconnect from the WebSocket.
|
||||
|
||||
This signals the handler to stop and cleanly close the connection.
|
||||
"""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
logger.info("Stopping trade stream handler...")
|
||||
self._running = False
|
||||
|
||||
if self._stop_event:
|
||||
self._stop_event.set()
|
||||
|
||||
await self._cleanup()
|
||||
|
||||
async def _cleanup(self) -> None:
|
||||
"""Clean up resources."""
|
||||
if self._ws:
|
||||
try:
|
||||
await self._ws.close()
|
||||
except Exception as e:
|
||||
logger.debug("Error closing WebSocket: %s", e)
|
||||
finally:
|
||||
self._ws = None
|
||||
|
||||
await self._set_state(ConnectionState.DISCONNECTED)
|
||||
logger.info("Trade stream handler stopped")
|
||||
|
||||
async def __aenter__(self) -> "TradeStreamHandler":
|
||||
"""Async context manager entry."""
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.stop()
|
||||
@@ -0,0 +1,480 @@
|
||||
"""Main pipeline orchestrator for Polymarket Insider Tracker.
|
||||
|
||||
This module provides the Pipeline class that wires together all detection
|
||||
components and manages the event flow from ingestion to alerting.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from polymarket_insider_tracker.alerter.channels.discord import DiscordChannel
|
||||
from polymarket_insider_tracker.alerter.channels.telegram import TelegramChannel
|
||||
from polymarket_insider_tracker.alerter.dispatcher import AlertChannel, AlertDispatcher
|
||||
from polymarket_insider_tracker.alerter.formatter import AlertFormatter
|
||||
from polymarket_insider_tracker.config import Settings, get_settings
|
||||
from polymarket_insider_tracker.detector.fresh_wallet import FreshWalletDetector
|
||||
from polymarket_insider_tracker.detector.scorer import RiskScorer, SignalBundle
|
||||
from polymarket_insider_tracker.detector.size_anomaly import SizeAnomalyDetector
|
||||
from polymarket_insider_tracker.ingestor.clob_client import ClobClient
|
||||
from polymarket_insider_tracker.ingestor.metadata_sync import MarketMetadataSync
|
||||
from polymarket_insider_tracker.ingestor.websocket import TradeStreamHandler
|
||||
from polymarket_insider_tracker.profiler.analyzer import WalletAnalyzer
|
||||
from polymarket_insider_tracker.profiler.chain import PolygonClient
|
||||
from polymarket_insider_tracker.storage.database import DatabaseManager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any
|
||||
|
||||
from polymarket_insider_tracker.detector.models import (
|
||||
FreshWalletSignal,
|
||||
SizeAnomalySignal,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PipelineState(str, Enum):
|
||||
"""Pipeline lifecycle states."""
|
||||
|
||||
STOPPED = "stopped"
|
||||
STARTING = "starting"
|
||||
RUNNING = "running"
|
||||
STOPPING = "stopping"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class PipelineStats:
|
||||
"""Statistics for the pipeline."""
|
||||
|
||||
started_at: datetime | None = None
|
||||
trades_processed: int = 0
|
||||
signals_generated: int = 0
|
||||
alerts_sent: int = 0
|
||||
errors: int = 0
|
||||
last_trade_time: datetime | None = None
|
||||
last_error: str | None = None
|
||||
|
||||
|
||||
class Pipeline:
|
||||
"""Main pipeline orchestrator for the Polymarket Insider Tracker.
|
||||
|
||||
This class wires together all detection components and manages the
|
||||
event flow from trade ingestion through profiling, detection, and alerting.
|
||||
|
||||
Pipeline flow:
|
||||
WebSocket Trade Stream → Wallet Profiler → Detectors → Risk Scorer → Alerter
|
||||
|
||||
Example:
|
||||
```python
|
||||
from polymarket_insider_tracker.config import get_settings
|
||||
from polymarket_insider_tracker.pipeline import Pipeline
|
||||
|
||||
settings = get_settings()
|
||||
pipeline = Pipeline(settings)
|
||||
|
||||
await pipeline.start()
|
||||
# Pipeline runs until stop() is called
|
||||
await pipeline.stop()
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
settings: Settings | None = None,
|
||||
*,
|
||||
dry_run: bool | None = None,
|
||||
) -> None:
|
||||
"""Initialize the pipeline.
|
||||
|
||||
Args:
|
||||
settings: Application settings. If not provided, uses get_settings().
|
||||
dry_run: If True, skip sending alerts. Overrides settings.dry_run.
|
||||
"""
|
||||
self._settings = settings or get_settings()
|
||||
self._dry_run = dry_run if dry_run is not None else self._settings.dry_run
|
||||
|
||||
self._state = PipelineState.STOPPED
|
||||
self._stats = PipelineStats()
|
||||
|
||||
# Components (initialized in start())
|
||||
self._redis: Redis | None = None
|
||||
self._db_manager: DatabaseManager | None = None
|
||||
self._polygon_client: PolygonClient | None = None
|
||||
self._clob_client: ClobClient | None = None
|
||||
self._metadata_sync: MarketMetadataSync | None = None
|
||||
self._wallet_analyzer: WalletAnalyzer | None = None
|
||||
self._fresh_wallet_detector: FreshWalletDetector | None = None
|
||||
self._size_anomaly_detector: SizeAnomalyDetector | None = None
|
||||
self._risk_scorer: RiskScorer | None = None
|
||||
self._alert_formatter: AlertFormatter | None = None
|
||||
self._alert_dispatcher: AlertDispatcher | None = None
|
||||
self._trade_stream: TradeStreamHandler | None = None
|
||||
|
||||
# Synchronization
|
||||
self._stop_event: asyncio.Event | None = None
|
||||
self._stream_task: asyncio.Task[None] | None = None
|
||||
|
||||
@property
|
||||
def state(self) -> PipelineState:
|
||||
"""Current pipeline state."""
|
||||
return self._state
|
||||
|
||||
@property
|
||||
def stats(self) -> PipelineStats:
|
||||
"""Current pipeline statistics."""
|
||||
return self._stats
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if pipeline is running."""
|
||||
return self._state == PipelineState.RUNNING
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the pipeline.
|
||||
|
||||
Initializes all components and begins processing trades.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If pipeline is already running.
|
||||
Exception: If any component fails to initialize.
|
||||
"""
|
||||
if self._state != PipelineState.STOPPED:
|
||||
raise RuntimeError(f"Cannot start pipeline in state {self._state}")
|
||||
|
||||
self._state = PipelineState.STARTING
|
||||
self._stop_event = asyncio.Event()
|
||||
logger.info("Starting pipeline...")
|
||||
|
||||
try:
|
||||
await self._initialize_components()
|
||||
await self._start_background_services()
|
||||
self._stats.started_at = datetime.now(UTC)
|
||||
self._state = PipelineState.RUNNING
|
||||
logger.info("Pipeline started successfully")
|
||||
except Exception as e:
|
||||
self._state = PipelineState.ERROR
|
||||
self._stats.last_error = str(e)
|
||||
logger.error("Failed to start pipeline: %s", e)
|
||||
await self._cleanup()
|
||||
raise
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the pipeline gracefully.
|
||||
|
||||
Stops all background services and cleans up resources.
|
||||
"""
|
||||
if self._state == PipelineState.STOPPED:
|
||||
return
|
||||
|
||||
self._state = PipelineState.STOPPING
|
||||
logger.info("Stopping pipeline...")
|
||||
|
||||
if self._stop_event:
|
||||
self._stop_event.set()
|
||||
|
||||
await self._stop_background_services()
|
||||
await self._cleanup()
|
||||
|
||||
self._state = PipelineState.STOPPED
|
||||
logger.info("Pipeline stopped")
|
||||
|
||||
async def _initialize_components(self) -> None:
|
||||
"""Initialize all pipeline components."""
|
||||
settings = self._settings
|
||||
|
||||
# Initialize Redis
|
||||
logger.debug("Initializing Redis connection...")
|
||||
self._redis = Redis.from_url(settings.redis.url)
|
||||
|
||||
# Initialize Database Manager
|
||||
logger.debug("Initializing database manager...")
|
||||
self._db_manager = DatabaseManager(
|
||||
settings.database.url,
|
||||
async_mode=True,
|
||||
)
|
||||
|
||||
# Initialize Polygon client
|
||||
logger.debug("Initializing Polygon client...")
|
||||
self._polygon_client = PolygonClient(
|
||||
settings.polygon.rpc_url,
|
||||
fallback_rpc_url=settings.polygon.fallback_rpc_url,
|
||||
redis=self._redis,
|
||||
)
|
||||
|
||||
# Initialize CLOB client
|
||||
logger.debug("Initializing CLOB client...")
|
||||
api_key = (
|
||||
settings.polymarket.api_key.get_secret_value() if settings.polymarket.api_key else None
|
||||
)
|
||||
self._clob_client = ClobClient(api_key=api_key)
|
||||
|
||||
# Initialize Market Metadata Sync
|
||||
logger.debug("Initializing market metadata sync...")
|
||||
self._metadata_sync = MarketMetadataSync(
|
||||
redis=self._redis,
|
||||
clob_client=self._clob_client,
|
||||
)
|
||||
|
||||
# Initialize Wallet Analyzer
|
||||
logger.debug("Initializing wallet analyzer...")
|
||||
self._wallet_analyzer = WalletAnalyzer(
|
||||
self._polygon_client,
|
||||
redis=self._redis,
|
||||
)
|
||||
|
||||
# Initialize Detectors
|
||||
logger.debug("Initializing detectors...")
|
||||
self._fresh_wallet_detector = FreshWalletDetector(self._wallet_analyzer)
|
||||
self._size_anomaly_detector = SizeAnomalyDetector(self._metadata_sync)
|
||||
|
||||
# Initialize Risk Scorer
|
||||
logger.debug("Initializing risk scorer...")
|
||||
self._risk_scorer = RiskScorer(self._redis)
|
||||
|
||||
# Initialize Alerting
|
||||
logger.debug("Initializing alerting components...")
|
||||
self._alert_formatter = AlertFormatter(verbosity="detailed")
|
||||
channels = self._build_alert_channels()
|
||||
self._alert_dispatcher = AlertDispatcher(channels)
|
||||
|
||||
# Initialize Trade Stream
|
||||
logger.debug("Initializing trade stream handler...")
|
||||
self._trade_stream = TradeStreamHandler(
|
||||
on_trade=self._on_trade,
|
||||
host=settings.polymarket.ws_url,
|
||||
)
|
||||
|
||||
logger.info("All components initialized")
|
||||
|
||||
def _build_alert_channels(self) -> list[AlertChannel]:
|
||||
"""Build list of enabled alert channels."""
|
||||
channels: list[AlertChannel] = []
|
||||
settings = self._settings
|
||||
|
||||
if settings.discord.enabled and settings.discord.webhook_url:
|
||||
webhook_url = settings.discord.webhook_url.get_secret_value()
|
||||
channels.append(DiscordChannel(webhook_url))
|
||||
logger.info("Discord channel enabled")
|
||||
|
||||
if settings.telegram.enabled:
|
||||
bot_token = settings.telegram.bot_token
|
||||
chat_id = settings.telegram.chat_id
|
||||
if bot_token and chat_id:
|
||||
channels.append(
|
||||
TelegramChannel(
|
||||
bot_token.get_secret_value(),
|
||||
chat_id,
|
||||
)
|
||||
)
|
||||
logger.info("Telegram channel enabled")
|
||||
|
||||
if not channels:
|
||||
logger.warning("No alert channels configured")
|
||||
|
||||
return channels
|
||||
|
||||
async def _start_background_services(self) -> None:
|
||||
"""Start background services."""
|
||||
# Start metadata sync
|
||||
if self._metadata_sync:
|
||||
logger.debug("Starting metadata sync service...")
|
||||
await self._metadata_sync.start()
|
||||
|
||||
# Start trade stream in background task
|
||||
if self._trade_stream:
|
||||
logger.debug("Starting trade stream...")
|
||||
self._stream_task = asyncio.create_task(self._run_trade_stream())
|
||||
|
||||
async def _run_trade_stream(self) -> None:
|
||||
"""Run the trade stream in a task."""
|
||||
if not self._trade_stream:
|
||||
return
|
||||
|
||||
try:
|
||||
await self._trade_stream.start()
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("Trade stream task cancelled")
|
||||
except Exception as e:
|
||||
logger.error("Trade stream error: %s", e)
|
||||
self._stats.last_error = str(e)
|
||||
self._stats.errors += 1
|
||||
|
||||
async def _stop_background_services(self) -> None:
|
||||
"""Stop background services."""
|
||||
# Stop trade stream
|
||||
if self._trade_stream:
|
||||
logger.debug("Stopping trade stream...")
|
||||
await self._trade_stream.stop()
|
||||
|
||||
# Cancel stream task
|
||||
if self._stream_task:
|
||||
self._stream_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._stream_task
|
||||
self._stream_task = None
|
||||
|
||||
# Stop metadata sync
|
||||
if self._metadata_sync:
|
||||
logger.debug("Stopping metadata sync...")
|
||||
await self._metadata_sync.stop()
|
||||
|
||||
async def _cleanup(self) -> None:
|
||||
"""Clean up resources."""
|
||||
# Close database connections
|
||||
if self._db_manager:
|
||||
await self._db_manager.dispose_async()
|
||||
self._db_manager = None
|
||||
|
||||
# Close Redis connection
|
||||
if self._redis:
|
||||
await self._redis.aclose()
|
||||
self._redis = None
|
||||
|
||||
logger.debug("Resources cleaned up")
|
||||
|
||||
async def _on_trade(self, trade: TradeEvent) -> None:
|
||||
"""Process a single trade event.
|
||||
|
||||
This is the main event handler that runs the detection pipeline:
|
||||
1. Run fresh wallet detection
|
||||
2. Run size anomaly detection
|
||||
3. Score the combined signals
|
||||
4. Send alert if threshold exceeded
|
||||
|
||||
Args:
|
||||
trade: The trade event from the WebSocket stream.
|
||||
"""
|
||||
self._stats.trades_processed += 1
|
||||
self._stats.last_trade_time = datetime.now(UTC)
|
||||
|
||||
try:
|
||||
# Run detectors in parallel
|
||||
fresh_signal, size_signal = await asyncio.gather(
|
||||
self._detect_fresh_wallet(trade),
|
||||
self._detect_size_anomaly(trade),
|
||||
)
|
||||
|
||||
# Bundle signals
|
||||
bundle = SignalBundle(
|
||||
trade_event=trade,
|
||||
fresh_wallet_signal=fresh_signal,
|
||||
size_anomaly_signal=size_signal,
|
||||
)
|
||||
|
||||
# Score and potentially alert
|
||||
if fresh_signal or size_signal:
|
||||
self._stats.signals_generated += 1
|
||||
await self._score_and_alert(bundle)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error processing trade %s: %s", trade.trade_id, e)
|
||||
self._stats.errors += 1
|
||||
self._stats.last_error = str(e)
|
||||
|
||||
async def _detect_fresh_wallet(self, trade: TradeEvent) -> FreshWalletSignal | None:
|
||||
"""Run fresh wallet detection."""
|
||||
if not self._fresh_wallet_detector:
|
||||
return None
|
||||
try:
|
||||
return await self._fresh_wallet_detector.analyze(trade)
|
||||
except Exception as e:
|
||||
logger.warning("Fresh wallet detection failed for %s: %s", trade.trade_id, e)
|
||||
return None
|
||||
|
||||
async def _detect_size_anomaly(self, trade: TradeEvent) -> SizeAnomalySignal | None:
|
||||
"""Run size anomaly detection."""
|
||||
if not self._size_anomaly_detector:
|
||||
return None
|
||||
try:
|
||||
return await self._size_anomaly_detector.analyze(trade)
|
||||
except Exception as e:
|
||||
logger.warning("Size anomaly detection failed for %s: %s", trade.trade_id, e)
|
||||
return None
|
||||
|
||||
async def _score_and_alert(self, bundle: SignalBundle) -> None:
|
||||
"""Score signals and send alert if threshold exceeded."""
|
||||
if not self._risk_scorer or not self._alert_formatter or not self._alert_dispatcher:
|
||||
return
|
||||
|
||||
# Get risk assessment
|
||||
assessment = await self._risk_scorer.assess(bundle)
|
||||
|
||||
if not assessment.should_alert:
|
||||
logger.debug(
|
||||
"Trade %s below alert threshold (score=%.2f)",
|
||||
bundle.trade_event.trade_id,
|
||||
assessment.weighted_score,
|
||||
)
|
||||
return
|
||||
|
||||
# Format and dispatch alert
|
||||
formatted_alert = self._alert_formatter.format(assessment)
|
||||
|
||||
if self._dry_run:
|
||||
logger.info(
|
||||
"[DRY RUN] Would send alert: wallet=%s, score=%.2f",
|
||||
assessment.wallet_address[:10] + "...",
|
||||
assessment.weighted_score,
|
||||
)
|
||||
return
|
||||
|
||||
result = await self._alert_dispatcher.dispatch(formatted_alert)
|
||||
|
||||
if result.all_succeeded:
|
||||
self._stats.alerts_sent += 1
|
||||
logger.info(
|
||||
"Alert sent successfully: wallet=%s, score=%.2f",
|
||||
assessment.wallet_address[:10] + "...",
|
||||
assessment.weighted_score,
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Alert partially failed: %d/%d channels succeeded",
|
||||
result.success_count,
|
||||
result.success_count + result.failure_count,
|
||||
)
|
||||
|
||||
async def run(self) -> None:
|
||||
"""Start the pipeline and run until interrupted.
|
||||
|
||||
This is a convenience method that starts the pipeline and
|
||||
blocks until a stop signal is received.
|
||||
|
||||
Example:
|
||||
```python
|
||||
pipeline = Pipeline()
|
||||
try:
|
||||
await pipeline.run()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
```
|
||||
"""
|
||||
await self.start()
|
||||
|
||||
try:
|
||||
if self._stop_event:
|
||||
await self._stop_event.wait()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
finally:
|
||||
await self.stop()
|
||||
|
||||
async def __aenter__(self) -> Pipeline:
|
||||
"""Async context manager entry."""
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.stop()
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Wallet profiler - Blockchain analysis for trader intelligence."""
|
||||
|
||||
from polymarket_insider_tracker.profiler.analyzer import (
|
||||
WalletAnalyzer,
|
||||
)
|
||||
from polymarket_insider_tracker.profiler.chain import (
|
||||
PolygonClient,
|
||||
PolygonClientError,
|
||||
RateLimitError,
|
||||
RPCError,
|
||||
)
|
||||
from polymarket_insider_tracker.profiler.entities import (
|
||||
EntityRegistry,
|
||||
)
|
||||
from polymarket_insider_tracker.profiler.entity_data import (
|
||||
EntityType,
|
||||
)
|
||||
from polymarket_insider_tracker.profiler.funding import (
|
||||
FundingTracer,
|
||||
)
|
||||
from polymarket_insider_tracker.profiler.models import (
|
||||
FundingChain,
|
||||
FundingTransfer,
|
||||
Transaction,
|
||||
WalletInfo,
|
||||
WalletProfile,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Analyzer
|
||||
"WalletAnalyzer",
|
||||
# Entity Registry
|
||||
"EntityRegistry",
|
||||
"EntityType",
|
||||
# Funding Tracer
|
||||
"FundingChain",
|
||||
"FundingTracer",
|
||||
"FundingTransfer",
|
||||
# Polygon Client
|
||||
"PolygonClient",
|
||||
"PolygonClientError",
|
||||
"RateLimitError",
|
||||
"RPCError",
|
||||
# Models
|
||||
"Transaction",
|
||||
"WalletInfo",
|
||||
"WalletProfile",
|
||||
]
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Wallet analysis for fresh wallet detection.
|
||||
|
||||
This module provides wallet analysis capabilities to identify potentially
|
||||
suspicious wallets based on their on-chain activity patterns.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from polymarket_insider_tracker.profiler.chain import PolygonClient
|
||||
from polymarket_insider_tracker.profiler.models import WalletProfile
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# USDC contract address on Polygon
|
||||
USDC_POLYGON_ADDRESS = "0x3c499c542cEF5E3811e1192ce70d8cC03d5c3359"
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_FRESH_THRESHOLD = 5 # Max nonce to be considered fresh
|
||||
DEFAULT_PROFILE_CACHE_TTL = 300 # 5 minutes
|
||||
|
||||
|
||||
class WalletAnalyzer:
|
||||
"""Analyzes wallets to detect fresh wallet patterns.
|
||||
|
||||
This class provides wallet analysis functionality including:
|
||||
- Fresh wallet detection based on transaction count
|
||||
- Wallet age calculation from first transaction
|
||||
- Balance queries for MATIC and USDC
|
||||
- Caching of analysis results
|
||||
|
||||
Example:
|
||||
```python
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=redis)
|
||||
analyzer = WalletAnalyzer(client, redis=redis)
|
||||
|
||||
# Full analysis
|
||||
profile = await analyzer.analyze("0x...")
|
||||
print(f"Fresh: {profile.is_fresh}, Score: {profile.freshness_score}")
|
||||
|
||||
# Quick check
|
||||
is_fresh = await analyzer.is_fresh("0x...")
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
polygon_client: PolygonClient,
|
||||
*,
|
||||
redis: Redis | None = None,
|
||||
fresh_threshold: int = DEFAULT_FRESH_THRESHOLD,
|
||||
cache_ttl_seconds: int = DEFAULT_PROFILE_CACHE_TTL,
|
||||
usdc_address: str = USDC_POLYGON_ADDRESS,
|
||||
) -> None:
|
||||
"""Initialize the wallet analyzer.
|
||||
|
||||
Args:
|
||||
polygon_client: PolygonClient for blockchain queries.
|
||||
redis: Optional Redis client for caching profiles.
|
||||
fresh_threshold: Maximum nonce to be considered fresh.
|
||||
cache_ttl_seconds: How long to cache analysis results.
|
||||
usdc_address: USDC token contract address on Polygon.
|
||||
"""
|
||||
self._client = polygon_client
|
||||
self._redis = redis
|
||||
self._fresh_threshold = fresh_threshold
|
||||
self._cache_ttl = cache_ttl_seconds
|
||||
self._usdc_address = usdc_address
|
||||
self._cache_prefix = "wallet_profile:"
|
||||
|
||||
def _cache_key(self, address: str) -> str:
|
||||
"""Generate cache key for wallet profile."""
|
||||
return f"{self._cache_prefix}{address.lower()}"
|
||||
|
||||
async def _get_cached_profile(self, address: str) -> WalletProfile | None:
|
||||
"""Get cached profile if available."""
|
||||
if not self._redis:
|
||||
return None
|
||||
|
||||
try:
|
||||
key = self._cache_key(address)
|
||||
cached = await self._redis.get(key)
|
||||
if cached is None:
|
||||
return None
|
||||
|
||||
data = json.loads(cached if isinstance(cached, str) else cached.decode())
|
||||
return WalletProfile(
|
||||
address=data["address"],
|
||||
nonce=data["nonce"],
|
||||
first_seen=datetime.fromisoformat(data["first_seen"])
|
||||
if data["first_seen"]
|
||||
else None,
|
||||
age_hours=data["age_hours"],
|
||||
is_fresh=data["is_fresh"],
|
||||
total_tx_count=data["total_tx_count"],
|
||||
matic_balance=Decimal(data["matic_balance"]),
|
||||
usdc_balance=Decimal(data["usdc_balance"]),
|
||||
analyzed_at=datetime.fromisoformat(data["analyzed_at"]),
|
||||
fresh_threshold=data["fresh_threshold"],
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to get cached profile for %s: %s", address, e)
|
||||
return None
|
||||
|
||||
async def _cache_profile(self, profile: WalletProfile) -> None:
|
||||
"""Cache a wallet profile."""
|
||||
if not self._redis:
|
||||
return
|
||||
|
||||
try:
|
||||
key = self._cache_key(profile.address)
|
||||
data = {
|
||||
"address": profile.address,
|
||||
"nonce": profile.nonce,
|
||||
"first_seen": profile.first_seen.isoformat() if profile.first_seen else None,
|
||||
"age_hours": profile.age_hours,
|
||||
"is_fresh": profile.is_fresh,
|
||||
"total_tx_count": profile.total_tx_count,
|
||||
"matic_balance": str(profile.matic_balance),
|
||||
"usdc_balance": str(profile.usdc_balance),
|
||||
"analyzed_at": profile.analyzed_at.isoformat(),
|
||||
"fresh_threshold": profile.fresh_threshold,
|
||||
}
|
||||
await self._redis.set(key, json.dumps(data), ex=self._cache_ttl)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to cache profile for %s: %s", profile.address, e)
|
||||
|
||||
async def analyze(
|
||||
self,
|
||||
address: str,
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> WalletProfile:
|
||||
"""Analyze a wallet and return its profile.
|
||||
|
||||
This method performs a comprehensive analysis of the wallet including:
|
||||
- Transaction count (nonce)
|
||||
- First transaction timestamp and wallet age
|
||||
- MATIC and USDC balances
|
||||
- Fresh wallet determination
|
||||
|
||||
Args:
|
||||
address: Wallet address to analyze.
|
||||
force_refresh: If True, bypass cache and re-analyze.
|
||||
|
||||
Returns:
|
||||
WalletProfile with analysis results.
|
||||
"""
|
||||
address = address.lower()
|
||||
|
||||
# Check cache unless force refresh
|
||||
if not force_refresh:
|
||||
cached = await self._get_cached_profile(address)
|
||||
if cached is not None:
|
||||
logger.debug("Using cached profile for %s", address)
|
||||
return cached
|
||||
|
||||
# Get wallet info from blockchain
|
||||
wallet_info = await self._client.get_wallet_info(address)
|
||||
|
||||
# Get USDC balance
|
||||
try:
|
||||
usdc_balance = await self._client.get_token_balance(address, self._usdc_address)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to get USDC balance for %s: %s", address, e)
|
||||
usdc_balance = Decimal(0)
|
||||
|
||||
# Calculate age from first transaction
|
||||
first_seen: datetime | None = None
|
||||
age_hours: float | None = None
|
||||
|
||||
if wallet_info.first_transaction is not None:
|
||||
first_seen = wallet_info.first_transaction.timestamp
|
||||
now = datetime.now(UTC)
|
||||
delta = now - first_seen
|
||||
age_hours = delta.total_seconds() / 3600
|
||||
|
||||
# Determine if fresh
|
||||
is_fresh = self._is_wallet_fresh(wallet_info.transaction_count, age_hours)
|
||||
|
||||
# Build profile
|
||||
profile = WalletProfile(
|
||||
address=address,
|
||||
nonce=wallet_info.transaction_count,
|
||||
first_seen=first_seen,
|
||||
age_hours=age_hours,
|
||||
is_fresh=is_fresh,
|
||||
total_tx_count=wallet_info.transaction_count,
|
||||
matic_balance=wallet_info.balance_wei,
|
||||
usdc_balance=usdc_balance,
|
||||
fresh_threshold=self._fresh_threshold,
|
||||
)
|
||||
|
||||
# Cache the result
|
||||
await self._cache_profile(profile)
|
||||
|
||||
return profile
|
||||
|
||||
def _is_wallet_fresh(self, nonce: int, age_hours: float | None) -> bool:
|
||||
"""Determine if wallet should be considered fresh.
|
||||
|
||||
A wallet is fresh if:
|
||||
- Transaction count (nonce) is below the threshold
|
||||
- AND either age is unknown OR age is less than 48 hours
|
||||
|
||||
Args:
|
||||
nonce: Transaction count.
|
||||
age_hours: Wallet age in hours, or None if unknown.
|
||||
|
||||
Returns:
|
||||
True if wallet is fresh.
|
||||
"""
|
||||
# Must have few transactions
|
||||
if nonce >= self._fresh_threshold:
|
||||
return False
|
||||
|
||||
# If age is known, must be recent (within 48 hours)
|
||||
return not (age_hours is not None and age_hours > 48)
|
||||
|
||||
async def is_fresh(self, address: str) -> bool:
|
||||
"""Quick check if wallet is fresh.
|
||||
|
||||
This is a convenience method that returns just the freshness status.
|
||||
It uses cached data if available.
|
||||
|
||||
Args:
|
||||
address: Wallet address to check.
|
||||
|
||||
Returns:
|
||||
True if wallet is fresh.
|
||||
"""
|
||||
profile = await self.analyze(address)
|
||||
return profile.is_fresh
|
||||
|
||||
async def analyze_batch(
|
||||
self,
|
||||
addresses: list[str],
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, WalletProfile]:
|
||||
"""Analyze multiple wallets.
|
||||
|
||||
Analyzes wallets in parallel for efficiency.
|
||||
|
||||
Args:
|
||||
addresses: List of wallet addresses to analyze.
|
||||
force_refresh: If True, bypass cache for all wallets.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping address (lowercase) to WalletProfile.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
results: dict[str, WalletProfile] = {}
|
||||
|
||||
# Analyze all in parallel
|
||||
tasks = [self.analyze(addr, force_refresh=force_refresh) for addr in addresses]
|
||||
profiles = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
for addr, profile in zip(addresses, profiles, strict=True):
|
||||
if isinstance(profile, BaseException):
|
||||
logger.warning("Failed to analyze %s: %s", addr, profile)
|
||||
continue
|
||||
results[addr.lower()] = profile
|
||||
|
||||
return results
|
||||
|
||||
async def get_fresh_wallets(
|
||||
self,
|
||||
addresses: list[str],
|
||||
) -> list[str]:
|
||||
"""Filter addresses to only return fresh wallets.
|
||||
|
||||
Args:
|
||||
addresses: List of wallet addresses to check.
|
||||
|
||||
Returns:
|
||||
List of addresses that are fresh wallets.
|
||||
"""
|
||||
profiles = await self.analyze_batch(addresses)
|
||||
return [addr for addr, profile in profiles.items() if profile.is_fresh]
|
||||
@@ -0,0 +1,525 @@
|
||||
"""Polygon blockchain client with connection pooling and caching.
|
||||
|
||||
This module provides a Polygon client for wallet data queries with:
|
||||
- Connection pooling for concurrent requests
|
||||
- Redis caching to avoid redundant RPC calls
|
||||
- Retry logic with exponential backoff
|
||||
- Rate limiting to respect provider limits
|
||||
- Failover to secondary RPC URL
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from decimal import Decimal
|
||||
from typing import Any, cast
|
||||
|
||||
from redis.asyncio import Redis
|
||||
from web3 import AsyncWeb3
|
||||
from web3.exceptions import Web3Exception
|
||||
from web3.providers import AsyncHTTPProvider
|
||||
|
||||
from polymarket_insider_tracker.profiler.models import Transaction, WalletInfo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_CACHE_TTL_SECONDS = 300 # 5 minutes
|
||||
DEFAULT_MAX_REQUESTS_PER_SECOND = 25
|
||||
DEFAULT_MAX_RETRIES = 3
|
||||
DEFAULT_RETRY_DELAY_SECONDS = 1.0
|
||||
DEFAULT_CONNECTION_POOL_SIZE = 10
|
||||
DEFAULT_REQUEST_TIMEOUT = 30
|
||||
|
||||
|
||||
class PolygonClientError(Exception):
|
||||
"""Base exception for Polygon client errors."""
|
||||
|
||||
|
||||
class RPCError(PolygonClientError):
|
||||
"""Raised when RPC call fails."""
|
||||
|
||||
|
||||
class RateLimitError(PolygonClientError):
|
||||
"""Raised when rate limit is exceeded."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class RateLimiter:
|
||||
"""Token bucket rate limiter."""
|
||||
|
||||
max_tokens: float
|
||||
refill_rate: float # tokens per second
|
||||
tokens: float
|
||||
last_refill: float
|
||||
|
||||
@classmethod
|
||||
def create(cls, max_requests_per_second: float) -> "RateLimiter":
|
||||
"""Create a rate limiter with specified max requests per second."""
|
||||
return cls(
|
||||
max_tokens=max_requests_per_second,
|
||||
refill_rate=max_requests_per_second,
|
||||
tokens=max_requests_per_second,
|
||||
last_refill=time.monotonic(),
|
||||
)
|
||||
|
||||
def _refill(self) -> None:
|
||||
"""Refill tokens based on elapsed time."""
|
||||
now = time.monotonic()
|
||||
elapsed = now - self.last_refill
|
||||
self.tokens = min(self.max_tokens, self.tokens + elapsed * self.refill_rate)
|
||||
self.last_refill = now
|
||||
|
||||
async def acquire(self, tokens: float = 1.0) -> None:
|
||||
"""Acquire tokens, waiting if necessary."""
|
||||
while True:
|
||||
self._refill()
|
||||
if self.tokens >= tokens:
|
||||
self.tokens -= tokens
|
||||
return
|
||||
# Wait for tokens to refill
|
||||
wait_time = (tokens - self.tokens) / self.refill_rate
|
||||
await asyncio.sleep(wait_time)
|
||||
|
||||
|
||||
class PolygonClient:
|
||||
"""Polygon blockchain client with caching and rate limiting.
|
||||
|
||||
Provides efficient access to wallet data with:
|
||||
- Connection pooling for concurrent requests
|
||||
- Redis caching with configurable TTL
|
||||
- Rate limiting to respect provider limits
|
||||
- Retry logic with exponential backoff
|
||||
- Failover to secondary RPC
|
||||
|
||||
Example:
|
||||
```python
|
||||
redis = Redis.from_url("redis://localhost:6379")
|
||||
client = PolygonClient(
|
||||
rpc_url="https://polygon-rpc.com",
|
||||
fallback_rpc_url="https://polygon-bor.publicnode.com",
|
||||
redis=redis,
|
||||
)
|
||||
|
||||
# Get single wallet info
|
||||
nonce = await client.get_transaction_count("0x...")
|
||||
|
||||
# Batch query multiple wallets
|
||||
nonces = await client.get_transaction_counts(["0x...", "0x..."])
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rpc_url: str,
|
||||
*,
|
||||
fallback_rpc_url: str | None = None,
|
||||
redis: Redis | None = None,
|
||||
cache_ttl_seconds: int = DEFAULT_CACHE_TTL_SECONDS,
|
||||
max_requests_per_second: float = DEFAULT_MAX_REQUESTS_PER_SECOND,
|
||||
max_retries: int = DEFAULT_MAX_RETRIES,
|
||||
retry_delay_seconds: float = DEFAULT_RETRY_DELAY_SECONDS,
|
||||
) -> None:
|
||||
"""Initialize the Polygon client.
|
||||
|
||||
Args:
|
||||
rpc_url: Primary Polygon RPC endpoint URL.
|
||||
fallback_rpc_url: Optional fallback RPC URL for failover.
|
||||
redis: Optional Redis client for caching.
|
||||
cache_ttl_seconds: Cache TTL in seconds.
|
||||
max_requests_per_second: Rate limit for RPC calls.
|
||||
max_retries: Maximum retry attempts on failure.
|
||||
retry_delay_seconds: Initial delay between retries.
|
||||
"""
|
||||
self._rpc_url = rpc_url
|
||||
self._fallback_rpc_url = fallback_rpc_url
|
||||
self._redis = redis
|
||||
self._cache_ttl = cache_ttl_seconds
|
||||
self._max_retries = max_retries
|
||||
self._retry_delay = retry_delay_seconds
|
||||
|
||||
# Create web3 instances
|
||||
self._w3 = AsyncWeb3(AsyncHTTPProvider(rpc_url))
|
||||
self._w3_fallback: AsyncWeb3[AsyncHTTPProvider] | None = None
|
||||
if fallback_rpc_url:
|
||||
self._w3_fallback = AsyncWeb3(AsyncHTTPProvider(fallback_rpc_url))
|
||||
|
||||
# Rate limiter
|
||||
self._rate_limiter = RateLimiter.create(max_requests_per_second)
|
||||
|
||||
# Track primary RPC health
|
||||
self._primary_healthy = True
|
||||
self._last_primary_check = 0.0
|
||||
self._primary_recovery_interval = 60.0 # Try primary again after 60s
|
||||
|
||||
# Cache key prefix
|
||||
self._cache_prefix = "polygon:"
|
||||
|
||||
def _cache_key(self, key_type: str, address: str) -> str:
|
||||
"""Generate a cache key."""
|
||||
return f"{self._cache_prefix}{key_type}:{address.lower()}"
|
||||
|
||||
async def _get_cached(self, key: str) -> str | None:
|
||||
"""Get value from cache."""
|
||||
if not self._redis:
|
||||
return None
|
||||
try:
|
||||
value = await self._redis.get(key)
|
||||
if isinstance(value, bytes):
|
||||
return value.decode()
|
||||
return str(value) if value is not None else None
|
||||
except Exception as e:
|
||||
logger.warning("Cache get failed: %s", e)
|
||||
return None
|
||||
|
||||
async def _set_cached(self, key: str, value: str, ttl: int | None = None) -> None:
|
||||
"""Set value in cache."""
|
||||
if not self._redis:
|
||||
return
|
||||
try:
|
||||
await self._redis.set(key, value, ex=ttl or self._cache_ttl)
|
||||
except Exception as e:
|
||||
logger.warning("Cache set failed: %s", e)
|
||||
|
||||
def _should_try_primary(self) -> bool:
|
||||
"""Check if we should try the primary RPC."""
|
||||
if self._primary_healthy:
|
||||
return True
|
||||
# Periodically retry primary
|
||||
now = time.monotonic()
|
||||
if now - self._last_primary_check > self._primary_recovery_interval:
|
||||
self._last_primary_check = now
|
||||
return True
|
||||
return False
|
||||
|
||||
async def _execute_with_retry(
|
||||
self,
|
||||
func_name: str,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Execute an RPC call with retry and failover logic.
|
||||
|
||||
Args:
|
||||
func_name: Name of the web3.eth method to call.
|
||||
*args: Positional arguments for the method.
|
||||
**kwargs: Keyword arguments for the method.
|
||||
|
||||
Returns:
|
||||
Result from the RPC call.
|
||||
|
||||
Raises:
|
||||
RPCError: If all retries and failover fail.
|
||||
"""
|
||||
await self._rate_limiter.acquire()
|
||||
|
||||
last_error: Exception | None = None
|
||||
delay = self._retry_delay
|
||||
|
||||
# Try primary RPC
|
||||
if self._should_try_primary():
|
||||
for attempt in range(self._max_retries):
|
||||
try:
|
||||
method = getattr(self._w3.eth, func_name)
|
||||
result = await method(*args, **kwargs)
|
||||
self._primary_healthy = True
|
||||
return result
|
||||
except Web3Exception as e:
|
||||
last_error = e
|
||||
logger.warning(
|
||||
"Primary RPC %s failed (attempt %d/%d): %s",
|
||||
func_name,
|
||||
attempt + 1,
|
||||
self._max_retries,
|
||||
e,
|
||||
)
|
||||
if attempt < self._max_retries - 1:
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2 # Exponential backoff
|
||||
|
||||
# Mark primary as unhealthy
|
||||
self._primary_healthy = False
|
||||
self._last_primary_check = time.monotonic()
|
||||
|
||||
# Try fallback RPC
|
||||
if self._w3_fallback:
|
||||
delay = self._retry_delay
|
||||
for attempt in range(self._max_retries):
|
||||
try:
|
||||
method = getattr(self._w3_fallback.eth, func_name)
|
||||
result = await method(*args, **kwargs)
|
||||
logger.info("Fallback RPC succeeded for %s", func_name)
|
||||
return result
|
||||
except Web3Exception as e:
|
||||
last_error = e
|
||||
logger.warning(
|
||||
"Fallback RPC %s failed (attempt %d/%d): %s",
|
||||
func_name,
|
||||
attempt + 1,
|
||||
self._max_retries,
|
||||
e,
|
||||
)
|
||||
if attempt < self._max_retries - 1:
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
|
||||
raise RPCError(f"RPC call {func_name} failed after all retries: {last_error}")
|
||||
|
||||
async def get_transaction_count(self, address: str) -> int:
|
||||
"""Get wallet transaction count (nonce).
|
||||
|
||||
Args:
|
||||
address: Wallet address.
|
||||
|
||||
Returns:
|
||||
Transaction count.
|
||||
"""
|
||||
cache_key = self._cache_key("nonce", address)
|
||||
|
||||
# Check cache
|
||||
cached = await self._get_cached(cache_key)
|
||||
if cached is not None:
|
||||
return int(cached)
|
||||
|
||||
# Query blockchain
|
||||
count = await self._execute_with_retry(
|
||||
"get_transaction_count",
|
||||
AsyncWeb3.to_checksum_address(address),
|
||||
)
|
||||
|
||||
# Cache result
|
||||
await self._set_cached(cache_key, str(count))
|
||||
|
||||
return int(count)
|
||||
|
||||
async def get_transaction_counts(
|
||||
self,
|
||||
addresses: Sequence[str],
|
||||
) -> dict[str, int]:
|
||||
"""Batch get transaction counts for multiple addresses.
|
||||
|
||||
Args:
|
||||
addresses: List of wallet addresses.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping address to transaction count.
|
||||
"""
|
||||
if not addresses:
|
||||
return {}
|
||||
|
||||
results: dict[str, int] = {}
|
||||
uncached: list[str] = []
|
||||
|
||||
# Check cache for each address
|
||||
for address in addresses:
|
||||
cache_key = self._cache_key("nonce", address)
|
||||
cached = await self._get_cached(cache_key)
|
||||
if cached is not None:
|
||||
results[address.lower()] = int(cached)
|
||||
else:
|
||||
uncached.append(address)
|
||||
|
||||
# Query uncached addresses concurrently
|
||||
if uncached:
|
||||
tasks = [self.get_transaction_count(addr) for addr in uncached]
|
||||
counts = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
for addr, count in zip(uncached, counts, strict=True):
|
||||
if isinstance(count, BaseException):
|
||||
logger.warning("Failed to get nonce for %s: %s", addr, count)
|
||||
results[addr.lower()] = 0
|
||||
else:
|
||||
results[addr.lower()] = count
|
||||
|
||||
return results
|
||||
|
||||
async def get_balance(self, address: str) -> Decimal:
|
||||
"""Get wallet MATIC balance in Wei.
|
||||
|
||||
Args:
|
||||
address: Wallet address.
|
||||
|
||||
Returns:
|
||||
Balance in Wei as Decimal.
|
||||
"""
|
||||
cache_key = self._cache_key("balance", address)
|
||||
|
||||
# Check cache
|
||||
cached = await self._get_cached(cache_key)
|
||||
if cached is not None:
|
||||
return Decimal(cached)
|
||||
|
||||
# Query blockchain
|
||||
balance = await self._execute_with_retry(
|
||||
"get_balance",
|
||||
AsyncWeb3.to_checksum_address(address),
|
||||
)
|
||||
|
||||
# Cache result
|
||||
await self._set_cached(cache_key, str(balance))
|
||||
|
||||
return Decimal(balance)
|
||||
|
||||
async def get_token_balance(
|
||||
self,
|
||||
address: str,
|
||||
token_address: str,
|
||||
) -> Decimal:
|
||||
"""Get ERC20 token balance.
|
||||
|
||||
Args:
|
||||
address: Wallet address.
|
||||
token_address: ERC20 token contract address.
|
||||
|
||||
Returns:
|
||||
Token balance in smallest unit as Decimal.
|
||||
"""
|
||||
cache_key = self._cache_key(f"token:{token_address.lower()}", address)
|
||||
|
||||
# Check cache
|
||||
cached = await self._get_cached(cache_key)
|
||||
if cached is not None:
|
||||
return Decimal(cached)
|
||||
|
||||
# ERC20 balanceOf ABI
|
||||
erc20_abi = [
|
||||
{
|
||||
"constant": True,
|
||||
"inputs": [{"name": "_owner", "type": "address"}],
|
||||
"name": "balanceOf",
|
||||
"outputs": [{"name": "balance", "type": "uint256"}],
|
||||
"type": "function",
|
||||
}
|
||||
]
|
||||
|
||||
await self._rate_limiter.acquire()
|
||||
|
||||
try:
|
||||
w3 = self._w3 if self._primary_healthy else (self._w3_fallback or self._w3)
|
||||
contract = w3.eth.contract(
|
||||
address=AsyncWeb3.to_checksum_address(token_address),
|
||||
abi=erc20_abi,
|
||||
)
|
||||
balance = await contract.functions.balanceOf(
|
||||
AsyncWeb3.to_checksum_address(address)
|
||||
).call()
|
||||
except Web3Exception as e:
|
||||
raise RPCError(f"Failed to get token balance: {e}") from e
|
||||
|
||||
# Cache result
|
||||
await self._set_cached(cache_key, str(balance))
|
||||
|
||||
return Decimal(balance)
|
||||
|
||||
async def get_block(self, block_number: int) -> dict[str, Any]:
|
||||
"""Get block by number.
|
||||
|
||||
Args:
|
||||
block_number: Block number.
|
||||
|
||||
Returns:
|
||||
Block data dictionary.
|
||||
"""
|
||||
cache_key = f"{self._cache_prefix}block:{block_number}"
|
||||
|
||||
# Check cache
|
||||
cached = await self._get_cached(cache_key)
|
||||
if cached is not None:
|
||||
return cast(dict[str, Any], json.loads(cached))
|
||||
|
||||
block = await self._execute_with_retry("get_block", block_number)
|
||||
|
||||
# Convert to serializable dict
|
||||
block_dict = dict(block)
|
||||
block_dict["timestamp"] = int(block_dict["timestamp"])
|
||||
|
||||
# Cache result (blocks are immutable, use longer TTL)
|
||||
await self._set_cached(cache_key, json.dumps(block_dict), ttl=3600)
|
||||
|
||||
return dict(block_dict)
|
||||
|
||||
async def get_first_transaction(self, address: str) -> Transaction | None:
|
||||
"""Get the first transaction for a wallet.
|
||||
|
||||
This is useful for determining wallet age. Note: This is an expensive
|
||||
operation as it may require scanning transaction history.
|
||||
|
||||
Args:
|
||||
address: Wallet address.
|
||||
|
||||
Returns:
|
||||
First transaction or None if no transactions.
|
||||
"""
|
||||
cache_key = self._cache_key("first_tx", address)
|
||||
|
||||
# Check cache
|
||||
cached = await self._get_cached(cache_key)
|
||||
if cached is not None:
|
||||
if cached == "null":
|
||||
return None
|
||||
data = json.loads(cached)
|
||||
return Transaction(
|
||||
hash=data["hash"],
|
||||
block_number=data["block_number"],
|
||||
timestamp=datetime.fromisoformat(data["timestamp"]),
|
||||
from_address=data["from_address"],
|
||||
to_address=data["to_address"],
|
||||
value=Decimal(data["value"]),
|
||||
gas_used=data["gas_used"],
|
||||
gas_price=Decimal(data["gas_price"]),
|
||||
)
|
||||
|
||||
# Check if wallet has any transactions
|
||||
nonce = await self.get_transaction_count(address)
|
||||
if nonce == 0:
|
||||
await self._set_cached(cache_key, "null", ttl=60) # Short TTL for empty
|
||||
return None
|
||||
|
||||
# Note: Getting the actual first transaction requires using an indexer
|
||||
# or scanning blocks, which is expensive. For now, we'll return None
|
||||
# and recommend using an indexer service for production.
|
||||
logger.warning(
|
||||
"get_first_transaction requires an indexer service for %s (nonce=%d)",
|
||||
address,
|
||||
nonce,
|
||||
)
|
||||
return None
|
||||
|
||||
async def get_wallet_info(self, address: str) -> WalletInfo:
|
||||
"""Get aggregated wallet information.
|
||||
|
||||
Args:
|
||||
address: Wallet address.
|
||||
|
||||
Returns:
|
||||
WalletInfo with transaction count, balance, and first transaction.
|
||||
"""
|
||||
# Fetch data concurrently
|
||||
nonce_task = self.get_transaction_count(address)
|
||||
balance_task = self.get_balance(address)
|
||||
first_tx_task = self.get_first_transaction(address)
|
||||
|
||||
nonce, balance, first_tx = await asyncio.gather(nonce_task, balance_task, first_tx_task)
|
||||
|
||||
return WalletInfo(
|
||||
address=address.lower(),
|
||||
transaction_count=nonce,
|
||||
balance_wei=balance,
|
||||
first_transaction=first_tx,
|
||||
)
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if the client can connect to the RPC.
|
||||
|
||||
Returns:
|
||||
True if healthy, False otherwise.
|
||||
"""
|
||||
try:
|
||||
await self._execute_with_retry("block_number")
|
||||
return True
|
||||
except RPCError:
|
||||
return False
|
||||
@@ -0,0 +1,276 @@
|
||||
"""Known entity registry for blockchain address classification.
|
||||
|
||||
This module provides the EntityRegistry class for classifying blockchain
|
||||
addresses as known entities (CEX hot wallets, bridges, DEX contracts, etc.)
|
||||
to support funding chain analysis and suspiciousness scoring.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from polymarket_insider_tracker.profiler.entity_data import (
|
||||
EntityType,
|
||||
get_all_known_entities,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EntityRegistry:
|
||||
"""Registry of known blockchain entities for address classification.
|
||||
|
||||
The registry contains mappings from blockchain addresses to known entity
|
||||
types (CEX, bridges, DEX, etc.). This is used to:
|
||||
- Terminate funding chain traces at known entities
|
||||
- Classify funding sources for suspiciousness scoring
|
||||
- Identify retail vs sophisticated wallet patterns
|
||||
|
||||
Attributes:
|
||||
_entities: Internal mapping of address to entity type.
|
||||
"""
|
||||
|
||||
# Entity types that should terminate funding chain traces
|
||||
TERMINAL_ENTITY_TYPES = frozenset(
|
||||
[
|
||||
EntityType.CEX_BINANCE,
|
||||
EntityType.CEX_COINBASE,
|
||||
EntityType.CEX_KRAKEN,
|
||||
EntityType.CEX_OKX,
|
||||
EntityType.CEX_KUCOIN,
|
||||
EntityType.CEX_BYBIT,
|
||||
EntityType.CEX_CRYPTO_COM,
|
||||
EntityType.CEX_OTHER,
|
||||
EntityType.BRIDGE_POLYGON,
|
||||
EntityType.BRIDGE_MULTICHAIN,
|
||||
EntityType.BRIDGE_STARGATE,
|
||||
EntityType.BRIDGE_HOP,
|
||||
EntityType.BRIDGE_OTHER,
|
||||
]
|
||||
)
|
||||
|
||||
# Entity types that indicate CEX origin
|
||||
CEX_ENTITY_TYPES = frozenset(
|
||||
[
|
||||
EntityType.CEX_BINANCE,
|
||||
EntityType.CEX_COINBASE,
|
||||
EntityType.CEX_KRAKEN,
|
||||
EntityType.CEX_OKX,
|
||||
EntityType.CEX_KUCOIN,
|
||||
EntityType.CEX_BYBIT,
|
||||
EntityType.CEX_CRYPTO_COM,
|
||||
EntityType.CEX_OTHER,
|
||||
]
|
||||
)
|
||||
|
||||
# Entity types that indicate bridge origin
|
||||
BRIDGE_ENTITY_TYPES = frozenset(
|
||||
[
|
||||
EntityType.BRIDGE_POLYGON,
|
||||
EntityType.BRIDGE_MULTICHAIN,
|
||||
EntityType.BRIDGE_STARGATE,
|
||||
EntityType.BRIDGE_HOP,
|
||||
EntityType.BRIDGE_OTHER,
|
||||
]
|
||||
)
|
||||
|
||||
# Entity types for DEX contracts
|
||||
DEX_ENTITY_TYPES = frozenset(
|
||||
[
|
||||
EntityType.DEX_UNISWAP,
|
||||
EntityType.DEX_SUSHISWAP,
|
||||
EntityType.DEX_QUICKSWAP,
|
||||
EntityType.DEX_1INCH,
|
||||
EntityType.DEX_OTHER,
|
||||
]
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
custom_entities: dict[str, EntityType] | None = None,
|
||||
*,
|
||||
include_defaults: bool = True,
|
||||
) -> None:
|
||||
"""Initialize the entity registry.
|
||||
|
||||
Args:
|
||||
custom_entities: Additional custom entity mappings to include.
|
||||
include_defaults: Whether to include default known entities.
|
||||
"""
|
||||
self._entities: dict[str, EntityType] = {}
|
||||
|
||||
if include_defaults:
|
||||
self._entities.update(get_all_known_entities())
|
||||
|
||||
if custom_entities:
|
||||
# Add custom entities (normalized to lowercase)
|
||||
for address, entity_type in custom_entities.items():
|
||||
self._entities[address.lower()] = entity_type
|
||||
|
||||
logger.info(f"EntityRegistry initialized with {len(self._entities)} known entities")
|
||||
|
||||
def classify(self, address: str) -> EntityType:
|
||||
"""Classify an address by its entity type.
|
||||
|
||||
Args:
|
||||
address: The blockchain address to classify.
|
||||
|
||||
Returns:
|
||||
The EntityType for the address, or UNKNOWN if not in registry.
|
||||
"""
|
||||
return self._entities.get(address.lower(), EntityType.UNKNOWN)
|
||||
|
||||
def is_known_entity(self, address: str) -> bool:
|
||||
"""Check if an address is a known entity.
|
||||
|
||||
Args:
|
||||
address: The blockchain address to check.
|
||||
|
||||
Returns:
|
||||
True if the address is in the registry, False otherwise.
|
||||
"""
|
||||
return address.lower() in self._entities
|
||||
|
||||
def is_cex(self, address: str) -> bool:
|
||||
"""Check if an address is a known CEX hot wallet.
|
||||
|
||||
Args:
|
||||
address: The blockchain address to check.
|
||||
|
||||
Returns:
|
||||
True if the address is a CEX hot wallet.
|
||||
"""
|
||||
return self.classify(address) in self.CEX_ENTITY_TYPES
|
||||
|
||||
def is_bridge(self, address: str) -> bool:
|
||||
"""Check if an address is a known bridge contract.
|
||||
|
||||
Args:
|
||||
address: The blockchain address to check.
|
||||
|
||||
Returns:
|
||||
True if the address is a bridge contract.
|
||||
"""
|
||||
return self.classify(address) in self.BRIDGE_ENTITY_TYPES
|
||||
|
||||
def is_dex(self, address: str) -> bool:
|
||||
"""Check if an address is a known DEX contract.
|
||||
|
||||
Args:
|
||||
address: The blockchain address to check.
|
||||
|
||||
Returns:
|
||||
True if the address is a DEX contract.
|
||||
"""
|
||||
return self.classify(address) in self.DEX_ENTITY_TYPES
|
||||
|
||||
def is_terminal(self, address: str) -> bool:
|
||||
"""Check if an address should terminate a funding chain trace.
|
||||
|
||||
Terminal entities are those where tracing further back provides
|
||||
diminishing returns (CEX, bridges). These indicate the practical
|
||||
origin of funds from the perspective of on-chain analysis.
|
||||
|
||||
Args:
|
||||
address: The blockchain address to check.
|
||||
|
||||
Returns:
|
||||
True if the address should terminate a funding trace.
|
||||
"""
|
||||
return self.classify(address) in self.TERMINAL_ENTITY_TYPES
|
||||
|
||||
def is_contract(self, address: str) -> bool:
|
||||
"""Check if an address is a known smart contract.
|
||||
|
||||
This includes DEX routers, token contracts, and DeFi protocols.
|
||||
|
||||
Args:
|
||||
address: The blockchain address to check.
|
||||
|
||||
Returns:
|
||||
True if the address is a known smart contract.
|
||||
"""
|
||||
entity_type = self.classify(address)
|
||||
contract_types = self.DEX_ENTITY_TYPES | {
|
||||
EntityType.TOKEN_USDC,
|
||||
EntityType.TOKEN_USDT,
|
||||
EntityType.TOKEN_WETH,
|
||||
EntityType.TOKEN_WMATIC,
|
||||
EntityType.DEFI_AAVE,
|
||||
EntityType.DEFI_COMPOUND,
|
||||
EntityType.DEFI_OTHER,
|
||||
EntityType.CONTRACT,
|
||||
}
|
||||
return entity_type in contract_types
|
||||
|
||||
def get_entity_category(self, address: str) -> str:
|
||||
"""Get a human-readable category for an address.
|
||||
|
||||
Args:
|
||||
address: The blockchain address to categorize.
|
||||
|
||||
Returns:
|
||||
Category string: "cex", "bridge", "dex", "token", "defi", "contract", or "unknown".
|
||||
"""
|
||||
entity_type = self.classify(address)
|
||||
|
||||
if entity_type in self.CEX_ENTITY_TYPES:
|
||||
return "cex"
|
||||
if entity_type in self.BRIDGE_ENTITY_TYPES:
|
||||
return "bridge"
|
||||
if entity_type in self.DEX_ENTITY_TYPES:
|
||||
return "dex"
|
||||
if entity_type in {
|
||||
EntityType.TOKEN_USDC,
|
||||
EntityType.TOKEN_USDT,
|
||||
EntityType.TOKEN_WETH,
|
||||
EntityType.TOKEN_WMATIC,
|
||||
}:
|
||||
return "token"
|
||||
if entity_type in {
|
||||
EntityType.DEFI_AAVE,
|
||||
EntityType.DEFI_COMPOUND,
|
||||
EntityType.DEFI_OTHER,
|
||||
}:
|
||||
return "defi"
|
||||
if entity_type == EntityType.CONTRACT:
|
||||
return "contract"
|
||||
|
||||
return "unknown"
|
||||
|
||||
def add_entity(self, address: str, entity_type: EntityType) -> None:
|
||||
"""Add or update an entity in the registry.
|
||||
|
||||
Args:
|
||||
address: The blockchain address.
|
||||
entity_type: The entity type to assign.
|
||||
"""
|
||||
self._entities[address.lower()] = entity_type
|
||||
logger.debug(f"Added entity: {address} -> {entity_type.value}")
|
||||
|
||||
def remove_entity(self, address: str) -> bool:
|
||||
"""Remove an entity from the registry.
|
||||
|
||||
Args:
|
||||
address: The blockchain address to remove.
|
||||
|
||||
Returns:
|
||||
True if the entity was removed, False if not found.
|
||||
"""
|
||||
normalized = address.lower()
|
||||
if normalized in self._entities:
|
||||
del self._entities[normalized]
|
||||
return True
|
||||
return False
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Return the number of entities in the registry."""
|
||||
return len(self._entities)
|
||||
|
||||
def __contains__(self, address: str) -> bool:
|
||||
"""Check if an address is in the registry."""
|
||||
return self.is_known_entity(address)
|
||||
@@ -0,0 +1,154 @@
|
||||
"""Known blockchain entity address mappings.
|
||||
|
||||
This module contains address-to-entity mappings for known blockchain
|
||||
entities on Polygon including CEX hot wallets, bridges, and DEX contracts.
|
||||
|
||||
Sources:
|
||||
- Etherscan labels
|
||||
- Arkham Intelligence
|
||||
- Official protocol documentation
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class EntityType(Enum):
|
||||
"""Classification of known blockchain entities."""
|
||||
|
||||
# Centralized Exchanges
|
||||
CEX_BINANCE = "cex_binance"
|
||||
CEX_COINBASE = "cex_coinbase"
|
||||
CEX_KRAKEN = "cex_kraken"
|
||||
CEX_OKX = "cex_okx"
|
||||
CEX_KUCOIN = "cex_kucoin"
|
||||
CEX_BYBIT = "cex_bybit"
|
||||
CEX_CRYPTO_COM = "cex_crypto_com"
|
||||
CEX_OTHER = "cex_other"
|
||||
|
||||
# Bridges
|
||||
BRIDGE_POLYGON = "bridge_polygon"
|
||||
BRIDGE_MULTICHAIN = "bridge_multichain"
|
||||
BRIDGE_STARGATE = "bridge_stargate"
|
||||
BRIDGE_HOP = "bridge_hop"
|
||||
BRIDGE_OTHER = "bridge_other"
|
||||
|
||||
# Decentralized Exchanges
|
||||
DEX_UNISWAP = "dex_uniswap"
|
||||
DEX_SUSHISWAP = "dex_sushiswap"
|
||||
DEX_QUICKSWAP = "dex_quickswap"
|
||||
DEX_1INCH = "dex_1inch"
|
||||
DEX_OTHER = "dex_other"
|
||||
|
||||
# Token Contracts
|
||||
TOKEN_USDC = "token_usdc"
|
||||
TOKEN_USDT = "token_usdt"
|
||||
TOKEN_WETH = "token_weth"
|
||||
TOKEN_WMATIC = "token_wmatic"
|
||||
|
||||
# Lending/DeFi
|
||||
DEFI_AAVE = "defi_aave"
|
||||
DEFI_COMPOUND = "defi_compound"
|
||||
DEFI_OTHER = "defi_other"
|
||||
|
||||
# Other
|
||||
CONTRACT = "contract"
|
||||
UNKNOWN = "unknown"
|
||||
|
||||
|
||||
# CEX hot wallet addresses on Polygon
|
||||
# Sources: Etherscan labels, Arkham Intelligence, public disclosures
|
||||
CEX_ADDRESSES: dict[str, EntityType] = {
|
||||
# Binance
|
||||
"0x28c6c06298d514db089934071355e5743bf21d60": EntityType.CEX_BINANCE,
|
||||
"0x21a31ee1afc51d94c2efccaa2092ad1028285549": EntityType.CEX_BINANCE,
|
||||
"0xf89d7b9c864f589bbf53a82105107622b35eaa40": EntityType.CEX_BINANCE,
|
||||
"0xdfd5293d8e347dfe59e90efd55b2956a1343963d": EntityType.CEX_BINANCE,
|
||||
# Coinbase
|
||||
"0x503828976d22510aad0339f595f37cc4e4645c80": EntityType.CEX_COINBASE,
|
||||
"0x71660c4005ba85c37ccec55d0c4493e66fe775d3": EntityType.CEX_COINBASE,
|
||||
"0xa9d1e08c7793af67e9d92fe308d5697fb81d3e43": EntityType.CEX_COINBASE,
|
||||
# Kraken
|
||||
"0x2910543af39aba0cd09dbb2d50200b3e800a63d2": EntityType.CEX_KRAKEN,
|
||||
"0x0a869d79a7052c7f1b55a8ebabbea3420f0d1e13": EntityType.CEX_KRAKEN,
|
||||
# OKX
|
||||
"0x5041ed759dd4afc3a72b8192c143f72f4724081a": EntityType.CEX_OKX,
|
||||
"0x6cc5f688a315f3dc28a7781717a9a798a59fda7b": EntityType.CEX_OKX,
|
||||
# KuCoin
|
||||
"0xf16e9b0d03470827a95cdfd0cb8a8a3b46969b91": EntityType.CEX_KUCOIN,
|
||||
"0xd6216fc19db775df9774a6e33526131da7d19a2c": EntityType.CEX_KUCOIN,
|
||||
# Bybit
|
||||
"0xf89e6d82be28f5cc97a9e6a94a16a17e5be73e78": EntityType.CEX_BYBIT,
|
||||
# Crypto.com
|
||||
"0x6262998ced04146fa42253a5c0af90ca02dfd2a3": EntityType.CEX_CRYPTO_COM,
|
||||
"0x46340b20830761efd32832a74d7169b29feb9758": EntityType.CEX_CRYPTO_COM,
|
||||
}
|
||||
|
||||
# Bridge contract addresses on Polygon
|
||||
BRIDGE_ADDRESSES: dict[str, EntityType] = {
|
||||
# Polygon PoS Bridge (RootChain / Plasma Bridge related)
|
||||
"0xa0c68c638235ee32657e8f720a23cec1bfc77c77": EntityType.BRIDGE_POLYGON,
|
||||
"0x401f6c983ea34274ec46f84d70b31c151321188b": EntityType.BRIDGE_POLYGON,
|
||||
# Multichain (formerly AnySwap)
|
||||
"0x4f3aff3a747fcade12598081e80c6605a8be192f": EntityType.BRIDGE_MULTICHAIN,
|
||||
# Stargate
|
||||
"0x45a01e4e04f14f7a4a6880d0cbaf2c3c1acfbed4": EntityType.BRIDGE_STARGATE,
|
||||
# Hop Protocol
|
||||
"0x76b22b8c1079a44f1211b0e72c5d26c5e3b3c3c9": EntityType.BRIDGE_HOP,
|
||||
}
|
||||
|
||||
# DEX router addresses on Polygon
|
||||
DEX_ADDRESSES: dict[str, EntityType] = {
|
||||
# Uniswap V3
|
||||
"0xe592427a0aece92de3edee1f18e0157c05861564": EntityType.DEX_UNISWAP,
|
||||
"0x68b3465833fb72a70ecdf485e0e4c7bd8665fc45": EntityType.DEX_UNISWAP, # SwapRouter02
|
||||
# SushiSwap
|
||||
"0x1b02da8cb0d097eb8d57a175b88c7d8b47997506": EntityType.DEX_SUSHISWAP,
|
||||
# QuickSwap
|
||||
"0xa5e0829caced8ffdd4de3c43696c57f7d7a678ff": EntityType.DEX_QUICKSWAP,
|
||||
# 1inch
|
||||
"0x1111111254eeb25477b68fb85ed929f73a960582": EntityType.DEX_1INCH,
|
||||
}
|
||||
|
||||
# Token contract addresses on Polygon
|
||||
TOKEN_ADDRESSES: dict[str, EntityType] = {
|
||||
# USDC (Bridged)
|
||||
"0x2791bca1f2de4661ed88a30c99a7a9449aa84174": EntityType.TOKEN_USDC,
|
||||
# USDC (Native)
|
||||
"0x3c499c542cef5e3811e1192ce70d8cc03d5c3359": EntityType.TOKEN_USDC,
|
||||
# USDT
|
||||
"0xc2132d05d31c914a87c6611c10748aeb04b58e8f": EntityType.TOKEN_USDT,
|
||||
# WETH
|
||||
"0x7ceb23fd6bc0add59e62ac25578270cff1b9f619": EntityType.TOKEN_WETH,
|
||||
# WMATIC
|
||||
"0x0d500b1d8e8ef31e21c99d1db9a6444d3adf1270": EntityType.TOKEN_WMATIC,
|
||||
}
|
||||
|
||||
# DeFi protocol addresses on Polygon
|
||||
DEFI_ADDRESSES: dict[str, EntityType] = {
|
||||
# Aave V3
|
||||
"0x794a61358d6845594f94dc1db02a252b5b4814ad": EntityType.DEFI_AAVE, # Pool
|
||||
"0x8145edddf43f50276641b55bd3ad95944510021e": EntityType.DEFI_AAVE, # PoolAddressesProvider
|
||||
}
|
||||
|
||||
|
||||
def get_all_known_entities() -> dict[str, EntityType]:
|
||||
"""Get all known entity addresses combined.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping lowercase addresses to their entity types.
|
||||
"""
|
||||
all_entities: dict[str, EntityType] = {}
|
||||
|
||||
for entities in [
|
||||
CEX_ADDRESSES,
|
||||
BRIDGE_ADDRESSES,
|
||||
DEX_ADDRESSES,
|
||||
TOKEN_ADDRESSES,
|
||||
DEFI_ADDRESSES,
|
||||
]:
|
||||
for address, entity_type in entities.items():
|
||||
all_entities[address.lower()] = entity_type
|
||||
|
||||
return all_entities
|
||||
@@ -0,0 +1,359 @@
|
||||
"""Funding chain tracer for wallet analysis.
|
||||
|
||||
This module provides the FundingTracer class for tracing the funding chain
|
||||
of wallets to identify where their USDC/MATIC originated from.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from web3 import AsyncWeb3
|
||||
|
||||
from polymarket_insider_tracker.profiler.entities import EntityRegistry
|
||||
from polymarket_insider_tracker.profiler.models import FundingChain, FundingTransfer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from polymarket_insider_tracker.profiler.chain import PolygonClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# USDC contract addresses on Polygon
|
||||
USDC_BRIDGED = "0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174"
|
||||
USDC_NATIVE = "0x3c499c542cEF5E3811e1192ce70d8cC03d5c3359"
|
||||
|
||||
# ERC20 Transfer event signature
|
||||
TRANSFER_EVENT_SIGNATURE = AsyncWeb3.keccak(text="Transfer(address,address,uint256)")
|
||||
|
||||
|
||||
class FundingTracer:
|
||||
"""Traces funding chains to identify wallet funding sources.
|
||||
|
||||
The tracer follows USDC transfers backwards from a target wallet
|
||||
to find where the funds originated, stopping at known entities
|
||||
(CEX hot wallets, bridges) or reaching the maximum hop count.
|
||||
|
||||
Attributes:
|
||||
polygon_client: Client for Polygon blockchain queries.
|
||||
entity_registry: Registry of known blockchain entities.
|
||||
max_hops: Maximum number of hops to trace (default 3).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
polygon_client: PolygonClient,
|
||||
entity_registry: EntityRegistry | None = None,
|
||||
*,
|
||||
max_hops: int = 3,
|
||||
usdc_addresses: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Initialize the funding tracer.
|
||||
|
||||
Args:
|
||||
polygon_client: Polygon blockchain client for queries.
|
||||
entity_registry: Registry for entity classification. Creates default if None.
|
||||
max_hops: Maximum hops to trace back (default 3).
|
||||
usdc_addresses: USDC contract addresses to track. Uses defaults if None.
|
||||
"""
|
||||
self.polygon_client = polygon_client
|
||||
self.entity_registry = entity_registry or EntityRegistry()
|
||||
self.max_hops = max_hops
|
||||
self._usdc_addresses = [
|
||||
addr.lower() for addr in (usdc_addresses or [USDC_BRIDGED, USDC_NATIVE])
|
||||
]
|
||||
|
||||
async def trace(
|
||||
self,
|
||||
address: str,
|
||||
max_hops: int | None = None,
|
||||
) -> FundingChain:
|
||||
"""Trace the funding chain for a wallet.
|
||||
|
||||
Follows the first USDC transfer into the wallet, then recursively
|
||||
traces the source wallet until reaching a known entity or max hops.
|
||||
|
||||
Args:
|
||||
address: Target wallet address to trace.
|
||||
max_hops: Override default max_hops for this trace.
|
||||
|
||||
Returns:
|
||||
FundingChain with the complete trace result.
|
||||
"""
|
||||
effective_max_hops = max_hops if max_hops is not None else self.max_hops
|
||||
normalized_address = address.lower()
|
||||
|
||||
chain: list[FundingTransfer] = []
|
||||
current_address = normalized_address
|
||||
origin_address = normalized_address
|
||||
origin_type = "unknown"
|
||||
|
||||
for hop in range(effective_max_hops):
|
||||
# Check if current address is a known entity
|
||||
if self.entity_registry.is_terminal(current_address):
|
||||
origin_address = current_address
|
||||
origin_type = self.entity_registry.classify(current_address).value
|
||||
logger.debug(
|
||||
"Trace terminated at known entity: %s (%s)",
|
||||
current_address,
|
||||
origin_type,
|
||||
)
|
||||
break
|
||||
|
||||
# Get first USDC transfer into this address
|
||||
transfer = await self.get_first_usdc_transfer(current_address)
|
||||
if transfer is None:
|
||||
logger.debug(
|
||||
"No USDC transfer found for %s at hop %d",
|
||||
current_address,
|
||||
hop,
|
||||
)
|
||||
origin_address = current_address
|
||||
break
|
||||
|
||||
chain.append(transfer)
|
||||
origin_address = transfer.from_address
|
||||
current_address = transfer.from_address
|
||||
|
||||
# Check if the source is a known entity
|
||||
if self.entity_registry.is_terminal(origin_address):
|
||||
origin_type = self.entity_registry.classify(origin_address).value
|
||||
logger.debug(
|
||||
"Trace found terminal entity: %s (%s)",
|
||||
origin_address,
|
||||
origin_type,
|
||||
)
|
||||
break
|
||||
|
||||
return FundingChain(
|
||||
target_address=normalized_address,
|
||||
chain=chain,
|
||||
origin_address=origin_address,
|
||||
origin_type=origin_type,
|
||||
hop_count=len(chain),
|
||||
traced_at=datetime.now(UTC),
|
||||
)
|
||||
|
||||
async def get_first_usdc_transfer(
|
||||
self,
|
||||
address: str,
|
||||
) -> FundingTransfer | None:
|
||||
"""Get the first USDC transfer into a wallet.
|
||||
|
||||
Queries the blockchain for ERC20 Transfer events to the target
|
||||
address for known USDC contracts.
|
||||
|
||||
Args:
|
||||
address: Target wallet address.
|
||||
|
||||
Returns:
|
||||
First FundingTransfer if found, None otherwise.
|
||||
"""
|
||||
normalized = address.lower()
|
||||
|
||||
# Query transfers for each USDC contract
|
||||
for usdc_address in self._usdc_addresses:
|
||||
transfer = await self._get_first_token_transfer(
|
||||
to_address=normalized,
|
||||
token_address=usdc_address,
|
||||
)
|
||||
if transfer is not None:
|
||||
return transfer
|
||||
|
||||
return None
|
||||
|
||||
async def _get_first_token_transfer(
|
||||
self,
|
||||
to_address: str,
|
||||
token_address: str,
|
||||
) -> FundingTransfer | None:
|
||||
"""Get the first ERC20 transfer to an address for a specific token.
|
||||
|
||||
Args:
|
||||
to_address: Recipient wallet address.
|
||||
token_address: ERC20 token contract address.
|
||||
|
||||
Returns:
|
||||
First FundingTransfer if found, None otherwise.
|
||||
"""
|
||||
try:
|
||||
logs = await self._get_transfer_logs(
|
||||
to_address=to_address,
|
||||
token_address=token_address,
|
||||
limit=1,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to get transfer logs for %s: %s",
|
||||
to_address,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
|
||||
if not logs:
|
||||
return None
|
||||
|
||||
log = logs[0]
|
||||
return await self._log_to_funding_transfer(log, token_address)
|
||||
|
||||
async def _get_transfer_logs(
|
||||
self,
|
||||
to_address: str,
|
||||
token_address: str,
|
||||
limit: int = 10,
|
||||
from_block: int | str = 0,
|
||||
to_block: int | str = "latest",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Get ERC20 Transfer event logs.
|
||||
|
||||
Args:
|
||||
to_address: Filter by recipient address.
|
||||
token_address: ERC20 token contract address.
|
||||
limit: Maximum logs to return.
|
||||
from_block: Starting block number.
|
||||
to_block: Ending block number.
|
||||
|
||||
Returns:
|
||||
List of log dictionaries.
|
||||
"""
|
||||
# Pad address to 32 bytes for topic filter
|
||||
padded_to = "0x" + to_address.lower().replace("0x", "").zfill(64)
|
||||
|
||||
await self.polygon_client._rate_limiter.acquire()
|
||||
|
||||
# Use the web3 instance from polygon client
|
||||
w3 = (
|
||||
self.polygon_client._w3
|
||||
if self.polygon_client._primary_healthy
|
||||
else (self.polygon_client._w3_fallback or self.polygon_client._w3)
|
||||
)
|
||||
|
||||
# Get logs with Transfer event filtering by recipient
|
||||
# Note: web3 typing is overly restrictive for block params
|
||||
logs = await w3.eth.get_logs(
|
||||
{
|
||||
"address": AsyncWeb3.to_checksum_address(token_address),
|
||||
"topics": [
|
||||
TRANSFER_EVENT_SIGNATURE.hex(), # Transfer event
|
||||
None, # from (any)
|
||||
padded_to, # to (target address)
|
||||
],
|
||||
"fromBlock": from_block, # type: ignore[typeddict-item]
|
||||
"toBlock": to_block, # type: ignore[typeddict-item]
|
||||
}
|
||||
)
|
||||
|
||||
# Convert to list of dicts and limit
|
||||
result = [dict(log) for log in logs[:limit]]
|
||||
return result
|
||||
|
||||
async def _log_to_funding_transfer(
|
||||
self,
|
||||
log: dict[str, Any],
|
||||
token_address: str,
|
||||
) -> FundingTransfer:
|
||||
"""Convert a log entry to a FundingTransfer.
|
||||
|
||||
Args:
|
||||
log: Log dictionary from get_logs.
|
||||
token_address: Token contract address.
|
||||
|
||||
Returns:
|
||||
FundingTransfer object.
|
||||
"""
|
||||
# Extract addresses from topics (padded to 32 bytes)
|
||||
from_address = "0x" + log["topics"][1].hex()[-40:]
|
||||
to_address = "0x" + log["topics"][2].hex()[-40:]
|
||||
|
||||
# Extract amount from data
|
||||
amount = int(log["data"].hex(), 16)
|
||||
|
||||
# Get block timestamp
|
||||
block_number = log["blockNumber"]
|
||||
try:
|
||||
block = await self.polygon_client.get_block(block_number)
|
||||
timestamp = datetime.fromtimestamp(block["timestamp"], tz=UTC)
|
||||
except Exception:
|
||||
timestamp = datetime.now(UTC)
|
||||
|
||||
# Determine token symbol
|
||||
token = "USDC" if token_address.lower() in self._usdc_addresses else "OTHER"
|
||||
|
||||
return FundingTransfer(
|
||||
from_address=from_address.lower(),
|
||||
to_address=to_address.lower(),
|
||||
amount=Decimal(amount),
|
||||
token=token,
|
||||
tx_hash=log["transactionHash"].hex(),
|
||||
block_number=block_number,
|
||||
timestamp=timestamp,
|
||||
)
|
||||
|
||||
async def get_funding_chains_batch(
|
||||
self,
|
||||
addresses: list[str],
|
||||
max_hops: int | None = None,
|
||||
) -> dict[str, FundingChain]:
|
||||
"""Trace funding chains for multiple addresses concurrently.
|
||||
|
||||
Args:
|
||||
addresses: List of wallet addresses to trace.
|
||||
max_hops: Override default max_hops for all traces.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping address to FundingChain.
|
||||
"""
|
||||
tasks = [self.trace(addr, max_hops=max_hops) for addr in addresses]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
chains: dict[str, FundingChain] = {}
|
||||
for addr, result in zip(addresses, results, strict=True):
|
||||
if isinstance(result, BaseException):
|
||||
logger.warning("Failed to trace %s: %s", addr, result)
|
||||
chains[addr.lower()] = FundingChain(
|
||||
target_address=addr.lower(),
|
||||
origin_type="error",
|
||||
)
|
||||
else:
|
||||
chains[addr.lower()] = result
|
||||
|
||||
return chains
|
||||
|
||||
def get_suspiciousness_score(self, chain: FundingChain) -> float:
|
||||
"""Calculate a suspiciousness score based on funding chain.
|
||||
|
||||
Higher scores indicate more suspicious funding patterns:
|
||||
- CEX origin: Lower suspicion (0.0-0.2)
|
||||
- Bridge origin: Low suspicion (0.2-0.4)
|
||||
- Unknown origin with few hops: High suspicion (0.8-1.0)
|
||||
- Unknown origin with many hops: Medium suspicion (0.5-0.8)
|
||||
|
||||
Args:
|
||||
chain: Funding chain to score.
|
||||
|
||||
Returns:
|
||||
Suspiciousness score from 0.0 to 1.0.
|
||||
"""
|
||||
if chain.is_cex_origin:
|
||||
# CEX origin is least suspicious
|
||||
return 0.1
|
||||
|
||||
if chain.is_bridge_origin:
|
||||
# Bridge origin is slightly more suspicious
|
||||
return 0.3
|
||||
|
||||
# Unknown origin
|
||||
if chain.hop_count == 0:
|
||||
# No transfers found - very suspicious (possible contract or new wallet)
|
||||
return 1.0
|
||||
|
||||
if chain.hop_count >= self.max_hops:
|
||||
# Max hops reached without finding known entity
|
||||
# More hops = more obfuscation = more suspicious
|
||||
return 0.7
|
||||
|
||||
# Some hops but didn't reach max - moderately suspicious
|
||||
return 0.5 + (0.3 * (1 - chain.hop_count / self.max_hops))
|
||||
@@ -0,0 +1,220 @@
|
||||
"""Data models for the profiler module."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Transaction:
|
||||
"""Represents a blockchain transaction."""
|
||||
|
||||
hash: str
|
||||
block_number: int
|
||||
timestamp: datetime
|
||||
from_address: str
|
||||
to_address: str | None
|
||||
value: Decimal # In Wei
|
||||
gas_used: int
|
||||
gas_price: Decimal # In Wei
|
||||
|
||||
@property
|
||||
def value_matic(self) -> Decimal:
|
||||
"""Return value in MATIC (10^18 Wei = 1 MATIC)."""
|
||||
return self.value / Decimal("1000000000000000000")
|
||||
|
||||
@property
|
||||
def gas_cost_wei(self) -> Decimal:
|
||||
"""Return total gas cost in Wei."""
|
||||
return Decimal(self.gas_used) * self.gas_price
|
||||
|
||||
@property
|
||||
def gas_cost_matic(self) -> Decimal:
|
||||
"""Return total gas cost in MATIC."""
|
||||
return self.gas_cost_wei / Decimal("1000000000000000000")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WalletInfo:
|
||||
"""Aggregated wallet information from blockchain queries."""
|
||||
|
||||
address: str
|
||||
transaction_count: int # Nonce
|
||||
balance_wei: Decimal
|
||||
first_transaction: Transaction | None = None
|
||||
|
||||
@property
|
||||
def balance_matic(self) -> Decimal:
|
||||
"""Return balance in MATIC."""
|
||||
return self.balance_wei / Decimal("1000000000000000000")
|
||||
|
||||
@property
|
||||
def is_fresh(self) -> bool:
|
||||
"""Return True if wallet has very few transactions (potential fresh wallet)."""
|
||||
return self.transaction_count < 10
|
||||
|
||||
@property
|
||||
def wallet_age_days(self) -> float | None:
|
||||
"""Return wallet age in days based on first transaction."""
|
||||
if self.first_transaction is None:
|
||||
return None
|
||||
delta = (
|
||||
datetime.now(tz=self.first_transaction.timestamp.tzinfo)
|
||||
- self.first_transaction.timestamp
|
||||
)
|
||||
return delta.total_seconds() / 86400
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WalletProfile:
|
||||
"""Complete wallet analysis profile.
|
||||
|
||||
This is the result of analyzing a wallet's on-chain activity to determine
|
||||
if it exhibits suspicious behavior patterns like fresh wallet trading.
|
||||
|
||||
Attributes:
|
||||
address: The wallet address (lowercase).
|
||||
nonce: Transaction count (number of outgoing transactions).
|
||||
first_seen: Timestamp of first transaction, if available.
|
||||
age_hours: Wallet age in hours since first transaction.
|
||||
is_fresh: True if wallet meets fresh wallet criteria.
|
||||
total_tx_count: Total number of transactions (same as nonce for now).
|
||||
matic_balance: MATIC balance in Wei.
|
||||
usdc_balance: USDC balance in smallest unit (6 decimals).
|
||||
analyzed_at: Timestamp when this profile was created.
|
||||
fresh_threshold: The threshold used to determine freshness.
|
||||
"""
|
||||
|
||||
address: str
|
||||
nonce: int
|
||||
first_seen: datetime | None
|
||||
age_hours: float | None
|
||||
is_fresh: bool
|
||||
total_tx_count: int
|
||||
matic_balance: Decimal
|
||||
usdc_balance: Decimal
|
||||
analyzed_at: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
fresh_threshold: int = 5
|
||||
|
||||
@property
|
||||
def age_days(self) -> float | None:
|
||||
"""Return wallet age in days."""
|
||||
if self.age_hours is None:
|
||||
return None
|
||||
return self.age_hours / 24.0
|
||||
|
||||
@property
|
||||
def matic_balance_formatted(self) -> Decimal:
|
||||
"""Return MATIC balance in human-readable format (18 decimals)."""
|
||||
return self.matic_balance / Decimal("1000000000000000000")
|
||||
|
||||
@property
|
||||
def usdc_balance_formatted(self) -> Decimal:
|
||||
"""Return USDC balance in human-readable format (6 decimals)."""
|
||||
return self.usdc_balance / Decimal("1000000")
|
||||
|
||||
@property
|
||||
def is_brand_new(self) -> bool:
|
||||
"""Return True if wallet has never transacted (nonce = 0)."""
|
||||
return self.nonce == 0
|
||||
|
||||
@property
|
||||
def freshness_score(self) -> float:
|
||||
"""Return a 0-1 score where 1 is maximally fresh.
|
||||
|
||||
Score is based on:
|
||||
- Nonce (fewer = fresher)
|
||||
- Age (younger = fresher)
|
||||
"""
|
||||
# Nonce component: 1.0 at 0, 0.0 at threshold or higher
|
||||
nonce_score = max(0.0, 1.0 - (self.nonce / self.fresh_threshold))
|
||||
|
||||
# Age component: 1.0 at 0 hours, 0.0 at 48 hours or more
|
||||
age_score = 1.0 if self.age_hours is None else max(0.0, 1.0 - self.age_hours / 48.0)
|
||||
|
||||
# Weighted average: nonce is slightly more important
|
||||
return 0.6 * nonce_score + 0.4 * age_score
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FundingTransfer:
|
||||
"""Represents an ERC20 token transfer for funding chain analysis.
|
||||
|
||||
Attributes:
|
||||
from_address: Source wallet address.
|
||||
to_address: Destination wallet address.
|
||||
amount: Transfer amount in token decimals.
|
||||
token: Token symbol (e.g., "USDC", "MATIC").
|
||||
tx_hash: Transaction hash.
|
||||
block_number: Block number of the transaction.
|
||||
timestamp: Timestamp of the transaction.
|
||||
"""
|
||||
|
||||
from_address: str
|
||||
to_address: str
|
||||
amount: Decimal
|
||||
token: str
|
||||
tx_hash: str
|
||||
block_number: int
|
||||
timestamp: datetime
|
||||
|
||||
@property
|
||||
def amount_formatted(self) -> Decimal:
|
||||
"""Return amount in human-readable format.
|
||||
|
||||
Assumes 6 decimals for USDC/USDT, 18 for others.
|
||||
"""
|
||||
if self.token in ("USDC", "USDT"):
|
||||
return self.amount / Decimal("1000000")
|
||||
return self.amount / Decimal("1000000000000000000")
|
||||
|
||||
|
||||
@dataclass
|
||||
class FundingChain:
|
||||
"""Result of funding chain analysis.
|
||||
|
||||
Represents the path of funds from origin to target wallet,
|
||||
tracing back through intermediate wallets.
|
||||
|
||||
Attributes:
|
||||
target_address: The wallet being analyzed.
|
||||
chain: Ordered list of transfers from target back to origin.
|
||||
origin_address: The ultimate source of funds.
|
||||
origin_type: Classification of the origin (cex, bridge, unknown, contract).
|
||||
hop_count: Number of hops from target to origin.
|
||||
traced_at: When the trace was performed.
|
||||
"""
|
||||
|
||||
target_address: str
|
||||
chain: list[FundingTransfer] = field(default_factory=list)
|
||||
origin_address: str = ""
|
||||
origin_type: str = "unknown"
|
||||
hop_count: int = 0
|
||||
traced_at: datetime = field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
@property
|
||||
def is_cex_origin(self) -> bool:
|
||||
"""Return True if funds originated from a CEX."""
|
||||
return self.origin_type.startswith("cex")
|
||||
|
||||
@property
|
||||
def is_bridge_origin(self) -> bool:
|
||||
"""Return True if funds came through a bridge."""
|
||||
return self.origin_type.startswith("bridge")
|
||||
|
||||
@property
|
||||
def is_unknown_origin(self) -> bool:
|
||||
"""Return True if origin could not be determined."""
|
||||
return self.origin_type == "unknown"
|
||||
|
||||
@property
|
||||
def total_amount(self) -> Decimal:
|
||||
"""Return total amount transferred in the chain."""
|
||||
if not self.chain:
|
||||
return Decimal("0")
|
||||
return self.chain[0].amount
|
||||
|
||||
@property
|
||||
def funding_depth(self) -> int:
|
||||
"""Return the funding depth (hops from known entity)."""
|
||||
return self.hop_count
|
||||
@@ -0,0 +1,325 @@
|
||||
"""Graceful shutdown handler for Polymarket Insider Tracker.
|
||||
|
||||
This module provides signal handling and graceful shutdown coordination
|
||||
for the async pipeline components.
|
||||
|
||||
Usage:
|
||||
```python
|
||||
async def main():
|
||||
shutdown = GracefulShutdown()
|
||||
|
||||
async with shutdown:
|
||||
pipeline = Pipeline(settings)
|
||||
await pipeline.start()
|
||||
|
||||
# Wait for shutdown signal
|
||||
await shutdown.wait()
|
||||
|
||||
# Graceful cleanup
|
||||
await pipeline.stop()
|
||||
```
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import signal
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from contextlib import suppress
|
||||
from types import FrameType
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default shutdown timeout in seconds
|
||||
DEFAULT_SHUTDOWN_TIMEOUT = 30.0
|
||||
|
||||
# Signals to trap for graceful shutdown
|
||||
SHUTDOWN_SIGNALS = (signal.SIGTERM, signal.SIGINT)
|
||||
|
||||
|
||||
class ShutdownTimeoutError(Exception):
|
||||
"""Raised when graceful shutdown exceeds timeout."""
|
||||
|
||||
|
||||
class GracefulShutdown:
|
||||
"""Graceful shutdown handler with signal trapping.
|
||||
|
||||
This class provides coordinated shutdown handling for async applications.
|
||||
It traps SIGTERM and SIGINT signals and provides an async event that
|
||||
can be awaited to detect shutdown requests.
|
||||
|
||||
Features:
|
||||
- SIGTERM and SIGINT signal trapping
|
||||
- Async event-based shutdown coordination
|
||||
- Configurable shutdown timeout
|
||||
- Async context manager support
|
||||
- Cleanup callback registration
|
||||
- Force exit on second signal or timeout
|
||||
|
||||
Example:
|
||||
```python
|
||||
shutdown = GracefulShutdown(timeout=30.0)
|
||||
|
||||
async with shutdown:
|
||||
await some_long_running_task()
|
||||
# Automatically handles cleanup on signals
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
timeout: float = DEFAULT_SHUTDOWN_TIMEOUT,
|
||||
*,
|
||||
exit_on_timeout: bool = True,
|
||||
) -> None:
|
||||
"""Initialize the shutdown handler.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time in seconds to wait for graceful shutdown.
|
||||
exit_on_timeout: If True, force exit when timeout is exceeded.
|
||||
"""
|
||||
self._timeout = timeout
|
||||
self._exit_on_timeout = exit_on_timeout
|
||||
|
||||
self._shutdown_event: asyncio.Event | None = None
|
||||
self._shutdown_requested = False
|
||||
self._force_exit_requested = False
|
||||
self._original_handlers: dict[signal.Signals, Any] = {}
|
||||
self._cleanup_callbacks: list[Callable[[], Any]] = []
|
||||
self._loop: asyncio.AbstractEventLoop | None = None
|
||||
|
||||
@property
|
||||
def timeout(self) -> float:
|
||||
"""Shutdown timeout in seconds."""
|
||||
return self._timeout
|
||||
|
||||
@property
|
||||
def is_shutdown_requested(self) -> bool:
|
||||
"""Check if shutdown has been requested."""
|
||||
return self._shutdown_requested
|
||||
|
||||
@property
|
||||
def is_force_exit_requested(self) -> bool:
|
||||
"""Check if force exit has been requested (second signal received)."""
|
||||
return self._force_exit_requested
|
||||
|
||||
def register_cleanup(self, callback: Callable[[], Any]) -> None:
|
||||
"""Register a cleanup callback to run during shutdown.
|
||||
|
||||
Args:
|
||||
callback: A callable (sync or async) to run during shutdown.
|
||||
"""
|
||||
self._cleanup_callbacks.append(callback)
|
||||
|
||||
def request_shutdown(self) -> None:
|
||||
"""Programmatically request shutdown.
|
||||
|
||||
This can be used to trigger shutdown from application code
|
||||
instead of waiting for a signal.
|
||||
"""
|
||||
if not self._shutdown_requested:
|
||||
self._shutdown_requested = True
|
||||
logger.info("Shutdown requested programmatically")
|
||||
if self._shutdown_event:
|
||||
self._shutdown_event.set()
|
||||
|
||||
async def wait(self) -> None:
|
||||
"""Wait for a shutdown signal.
|
||||
|
||||
This coroutine blocks until a shutdown signal is received
|
||||
or request_shutdown() is called.
|
||||
"""
|
||||
if self._shutdown_event is None:
|
||||
self._shutdown_event = asyncio.Event()
|
||||
|
||||
await self._shutdown_event.wait()
|
||||
|
||||
async def wait_with_timeout(self) -> bool:
|
||||
"""Wait for shutdown with timeout.
|
||||
|
||||
Returns:
|
||||
True if shutdown was requested, False if timeout occurred.
|
||||
"""
|
||||
if self._shutdown_event is None:
|
||||
self._shutdown_event = asyncio.Event()
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._shutdown_event.wait(),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
return True
|
||||
except TimeoutError:
|
||||
return False
|
||||
|
||||
def install_signal_handlers(self) -> None:
|
||||
"""Install signal handlers for graceful shutdown.
|
||||
|
||||
Traps SIGTERM and SIGINT to trigger graceful shutdown.
|
||||
On Windows, only SIGINT is trapped as SIGTERM is not available.
|
||||
"""
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._shutdown_event = asyncio.Event()
|
||||
|
||||
# Platform-specific signal handling
|
||||
if sys.platform == "win32":
|
||||
# Windows: use signal.signal for SIGINT only
|
||||
self._install_windows_handlers()
|
||||
else:
|
||||
# Unix: use loop.add_signal_handler for both signals
|
||||
self._install_unix_handlers()
|
||||
|
||||
logger.debug("Signal handlers installed")
|
||||
|
||||
def _install_unix_handlers(self) -> None:
|
||||
"""Install Unix signal handlers using the event loop."""
|
||||
if self._loop is None:
|
||||
return
|
||||
|
||||
for sig in SHUTDOWN_SIGNALS:
|
||||
try:
|
||||
self._loop.add_signal_handler(
|
||||
sig,
|
||||
self._handle_signal,
|
||||
sig,
|
||||
)
|
||||
logger.debug("Installed handler for %s", sig.name)
|
||||
except (ValueError, OSError) as e:
|
||||
logger.warning("Could not install handler for %s: %s", sig.name, e)
|
||||
|
||||
def _install_windows_handlers(self) -> None:
|
||||
"""Install Windows signal handlers using signal.signal."""
|
||||
for sig in SHUTDOWN_SIGNALS:
|
||||
try:
|
||||
self._original_handlers[sig] = signal.signal(
|
||||
sig,
|
||||
self._handle_signal_sync,
|
||||
)
|
||||
logger.debug("Installed handler for %s", sig.name)
|
||||
except (ValueError, OSError) as e:
|
||||
logger.warning("Could not install handler for %s: %s", sig.name, e)
|
||||
|
||||
def remove_signal_handlers(self) -> None:
|
||||
"""Remove installed signal handlers and restore originals."""
|
||||
if sys.platform == "win32":
|
||||
self._remove_windows_handlers()
|
||||
else:
|
||||
self._remove_unix_handlers()
|
||||
|
||||
logger.debug("Signal handlers removed")
|
||||
|
||||
def _remove_unix_handlers(self) -> None:
|
||||
"""Remove Unix signal handlers."""
|
||||
if self._loop is None:
|
||||
return
|
||||
|
||||
for sig in SHUTDOWN_SIGNALS:
|
||||
with suppress(ValueError, OSError):
|
||||
self._loop.remove_signal_handler(sig)
|
||||
|
||||
def _remove_windows_handlers(self) -> None:
|
||||
"""Remove Windows signal handlers and restore originals."""
|
||||
for sig, original in self._original_handlers.items():
|
||||
with suppress(ValueError, OSError):
|
||||
signal.signal(sig, original)
|
||||
self._original_handlers.clear()
|
||||
|
||||
def _handle_signal(self, sig: signal.Signals) -> None:
|
||||
"""Handle shutdown signal (Unix version).
|
||||
|
||||
Args:
|
||||
sig: The signal that was received.
|
||||
"""
|
||||
if self._shutdown_requested:
|
||||
# Second signal - force exit
|
||||
self._force_exit_requested = True
|
||||
logger.warning("Received %s again - forcing exit!", sig.name)
|
||||
sys.exit(128 + sig.value)
|
||||
else:
|
||||
self._shutdown_requested = True
|
||||
logger.info("Received %s - initiating graceful shutdown...", sig.name)
|
||||
if self._shutdown_event:
|
||||
self._shutdown_event.set()
|
||||
|
||||
def _handle_signal_sync(self, sig: int, _frame: FrameType | None) -> None:
|
||||
"""Handle shutdown signal (Windows version).
|
||||
|
||||
Args:
|
||||
sig: The signal number that was received.
|
||||
_frame: The current stack frame (unused).
|
||||
"""
|
||||
sig_enum = signal.Signals(sig)
|
||||
self._handle_signal(sig_enum)
|
||||
|
||||
async def run_cleanup_callbacks(self) -> None:
|
||||
"""Run all registered cleanup callbacks."""
|
||||
for callback in self._cleanup_callbacks:
|
||||
try:
|
||||
result = callback()
|
||||
if asyncio.iscoroutine(result):
|
||||
await result
|
||||
except Exception as e:
|
||||
logger.error("Cleanup callback failed: %s", e)
|
||||
|
||||
async def __aenter__(self) -> GracefulShutdown:
|
||||
"""Async context manager entry - install signal handlers."""
|
||||
self.install_signal_handlers()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args: Any) -> None:
|
||||
"""Async context manager exit - cleanup."""
|
||||
self.remove_signal_handlers()
|
||||
await self.run_cleanup_callbacks()
|
||||
|
||||
|
||||
async def run_with_graceful_shutdown(
|
||||
coro: Any,
|
||||
*,
|
||||
timeout: float = DEFAULT_SHUTDOWN_TIMEOUT,
|
||||
) -> None:
|
||||
"""Run an async coroutine with graceful shutdown handling.
|
||||
|
||||
This is a convenience function that wraps a coroutine with
|
||||
signal handling and timeout-based cleanup.
|
||||
|
||||
Args:
|
||||
coro: The coroutine to run.
|
||||
timeout: Maximum time to wait for graceful shutdown.
|
||||
|
||||
Example:
|
||||
```python
|
||||
async def my_app():
|
||||
await asyncio.sleep(3600) # Run for an hour
|
||||
|
||||
# Will handle SIGTERM/SIGINT gracefully
|
||||
await run_with_graceful_shutdown(my_app())
|
||||
```
|
||||
"""
|
||||
shutdown = GracefulShutdown(timeout=timeout)
|
||||
|
||||
async with shutdown:
|
||||
task = asyncio.create_task(coro)
|
||||
|
||||
# Wait for either task completion or shutdown signal
|
||||
shutdown_wait = asyncio.create_task(shutdown.wait())
|
||||
|
||||
done, pending = await asyncio.wait(
|
||||
[task, shutdown_wait],
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
|
||||
# Cancel pending tasks
|
||||
for pending_task in pending:
|
||||
pending_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await pending_task
|
||||
|
||||
# If the main task is in done, get any exception
|
||||
if task in done:
|
||||
task.result()
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Storage layer - Database schemas and repositories."""
|
||||
|
||||
from polymarket_insider_tracker.storage.database import (
|
||||
DatabaseManager,
|
||||
create_async_db_engine,
|
||||
create_async_session_factory,
|
||||
create_sync_engine,
|
||||
create_sync_session_factory,
|
||||
init_async_db,
|
||||
init_db,
|
||||
)
|
||||
from polymarket_insider_tracker.storage.models import (
|
||||
Base,
|
||||
FundingTransferModel,
|
||||
WalletProfileModel,
|
||||
WalletRelationshipModel,
|
||||
)
|
||||
from polymarket_insider_tracker.storage.repos import (
|
||||
FundingRepository,
|
||||
FundingTransferDTO,
|
||||
RelationshipRepository,
|
||||
WalletProfileDTO,
|
||||
WalletRelationshipDTO,
|
||||
WalletRepository,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Base",
|
||||
"DatabaseManager",
|
||||
"FundingRepository",
|
||||
"FundingTransferDTO",
|
||||
"FundingTransferModel",
|
||||
"RelationshipRepository",
|
||||
"WalletProfileDTO",
|
||||
"WalletProfileModel",
|
||||
"WalletRelationshipDTO",
|
||||
"WalletRelationshipModel",
|
||||
"WalletRepository",
|
||||
"create_async_db_engine",
|
||||
"create_async_session_factory",
|
||||
"create_sync_engine",
|
||||
"create_sync_session_factory",
|
||||
"init_async_db",
|
||||
"init_db",
|
||||
]
|
||||
@@ -0,0 +1,209 @@
|
||||
"""Database connection and session management.
|
||||
|
||||
This module provides the database engine, session factory, and
|
||||
async session support for the storage layer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from polymarket_insider_tracker.storage.models import Base
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy import Engine
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_sync_engine(database_url: str, **kwargs: Any) -> Engine:
|
||||
"""Create a synchronous SQLAlchemy engine.
|
||||
|
||||
Args:
|
||||
database_url: Database connection URL (e.g., postgresql://...).
|
||||
**kwargs: Additional engine options.
|
||||
|
||||
Returns:
|
||||
SQLAlchemy Engine instance.
|
||||
"""
|
||||
return create_engine(database_url, **kwargs)
|
||||
|
||||
|
||||
def create_async_db_engine(database_url: str, **kwargs: Any) -> AsyncEngine:
|
||||
"""Create an asynchronous SQLAlchemy engine.
|
||||
|
||||
Args:
|
||||
database_url: Database connection URL (e.g., postgresql+asyncpg://...).
|
||||
**kwargs: Additional engine options.
|
||||
|
||||
Returns:
|
||||
SQLAlchemy AsyncEngine instance.
|
||||
"""
|
||||
return create_async_engine(database_url, **kwargs)
|
||||
|
||||
|
||||
def create_sync_session_factory(engine: Engine) -> sessionmaker[Session]:
|
||||
"""Create a synchronous session factory.
|
||||
|
||||
Args:
|
||||
engine: SQLAlchemy Engine instance.
|
||||
|
||||
Returns:
|
||||
Session factory.
|
||||
"""
|
||||
return sessionmaker(bind=engine, expire_on_commit=False)
|
||||
|
||||
|
||||
def create_async_session_factory(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
|
||||
"""Create an asynchronous session factory.
|
||||
|
||||
Args:
|
||||
engine: SQLAlchemy AsyncEngine instance.
|
||||
|
||||
Returns:
|
||||
Async session factory.
|
||||
"""
|
||||
return async_sessionmaker(bind=engine, expire_on_commit=False)
|
||||
|
||||
|
||||
def init_db(engine: Engine) -> None:
|
||||
"""Initialize the database schema.
|
||||
|
||||
Creates all tables defined in the models.
|
||||
|
||||
Args:
|
||||
engine: SQLAlchemy Engine instance.
|
||||
"""
|
||||
Base.metadata.create_all(engine)
|
||||
logger.info("Database schema initialized")
|
||||
|
||||
|
||||
async def init_async_db(engine: AsyncEngine) -> None:
|
||||
"""Initialize the database schema asynchronously.
|
||||
|
||||
Creates all tables defined in the models.
|
||||
|
||||
Args:
|
||||
engine: SQLAlchemy AsyncEngine instance.
|
||||
"""
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
logger.info("Database schema initialized (async)")
|
||||
|
||||
|
||||
class DatabaseManager:
|
||||
"""Manages database connections and sessions.
|
||||
|
||||
Provides a unified interface for both sync and async database operations.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
database_url: str,
|
||||
*,
|
||||
async_mode: bool = True,
|
||||
pool_size: int = 5,
|
||||
max_overflow: int = 10,
|
||||
echo: bool = False,
|
||||
) -> None:
|
||||
"""Initialize database manager.
|
||||
|
||||
Args:
|
||||
database_url: Database connection URL.
|
||||
async_mode: Use async engine/sessions if True.
|
||||
pool_size: Connection pool size.
|
||||
max_overflow: Maximum overflow connections.
|
||||
echo: Echo SQL statements for debugging.
|
||||
"""
|
||||
self.database_url = database_url
|
||||
self.async_mode = async_mode
|
||||
self._pool_size = pool_size
|
||||
self._max_overflow = max_overflow
|
||||
self._echo = echo
|
||||
|
||||
self._sync_engine: Engine | None = None
|
||||
self._async_engine: AsyncEngine | None = None
|
||||
self._sync_session_factory: sessionmaker[Session] | None = None
|
||||
self._async_session_factory: async_sessionmaker[AsyncSession] | None = None
|
||||
|
||||
def _get_sync_engine(self) -> Engine:
|
||||
"""Get or create the synchronous engine."""
|
||||
if self._sync_engine is None:
|
||||
self._sync_engine = create_sync_engine(
|
||||
self.database_url,
|
||||
pool_size=self._pool_size,
|
||||
max_overflow=self._max_overflow,
|
||||
echo=self._echo,
|
||||
)
|
||||
return self._sync_engine
|
||||
|
||||
def _get_async_engine(self) -> AsyncEngine:
|
||||
"""Get or create the asynchronous engine."""
|
||||
if self._async_engine is None:
|
||||
self._async_engine = create_async_db_engine(
|
||||
self.database_url,
|
||||
pool_size=self._pool_size,
|
||||
max_overflow=self._max_overflow,
|
||||
echo=self._echo,
|
||||
)
|
||||
return self._async_engine
|
||||
|
||||
def get_sync_session(self) -> Session:
|
||||
"""Get a new synchronous session.
|
||||
|
||||
Returns:
|
||||
SQLAlchemy Session instance.
|
||||
"""
|
||||
if self._sync_session_factory is None:
|
||||
self._sync_session_factory = create_sync_session_factory(self._get_sync_engine())
|
||||
return self._sync_session_factory()
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_async_session(self) -> AsyncGenerator[AsyncSession, None]:
|
||||
"""Get an asynchronous session as a context manager.
|
||||
|
||||
Yields:
|
||||
SQLAlchemy AsyncSession instance.
|
||||
"""
|
||||
if self._async_session_factory is None:
|
||||
self._async_session_factory = create_async_session_factory(self._get_async_engine())
|
||||
|
||||
session = self._async_session_factory()
|
||||
try:
|
||||
yield session
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
def init_schema(self) -> None:
|
||||
"""Initialize database schema synchronously."""
|
||||
init_db(self._get_sync_engine())
|
||||
|
||||
async def init_schema_async(self) -> None:
|
||||
"""Initialize database schema asynchronously."""
|
||||
await init_async_db(self._get_async_engine())
|
||||
|
||||
def dispose(self) -> None:
|
||||
"""Dispose of all database connections."""
|
||||
if self._sync_engine is not None:
|
||||
self._sync_engine.dispose()
|
||||
self._sync_engine = None
|
||||
logger.info("Database connections disposed")
|
||||
|
||||
async def dispose_async(self) -> None:
|
||||
"""Dispose of all async database connections."""
|
||||
if self._async_engine is not None:
|
||||
await self._async_engine.dispose()
|
||||
self._async_engine = None
|
||||
logger.info("Async database connections disposed")
|
||||
@@ -0,0 +1,117 @@
|
||||
"""SQLAlchemy models for persistent storage.
|
||||
|
||||
This module defines the database schema for storing wallet profiles,
|
||||
funding transfers, and wallet relationships.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import (
|
||||
Boolean,
|
||||
DateTime,
|
||||
Index,
|
||||
Integer,
|
||||
Numeric,
|
||||
String,
|
||||
UniqueConstraint,
|
||||
)
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
"""Base class for all SQLAlchemy models."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WalletProfileModel(Base):
|
||||
"""SQLAlchemy model for wallet profiles.
|
||||
|
||||
Stores analyzed wallet information including age, transaction count,
|
||||
balances, and freshness classification.
|
||||
"""
|
||||
|
||||
__tablename__ = "wallet_profiles"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
address: Mapped[str] = mapped_column(String(42), unique=True, nullable=False)
|
||||
nonce: Mapped[int] = mapped_column(Integer, nullable=False)
|
||||
first_seen_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
is_fresh: Mapped[bool] = mapped_column(Boolean, nullable=False)
|
||||
matic_balance: Mapped[Decimal | None] = mapped_column(Numeric(30, 0), nullable=True)
|
||||
usdc_balance: Mapped[Decimal | None] = mapped_column(Numeric(20, 6), nullable=True)
|
||||
analyzed_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), nullable=False, default=lambda: datetime.now(UTC)
|
||||
)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), nullable=False, default=lambda: datetime.now(UTC)
|
||||
)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
nullable=False,
|
||||
default=lambda: datetime.now(UTC),
|
||||
onupdate=lambda: datetime.now(UTC),
|
||||
)
|
||||
|
||||
__table_args__ = (Index("idx_wallet_profiles_address", "address"),)
|
||||
|
||||
|
||||
class FundingTransferModel(Base):
|
||||
"""SQLAlchemy model for funding transfers.
|
||||
|
||||
Stores ERC20 transfer events to track wallet funding sources.
|
||||
"""
|
||||
|
||||
__tablename__ = "funding_transfers"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
from_address: Mapped[str] = mapped_column(String(42), nullable=False)
|
||||
to_address: Mapped[str] = mapped_column(String(42), nullable=False)
|
||||
amount: Mapped[Decimal] = mapped_column(Numeric(30, 6), nullable=False)
|
||||
token: Mapped[str] = mapped_column(String(10), nullable=False)
|
||||
tx_hash: Mapped[str] = mapped_column(String(66), unique=True, nullable=False)
|
||||
block_number: Mapped[int] = mapped_column(Integer, nullable=False)
|
||||
timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), nullable=False, default=lambda: datetime.now(UTC)
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
Index("idx_funding_transfers_to", "to_address"),
|
||||
Index("idx_funding_transfers_from", "from_address"),
|
||||
Index("idx_funding_transfers_block", "block_number"),
|
||||
)
|
||||
|
||||
|
||||
class WalletRelationshipModel(Base):
|
||||
"""SQLAlchemy model for wallet relationships.
|
||||
|
||||
Stores graph edges between wallets representing funding relationships
|
||||
or entity linkages.
|
||||
"""
|
||||
|
||||
__tablename__ = "wallet_relationships"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
wallet_a: Mapped[str] = mapped_column(String(42), nullable=False)
|
||||
wallet_b: Mapped[str] = mapped_column(String(42), nullable=False)
|
||||
relationship_type: Mapped[str] = mapped_column(String(20), nullable=False)
|
||||
confidence: Mapped[Decimal] = mapped_column(Numeric(3, 2), nullable=False)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), nullable=False, default=lambda: datetime.now(UTC)
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint(
|
||||
"wallet_a", "wallet_b", "relationship_type", name="uq_wallet_relationship"
|
||||
),
|
||||
Index("idx_wallet_relationships_a", "wallet_a"),
|
||||
Index("idx_wallet_relationships_b", "wallet_b"),
|
||||
)
|
||||
@@ -0,0 +1,512 @@
|
||||
"""Repository pattern implementations for data access.
|
||||
|
||||
This module provides clean data access abstractions for wallet profiles,
|
||||
funding transfers, and wallet relationships.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import delete, select, update
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
|
||||
from polymarket_insider_tracker.storage.models import (
|
||||
FundingTransferModel,
|
||||
WalletProfileModel,
|
||||
WalletRelationshipModel,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class WalletProfileDTO:
|
||||
"""Data transfer object for wallet profiles."""
|
||||
|
||||
address: str
|
||||
nonce: int
|
||||
first_seen_at: datetime | None
|
||||
is_fresh: bool
|
||||
matic_balance: Decimal | None
|
||||
usdc_balance: Decimal | None
|
||||
analyzed_at: datetime
|
||||
created_at: datetime | None = None
|
||||
updated_at: datetime | None = None
|
||||
|
||||
@classmethod
|
||||
def from_model(cls, model: WalletProfileModel) -> WalletProfileDTO:
|
||||
"""Create DTO from SQLAlchemy model."""
|
||||
return cls(
|
||||
address=model.address,
|
||||
nonce=model.nonce,
|
||||
first_seen_at=model.first_seen_at,
|
||||
is_fresh=model.is_fresh,
|
||||
matic_balance=model.matic_balance,
|
||||
usdc_balance=model.usdc_balance,
|
||||
analyzed_at=model.analyzed_at,
|
||||
created_at=model.created_at,
|
||||
updated_at=model.updated_at,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FundingTransferDTO:
|
||||
"""Data transfer object for funding transfers."""
|
||||
|
||||
from_address: str
|
||||
to_address: str
|
||||
amount: Decimal
|
||||
token: str
|
||||
tx_hash: str
|
||||
block_number: int
|
||||
timestamp: datetime
|
||||
created_at: datetime | None = None
|
||||
|
||||
@classmethod
|
||||
def from_model(cls, model: FundingTransferModel) -> FundingTransferDTO:
|
||||
"""Create DTO from SQLAlchemy model."""
|
||||
return cls(
|
||||
from_address=model.from_address,
|
||||
to_address=model.to_address,
|
||||
amount=model.amount,
|
||||
token=model.token,
|
||||
tx_hash=model.tx_hash,
|
||||
block_number=model.block_number,
|
||||
timestamp=model.timestamp,
|
||||
created_at=model.created_at,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class WalletRelationshipDTO:
|
||||
"""Data transfer object for wallet relationships."""
|
||||
|
||||
wallet_a: str
|
||||
wallet_b: str
|
||||
relationship_type: str
|
||||
confidence: Decimal
|
||||
created_at: datetime | None = None
|
||||
|
||||
@classmethod
|
||||
def from_model(cls, model: WalletRelationshipModel) -> WalletRelationshipDTO:
|
||||
"""Create DTO from SQLAlchemy model."""
|
||||
return cls(
|
||||
wallet_a=model.wallet_a,
|
||||
wallet_b=model.wallet_b,
|
||||
relationship_type=model.relationship_type,
|
||||
confidence=model.confidence,
|
||||
created_at=model.created_at,
|
||||
)
|
||||
|
||||
|
||||
class WalletRepository:
|
||||
"""Repository for wallet profile data access.
|
||||
|
||||
Provides CRUD operations for wallet profiles with async support.
|
||||
"""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
"""Initialize repository with database session.
|
||||
|
||||
Args:
|
||||
session: SQLAlchemy async session.
|
||||
"""
|
||||
self.session = session
|
||||
|
||||
async def get_by_address(self, address: str) -> WalletProfileDTO | None:
|
||||
"""Get wallet profile by address.
|
||||
|
||||
Args:
|
||||
address: Wallet address (lowercase).
|
||||
|
||||
Returns:
|
||||
WalletProfileDTO if found, None otherwise.
|
||||
"""
|
||||
result = await self.session.execute(
|
||||
select(WalletProfileModel).where(WalletProfileModel.address == address.lower())
|
||||
)
|
||||
model = result.scalar_one_or_none()
|
||||
return WalletProfileDTO.from_model(model) if model else None
|
||||
|
||||
async def get_many(self, addresses: list[str]) -> list[WalletProfileDTO]:
|
||||
"""Get multiple wallet profiles by addresses.
|
||||
|
||||
Args:
|
||||
addresses: List of wallet addresses.
|
||||
|
||||
Returns:
|
||||
List of WalletProfileDTOs for found addresses.
|
||||
"""
|
||||
normalized = [addr.lower() for addr in addresses]
|
||||
result = await self.session.execute(
|
||||
select(WalletProfileModel).where(WalletProfileModel.address.in_(normalized))
|
||||
)
|
||||
return [WalletProfileDTO.from_model(m) for m in result.scalars().all()]
|
||||
|
||||
async def get_fresh_wallets(self, limit: int = 100) -> list[WalletProfileDTO]:
|
||||
"""Get recent fresh wallets.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of results.
|
||||
|
||||
Returns:
|
||||
List of WalletProfileDTOs marked as fresh.
|
||||
"""
|
||||
result = await self.session.execute(
|
||||
select(WalletProfileModel)
|
||||
.where(WalletProfileModel.is_fresh.is_(True))
|
||||
.order_by(WalletProfileModel.analyzed_at.desc())
|
||||
.limit(limit)
|
||||
)
|
||||
return [WalletProfileDTO.from_model(m) for m in result.scalars().all()]
|
||||
|
||||
async def upsert(self, dto: WalletProfileDTO) -> WalletProfileDTO:
|
||||
"""Insert or update wallet profile.
|
||||
|
||||
Args:
|
||||
dto: Wallet profile data.
|
||||
|
||||
Returns:
|
||||
Updated WalletProfileDTO.
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
values = {
|
||||
"address": dto.address.lower(),
|
||||
"nonce": dto.nonce,
|
||||
"first_seen_at": dto.first_seen_at,
|
||||
"is_fresh": dto.is_fresh,
|
||||
"matic_balance": dto.matic_balance,
|
||||
"usdc_balance": dto.usdc_balance,
|
||||
"analyzed_at": dto.analyzed_at,
|
||||
"updated_at": now,
|
||||
}
|
||||
|
||||
# Try PostgreSQL upsert first, fall back to SQLite for testing
|
||||
try:
|
||||
stmt = pg_insert(WalletProfileModel).values(**values, created_at=now)
|
||||
stmt = stmt.on_conflict_do_update(
|
||||
index_elements=["address"],
|
||||
set_={
|
||||
"nonce": stmt.excluded.nonce,
|
||||
"first_seen_at": stmt.excluded.first_seen_at,
|
||||
"is_fresh": stmt.excluded.is_fresh,
|
||||
"matic_balance": stmt.excluded.matic_balance,
|
||||
"usdc_balance": stmt.excluded.usdc_balance,
|
||||
"analyzed_at": stmt.excluded.analyzed_at,
|
||||
"updated_at": stmt.excluded.updated_at,
|
||||
},
|
||||
)
|
||||
await self.session.execute(stmt)
|
||||
except Exception:
|
||||
# Fall back to SQLite upsert for testing
|
||||
sqlite_stmt = sqlite_insert(WalletProfileModel).values(**values, created_at=now)
|
||||
sqlite_stmt = sqlite_stmt.on_conflict_do_update(
|
||||
index_elements=["address"],
|
||||
set_={
|
||||
"nonce": sqlite_stmt.excluded.nonce,
|
||||
"first_seen_at": sqlite_stmt.excluded.first_seen_at,
|
||||
"is_fresh": sqlite_stmt.excluded.is_fresh,
|
||||
"matic_balance": sqlite_stmt.excluded.matic_balance,
|
||||
"usdc_balance": sqlite_stmt.excluded.usdc_balance,
|
||||
"analyzed_at": sqlite_stmt.excluded.analyzed_at,
|
||||
"updated_at": sqlite_stmt.excluded.updated_at,
|
||||
},
|
||||
)
|
||||
await self.session.execute(sqlite_stmt)
|
||||
|
||||
await self.session.flush()
|
||||
return dto
|
||||
|
||||
async def delete(self, address: str) -> bool:
|
||||
"""Delete wallet profile by address.
|
||||
|
||||
Args:
|
||||
address: Wallet address.
|
||||
|
||||
Returns:
|
||||
True if deleted, False if not found.
|
||||
"""
|
||||
result = await self.session.execute(
|
||||
delete(WalletProfileModel).where(WalletProfileModel.address == address.lower())
|
||||
)
|
||||
# SQLAlchemy Result does have rowcount but typing doesn't reflect it
|
||||
return (result.rowcount or 0) > 0 # type: ignore[attr-defined]
|
||||
|
||||
async def mark_stale(self, address: str) -> bool:
|
||||
"""Mark a wallet profile as stale (soft delete).
|
||||
|
||||
Sets analyzed_at to a very old date to trigger re-analysis.
|
||||
|
||||
Args:
|
||||
address: Wallet address.
|
||||
|
||||
Returns:
|
||||
True if updated, False if not found.
|
||||
"""
|
||||
stale_time = datetime(2000, 1, 1, tzinfo=UTC)
|
||||
result = await self.session.execute(
|
||||
update(WalletProfileModel)
|
||||
.where(WalletProfileModel.address == address.lower())
|
||||
.values(analyzed_at=stale_time, updated_at=datetime.now(UTC))
|
||||
)
|
||||
# SQLAlchemy Result does have rowcount but typing doesn't reflect it
|
||||
return (result.rowcount or 0) > 0 # type: ignore[attr-defined]
|
||||
|
||||
|
||||
class FundingRepository:
|
||||
"""Repository for funding transfer data access.
|
||||
|
||||
Provides CRUD operations for funding transfers with async support.
|
||||
"""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
"""Initialize repository with database session.
|
||||
|
||||
Args:
|
||||
session: SQLAlchemy async session.
|
||||
"""
|
||||
self.session = session
|
||||
|
||||
async def get_transfers_to(self, address: str, limit: int = 100) -> list[FundingTransferDTO]:
|
||||
"""Get transfers to a wallet address.
|
||||
|
||||
Args:
|
||||
address: Destination wallet address.
|
||||
limit: Maximum number of results.
|
||||
|
||||
Returns:
|
||||
List of FundingTransferDTOs ordered by timestamp.
|
||||
"""
|
||||
result = await self.session.execute(
|
||||
select(FundingTransferModel)
|
||||
.where(FundingTransferModel.to_address == address.lower())
|
||||
.order_by(FundingTransferModel.timestamp.asc())
|
||||
.limit(limit)
|
||||
)
|
||||
return [FundingTransferDTO.from_model(m) for m in result.scalars().all()]
|
||||
|
||||
async def get_transfers_from(self, address: str, limit: int = 100) -> list[FundingTransferDTO]:
|
||||
"""Get transfers from a wallet address.
|
||||
|
||||
Args:
|
||||
address: Source wallet address.
|
||||
limit: Maximum number of results.
|
||||
|
||||
Returns:
|
||||
List of FundingTransferDTOs ordered by timestamp.
|
||||
"""
|
||||
result = await self.session.execute(
|
||||
select(FundingTransferModel)
|
||||
.where(FundingTransferModel.from_address == address.lower())
|
||||
.order_by(FundingTransferModel.timestamp.asc())
|
||||
.limit(limit)
|
||||
)
|
||||
return [FundingTransferDTO.from_model(m) for m in result.scalars().all()]
|
||||
|
||||
async def get_first_transfer_to(self, address: str) -> FundingTransferDTO | None:
|
||||
"""Get the first transfer to a wallet.
|
||||
|
||||
Args:
|
||||
address: Wallet address.
|
||||
|
||||
Returns:
|
||||
First FundingTransferDTO if found, None otherwise.
|
||||
"""
|
||||
result = await self.session.execute(
|
||||
select(FundingTransferModel)
|
||||
.where(FundingTransferModel.to_address == address.lower())
|
||||
.order_by(FundingTransferModel.timestamp.asc())
|
||||
.limit(1)
|
||||
)
|
||||
model = result.scalar_one_or_none()
|
||||
return FundingTransferDTO.from_model(model) if model else None
|
||||
|
||||
async def get_by_tx_hash(self, tx_hash: str) -> FundingTransferDTO | None:
|
||||
"""Get transfer by transaction hash.
|
||||
|
||||
Args:
|
||||
tx_hash: Transaction hash.
|
||||
|
||||
Returns:
|
||||
FundingTransferDTO if found, None otherwise.
|
||||
"""
|
||||
result = await self.session.execute(
|
||||
select(FundingTransferModel).where(FundingTransferModel.tx_hash == tx_hash.lower())
|
||||
)
|
||||
model = result.scalar_one_or_none()
|
||||
return FundingTransferDTO.from_model(model) if model else None
|
||||
|
||||
async def insert(self, dto: FundingTransferDTO) -> FundingTransferDTO:
|
||||
"""Insert a new funding transfer.
|
||||
|
||||
Args:
|
||||
dto: Funding transfer data.
|
||||
|
||||
Returns:
|
||||
Inserted FundingTransferDTO.
|
||||
|
||||
Raises:
|
||||
IntegrityError if tx_hash already exists.
|
||||
"""
|
||||
model = FundingTransferModel(
|
||||
from_address=dto.from_address.lower(),
|
||||
to_address=dto.to_address.lower(),
|
||||
amount=dto.amount,
|
||||
token=dto.token,
|
||||
tx_hash=dto.tx_hash.lower(),
|
||||
block_number=dto.block_number,
|
||||
timestamp=dto.timestamp,
|
||||
)
|
||||
self.session.add(model)
|
||||
await self.session.flush()
|
||||
return dto
|
||||
|
||||
async def insert_many(self, dtos: list[FundingTransferDTO]) -> int:
|
||||
"""Insert multiple funding transfers.
|
||||
|
||||
Skips duplicates silently.
|
||||
|
||||
Args:
|
||||
dtos: List of funding transfer data.
|
||||
|
||||
Returns:
|
||||
Number of transfers inserted.
|
||||
"""
|
||||
inserted = 0
|
||||
for dto in dtos:
|
||||
try:
|
||||
await self.insert(dto)
|
||||
inserted += 1
|
||||
except Exception as e:
|
||||
# Skip duplicates
|
||||
if "UNIQUE constraint" in str(e) or "duplicate key" in str(e).lower():
|
||||
continue
|
||||
raise
|
||||
return inserted
|
||||
|
||||
|
||||
class RelationshipRepository:
|
||||
"""Repository for wallet relationship data access.
|
||||
|
||||
Provides CRUD operations for wallet relationships with async support.
|
||||
"""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
"""Initialize repository with database session.
|
||||
|
||||
Args:
|
||||
session: SQLAlchemy async session.
|
||||
"""
|
||||
self.session = session
|
||||
|
||||
async def get_relationships(
|
||||
self, wallet: str, relationship_type: str | None = None
|
||||
) -> list[WalletRelationshipDTO]:
|
||||
"""Get relationships for a wallet.
|
||||
|
||||
Args:
|
||||
wallet: Wallet address.
|
||||
relationship_type: Optional filter by type.
|
||||
|
||||
Returns:
|
||||
List of WalletRelationshipDTOs.
|
||||
"""
|
||||
stmt = select(WalletRelationshipModel).where(
|
||||
(WalletRelationshipModel.wallet_a == wallet.lower())
|
||||
| (WalletRelationshipModel.wallet_b == wallet.lower())
|
||||
)
|
||||
if relationship_type:
|
||||
stmt = stmt.where(WalletRelationshipModel.relationship_type == relationship_type)
|
||||
|
||||
result = await self.session.execute(stmt)
|
||||
return [WalletRelationshipDTO.from_model(m) for m in result.scalars().all()]
|
||||
|
||||
async def get_related_wallets(
|
||||
self, wallet: str, relationship_type: str | None = None
|
||||
) -> list[str]:
|
||||
"""Get addresses of related wallets.
|
||||
|
||||
Args:
|
||||
wallet: Wallet address.
|
||||
relationship_type: Optional filter by type.
|
||||
|
||||
Returns:
|
||||
List of related wallet addresses.
|
||||
"""
|
||||
relationships = await self.get_relationships(wallet, relationship_type)
|
||||
related = set()
|
||||
normalized = wallet.lower()
|
||||
for rel in relationships:
|
||||
if rel.wallet_a == normalized:
|
||||
related.add(rel.wallet_b)
|
||||
else:
|
||||
related.add(rel.wallet_a)
|
||||
return list(related)
|
||||
|
||||
async def upsert(self, dto: WalletRelationshipDTO) -> WalletRelationshipDTO:
|
||||
"""Insert or update wallet relationship.
|
||||
|
||||
Args:
|
||||
dto: Wallet relationship data.
|
||||
|
||||
Returns:
|
||||
Updated WalletRelationshipDTO.
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
values = {
|
||||
"wallet_a": dto.wallet_a.lower(),
|
||||
"wallet_b": dto.wallet_b.lower(),
|
||||
"relationship_type": dto.relationship_type,
|
||||
"confidence": dto.confidence,
|
||||
"created_at": now,
|
||||
}
|
||||
|
||||
# Try PostgreSQL upsert first, fall back to SQLite for testing
|
||||
try:
|
||||
stmt = pg_insert(WalletRelationshipModel).values(**values)
|
||||
stmt = stmt.on_conflict_do_update(
|
||||
constraint="uq_wallet_relationship",
|
||||
set_={"confidence": stmt.excluded.confidence},
|
||||
)
|
||||
await self.session.execute(stmt)
|
||||
except Exception:
|
||||
# Fall back to SQLite upsert for testing
|
||||
sqlite_stmt = sqlite_insert(WalletRelationshipModel).values(**values)
|
||||
sqlite_stmt = sqlite_stmt.on_conflict_do_update(
|
||||
index_elements=["wallet_a", "wallet_b", "relationship_type"],
|
||||
set_={"confidence": sqlite_stmt.excluded.confidence},
|
||||
)
|
||||
await self.session.execute(sqlite_stmt)
|
||||
|
||||
await self.session.flush()
|
||||
return dto
|
||||
|
||||
async def delete(self, wallet_a: str, wallet_b: str, relationship_type: str) -> bool:
|
||||
"""Delete a specific relationship.
|
||||
|
||||
Args:
|
||||
wallet_a: First wallet address.
|
||||
wallet_b: Second wallet address.
|
||||
relationship_type: Type of relationship.
|
||||
|
||||
Returns:
|
||||
True if deleted, False if not found.
|
||||
"""
|
||||
result = await self.session.execute(
|
||||
delete(WalletRelationshipModel).where(
|
||||
WalletRelationshipModel.wallet_a == wallet_a.lower(),
|
||||
WalletRelationshipModel.wallet_b == wallet_b.lower(),
|
||||
WalletRelationshipModel.relationship_type == relationship_type,
|
||||
)
|
||||
)
|
||||
# SQLAlchemy Result does have rowcount but typing doesn't reflect it
|
||||
return (result.rowcount or 0) > 0 # type: ignore[attr-defined]
|
||||
@@ -0,0 +1 @@
|
||||
"""Test suite for Polymarket Insider Tracker."""
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for the alerter module."""
|
||||
@@ -0,0 +1,492 @@
|
||||
"""Tests for alert dispatcher and channels."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.alerter.channels.discord import DiscordChannel
|
||||
from polymarket_insider_tracker.alerter.channels.telegram import TelegramChannel
|
||||
from polymarket_insider_tracker.alerter.dispatcher import (
|
||||
AlertDispatcher,
|
||||
CircuitBreakerState,
|
||||
DispatchResult,
|
||||
)
|
||||
from polymarket_insider_tracker.alerter.models import FormattedAlert
|
||||
|
||||
# ============================================================================
|
||||
# Fixtures
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_alert() -> FormattedAlert:
|
||||
"""Create a sample formatted alert."""
|
||||
return FormattedAlert(
|
||||
title="Test Alert",
|
||||
body="Test body",
|
||||
discord_embed={
|
||||
"title": "Test",
|
||||
"color": 15158332,
|
||||
"fields": [],
|
||||
},
|
||||
telegram_markdown="*Test Alert*\nTest body",
|
||||
plain_text="TEST ALERT\nTest body",
|
||||
links={"market": "https://polymarket.com/test"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_discord_channel() -> MagicMock:
|
||||
"""Create a mock Discord channel."""
|
||||
channel = MagicMock()
|
||||
channel.name = "discord"
|
||||
channel.send = AsyncMock(return_value=True)
|
||||
return channel
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_telegram_channel() -> MagicMock:
|
||||
"""Create a mock Telegram channel."""
|
||||
channel = MagicMock()
|
||||
channel.name = "telegram"
|
||||
channel.send = AsyncMock(return_value=True)
|
||||
return channel
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# DiscordChannel Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestDiscordChannel:
|
||||
"""Tests for Discord channel."""
|
||||
|
||||
def test_init(self) -> None:
|
||||
"""Test channel initialization."""
|
||||
channel = DiscordChannel(
|
||||
webhook_url="https://discord.com/api/webhooks/123/abc",
|
||||
rate_limit_per_minute=30,
|
||||
)
|
||||
assert channel.webhook_url == "https://discord.com/api/webhooks/123/abc"
|
||||
assert channel.rate_limit_per_minute == 30
|
||||
assert channel.name == "discord"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_success(self, sample_alert: FormattedAlert) -> None:
|
||||
"""Test successful Discord message send."""
|
||||
channel = DiscordChannel(webhook_url="https://discord.com/api/webhooks/123/abc")
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 204
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
mock_client.__aenter__.return_value = mock_client
|
||||
mock_client.__aexit__.return_value = None
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
result = await channel.send(sample_alert)
|
||||
|
||||
assert result is True
|
||||
mock_client.post.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_rate_limited(self, sample_alert: FormattedAlert) -> None:
|
||||
"""Test Discord rate limit handling."""
|
||||
channel = DiscordChannel(
|
||||
webhook_url="https://discord.com/api/webhooks/123/abc",
|
||||
max_retries=2,
|
||||
retry_delay=0.01,
|
||||
)
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_response_429 = MagicMock()
|
||||
mock_response_429.status_code = 429
|
||||
mock_response_429.json.return_value = {"retry_after": 0.01}
|
||||
|
||||
mock_response_success = MagicMock()
|
||||
mock_response_success.status_code = 204
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.side_effect = [mock_response_429, mock_response_success]
|
||||
mock_client.__aenter__.return_value = mock_client
|
||||
mock_client.__aexit__.return_value = None
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
result = await channel.send(sample_alert)
|
||||
|
||||
assert result is True
|
||||
assert mock_client.post.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_failure(self, sample_alert: FormattedAlert) -> None:
|
||||
"""Test Discord send failure after retries."""
|
||||
channel = DiscordChannel(
|
||||
webhook_url="https://discord.com/api/webhooks/123/abc",
|
||||
max_retries=2,
|
||||
retry_delay=0.01,
|
||||
)
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_response.text = "Internal Server Error"
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
mock_client.__aenter__.return_value = mock_client
|
||||
mock_client.__aexit__.return_value = None
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
result = await channel.send(sample_alert)
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# TelegramChannel Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestTelegramChannel:
|
||||
"""Tests for Telegram channel."""
|
||||
|
||||
def test_init(self) -> None:
|
||||
"""Test channel initialization."""
|
||||
channel = TelegramChannel(
|
||||
bot_token="123456:ABC-DEF",
|
||||
chat_id="-1001234567890",
|
||||
rate_limit_per_minute=20,
|
||||
)
|
||||
assert channel.bot_token == "123456:ABC-DEF"
|
||||
assert channel.chat_id == "-1001234567890"
|
||||
assert channel.name == "telegram"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_success(self, sample_alert: FormattedAlert) -> None:
|
||||
"""Test successful Telegram message send."""
|
||||
channel = TelegramChannel(
|
||||
bot_token="123456:ABC-DEF",
|
||||
chat_id="-1001234567890",
|
||||
)
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"ok": True}
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
mock_client.__aenter__.return_value = mock_client
|
||||
mock_client.__aexit__.return_value = None
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
result = await channel.send(sample_alert)
|
||||
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_rate_limited(self, sample_alert: FormattedAlert) -> None:
|
||||
"""Test Telegram rate limit handling."""
|
||||
channel = TelegramChannel(
|
||||
bot_token="123456:ABC-DEF",
|
||||
chat_id="-1001234567890",
|
||||
max_retries=2,
|
||||
retry_delay=0.01,
|
||||
)
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_response_429 = MagicMock()
|
||||
mock_response_429.json.return_value = {
|
||||
"ok": False,
|
||||
"error_code": 429,
|
||||
"parameters": {"retry_after": 0.01},
|
||||
}
|
||||
|
||||
mock_response_success = MagicMock()
|
||||
mock_response_success.json.return_value = {"ok": True}
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.side_effect = [mock_response_429, mock_response_success]
|
||||
mock_client.__aenter__.return_value = mock_client
|
||||
mock_client.__aexit__.return_value = None
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
result = await channel.send(sample_alert)
|
||||
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_failure(self, sample_alert: FormattedAlert) -> None:
|
||||
"""Test Telegram send failure."""
|
||||
channel = TelegramChannel(
|
||||
bot_token="123456:ABC-DEF",
|
||||
chat_id="-1001234567890",
|
||||
max_retries=2,
|
||||
retry_delay=0.01,
|
||||
)
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client_class:
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"ok": False,
|
||||
"error_code": 400,
|
||||
"description": "Bad Request",
|
||||
}
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
mock_client.__aenter__.return_value = mock_client
|
||||
mock_client.__aexit__.return_value = None
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
result = await channel.send(sample_alert)
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# CircuitBreakerState Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestCircuitBreakerState:
|
||||
"""Tests for circuit breaker state."""
|
||||
|
||||
def test_default_state(self) -> None:
|
||||
"""Test default circuit breaker state."""
|
||||
state = CircuitBreakerState()
|
||||
assert state.failure_count == 0
|
||||
assert state.is_open is False
|
||||
assert state.half_open_attempts == 0
|
||||
assert state.last_failure_time is None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# DispatchResult Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestDispatchResult:
|
||||
"""Tests for dispatch result."""
|
||||
|
||||
def test_all_succeeded(self) -> None:
|
||||
"""Test all_succeeded property."""
|
||||
result = DispatchResult(
|
||||
success_count=2,
|
||||
failure_count=0,
|
||||
channel_results={"discord": True, "telegram": True},
|
||||
)
|
||||
assert result.all_succeeded is True
|
||||
|
||||
def test_partial_success(self) -> None:
|
||||
"""Test partial success."""
|
||||
result = DispatchResult(
|
||||
success_count=1,
|
||||
failure_count=1,
|
||||
channel_results={"discord": True, "telegram": False},
|
||||
)
|
||||
assert result.all_succeeded is False
|
||||
|
||||
def test_empty_channels(self) -> None:
|
||||
"""Test with no channels."""
|
||||
result = DispatchResult(success_count=0, failure_count=0)
|
||||
assert result.all_succeeded is False
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# AlertDispatcher Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestAlertDispatcher:
|
||||
"""Tests for alert dispatcher."""
|
||||
|
||||
def test_init(
|
||||
self,
|
||||
mock_discord_channel: MagicMock,
|
||||
mock_telegram_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test dispatcher initialization."""
|
||||
dispatcher = AlertDispatcher(channels=[mock_discord_channel, mock_telegram_channel])
|
||||
assert len(dispatcher.channels) == 2
|
||||
assert "discord" in dispatcher._circuit_state
|
||||
assert "telegram" in dispatcher._circuit_state
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_all_success(
|
||||
self,
|
||||
sample_alert: FormattedAlert,
|
||||
mock_discord_channel: MagicMock,
|
||||
mock_telegram_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test successful dispatch to all channels."""
|
||||
dispatcher = AlertDispatcher(channels=[mock_discord_channel, mock_telegram_channel])
|
||||
|
||||
result = await dispatcher.dispatch(sample_alert)
|
||||
|
||||
assert result.success_count == 2
|
||||
assert result.failure_count == 0
|
||||
assert result.all_succeeded is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_partial_failure(
|
||||
self,
|
||||
sample_alert: FormattedAlert,
|
||||
mock_discord_channel: MagicMock,
|
||||
mock_telegram_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test dispatch with one channel failing."""
|
||||
mock_telegram_channel.send.return_value = False
|
||||
|
||||
dispatcher = AlertDispatcher(channels=[mock_discord_channel, mock_telegram_channel])
|
||||
|
||||
result = await dispatcher.dispatch(sample_alert)
|
||||
|
||||
assert result.success_count == 1
|
||||
assert result.failure_count == 1
|
||||
assert result.channel_results["discord"] is True
|
||||
assert result.channel_results["telegram"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_no_channels(self, sample_alert: FormattedAlert) -> None:
|
||||
"""Test dispatch with no channels configured."""
|
||||
dispatcher = AlertDispatcher(channels=[])
|
||||
|
||||
result = await dispatcher.dispatch(sample_alert)
|
||||
|
||||
assert result.success_count == 0
|
||||
assert result.failure_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_circuit_opens_after_failures(
|
||||
self,
|
||||
sample_alert: FormattedAlert,
|
||||
mock_discord_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test circuit breaker opens after threshold failures."""
|
||||
mock_discord_channel.send.return_value = False
|
||||
|
||||
dispatcher = AlertDispatcher(
|
||||
channels=[mock_discord_channel],
|
||||
failure_threshold=3,
|
||||
)
|
||||
|
||||
# First 3 failures
|
||||
for _ in range(3):
|
||||
await dispatcher.dispatch(sample_alert)
|
||||
|
||||
# Circuit should be open now
|
||||
assert dispatcher._circuit_state["discord"].is_open is True
|
||||
assert dispatcher._circuit_state["discord"].failure_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_circuit_skips_when_open(
|
||||
self,
|
||||
sample_alert: FormattedAlert,
|
||||
mock_discord_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test that open circuit skips delivery."""
|
||||
dispatcher = AlertDispatcher(
|
||||
channels=[mock_discord_channel],
|
||||
failure_threshold=3,
|
||||
recovery_timeout_seconds=3600, # Long timeout
|
||||
)
|
||||
|
||||
# Manually open the circuit
|
||||
dispatcher._circuit_state["discord"].is_open = True
|
||||
dispatcher._circuit_state["discord"].last_failure_time = datetime.now(UTC)
|
||||
|
||||
result = await dispatcher.dispatch(sample_alert)
|
||||
|
||||
assert result.channel_results["discord"] is False
|
||||
# send() should not be called
|
||||
mock_discord_channel.send.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_circuit_closes_on_success(
|
||||
self,
|
||||
sample_alert: FormattedAlert,
|
||||
mock_discord_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test circuit closes on successful delivery."""
|
||||
mock_discord_channel.send.return_value = False
|
||||
|
||||
dispatcher = AlertDispatcher(
|
||||
channels=[mock_discord_channel],
|
||||
failure_threshold=2,
|
||||
)
|
||||
|
||||
# Cause failures to open circuit
|
||||
await dispatcher.dispatch(sample_alert)
|
||||
await dispatcher.dispatch(sample_alert)
|
||||
assert dispatcher._circuit_state["discord"].is_open is True
|
||||
|
||||
# Now succeed
|
||||
mock_discord_channel.send.return_value = True
|
||||
# Force half-open by resetting last_failure to past
|
||||
dispatcher._circuit_state["discord"].last_failure_time = datetime(2020, 1, 1, tzinfo=UTC)
|
||||
|
||||
result = await dispatcher.dispatch(sample_alert)
|
||||
|
||||
assert result.channel_results["discord"] is True
|
||||
assert dispatcher._circuit_state["discord"].is_open is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_batch(
|
||||
self,
|
||||
sample_alert: FormattedAlert,
|
||||
mock_discord_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test batch dispatch."""
|
||||
dispatcher = AlertDispatcher(channels=[mock_discord_channel])
|
||||
|
||||
alerts = [sample_alert, sample_alert, sample_alert]
|
||||
results = await dispatcher.dispatch_batch(alerts)
|
||||
|
||||
assert len(results) == 3
|
||||
assert all(r.success_count == 1 for r in results)
|
||||
|
||||
def test_get_circuit_status(
|
||||
self,
|
||||
mock_discord_channel: MagicMock,
|
||||
mock_telegram_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test getting circuit status."""
|
||||
dispatcher = AlertDispatcher(channels=[mock_discord_channel, mock_telegram_channel])
|
||||
|
||||
status = dispatcher.get_circuit_status()
|
||||
|
||||
assert "discord" in status
|
||||
assert "telegram" in status
|
||||
assert status["discord"]["is_open"] is False
|
||||
assert status["discord"]["failure_count"] == 0
|
||||
|
||||
def test_reset_circuit(
|
||||
self,
|
||||
mock_discord_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test manual circuit reset."""
|
||||
dispatcher = AlertDispatcher(channels=[mock_discord_channel])
|
||||
|
||||
# Set up failure state
|
||||
dispatcher._circuit_state["discord"].failure_count = 5
|
||||
dispatcher._circuit_state["discord"].is_open = True
|
||||
|
||||
# Reset
|
||||
result = dispatcher.reset_circuit("discord")
|
||||
|
||||
assert result is True
|
||||
assert dispatcher._circuit_state["discord"].failure_count == 0
|
||||
assert dispatcher._circuit_state["discord"].is_open is False
|
||||
|
||||
def test_reset_circuit_unknown_channel(
|
||||
self,
|
||||
mock_discord_channel: MagicMock,
|
||||
) -> None:
|
||||
"""Test reset with unknown channel name."""
|
||||
dispatcher = AlertDispatcher(channels=[mock_discord_channel])
|
||||
|
||||
result = dispatcher.reset_circuit("unknown")
|
||||
|
||||
assert result is False
|
||||
@@ -0,0 +1,587 @@
|
||||
"""Tests for alert message formatter."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.alerter.formatter import (
|
||||
COLOR_HIGH_RISK,
|
||||
COLOR_LOW_RISK,
|
||||
COLOR_MEDIUM_RISK,
|
||||
AlertFormatter,
|
||||
format_usdc,
|
||||
get_risk_color,
|
||||
get_risk_level,
|
||||
get_triggered_signals,
|
||||
truncate_address,
|
||||
)
|
||||
from polymarket_insider_tracker.alerter.models import FormattedAlert
|
||||
from polymarket_insider_tracker.detector.models import (
|
||||
FreshWalletSignal,
|
||||
RiskAssessment,
|
||||
SizeAnomalySignal,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.models import MarketMetadata, Token, TradeEvent
|
||||
from polymarket_insider_tracker.profiler.models import WalletProfile
|
||||
|
||||
# ============================================================================
|
||||
# Fixtures
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_trade() -> TradeEvent:
|
||||
"""Create a sample trade event."""
|
||||
return TradeEvent(
|
||||
market_id="market_abc123",
|
||||
trade_id="tx_001",
|
||||
wallet_address="0x1234567890abcdef1234567890abcdef12345678",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.075"),
|
||||
size=Decimal("200000"),
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_123",
|
||||
market_slug="will-x-happen",
|
||||
event_title="Will X happen by Y?",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_wallet_profile() -> WalletProfile:
|
||||
"""Create a sample wallet profile."""
|
||||
return WalletProfile(
|
||||
address="0x1234567890abcdef1234567890abcdef12345678",
|
||||
nonce=2,
|
||||
first_seen=datetime.now(UTC),
|
||||
age_hours=2.0,
|
||||
is_fresh=True,
|
||||
total_tx_count=2,
|
||||
matic_balance=Decimal("1000000000000000000"),
|
||||
usdc_balance=Decimal("1000000"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_metadata() -> MarketMetadata:
|
||||
"""Create sample market metadata."""
|
||||
return MarketMetadata(
|
||||
condition_id="market_abc123",
|
||||
question="Will X happen by Y?",
|
||||
description="Test market description",
|
||||
tokens=(Token(token_id="token_123", outcome="Yes", price=Decimal("0.075")),),
|
||||
category="other",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_wallet_signal(
|
||||
sample_trade: TradeEvent, sample_wallet_profile: WalletProfile
|
||||
) -> FreshWalletSignal:
|
||||
"""Create a sample fresh wallet signal."""
|
||||
return FreshWalletSignal(
|
||||
trade_event=sample_trade,
|
||||
wallet_profile=sample_wallet_profile,
|
||||
confidence=0.8,
|
||||
factors={"base": 0.5, "brand_new_bonus": 0.2, "large_trade_bonus": 0.1},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def size_anomaly_signal(
|
||||
sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> SizeAnomalySignal:
|
||||
"""Create a sample size anomaly signal."""
|
||||
return SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.10,
|
||||
book_impact=0.15,
|
||||
is_niche_market=True,
|
||||
confidence=0.7,
|
||||
factors={"volume_impact": 0.4, "book_impact": 0.3},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def high_risk_assessment(
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
size_anomaly_signal: SizeAnomalySignal,
|
||||
) -> RiskAssessment:
|
||||
"""Create a high-risk assessment with multiple signals."""
|
||||
return RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
size_anomaly_signal=size_anomaly_signal,
|
||||
signals_triggered=2,
|
||||
weighted_score=0.82,
|
||||
should_alert=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def medium_risk_assessment(
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
) -> RiskAssessment:
|
||||
"""Create a medium-risk assessment with one signal."""
|
||||
return RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=1,
|
||||
weighted_score=0.55,
|
||||
should_alert=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def low_risk_assessment(sample_trade: TradeEvent) -> RiskAssessment:
|
||||
"""Create a low-risk assessment with no signals."""
|
||||
return RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=None,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=0,
|
||||
weighted_score=0.25,
|
||||
should_alert=False,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helper Function Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestTruncateAddress:
|
||||
"""Tests for truncate_address helper."""
|
||||
|
||||
def test_truncate_standard_address(self) -> None:
|
||||
"""Test truncating a standard Ethereum address."""
|
||||
address = "0x1234567890abcdef1234567890abcdef12345678"
|
||||
result = truncate_address(address)
|
||||
assert result == "0x1234...5678"
|
||||
|
||||
def test_truncate_with_custom_length(self) -> None:
|
||||
"""Test truncating with custom character count."""
|
||||
address = "0x1234567890abcdef1234567890abcdef12345678"
|
||||
result = truncate_address(address, chars=6)
|
||||
assert result == "0x123456...345678"
|
||||
|
||||
def test_short_address_not_truncated(self) -> None:
|
||||
"""Test that short addresses are not truncated."""
|
||||
address = "0x1234"
|
||||
result = truncate_address(address)
|
||||
assert result == "0x1234"
|
||||
|
||||
|
||||
class TestFormatUsdc:
|
||||
"""Tests for format_usdc helper."""
|
||||
|
||||
def test_format_whole_dollars(self) -> None:
|
||||
"""Test formatting whole dollar amounts."""
|
||||
result = format_usdc(Decimal("15000"))
|
||||
assert result == "$15,000.00"
|
||||
|
||||
def test_format_with_cents(self) -> None:
|
||||
"""Test formatting with decimal places."""
|
||||
result = format_usdc(Decimal("1234.56"))
|
||||
assert result == "$1,234.56"
|
||||
|
||||
def test_format_large_amount(self) -> None:
|
||||
"""Test formatting large amounts."""
|
||||
result = format_usdc(Decimal("1000000"))
|
||||
assert result == "$1,000,000.00"
|
||||
|
||||
|
||||
class TestGetRiskLevel:
|
||||
"""Tests for get_risk_level helper."""
|
||||
|
||||
def test_high_risk(self) -> None:
|
||||
"""Test high risk threshold."""
|
||||
assert get_risk_level(0.85) == "HIGH"
|
||||
assert get_risk_level(0.70) == "HIGH"
|
||||
|
||||
def test_medium_risk(self) -> None:
|
||||
"""Test medium risk threshold."""
|
||||
assert get_risk_level(0.65) == "MEDIUM"
|
||||
assert get_risk_level(0.50) == "MEDIUM"
|
||||
|
||||
def test_low_risk(self) -> None:
|
||||
"""Test low risk threshold."""
|
||||
assert get_risk_level(0.40) == "LOW"
|
||||
assert get_risk_level(0.10) == "LOW"
|
||||
|
||||
|
||||
class TestGetRiskColor:
|
||||
"""Tests for get_risk_color helper."""
|
||||
|
||||
def test_high_risk_color(self) -> None:
|
||||
"""Test high risk returns red color."""
|
||||
assert get_risk_color(0.85) == COLOR_HIGH_RISK
|
||||
assert get_risk_color(0.70) == COLOR_HIGH_RISK
|
||||
|
||||
def test_medium_risk_color(self) -> None:
|
||||
"""Test medium risk returns orange color."""
|
||||
assert get_risk_color(0.65) == COLOR_MEDIUM_RISK
|
||||
assert get_risk_color(0.50) == COLOR_MEDIUM_RISK
|
||||
|
||||
def test_low_risk_color(self) -> None:
|
||||
"""Test low risk returns yellow color."""
|
||||
assert get_risk_color(0.40) == COLOR_LOW_RISK
|
||||
|
||||
|
||||
class TestGetTriggeredSignals:
|
||||
"""Tests for get_triggered_signals helper."""
|
||||
|
||||
def test_no_signals(self, low_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test assessment with no signals."""
|
||||
signals = get_triggered_signals(low_risk_assessment)
|
||||
assert signals == []
|
||||
|
||||
def test_fresh_wallet_only(self, medium_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test assessment with only fresh wallet signal."""
|
||||
signals = get_triggered_signals(medium_risk_assessment)
|
||||
assert "Fresh Wallet" in signals
|
||||
assert "Large Position" not in signals
|
||||
|
||||
def test_both_signals(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test assessment with both signals."""
|
||||
signals = get_triggered_signals(high_risk_assessment)
|
||||
assert "Fresh Wallet" in signals
|
||||
assert "Large Position" in signals
|
||||
assert "Niche Market" in signals # From size anomaly with is_niche_market=True
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# AlertFormatter Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestAlertFormatterInit:
|
||||
"""Tests for AlertFormatter initialization."""
|
||||
|
||||
def test_default_verbosity(self) -> None:
|
||||
"""Test default verbosity is detailed."""
|
||||
formatter = AlertFormatter()
|
||||
assert formatter.verbosity == "detailed"
|
||||
|
||||
def test_compact_verbosity(self) -> None:
|
||||
"""Test setting compact verbosity."""
|
||||
formatter = AlertFormatter(verbosity="compact")
|
||||
assert formatter.verbosity == "compact"
|
||||
|
||||
|
||||
class TestAlertFormatterFormat:
|
||||
"""Tests for AlertFormatter.format method."""
|
||||
|
||||
def test_format_returns_formatted_alert(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that format returns a FormattedAlert."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert isinstance(result, FormattedAlert)
|
||||
|
||||
def test_format_includes_all_fields(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that all fields are populated."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
|
||||
assert result.title != ""
|
||||
assert result.body != ""
|
||||
assert result.discord_embed != {}
|
||||
assert result.telegram_markdown != ""
|
||||
assert result.plain_text != ""
|
||||
assert result.links != {}
|
||||
|
||||
def test_format_title_includes_risk_level(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that title includes risk level."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "HIGH" in result.title
|
||||
|
||||
def test_format_includes_wallet_link(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that wallet explorer link is included."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "wallet" in result.links
|
||||
assert "polygonscan.com" in result.links["wallet"]
|
||||
|
||||
def test_format_includes_market_link(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that market link is included when slug available."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "market" in result.links
|
||||
assert "polymarket.com" in result.links["market"]
|
||||
|
||||
|
||||
class TestDiscordEmbed:
|
||||
"""Tests for Discord embed format."""
|
||||
|
||||
def test_embed_has_required_fields(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that embed has required Discord fields."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
embed = result.discord_embed
|
||||
|
||||
assert "title" in embed
|
||||
assert "color" in embed
|
||||
assert "fields" in embed
|
||||
assert "footer" in embed
|
||||
|
||||
def test_embed_color_reflects_risk(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that embed color matches risk level."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert result.discord_embed["color"] == COLOR_HIGH_RISK
|
||||
|
||||
def test_embed_includes_wallet_field(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that embed includes wallet field."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
fields = result.discord_embed["fields"]
|
||||
|
||||
wallet_field = next((f for f in fields if f["name"] == "Wallet"), None)
|
||||
assert wallet_field is not None
|
||||
assert "0x1234" in wallet_field["value"]
|
||||
|
||||
def test_embed_includes_wallet_age(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that wallet age is shown when fresh wallet signal present."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
fields = result.discord_embed["fields"]
|
||||
|
||||
wallet_field = next((f for f in fields if f["name"] == "Wallet"), None)
|
||||
assert "Age:" in wallet_field["value"]
|
||||
|
||||
def test_embed_includes_trade_details(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that trade details are in embed."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
fields = result.discord_embed["fields"]
|
||||
|
||||
trade_field = next((f for f in fields if f["name"] == "Trade"), None)
|
||||
assert trade_field is not None
|
||||
assert "BUY" in trade_field["value"]
|
||||
assert "Yes" in trade_field["value"]
|
||||
|
||||
def test_embed_includes_signals_field(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that signals are listed in embed."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
fields = result.discord_embed["fields"]
|
||||
|
||||
signals_field = next((f for f in fields if f["name"] == "Signals"), None)
|
||||
assert signals_field is not None
|
||||
assert "Fresh Wallet" in signals_field["value"]
|
||||
|
||||
def test_detailed_embed_includes_confidence(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that detailed mode includes confidence breakdown."""
|
||||
formatter = AlertFormatter(verbosity="detailed")
|
||||
result = formatter.format(high_risk_assessment)
|
||||
fields = result.discord_embed["fields"]
|
||||
|
||||
conf_field = next((f for f in fields if f["name"] == "Confidence"), None)
|
||||
assert conf_field is not None
|
||||
|
||||
|
||||
class TestTelegramMarkdown:
|
||||
"""Tests for Telegram markdown format."""
|
||||
|
||||
def test_telegram_includes_header(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that Telegram message has header."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "*Suspicious Activity Detected*" in result.telegram_markdown
|
||||
|
||||
def test_telegram_includes_wallet(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that Telegram message includes wallet."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "`0x1234...5678`" in result.telegram_markdown
|
||||
|
||||
def test_telegram_includes_risk_score(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that Telegram message includes risk score."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "0.82" in result.telegram_markdown
|
||||
assert "HIGH" in result.telegram_markdown
|
||||
|
||||
def test_telegram_includes_links(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that Telegram message includes links."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "[View Wallet]" in result.telegram_markdown
|
||||
assert "[View Market]" in result.telegram_markdown
|
||||
|
||||
|
||||
class TestPlainText:
|
||||
"""Tests for plain text format."""
|
||||
|
||||
def test_plain_text_header(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test plain text has header."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "SUSPICIOUS ACTIVITY DETECTED" in result.plain_text
|
||||
|
||||
def test_plain_text_wallet(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test plain text includes wallet."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "Wallet:" in result.plain_text
|
||||
assert "0x1234...5678" in result.plain_text
|
||||
|
||||
def test_plain_text_trade(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test plain text includes trade details."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "Trade:" in result.plain_text
|
||||
assert "BUY" in result.plain_text
|
||||
|
||||
def test_plain_text_signals(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test plain text includes signals."""
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(high_risk_assessment)
|
||||
assert "Signals:" in result.plain_text
|
||||
assert "Fresh Wallet" in result.plain_text
|
||||
|
||||
|
||||
class TestCompactVerbosity:
|
||||
"""Tests for compact verbosity mode."""
|
||||
|
||||
def test_compact_body_is_shorter(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test that compact mode produces shorter body."""
|
||||
detailed_formatter = AlertFormatter(verbosity="detailed")
|
||||
compact_formatter = AlertFormatter(verbosity="compact")
|
||||
|
||||
detailed_result = detailed_formatter.format(high_risk_assessment)
|
||||
compact_result = compact_formatter.format(high_risk_assessment)
|
||||
|
||||
assert len(compact_result.body) < len(detailed_result.body)
|
||||
|
||||
def test_compact_body_includes_essential_info(
|
||||
self, high_risk_assessment: RiskAssessment
|
||||
) -> None:
|
||||
"""Test that compact mode still has essential info."""
|
||||
formatter = AlertFormatter(verbosity="compact")
|
||||
result = formatter.format(high_risk_assessment)
|
||||
|
||||
assert "0x1234...5678" in result.body
|
||||
assert "0.82" in result.body
|
||||
assert "HIGH" in result.body
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Tests for edge cases and special scenarios."""
|
||||
|
||||
def test_no_market_slug(self, sample_trade: TradeEvent) -> None:
|
||||
"""Test formatting when market slug is empty."""
|
||||
trade = TradeEvent(
|
||||
market_id=sample_trade.market_id,
|
||||
trade_id=sample_trade.trade_id,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
side=sample_trade.side,
|
||||
outcome=sample_trade.outcome,
|
||||
outcome_index=sample_trade.outcome_index,
|
||||
price=sample_trade.price,
|
||||
size=sample_trade.size,
|
||||
timestamp=sample_trade.timestamp,
|
||||
asset_id=sample_trade.asset_id,
|
||||
market_slug="", # Empty slug
|
||||
event_title="", # Empty title
|
||||
)
|
||||
assessment = RiskAssessment(
|
||||
trade_event=trade,
|
||||
wallet_address=trade.wallet_address,
|
||||
market_id=trade.market_id,
|
||||
fresh_wallet_signal=None,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=0,
|
||||
weighted_score=0.5,
|
||||
should_alert=False,
|
||||
)
|
||||
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(assessment)
|
||||
|
||||
# Should not have market link
|
||||
assert "market" not in result.links
|
||||
# Should have fallback text
|
||||
assert "Unknown Market" in result.plain_text
|
||||
|
||||
def test_very_short_wallet_age(
|
||||
self,
|
||||
sample_trade: TradeEvent,
|
||||
sample_wallet_profile: WalletProfile,
|
||||
) -> None:
|
||||
"""Test formatting with wallet age less than 1 hour."""
|
||||
profile = WalletProfile(
|
||||
address=sample_wallet_profile.address,
|
||||
nonce=sample_wallet_profile.nonce,
|
||||
first_seen=datetime.now(UTC),
|
||||
age_hours=0.5, # 30 minutes
|
||||
is_fresh=True,
|
||||
total_tx_count=1,
|
||||
matic_balance=sample_wallet_profile.matic_balance,
|
||||
usdc_balance=sample_wallet_profile.usdc_balance,
|
||||
)
|
||||
signal = FreshWalletSignal(
|
||||
trade_event=sample_trade,
|
||||
wallet_profile=profile,
|
||||
confidence=0.9,
|
||||
factors={},
|
||||
)
|
||||
assessment = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=signal,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=1,
|
||||
weighted_score=0.75,
|
||||
should_alert=True,
|
||||
)
|
||||
|
||||
formatter = AlertFormatter()
|
||||
result = formatter.format(assessment)
|
||||
|
||||
# Should show age in minutes
|
||||
assert "30m" in result.plain_text or "Age: 30m" in result.plain_text
|
||||
|
||||
def test_size_anomaly_without_niche(
|
||||
self,
|
||||
sample_trade: TradeEvent,
|
||||
sample_metadata: MarketMetadata,
|
||||
) -> None:
|
||||
"""Test size anomaly signal without niche market flag."""
|
||||
signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.10,
|
||||
book_impact=0.15,
|
||||
is_niche_market=False, # Not niche
|
||||
confidence=0.7,
|
||||
factors={},
|
||||
)
|
||||
assessment = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=None,
|
||||
size_anomaly_signal=signal,
|
||||
signals_triggered=1,
|
||||
weighted_score=0.6,
|
||||
should_alert=True,
|
||||
)
|
||||
|
||||
signals = get_triggered_signals(assessment)
|
||||
assert "Large Position" in signals
|
||||
assert "Niche Market" not in signals
|
||||
@@ -0,0 +1,534 @@
|
||||
"""Tests for alert history and deduplication."""
|
||||
|
||||
import json
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.alerter.history import (
|
||||
AlertHistory,
|
||||
AlertRecord,
|
||||
_generate_dedup_key,
|
||||
_get_signals_from_assessment,
|
||||
)
|
||||
from polymarket_insider_tracker.detector.models import (
|
||||
FreshWalletSignal,
|
||||
RiskAssessment,
|
||||
SizeAnomalySignal,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.models import MarketMetadata, Token, TradeEvent
|
||||
from polymarket_insider_tracker.profiler.models import WalletProfile
|
||||
|
||||
# ============================================================================
|
||||
# Fixtures
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_trade() -> TradeEvent:
|
||||
"""Create a sample trade event."""
|
||||
return TradeEvent(
|
||||
market_id="market_abc123",
|
||||
trade_id="tx_001",
|
||||
wallet_address="0x1234567890abcdef1234567890abcdef12345678",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.65"),
|
||||
size=Decimal("10000"),
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_123",
|
||||
event_title="Test Market",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_wallet_profile() -> WalletProfile:
|
||||
"""Create a sample wallet profile."""
|
||||
return WalletProfile(
|
||||
address="0x1234567890abcdef1234567890abcdef12345678",
|
||||
nonce=2,
|
||||
first_seen=datetime.now(UTC),
|
||||
age_hours=1.0,
|
||||
is_fresh=True,
|
||||
total_tx_count=2,
|
||||
matic_balance=Decimal("1000000000000000000"),
|
||||
usdc_balance=Decimal("1000000"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_metadata() -> MarketMetadata:
|
||||
"""Create sample market metadata."""
|
||||
return MarketMetadata(
|
||||
condition_id="market_abc123",
|
||||
question="Test market?",
|
||||
description="Test",
|
||||
tokens=(Token(token_id="token_123", outcome="Yes", price=Decimal("0.65")),),
|
||||
category="other",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_wallet_signal(
|
||||
sample_trade: TradeEvent, sample_wallet_profile: WalletProfile
|
||||
) -> FreshWalletSignal:
|
||||
"""Create a sample fresh wallet signal."""
|
||||
return FreshWalletSignal(
|
||||
trade_event=sample_trade,
|
||||
wallet_profile=sample_wallet_profile,
|
||||
confidence=0.8,
|
||||
factors={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def size_anomaly_signal(
|
||||
sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> SizeAnomalySignal:
|
||||
"""Create a sample size anomaly signal."""
|
||||
return SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.10,
|
||||
book_impact=0.15,
|
||||
is_niche_market=True,
|
||||
confidence=0.7,
|
||||
factors={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def high_risk_assessment(
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
size_anomaly_signal: SizeAnomalySignal,
|
||||
) -> RiskAssessment:
|
||||
"""Create a high-risk assessment."""
|
||||
return RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
size_anomaly_signal=size_anomaly_signal,
|
||||
signals_triggered=2,
|
||||
weighted_score=0.82,
|
||||
should_alert=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis() -> MagicMock:
|
||||
"""Create a mock Redis client."""
|
||||
redis = MagicMock()
|
||||
|
||||
# Make async methods return AsyncMock
|
||||
redis.exists = AsyncMock(return_value=0) # Key doesn't exist (not duplicate)
|
||||
redis.get = AsyncMock(return_value=None)
|
||||
redis.set = AsyncMock(return_value=True)
|
||||
redis.ttl = AsyncMock(return_value=3600)
|
||||
redis.zadd = AsyncMock(return_value=1)
|
||||
redis.expire = AsyncMock(return_value=True)
|
||||
redis.zrangebyscore = AsyncMock(return_value=[])
|
||||
redis.zcount = AsyncMock(return_value=0)
|
||||
redis.zremrangebyscore = AsyncMock(return_value=0)
|
||||
|
||||
# Mock pipeline - async context manager
|
||||
pipeline = MagicMock()
|
||||
pipeline.__aenter__ = AsyncMock(return_value=pipeline)
|
||||
pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
pipeline.set.return_value = pipeline
|
||||
pipeline.zadd.return_value = pipeline
|
||||
pipeline.expire.return_value = pipeline
|
||||
pipeline.execute = AsyncMock(return_value=[True, True, True, True, True, True])
|
||||
redis.pipeline.return_value = pipeline
|
||||
|
||||
return redis
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# AlertRecord Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestAlertRecord:
|
||||
"""Tests for AlertRecord dataclass."""
|
||||
|
||||
def test_to_dict(self) -> None:
|
||||
"""Test serialization to dict."""
|
||||
now = datetime.now(UTC)
|
||||
record = AlertRecord(
|
||||
alert_id="test-123",
|
||||
wallet_address="0x1234",
|
||||
market_id="market_abc",
|
||||
risk_score=0.75,
|
||||
signals_triggered=["fresh_wallet"],
|
||||
channels_attempted=["discord", "telegram"],
|
||||
channels_succeeded=["discord"],
|
||||
dedup_key="0x1234:market_abc:2026010416",
|
||||
feedback_useful=True,
|
||||
created_at=now,
|
||||
)
|
||||
|
||||
data = record.to_dict()
|
||||
|
||||
assert data["alert_id"] == "test-123"
|
||||
assert data["risk_score"] == 0.75
|
||||
assert data["feedback_useful"] is True
|
||||
assert data["created_at"] == now.isoformat()
|
||||
|
||||
def test_from_dict(self) -> None:
|
||||
"""Test deserialization from dict."""
|
||||
data = {
|
||||
"alert_id": "test-456",
|
||||
"wallet_address": "0x5678",
|
||||
"market_id": "market_xyz",
|
||||
"risk_score": 0.82,
|
||||
"signals_triggered": ["size_anomaly"],
|
||||
"channels_attempted": ["discord"],
|
||||
"channels_succeeded": ["discord"],
|
||||
"dedup_key": "0x5678:market_xyz:2026010416",
|
||||
"feedback_useful": None,
|
||||
"created_at": "2026-01-04T16:00:00+00:00",
|
||||
}
|
||||
|
||||
record = AlertRecord.from_dict(data)
|
||||
|
||||
assert record.alert_id == "test-456"
|
||||
assert record.risk_score == 0.82
|
||||
assert record.feedback_useful is None
|
||||
|
||||
def test_from_dict_missing_optional(self) -> None:
|
||||
"""Test deserialization with missing optional fields."""
|
||||
data = {
|
||||
"alert_id": "test-789",
|
||||
"wallet_address": "0x9999",
|
||||
"market_id": "market_aaa",
|
||||
"risk_score": "0.5", # Test string conversion
|
||||
"dedup_key": "key",
|
||||
}
|
||||
|
||||
record = AlertRecord.from_dict(data)
|
||||
|
||||
assert record.signals_triggered == []
|
||||
assert record.channels_attempted == []
|
||||
assert record.feedback_useful is None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helper Function Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestGenerateDedupKey:
|
||||
"""Tests for dedup key generation."""
|
||||
|
||||
def test_basic_key(self) -> None:
|
||||
"""Test basic dedup key generation."""
|
||||
hour = datetime(2026, 1, 4, 16, 30, 0, tzinfo=UTC)
|
||||
key = _generate_dedup_key("0x1234", "market_abc", hour)
|
||||
assert key == "0x1234:market_abc:2026010416"
|
||||
|
||||
def test_different_hours(self) -> None:
|
||||
"""Test that different hours produce different keys."""
|
||||
hour1 = datetime(2026, 1, 4, 16, 0, 0, tzinfo=UTC)
|
||||
hour2 = datetime(2026, 1, 4, 17, 0, 0, tzinfo=UTC)
|
||||
|
||||
key1 = _generate_dedup_key("0x1234", "market_abc", hour1)
|
||||
key2 = _generate_dedup_key("0x1234", "market_abc", hour2)
|
||||
|
||||
assert key1 != key2
|
||||
|
||||
|
||||
class TestGetSignalsFromAssessment:
|
||||
"""Tests for signal extraction."""
|
||||
|
||||
def test_no_signals(self, sample_trade: TradeEvent) -> None:
|
||||
"""Test extraction with no signals."""
|
||||
assessment = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=None,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=0,
|
||||
weighted_score=0.0,
|
||||
should_alert=False,
|
||||
)
|
||||
|
||||
signals = _get_signals_from_assessment(assessment)
|
||||
assert signals == []
|
||||
|
||||
def test_fresh_wallet_only(
|
||||
self,
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
) -> None:
|
||||
"""Test extraction with fresh wallet signal."""
|
||||
assessment = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=1,
|
||||
weighted_score=0.5,
|
||||
should_alert=True,
|
||||
)
|
||||
|
||||
signals = _get_signals_from_assessment(assessment)
|
||||
assert "fresh_wallet" in signals
|
||||
assert "size_anomaly" not in signals
|
||||
|
||||
def test_all_signals(self, high_risk_assessment: RiskAssessment) -> None:
|
||||
"""Test extraction with all signals."""
|
||||
signals = _get_signals_from_assessment(high_risk_assessment)
|
||||
assert "fresh_wallet" in signals
|
||||
assert "size_anomaly" in signals
|
||||
assert "niche_market" in signals
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# AlertHistory Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestAlertHistoryInit:
|
||||
"""Tests for AlertHistory initialization."""
|
||||
|
||||
def test_default_settings(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test default configuration."""
|
||||
history = AlertHistory(mock_redis)
|
||||
|
||||
assert history.dedup_window_hours == 1
|
||||
assert history.retention_days == 30
|
||||
assert history._dedup_ttl == 3600
|
||||
assert history._retention_ttl == 30 * 86400
|
||||
|
||||
def test_custom_settings(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test custom configuration."""
|
||||
history = AlertHistory(
|
||||
mock_redis,
|
||||
dedup_window_hours=2,
|
||||
retention_days=7,
|
||||
)
|
||||
|
||||
assert history.dedup_window_hours == 2
|
||||
assert history._dedup_ttl == 7200
|
||||
|
||||
|
||||
class TestShouldSend:
|
||||
"""Tests for should_send method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_not_duplicate(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
high_risk_assessment: RiskAssessment,
|
||||
) -> None:
|
||||
"""Test that non-duplicate returns True."""
|
||||
mock_redis.exists.return_value = 0
|
||||
history = AlertHistory(mock_redis)
|
||||
|
||||
result = await history.should_send(high_risk_assessment)
|
||||
|
||||
assert result is True
|
||||
mock_redis.exists.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_duplicate(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
high_risk_assessment: RiskAssessment,
|
||||
) -> None:
|
||||
"""Test that duplicate returns False."""
|
||||
mock_redis.exists.return_value = 1
|
||||
history = AlertHistory(mock_redis)
|
||||
|
||||
result = await history.should_send(high_risk_assessment)
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestRecordSent:
|
||||
"""Tests for record_sent method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_success(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
high_risk_assessment: RiskAssessment,
|
||||
) -> None:
|
||||
"""Test recording a sent alert."""
|
||||
history = AlertHistory(mock_redis)
|
||||
|
||||
alert_id = await history.record_sent(
|
||||
high_risk_assessment,
|
||||
channels_attempted=["discord", "telegram"],
|
||||
channels_succeeded={"discord": True, "telegram": False},
|
||||
)
|
||||
|
||||
assert alert_id is not None
|
||||
assert len(alert_id) == 36 # UUID length
|
||||
|
||||
# Verify pipeline was used
|
||||
mock_redis.pipeline.assert_called_once()
|
||||
|
||||
|
||||
class TestRecordFeedback:
|
||||
"""Tests for record_feedback method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_feedback_success(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test recording feedback for existing alert."""
|
||||
existing_record = {
|
||||
"alert_id": "test-123",
|
||||
"wallet_address": "0x1234",
|
||||
"market_id": "market_abc",
|
||||
"risk_score": 0.75,
|
||||
"signals_triggered": [],
|
||||
"channels_attempted": [],
|
||||
"channels_succeeded": [],
|
||||
"dedup_key": "key",
|
||||
"feedback_useful": None,
|
||||
}
|
||||
mock_redis.get.return_value = json.dumps(existing_record)
|
||||
mock_redis.ttl.return_value = 3600
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
result = await history.record_feedback("test-123", useful=True)
|
||||
|
||||
assert result is True
|
||||
mock_redis.set.assert_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_feedback_not_found(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test feedback for non-existent alert."""
|
||||
mock_redis.get.return_value = None
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
result = await history.record_feedback("nonexistent", useful=True)
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestGetAlert:
|
||||
"""Tests for get_alert method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_existing(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting existing alert."""
|
||||
existing_record = {
|
||||
"alert_id": "test-123",
|
||||
"wallet_address": "0x1234",
|
||||
"market_id": "market_abc",
|
||||
"risk_score": 0.75,
|
||||
"signals_triggered": ["fresh_wallet"],
|
||||
"channels_attempted": ["discord"],
|
||||
"channels_succeeded": ["discord"],
|
||||
"dedup_key": "key",
|
||||
"created_at": "2026-01-04T16:00:00+00:00",
|
||||
}
|
||||
mock_redis.get.return_value = json.dumps(existing_record)
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
record = await history.get_alert("test-123")
|
||||
|
||||
assert record is not None
|
||||
assert record.alert_id == "test-123"
|
||||
assert record.risk_score == 0.75
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nonexistent(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting non-existent alert."""
|
||||
mock_redis.get.return_value = None
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
record = await history.get_alert("nonexistent")
|
||||
|
||||
assert record is None
|
||||
|
||||
|
||||
class TestGetAlerts:
|
||||
"""Tests for get_alerts query method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_results(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test query with no results."""
|
||||
mock_redis.zrangebyscore.return_value = []
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
results = await history.get_alerts(
|
||||
start=datetime.now(UTC) - timedelta(hours=24),
|
||||
end=datetime.now(UTC),
|
||||
)
|
||||
|
||||
assert results == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_wallet_filter(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test query with wallet filter uses correct index."""
|
||||
mock_redis.zrangebyscore.return_value = []
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
await history.get_alerts(
|
||||
start=datetime.now(UTC) - timedelta(hours=24),
|
||||
end=datetime.now(UTC),
|
||||
wallet="0x1234",
|
||||
)
|
||||
|
||||
# Verify correct index was used
|
||||
call_args = mock_redis.zrangebyscore.call_args
|
||||
assert "wallet:0x1234" in call_args[0][0]
|
||||
|
||||
|
||||
class TestGetRecentCount:
|
||||
"""Tests for get_recent_count method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_all(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test counting all recent alerts."""
|
||||
mock_redis.zcount.return_value = 42
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
count = await history.get_recent_count(hours=24)
|
||||
|
||||
assert count == 42
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_by_wallet(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test counting alerts for specific wallet."""
|
||||
mock_redis.zcount.return_value = 5
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
count = await history.get_recent_count(hours=24, wallet="0x1234")
|
||||
|
||||
assert count == 5
|
||||
|
||||
|
||||
class TestCleanupOldAlerts:
|
||||
"""Tests for cleanup_old_alerts method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_empty(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test cleanup with no old alerts."""
|
||||
mock_redis.zrangebyscore.return_value = []
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
removed = await history.cleanup_old_alerts()
|
||||
|
||||
assert removed == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_removes_old(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test cleanup removes old alerts."""
|
||||
mock_redis.zrangebyscore.return_value = [b"alert-1", b"alert-2"]
|
||||
mock_redis.zremrangebyscore.return_value = 2
|
||||
|
||||
history = AlertHistory(mock_redis)
|
||||
removed = await history.cleanup_old_alerts()
|
||||
|
||||
assert removed == 2
|
||||
mock_redis.zremrangebyscore.assert_called_once()
|
||||
@@ -0,0 +1,9 @@
|
||||
"""Pytest configuration and fixtures."""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_market_id() -> str:
|
||||
"""Sample market ID for testing."""
|
||||
return "0x1234567890abcdef1234567890abcdef12345678"
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for detector module."""
|
||||
@@ -0,0 +1,608 @@
|
||||
"""Tests for the fresh wallet detector."""
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.detector.fresh_wallet import (
|
||||
BASE_CONFIDENCE,
|
||||
BRAND_NEW_BONUS,
|
||||
DEFAULT_MAX_AGE_HOURS,
|
||||
DEFAULT_MAX_NONCE,
|
||||
DEFAULT_MIN_TRADE_SIZE,
|
||||
LARGE_TRADE_BONUS,
|
||||
VERY_YOUNG_BONUS,
|
||||
FreshWalletDetector,
|
||||
)
|
||||
from polymarket_insider_tracker.detector.models import FreshWalletSignal
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
from polymarket_insider_tracker.profiler.models import WalletProfile
|
||||
|
||||
|
||||
# Test fixtures
|
||||
@pytest.fixture
|
||||
def mock_wallet_analyzer():
|
||||
"""Create a mock WalletAnalyzer."""
|
||||
return AsyncMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def detector(mock_wallet_analyzer):
|
||||
"""Create a FreshWalletDetector with mocked analyzer."""
|
||||
return FreshWalletDetector(mock_wallet_analyzer)
|
||||
|
||||
|
||||
def create_trade_event(
|
||||
*,
|
||||
wallet_address: str = "0x1234567890123456789012345678901234567890",
|
||||
market_id: str = "market123",
|
||||
trade_id: str = "trade123",
|
||||
price: Decimal = Decimal("0.5"),
|
||||
size: Decimal = Decimal("2000"),
|
||||
side: str = "BUY",
|
||||
) -> TradeEvent:
|
||||
"""Create a TradeEvent for testing."""
|
||||
return TradeEvent(
|
||||
market_id=market_id,
|
||||
trade_id=trade_id,
|
||||
wallet_address=wallet_address,
|
||||
side=side,
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=price,
|
||||
size=size,
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="asset123",
|
||||
)
|
||||
|
||||
|
||||
def create_wallet_profile(
|
||||
*,
|
||||
address: str = "0x1234567890123456789012345678901234567890",
|
||||
nonce: int = 2,
|
||||
age_hours: float | None = 24.0,
|
||||
is_fresh: bool = True,
|
||||
) -> WalletProfile:
|
||||
"""Create a WalletProfile for testing."""
|
||||
first_seen = None
|
||||
if age_hours is not None:
|
||||
first_seen = datetime.now(UTC) - timedelta(hours=age_hours)
|
||||
|
||||
return WalletProfile(
|
||||
address=address.lower(),
|
||||
nonce=nonce,
|
||||
first_seen=first_seen,
|
||||
age_hours=age_hours,
|
||||
is_fresh=is_fresh,
|
||||
total_tx_count=nonce,
|
||||
matic_balance=Decimal("1000000000000000000"), # 1 MATIC
|
||||
usdc_balance=Decimal("1000000000"), # 1000 USDC
|
||||
fresh_threshold=5,
|
||||
)
|
||||
|
||||
|
||||
# === Test FreshWalletSignal Model ===
|
||||
|
||||
|
||||
class TestFreshWalletSignal:
|
||||
"""Tests for FreshWalletSignal dataclass."""
|
||||
|
||||
def test_wallet_address_property(self):
|
||||
"""Test wallet_address property returns trade wallet."""
|
||||
trade = create_trade_event(wallet_address="0xabc123")
|
||||
profile = create_wallet_profile(address="0xabc123")
|
||||
signal = FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=profile,
|
||||
confidence=0.7,
|
||||
factors={"base": 0.5},
|
||||
)
|
||||
assert signal.wallet_address == "0xabc123"
|
||||
|
||||
def test_market_id_property(self):
|
||||
"""Test market_id property returns trade market."""
|
||||
trade = create_trade_event(market_id="market456")
|
||||
profile = create_wallet_profile()
|
||||
signal = FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=profile,
|
||||
confidence=0.7,
|
||||
factors={"base": 0.5},
|
||||
)
|
||||
assert signal.market_id == "market456"
|
||||
|
||||
def test_trade_size_usdc_property(self):
|
||||
"""Test trade_size_usdc returns notional value."""
|
||||
trade = create_trade_event(price=Decimal("0.5"), size=Decimal("2000"))
|
||||
profile = create_wallet_profile()
|
||||
signal = FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=profile,
|
||||
confidence=0.7,
|
||||
factors={"base": 0.5},
|
||||
)
|
||||
assert signal.trade_size_usdc == Decimal("1000")
|
||||
|
||||
def test_is_high_confidence(self):
|
||||
"""Test is_high_confidence threshold."""
|
||||
trade = create_trade_event()
|
||||
profile = create_wallet_profile()
|
||||
|
||||
# At threshold
|
||||
signal = FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=profile,
|
||||
confidence=0.7,
|
||||
factors={},
|
||||
)
|
||||
assert signal.is_high_confidence is True
|
||||
|
||||
# Below threshold
|
||||
signal_low = FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=profile,
|
||||
confidence=0.69,
|
||||
factors={},
|
||||
)
|
||||
assert signal_low.is_high_confidence is False
|
||||
|
||||
def test_is_very_high_confidence(self):
|
||||
"""Test is_very_high_confidence threshold."""
|
||||
trade = create_trade_event()
|
||||
profile = create_wallet_profile()
|
||||
|
||||
# At threshold
|
||||
signal = FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=profile,
|
||||
confidence=0.85,
|
||||
factors={},
|
||||
)
|
||||
assert signal.is_very_high_confidence is True
|
||||
|
||||
# Below threshold
|
||||
signal_low = FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=profile,
|
||||
confidence=0.84,
|
||||
factors={},
|
||||
)
|
||||
assert signal_low.is_very_high_confidence is False
|
||||
|
||||
def test_to_dict(self):
|
||||
"""Test serialization to dictionary."""
|
||||
trade = create_trade_event(
|
||||
wallet_address="0xtest",
|
||||
market_id="market1",
|
||||
trade_id="tx1",
|
||||
price=Decimal("0.6"),
|
||||
size=Decimal("5000"),
|
||||
side="BUY",
|
||||
)
|
||||
profile = create_wallet_profile(address="0xtest", nonce=0, age_hours=1.0)
|
||||
signal = FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=profile,
|
||||
confidence=0.8,
|
||||
factors={"base": 0.5, "brand_new": 0.2},
|
||||
)
|
||||
|
||||
result = signal.to_dict()
|
||||
|
||||
assert result["wallet_address"] == "0xtest"
|
||||
assert result["market_id"] == "market1"
|
||||
assert result["trade_id"] == "tx1"
|
||||
assert result["trade_size"] == "3000.0"
|
||||
assert result["trade_side"] == "BUY"
|
||||
assert result["trade_price"] == "0.6"
|
||||
assert result["wallet_nonce"] == 0
|
||||
assert result["wallet_age_hours"] == 1.0
|
||||
assert result["wallet_is_fresh"] is True
|
||||
assert result["confidence"] == 0.8
|
||||
assert result["factors"] == {"base": 0.5, "brand_new": 0.2}
|
||||
assert "timestamp" in result
|
||||
|
||||
|
||||
# === Test FreshWalletDetector Initialization ===
|
||||
|
||||
|
||||
class TestFreshWalletDetectorInit:
|
||||
"""Tests for FreshWalletDetector initialization."""
|
||||
|
||||
def test_default_config(self, mock_wallet_analyzer):
|
||||
"""Test detector initializes with default config."""
|
||||
detector = FreshWalletDetector(mock_wallet_analyzer)
|
||||
assert detector._min_trade_size == DEFAULT_MIN_TRADE_SIZE
|
||||
assert detector._max_nonce == DEFAULT_MAX_NONCE
|
||||
assert detector._max_age_hours == DEFAULT_MAX_AGE_HOURS
|
||||
|
||||
def test_custom_config(self, mock_wallet_analyzer):
|
||||
"""Test detector accepts custom config."""
|
||||
detector = FreshWalletDetector(
|
||||
mock_wallet_analyzer,
|
||||
min_trade_size=Decimal("5000"),
|
||||
max_nonce=10,
|
||||
max_age_hours=72.0,
|
||||
)
|
||||
assert detector._min_trade_size == Decimal("5000")
|
||||
assert detector._max_nonce == 10
|
||||
assert detector._max_age_hours == 72.0
|
||||
|
||||
|
||||
# === Test FreshWalletDetector.analyze ===
|
||||
|
||||
|
||||
class TestFreshWalletDetectorAnalyze:
|
||||
"""Tests for FreshWalletDetector.analyze method."""
|
||||
|
||||
async def test_filters_small_trades(self, detector, mock_wallet_analyzer):
|
||||
"""Test that trades below minimum size are filtered out."""
|
||||
trade = create_trade_event(
|
||||
price=Decimal("0.1"),
|
||||
size=Decimal("100"), # notional = $10
|
||||
)
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
assert result is None
|
||||
mock_wallet_analyzer.analyze.assert_not_called()
|
||||
|
||||
async def test_detects_fresh_wallet(self, detector, mock_wallet_analyzer):
|
||||
"""Test detection of fresh wallet trade."""
|
||||
trade = create_trade_event(
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("2000"), # notional = $1000
|
||||
)
|
||||
profile = create_wallet_profile(nonce=2, age_hours=24.0, is_fresh=True)
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
assert result is not None
|
||||
assert isinstance(result, FreshWalletSignal)
|
||||
assert result.trade_event == trade
|
||||
assert result.wallet_profile == profile
|
||||
assert result.confidence >= BASE_CONFIDENCE
|
||||
|
||||
async def test_filters_non_fresh_wallet(self, detector, mock_wallet_analyzer):
|
||||
"""Test that non-fresh wallets are filtered out."""
|
||||
trade = create_trade_event(
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("4000"), # notional = $2000
|
||||
)
|
||||
profile = create_wallet_profile(nonce=10, age_hours=100.0, is_fresh=False)
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
assert result is None
|
||||
|
||||
async def test_handles_analyzer_error(self, detector, mock_wallet_analyzer):
|
||||
"""Test graceful handling of analyzer errors."""
|
||||
trade = create_trade_event()
|
||||
mock_wallet_analyzer.analyze.side_effect = Exception("RPC error")
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
assert result is None
|
||||
|
||||
async def test_wallet_at_nonce_threshold(self, detector, mock_wallet_analyzer):
|
||||
"""Test wallet exactly at max_nonce threshold."""
|
||||
trade = create_trade_event(size=Decimal("3000"))
|
||||
profile = create_wallet_profile(nonce=5, age_hours=24.0) # At threshold
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
assert result is not None
|
||||
|
||||
async def test_wallet_above_nonce_threshold(self, detector, mock_wallet_analyzer):
|
||||
"""Test wallet above max_nonce threshold."""
|
||||
trade = create_trade_event(size=Decimal("3000"))
|
||||
profile = create_wallet_profile(nonce=6, age_hours=24.0) # Above threshold
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
assert result is None
|
||||
|
||||
async def test_wallet_at_age_threshold(self, detector, mock_wallet_analyzer):
|
||||
"""Test wallet exactly at max_age_hours threshold."""
|
||||
trade = create_trade_event(size=Decimal("3000"))
|
||||
profile = create_wallet_profile(nonce=2, age_hours=48.0) # At threshold
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
assert result is not None
|
||||
|
||||
async def test_wallet_above_age_threshold(self, detector, mock_wallet_analyzer):
|
||||
"""Test wallet above max_age_hours threshold."""
|
||||
trade = create_trade_event(size=Decimal("3000"))
|
||||
profile = create_wallet_profile(nonce=2, age_hours=49.0) # Above threshold
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
assert result is None
|
||||
|
||||
async def test_wallet_with_unknown_age(self, detector, mock_wallet_analyzer):
|
||||
"""Test wallet with unknown age (None)."""
|
||||
trade = create_trade_event(size=Decimal("3000"))
|
||||
profile = create_wallet_profile(nonce=2, age_hours=None)
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
# Should pass since age is unknown
|
||||
assert result is not None
|
||||
|
||||
|
||||
# === Test Confidence Scoring ===
|
||||
|
||||
|
||||
class TestConfidenceScoring:
|
||||
"""Tests for confidence score calculation."""
|
||||
|
||||
def test_base_confidence_only(self, detector):
|
||||
"""Test base confidence for fresh wallet."""
|
||||
profile = create_wallet_profile(nonce=3, age_hours=24.0)
|
||||
trade = create_trade_event(size=Decimal("2000")) # $1000 notional
|
||||
|
||||
confidence, factors = detector.calculate_confidence(profile, trade)
|
||||
|
||||
assert confidence == BASE_CONFIDENCE
|
||||
assert factors == {"base": BASE_CONFIDENCE}
|
||||
|
||||
def test_brand_new_bonus(self, detector):
|
||||
"""Test bonus for brand new wallet (nonce=0)."""
|
||||
profile = create_wallet_profile(nonce=0, age_hours=24.0)
|
||||
trade = create_trade_event(size=Decimal("2000"))
|
||||
|
||||
confidence, factors = detector.calculate_confidence(profile, trade)
|
||||
|
||||
expected = BASE_CONFIDENCE + BRAND_NEW_BONUS
|
||||
assert confidence == expected
|
||||
assert factors["brand_new"] == BRAND_NEW_BONUS
|
||||
|
||||
def test_very_young_bonus(self, detector):
|
||||
"""Test bonus for very young wallet (age < 2 hours)."""
|
||||
profile = create_wallet_profile(nonce=3, age_hours=1.5)
|
||||
trade = create_trade_event(size=Decimal("2000"))
|
||||
|
||||
confidence, factors = detector.calculate_confidence(profile, trade)
|
||||
|
||||
expected = BASE_CONFIDENCE + VERY_YOUNG_BONUS
|
||||
assert confidence == expected
|
||||
assert factors["very_young"] == VERY_YOUNG_BONUS
|
||||
|
||||
def test_large_trade_bonus(self, detector):
|
||||
"""Test bonus for large trade (> $10,000)."""
|
||||
profile = create_wallet_profile(nonce=3, age_hours=24.0)
|
||||
trade = create_trade_event(
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("25000"), # $12,500 notional
|
||||
)
|
||||
|
||||
confidence, factors = detector.calculate_confidence(profile, trade)
|
||||
|
||||
expected = BASE_CONFIDENCE + LARGE_TRADE_BONUS
|
||||
assert confidence == expected
|
||||
assert factors["large_trade"] == LARGE_TRADE_BONUS
|
||||
|
||||
def test_all_bonuses_combined(self, detector):
|
||||
"""Test all bonuses stacking."""
|
||||
profile = create_wallet_profile(nonce=0, age_hours=0.5)
|
||||
trade = create_trade_event(
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("30000"), # $15,000 notional
|
||||
)
|
||||
|
||||
confidence, factors = detector.calculate_confidence(profile, trade)
|
||||
|
||||
expected = BASE_CONFIDENCE + BRAND_NEW_BONUS + VERY_YOUNG_BONUS + LARGE_TRADE_BONUS
|
||||
assert confidence == expected
|
||||
assert factors["brand_new"] == BRAND_NEW_BONUS
|
||||
assert factors["very_young"] == VERY_YOUNG_BONUS
|
||||
assert factors["large_trade"] == LARGE_TRADE_BONUS
|
||||
|
||||
def test_confidence_clamped_to_max(self, detector):
|
||||
"""Test confidence is clamped to 1.0 max."""
|
||||
# Create scenario where total would exceed 1.0
|
||||
# BASE=0.5 + BRAND_NEW=0.2 + VERY_YOUNG=0.1 + LARGE_TRADE=0.1 = 0.9
|
||||
# This is less than 1.0, so let's verify clamping works
|
||||
profile = create_wallet_profile(nonce=0, age_hours=0.5)
|
||||
trade = create_trade_event(size=Decimal("30000"))
|
||||
|
||||
confidence, _ = detector.calculate_confidence(profile, trade)
|
||||
|
||||
assert confidence <= 1.0
|
||||
|
||||
def test_no_bonus_at_age_boundary(self, detector):
|
||||
"""Test no bonus when age exactly at 2 hours."""
|
||||
profile = create_wallet_profile(nonce=3, age_hours=2.0)
|
||||
trade = create_trade_event(size=Decimal("2000"))
|
||||
|
||||
confidence, factors = detector.calculate_confidence(profile, trade)
|
||||
|
||||
assert "very_young" not in factors
|
||||
assert confidence == BASE_CONFIDENCE
|
||||
|
||||
def test_no_bonus_at_trade_size_boundary(self, detector):
|
||||
"""Test no bonus when trade exactly at $10,000."""
|
||||
profile = create_wallet_profile(nonce=3, age_hours=24.0)
|
||||
trade = create_trade_event(
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("20000"), # $10,000 notional exactly
|
||||
)
|
||||
|
||||
confidence, factors = detector.calculate_confidence(profile, trade)
|
||||
|
||||
# Exactly at threshold should NOT get bonus
|
||||
assert "large_trade" not in factors
|
||||
assert confidence == BASE_CONFIDENCE
|
||||
|
||||
def test_unknown_age_no_young_bonus(self, detector):
|
||||
"""Test no young bonus when age is None."""
|
||||
profile = create_wallet_profile(nonce=3, age_hours=None)
|
||||
trade = create_trade_event(size=Decimal("2000"))
|
||||
|
||||
confidence, factors = detector.calculate_confidence(profile, trade)
|
||||
|
||||
assert "very_young" not in factors
|
||||
assert confidence == BASE_CONFIDENCE
|
||||
|
||||
|
||||
# === Test Batch Analysis ===
|
||||
|
||||
|
||||
class TestBatchAnalysis:
|
||||
"""Tests for batch trade analysis."""
|
||||
|
||||
async def test_analyze_batch_success(self, detector, mock_wallet_analyzer):
|
||||
"""Test analyzing multiple trades."""
|
||||
trades = [
|
||||
create_trade_event(trade_id="trade1", size=Decimal("3000")),
|
||||
create_trade_event(trade_id="trade2", size=Decimal("4000")),
|
||||
]
|
||||
profile = create_wallet_profile(nonce=2, age_hours=24.0)
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
results = await detector.analyze_batch(trades)
|
||||
|
||||
assert len(results) == 2
|
||||
assert all(isinstance(r, FreshWalletSignal) for r in results)
|
||||
|
||||
async def test_analyze_batch_filters_small(self, detector, mock_wallet_analyzer):
|
||||
"""Test batch analysis filters small trades."""
|
||||
trades = [
|
||||
create_trade_event(trade_id="trade1", size=Decimal("3000")), # Above threshold
|
||||
create_trade_event(trade_id="trade2", size=Decimal("100")), # Below threshold
|
||||
]
|
||||
profile = create_wallet_profile(nonce=2, age_hours=24.0)
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
results = await detector.analyze_batch(trades)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].trade_event.trade_id == "trade1"
|
||||
|
||||
async def test_analyze_batch_handles_errors(self, detector, mock_wallet_analyzer):
|
||||
"""Test batch analysis handles individual errors gracefully."""
|
||||
trades = [
|
||||
create_trade_event(trade_id="trade1", size=Decimal("3000")),
|
||||
create_trade_event(trade_id="trade2", size=Decimal("4000")),
|
||||
]
|
||||
|
||||
# First call succeeds, second fails
|
||||
profile = create_wallet_profile(nonce=2, age_hours=24.0)
|
||||
mock_wallet_analyzer.analyze.side_effect = [profile, Exception("Error")]
|
||||
|
||||
results = await detector.analyze_batch(trades)
|
||||
|
||||
assert len(results) == 1
|
||||
|
||||
async def test_analyze_batch_empty_list(self, detector):
|
||||
"""Test batch analysis with empty list."""
|
||||
results = await detector.analyze_batch([])
|
||||
|
||||
assert results == []
|
||||
|
||||
async def test_analyze_batch_all_filtered(self, detector, mock_wallet_analyzer):
|
||||
"""Test batch analysis when all trades are filtered."""
|
||||
trades = [
|
||||
create_trade_event(trade_id="trade1", size=Decimal("100")),
|
||||
create_trade_event(trade_id="trade2", size=Decimal("50")),
|
||||
]
|
||||
|
||||
results = await detector.analyze_batch(trades)
|
||||
|
||||
assert results == []
|
||||
mock_wallet_analyzer.analyze.assert_not_called()
|
||||
|
||||
|
||||
# === Integration-Style Tests ===
|
||||
|
||||
|
||||
class TestDetectorIntegration:
|
||||
"""Integration-style tests for complete detector flow."""
|
||||
|
||||
async def test_full_detection_flow(self, mock_wallet_analyzer):
|
||||
"""Test complete detection flow from trade to signal."""
|
||||
detector = FreshWalletDetector(mock_wallet_analyzer)
|
||||
|
||||
# Create a suspicious trade
|
||||
trade = create_trade_event(
|
||||
wallet_address="0xfresh123",
|
||||
market_id="suspicious_market",
|
||||
price=Decimal("0.65"),
|
||||
size=Decimal("20000"), # $13,000 notional
|
||||
)
|
||||
|
||||
# Create a fresh wallet profile
|
||||
profile = create_wallet_profile(
|
||||
address="0xfresh123",
|
||||
nonce=0, # Brand new
|
||||
age_hours=0.5, # Very young
|
||||
)
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
# Analyze
|
||||
signal = await detector.analyze(trade)
|
||||
|
||||
# Verify signal
|
||||
assert signal is not None
|
||||
assert signal.wallet_address == "0xfresh123"
|
||||
assert signal.market_id == "suspicious_market"
|
||||
assert signal.is_high_confidence is True
|
||||
assert "brand_new" in signal.factors
|
||||
assert "very_young" in signal.factors
|
||||
assert "large_trade" in signal.factors
|
||||
|
||||
async def test_custom_thresholds(self, mock_wallet_analyzer):
|
||||
"""Test detection with custom thresholds."""
|
||||
detector = FreshWalletDetector(
|
||||
mock_wallet_analyzer,
|
||||
min_trade_size=Decimal("5000"),
|
||||
max_nonce=3,
|
||||
max_age_hours=24.0,
|
||||
)
|
||||
|
||||
# Trade that would pass default but fails custom thresholds
|
||||
trade = create_trade_event(size=Decimal("8000")) # $4000 notional
|
||||
profile = create_wallet_profile(nonce=4, age_hours=30.0)
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
# Should fail min_trade_size check
|
||||
result = await detector.analyze(trade)
|
||||
assert result is None
|
||||
|
||||
async def test_edge_case_exact_min_trade_size(self, detector, mock_wallet_analyzer):
|
||||
"""Test trade exactly at minimum size threshold."""
|
||||
trade = create_trade_event(
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("2000"), # $1000 notional exactly at default
|
||||
)
|
||||
profile = create_wallet_profile(nonce=2)
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
# At threshold should pass
|
||||
assert result is not None
|
||||
|
||||
async def test_edge_case_below_min_trade_size(self, detector, mock_wallet_analyzer):
|
||||
"""Test trade just below minimum size threshold."""
|
||||
trade = create_trade_event(
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("1999"), # $999.50 - just under $1000
|
||||
)
|
||||
profile = create_wallet_profile(nonce=2)
|
||||
mock_wallet_analyzer.analyze.return_value = profile
|
||||
|
||||
result = await detector.analyze(trade)
|
||||
|
||||
# Below threshold should fail
|
||||
assert result is None
|
||||
@@ -0,0 +1,677 @@
|
||||
"""Tests for composite risk scorer."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.detector.models import (
|
||||
FreshWalletSignal,
|
||||
RiskAssessment,
|
||||
SizeAnomalySignal,
|
||||
)
|
||||
from polymarket_insider_tracker.detector.scorer import (
|
||||
DEFAULT_ALERT_THRESHOLD,
|
||||
DEFAULT_WEIGHTS,
|
||||
MULTI_SIGNAL_BONUS_2,
|
||||
RiskScorer,
|
||||
SignalBundle,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.models import MarketMetadata, Token, TradeEvent
|
||||
from polymarket_insider_tracker.profiler.models import WalletProfile
|
||||
|
||||
# ============================================================================
|
||||
# Fixtures
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis() -> AsyncMock:
|
||||
"""Create a mock Redis client."""
|
||||
mock = AsyncMock()
|
||||
# Default: key doesn't exist (not a duplicate)
|
||||
mock.set.return_value = True
|
||||
mock.delete.return_value = 1
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_trade() -> TradeEvent:
|
||||
"""Create a sample trade event."""
|
||||
return TradeEvent(
|
||||
market_id="market_abc123",
|
||||
trade_id="tx_001",
|
||||
wallet_address="0x1234567890abcdef1234567890abcdef12345678",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.65"),
|
||||
size=Decimal("10000"),
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_123",
|
||||
event_title="Test Market",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_wallet_profile() -> WalletProfile:
|
||||
"""Create a sample wallet profile."""
|
||||
return WalletProfile(
|
||||
address="0x1234567890abcdef1234567890abcdef12345678",
|
||||
nonce=2,
|
||||
first_seen=datetime.now(UTC),
|
||||
age_hours=1.0,
|
||||
is_fresh=True,
|
||||
total_tx_count=2,
|
||||
matic_balance=Decimal("1000000000000000000"), # 1 MATIC
|
||||
usdc_balance=Decimal("1000000"), # 1 USDC
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_metadata() -> MarketMetadata:
|
||||
"""Create sample market metadata."""
|
||||
return MarketMetadata(
|
||||
condition_id="market_abc123",
|
||||
question="Will it rain tomorrow?",
|
||||
description="Weather prediction market",
|
||||
tokens=(Token(token_id="token_123", outcome="Yes", price=Decimal("0.65")),),
|
||||
category="science",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_wallet_signal(
|
||||
sample_trade: TradeEvent, sample_wallet_profile: WalletProfile
|
||||
) -> FreshWalletSignal:
|
||||
"""Create a sample fresh wallet signal."""
|
||||
return FreshWalletSignal(
|
||||
trade_event=sample_trade,
|
||||
wallet_profile=sample_wallet_profile,
|
||||
confidence=0.8,
|
||||
factors={"base": 0.5, "brand_new_bonus": 0.2, "large_trade_bonus": 0.1},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def size_anomaly_signal(
|
||||
sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> SizeAnomalySignal:
|
||||
"""Create a sample size anomaly signal."""
|
||||
return SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.10,
|
||||
book_impact=0.15,
|
||||
is_niche_market=True,
|
||||
confidence=0.7,
|
||||
factors={"volume_impact": 0.4, "book_impact": 0.3},
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SignalBundle Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestSignalBundle:
|
||||
"""Tests for the SignalBundle dataclass."""
|
||||
|
||||
def test_bundle_with_no_signals(self, sample_trade: TradeEvent) -> None:
|
||||
"""Test bundle with only trade, no signals."""
|
||||
bundle = SignalBundle(trade_event=sample_trade)
|
||||
|
||||
assert bundle.trade_event == sample_trade
|
||||
assert bundle.fresh_wallet_signal is None
|
||||
assert bundle.size_anomaly_signal is None
|
||||
assert bundle.wallet_address == sample_trade.wallet_address
|
||||
assert bundle.market_id == sample_trade.market_id
|
||||
|
||||
def test_bundle_with_fresh_wallet_signal(
|
||||
self,
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
) -> None:
|
||||
"""Test bundle with fresh wallet signal."""
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
)
|
||||
|
||||
assert bundle.fresh_wallet_signal == fresh_wallet_signal
|
||||
assert bundle.size_anomaly_signal is None
|
||||
|
||||
def test_bundle_with_all_signals(
|
||||
self,
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
size_anomaly_signal: SizeAnomalySignal,
|
||||
) -> None:
|
||||
"""Test bundle with all signal types."""
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
size_anomaly_signal=size_anomaly_signal,
|
||||
)
|
||||
|
||||
assert bundle.fresh_wallet_signal == fresh_wallet_signal
|
||||
assert bundle.size_anomaly_signal == size_anomaly_signal
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# RiskAssessment Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestRiskAssessment:
|
||||
"""Tests for the RiskAssessment dataclass."""
|
||||
|
||||
def test_assessment_creation(self, sample_trade: TradeEvent) -> None:
|
||||
"""Test basic assessment creation."""
|
||||
assessment = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=None,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=0,
|
||||
weighted_score=0.0,
|
||||
should_alert=False,
|
||||
)
|
||||
|
||||
assert assessment.trade_event == sample_trade
|
||||
assert assessment.signals_triggered == 0
|
||||
assert assessment.weighted_score == 0.0
|
||||
assert assessment.should_alert is False
|
||||
assert assessment.assessment_id is not None
|
||||
assert assessment.timestamp is not None
|
||||
|
||||
def test_is_high_risk(self, sample_trade: TradeEvent) -> None:
|
||||
"""Test is_high_risk property."""
|
||||
high_risk = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=None,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=1,
|
||||
weighted_score=0.70,
|
||||
should_alert=True,
|
||||
)
|
||||
low_risk = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=None,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=1,
|
||||
weighted_score=0.69,
|
||||
should_alert=True,
|
||||
)
|
||||
|
||||
assert high_risk.is_high_risk is True
|
||||
assert low_risk.is_high_risk is False
|
||||
|
||||
def test_is_very_high_risk(self, sample_trade: TradeEvent) -> None:
|
||||
"""Test is_very_high_risk property."""
|
||||
very_high = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=None,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=2,
|
||||
weighted_score=0.85,
|
||||
should_alert=True,
|
||||
)
|
||||
high = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=None,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=2,
|
||||
weighted_score=0.84,
|
||||
should_alert=True,
|
||||
)
|
||||
|
||||
assert very_high.is_very_high_risk is True
|
||||
assert high.is_very_high_risk is False
|
||||
|
||||
def test_to_dict(
|
||||
self,
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
) -> None:
|
||||
"""Test to_dict serialization."""
|
||||
assessment = RiskAssessment(
|
||||
trade_event=sample_trade,
|
||||
wallet_address=sample_trade.wallet_address,
|
||||
market_id=sample_trade.market_id,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
size_anomaly_signal=None,
|
||||
signals_triggered=1,
|
||||
weighted_score=0.65,
|
||||
should_alert=True,
|
||||
)
|
||||
|
||||
result = assessment.to_dict()
|
||||
|
||||
assert result["wallet_address"] == sample_trade.wallet_address
|
||||
assert result["market_id"] == sample_trade.market_id
|
||||
assert result["signals_triggered"] == 1
|
||||
assert result["weighted_score"] == 0.65
|
||||
assert result["should_alert"] is True
|
||||
assert result["has_fresh_wallet_signal"] is True
|
||||
assert result["has_size_anomaly_signal"] is False
|
||||
assert result["fresh_wallet_confidence"] == 0.8
|
||||
assert result["size_anomaly_confidence"] is None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# RiskScorer Initialization Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestRiskScorerInit:
|
||||
"""Tests for RiskScorer initialization."""
|
||||
|
||||
def test_default_initialization(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test scorer initializes with default values."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
|
||||
assert scorer._alert_threshold == DEFAULT_ALERT_THRESHOLD
|
||||
assert scorer._weights == DEFAULT_WEIGHTS
|
||||
assert scorer._dedup_window == 3600
|
||||
|
||||
def test_custom_configuration(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test scorer with custom configuration."""
|
||||
custom_weights = {"fresh_wallet": 0.5, "size_anomaly": 0.5}
|
||||
scorer = RiskScorer(
|
||||
mock_redis,
|
||||
weights=custom_weights,
|
||||
alert_threshold=0.7,
|
||||
dedup_window_seconds=1800,
|
||||
)
|
||||
|
||||
assert scorer._alert_threshold == 0.7
|
||||
assert scorer._weights == custom_weights
|
||||
assert scorer._dedup_window == 1800
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Weighted Score Calculation Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestWeightedScoreCalculation:
|
||||
"""Tests for weighted score calculation."""
|
||||
|
||||
def test_no_signals_zero_score(self, mock_redis: AsyncMock, sample_trade: TradeEvent) -> None:
|
||||
"""Test score is zero when no signals present."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(trade_event=sample_trade)
|
||||
|
||||
score, count = scorer.calculate_weighted_score(bundle)
|
||||
|
||||
assert score == 0.0
|
||||
assert count == 0
|
||||
|
||||
def test_fresh_wallet_only(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
) -> None:
|
||||
"""Test score with only fresh wallet signal."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
)
|
||||
|
||||
score, count = scorer.calculate_weighted_score(bundle)
|
||||
|
||||
# 0.8 confidence * 0.4 weight = 0.32
|
||||
expected = 0.8 * DEFAULT_WEIGHTS["fresh_wallet"]
|
||||
assert score == pytest.approx(expected)
|
||||
assert count == 1
|
||||
|
||||
def test_size_anomaly_only(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
size_anomaly_signal: SizeAnomalySignal,
|
||||
) -> None:
|
||||
"""Test score with only size anomaly signal."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
size_anomaly_signal=size_anomaly_signal,
|
||||
)
|
||||
|
||||
score, count = scorer.calculate_weighted_score(bundle)
|
||||
|
||||
# 0.7 confidence * 0.35 weight + 0.7 * 0.25 niche weight = 0.42
|
||||
expected = 0.7 * DEFAULT_WEIGHTS["size_anomaly"] + 0.7 * DEFAULT_WEIGHTS["niche_market"]
|
||||
assert score == pytest.approx(expected)
|
||||
assert count == 1
|
||||
|
||||
def test_size_anomaly_non_niche(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
sample_metadata: MarketMetadata,
|
||||
) -> None:
|
||||
"""Test size anomaly without niche bonus."""
|
||||
signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.10,
|
||||
book_impact=0.15,
|
||||
is_niche_market=False,
|
||||
confidence=0.7,
|
||||
factors={},
|
||||
)
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
size_anomaly_signal=signal,
|
||||
)
|
||||
|
||||
score, count = scorer.calculate_weighted_score(bundle)
|
||||
|
||||
# 0.7 * 0.35 = 0.245 (no niche bonus)
|
||||
expected = 0.7 * DEFAULT_WEIGHTS["size_anomaly"]
|
||||
assert score == pytest.approx(expected)
|
||||
|
||||
def test_multi_signal_bonus_two_signals(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
size_anomaly_signal: SizeAnomalySignal,
|
||||
) -> None:
|
||||
"""Test 20% bonus for two signals."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
size_anomaly_signal=size_anomaly_signal,
|
||||
)
|
||||
|
||||
score, count = scorer.calculate_weighted_score(bundle)
|
||||
|
||||
# Calculate base score
|
||||
base = (
|
||||
0.8 * DEFAULT_WEIGHTS["fresh_wallet"]
|
||||
+ 0.7 * DEFAULT_WEIGHTS["size_anomaly"]
|
||||
+ 0.7 * DEFAULT_WEIGHTS["niche_market"]
|
||||
)
|
||||
expected = base * MULTI_SIGNAL_BONUS_2
|
||||
assert score == pytest.approx(expected)
|
||||
assert count == 2
|
||||
|
||||
def test_score_capped_at_one(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
sample_wallet_profile: WalletProfile,
|
||||
sample_metadata: MarketMetadata,
|
||||
) -> None:
|
||||
"""Test score is capped at 1.0."""
|
||||
# Create high confidence signals
|
||||
fresh_signal = FreshWalletSignal(
|
||||
trade_event=sample_trade,
|
||||
wallet_profile=sample_wallet_profile,
|
||||
confidence=1.0,
|
||||
factors={},
|
||||
)
|
||||
size_signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.10,
|
||||
book_impact=0.15,
|
||||
is_niche_market=True,
|
||||
confidence=1.0,
|
||||
factors={},
|
||||
)
|
||||
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
fresh_wallet_signal=fresh_signal,
|
||||
size_anomaly_signal=size_signal,
|
||||
)
|
||||
|
||||
score, count = scorer.calculate_weighted_score(bundle)
|
||||
|
||||
assert score == 1.0 # Capped
|
||||
assert count == 2
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Assess Method Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestAssessMethod:
|
||||
"""Tests for the assess method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assess_triggers_alert(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
size_anomaly_signal: SizeAnomalySignal,
|
||||
) -> None:
|
||||
"""Test assess triggers alert for high-risk trades."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
size_anomaly_signal=size_anomaly_signal,
|
||||
)
|
||||
|
||||
assessment = await scorer.assess(bundle)
|
||||
|
||||
assert assessment.should_alert is True
|
||||
assert assessment.signals_triggered == 2
|
||||
assert assessment.weighted_score >= DEFAULT_ALERT_THRESHOLD
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assess_no_alert_below_threshold(
|
||||
self, mock_redis: AsyncMock, sample_trade: TradeEvent
|
||||
) -> None:
|
||||
"""Test assess does not alert for low-risk trades."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(trade_event=sample_trade)
|
||||
|
||||
assessment = await scorer.assess(bundle)
|
||||
|
||||
assert assessment.should_alert is False
|
||||
assert assessment.signals_triggered == 0
|
||||
assert assessment.weighted_score == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assess_deduplication(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
size_anomaly_signal: SizeAnomalySignal,
|
||||
) -> None:
|
||||
"""Test assess deduplicates repeated alerts."""
|
||||
# First call: key doesn't exist (returns True)
|
||||
# Second call: key exists (returns False/None)
|
||||
mock_redis.set.side_effect = [True, False]
|
||||
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
size_anomaly_signal=size_anomaly_signal,
|
||||
)
|
||||
|
||||
# First assessment should alert
|
||||
assessment1 = await scorer.assess(bundle)
|
||||
# Second assessment should be deduplicated
|
||||
assessment2 = await scorer.assess(bundle)
|
||||
|
||||
assert assessment1.should_alert is True
|
||||
assert assessment2.should_alert is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assess_preserves_signals(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
fresh_wallet_signal: FreshWalletSignal,
|
||||
) -> None:
|
||||
"""Test assess preserves original signals in assessment."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade,
|
||||
fresh_wallet_signal=fresh_wallet_signal,
|
||||
)
|
||||
|
||||
assessment = await scorer.assess(bundle)
|
||||
|
||||
assert assessment.fresh_wallet_signal == fresh_wallet_signal
|
||||
assert assessment.size_anomaly_signal is None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Deduplication Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestDeduplication:
|
||||
"""Tests for deduplication functionality."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_and_set_dedup_new_key(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test dedup returns False for new key."""
|
||||
mock_redis.set.return_value = True
|
||||
|
||||
scorer = RiskScorer(mock_redis)
|
||||
is_dup = await scorer._check_and_set_dedup("0xwallet", "market123")
|
||||
|
||||
assert is_dup is False
|
||||
mock_redis.set.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_and_set_dedup_existing_key(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test dedup returns True for existing key."""
|
||||
mock_redis.set.return_value = False # Key exists, NX failed
|
||||
|
||||
scorer = RiskScorer(mock_redis)
|
||||
is_dup = await scorer._check_and_set_dedup("0xwallet", "market123")
|
||||
|
||||
assert is_dup is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_dedup(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test clearing dedup key."""
|
||||
mock_redis.delete.return_value = 1
|
||||
|
||||
scorer = RiskScorer(mock_redis)
|
||||
cleared = await scorer.clear_dedup("0xwallet", "market123")
|
||||
|
||||
assert cleared is True
|
||||
mock_redis.delete.assert_called_once()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Batch Analysis Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestBatchAnalysis:
|
||||
"""Tests for batch assessment."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assess_batch(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
sample_wallet_profile: WalletProfile,
|
||||
) -> None:
|
||||
"""Test batch assessment returns assessments for all bundles."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
|
||||
bundles = []
|
||||
for i in range(3):
|
||||
trade = TradeEvent(
|
||||
market_id=f"market_{i}",
|
||||
trade_id=f"tx_{i}",
|
||||
wallet_address=f"0xwallet{i}",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.50"),
|
||||
size=Decimal("10000"),
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_123",
|
||||
)
|
||||
signal = FreshWalletSignal(
|
||||
trade_event=trade,
|
||||
wallet_profile=sample_wallet_profile,
|
||||
confidence=0.8,
|
||||
factors={},
|
||||
)
|
||||
bundles.append(SignalBundle(trade_event=trade, fresh_wallet_signal=signal))
|
||||
|
||||
assessments = await scorer.assess_batch(bundles)
|
||||
|
||||
assert len(assessments) == 3
|
||||
assert all(isinstance(a, RiskAssessment) for a in assessments)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assess_batch_empty(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test batch assessment with empty list."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
|
||||
assessments = await scorer.assess_batch([])
|
||||
|
||||
assert assessments == []
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Weight Management Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestWeightManagement:
|
||||
"""Tests for weight get/set functionality."""
|
||||
|
||||
def test_get_weights(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting weights returns a copy."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
|
||||
weights = scorer.get_weights()
|
||||
|
||||
assert weights == DEFAULT_WEIGHTS
|
||||
# Verify it's a copy, not the original
|
||||
weights["fresh_wallet"] = 999
|
||||
assert scorer._weights["fresh_wallet"] != 999
|
||||
|
||||
def test_set_weights(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test setting new weights."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
new_weights = {"fresh_wallet": 0.5, "size_anomaly": 0.5}
|
||||
|
||||
scorer.set_weights(new_weights)
|
||||
|
||||
assert scorer._weights == new_weights
|
||||
|
||||
def test_set_weights_makes_copy(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test set_weights makes a copy of the input."""
|
||||
scorer = RiskScorer(mock_redis)
|
||||
new_weights = {"fresh_wallet": 0.5, "size_anomaly": 0.5}
|
||||
|
||||
scorer.set_weights(new_weights)
|
||||
new_weights["fresh_wallet"] = 999
|
||||
|
||||
assert scorer._weights["fresh_wallet"] == 0.5
|
||||
@@ -0,0 +1,821 @@
|
||||
"""Tests for position size anomaly detection."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.detector.models import SizeAnomalySignal
|
||||
from polymarket_insider_tracker.detector.size_anomaly import (
|
||||
DEFAULT_BOOK_THRESHOLD,
|
||||
DEFAULT_NICHE_VOLUME_THRESHOLD,
|
||||
DEFAULT_VOLUME_THRESHOLD,
|
||||
NICHE_PRONE_CATEGORIES,
|
||||
SizeAnomalyDetector,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.metadata_sync import MarketMetadataSync
|
||||
from polymarket_insider_tracker.ingestor.models import MarketMetadata, Token, TradeEvent
|
||||
|
||||
# ============================================================================
|
||||
# Fixtures
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_metadata_sync() -> AsyncMock:
|
||||
"""Create a mock MarketMetadataSync."""
|
||||
return AsyncMock(spec=MarketMetadataSync)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_token() -> Token:
|
||||
"""Create a sample token."""
|
||||
return Token(
|
||||
token_id="token_123",
|
||||
outcome="Yes",
|
||||
price=Decimal("0.65"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_metadata(sample_token: Token) -> MarketMetadata:
|
||||
"""Create sample market metadata."""
|
||||
return MarketMetadata(
|
||||
condition_id="market_abc123",
|
||||
question="Will it rain tomorrow?",
|
||||
description="Weather prediction market",
|
||||
tokens=(sample_token,),
|
||||
category="science",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_trade() -> TradeEvent:
|
||||
"""Create a sample trade event."""
|
||||
return TradeEvent(
|
||||
market_id="market_abc123",
|
||||
trade_id="tx_001",
|
||||
wallet_address="0x1234567890abcdef",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.65"),
|
||||
size=Decimal("10000"), # $6,500 notional
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_123",
|
||||
event_title="Weather Market",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def large_trade() -> TradeEvent:
|
||||
"""Create a large trade event."""
|
||||
return TradeEvent(
|
||||
market_id="market_abc123",
|
||||
trade_id="tx_002",
|
||||
wallet_address="0xlargewallet",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.50"),
|
||||
size=Decimal("100000"), # $50,000 notional
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_123",
|
||||
event_title="Big Market",
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SizeAnomalySignal Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestSizeAnomalySignal:
|
||||
"""Tests for the SizeAnomalySignal dataclass."""
|
||||
|
||||
def test_signal_creation(
|
||||
self, sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test basic signal creation."""
|
||||
signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.05,
|
||||
book_impact=0.10,
|
||||
is_niche_market=True,
|
||||
confidence=0.75,
|
||||
factors={"volume_impact": 0.4, "niche_multiplier": 1.5},
|
||||
)
|
||||
|
||||
assert signal.trade_event == sample_trade
|
||||
assert signal.market_metadata == sample_metadata
|
||||
assert signal.volume_impact == 0.05
|
||||
assert signal.book_impact == 0.10
|
||||
assert signal.is_niche_market is True
|
||||
assert signal.confidence == 0.75
|
||||
assert "volume_impact" in signal.factors
|
||||
|
||||
def test_wallet_address_property(
|
||||
self, sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test wallet_address property."""
|
||||
signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.05,
|
||||
book_impact=0.10,
|
||||
is_niche_market=False,
|
||||
confidence=0.5,
|
||||
factors={},
|
||||
)
|
||||
|
||||
assert signal.wallet_address == sample_trade.wallet_address
|
||||
|
||||
def test_market_id_property(
|
||||
self, sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test market_id property."""
|
||||
signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.05,
|
||||
book_impact=0.10,
|
||||
is_niche_market=False,
|
||||
confidence=0.5,
|
||||
factors={},
|
||||
)
|
||||
|
||||
assert signal.market_id == sample_trade.market_id
|
||||
|
||||
def test_trade_size_usdc_property(
|
||||
self, sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test trade_size_usdc property returns notional value."""
|
||||
signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.05,
|
||||
book_impact=0.10,
|
||||
is_niche_market=False,
|
||||
confidence=0.5,
|
||||
factors={},
|
||||
)
|
||||
|
||||
# notional = price * size = 0.65 * 10000 = 6500
|
||||
assert signal.trade_size_usdc == Decimal("6500.00")
|
||||
|
||||
def test_is_high_confidence(
|
||||
self, sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test is_high_confidence threshold."""
|
||||
high_signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.05,
|
||||
book_impact=0.10,
|
||||
is_niche_market=False,
|
||||
confidence=0.70,
|
||||
factors={},
|
||||
)
|
||||
low_signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.05,
|
||||
book_impact=0.10,
|
||||
is_niche_market=False,
|
||||
confidence=0.69,
|
||||
factors={},
|
||||
)
|
||||
|
||||
assert high_signal.is_high_confidence is True
|
||||
assert low_signal.is_high_confidence is False
|
||||
|
||||
def test_is_very_high_confidence(
|
||||
self, sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test is_very_high_confidence threshold."""
|
||||
very_high = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.05,
|
||||
book_impact=0.10,
|
||||
is_niche_market=False,
|
||||
confidence=0.85,
|
||||
factors={},
|
||||
)
|
||||
high = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.05,
|
||||
book_impact=0.10,
|
||||
is_niche_market=False,
|
||||
confidence=0.84,
|
||||
factors={},
|
||||
)
|
||||
|
||||
assert very_high.is_very_high_confidence is True
|
||||
assert high.is_very_high_confidence is False
|
||||
|
||||
def test_to_dict_serialization(
|
||||
self, sample_trade: TradeEvent, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test to_dict produces valid serialization."""
|
||||
signal = SizeAnomalySignal(
|
||||
trade_event=sample_trade,
|
||||
market_metadata=sample_metadata,
|
||||
volume_impact=0.05,
|
||||
book_impact=0.10,
|
||||
is_niche_market=True,
|
||||
confidence=0.75,
|
||||
factors={"volume_impact": 0.5},
|
||||
)
|
||||
|
||||
result = signal.to_dict()
|
||||
|
||||
assert result["wallet_address"] == sample_trade.wallet_address
|
||||
assert result["market_id"] == sample_trade.market_id
|
||||
assert result["trade_id"] == sample_trade.trade_id
|
||||
assert result["trade_size"] == "6500.00"
|
||||
assert result["trade_side"] == "BUY"
|
||||
assert result["market_category"] == "science"
|
||||
assert result["volume_impact"] == 0.05
|
||||
assert result["book_impact"] == 0.10
|
||||
assert result["is_niche_market"] is True
|
||||
assert result["confidence"] == 0.75
|
||||
assert result["factors"] == {"volume_impact": 0.5}
|
||||
assert "timestamp" in result
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SizeAnomalyDetector Initialization Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestSizeAnomalyDetectorInit:
|
||||
"""Tests for SizeAnomalyDetector initialization."""
|
||||
|
||||
def test_default_initialization(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test detector initializes with default values."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
assert detector._volume_threshold == DEFAULT_VOLUME_THRESHOLD
|
||||
assert detector._book_threshold == DEFAULT_BOOK_THRESHOLD
|
||||
assert detector._niche_volume_threshold == DEFAULT_NICHE_VOLUME_THRESHOLD
|
||||
|
||||
def test_custom_thresholds(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test detector with custom thresholds."""
|
||||
detector = SizeAnomalyDetector(
|
||||
mock_metadata_sync,
|
||||
volume_threshold=0.05,
|
||||
book_threshold=0.10,
|
||||
niche_volume_threshold=Decimal("100000"),
|
||||
)
|
||||
|
||||
assert detector._volume_threshold == 0.05
|
||||
assert detector._book_threshold == 0.10
|
||||
assert detector._niche_volume_threshold == Decimal("100000")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Volume Impact Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestVolumeImpactCalculation:
|
||||
"""Tests for volume impact calculation."""
|
||||
|
||||
def test_volume_impact_calculation(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test correct volume impact calculation."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Trade size $1000, daily volume $50000 = 2% impact
|
||||
impact = detector._calculate_volume_impact(Decimal("1000"), Decimal("50000"))
|
||||
assert impact == pytest.approx(0.02)
|
||||
|
||||
def test_volume_impact_none_volume(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test volume impact returns 0 when volume is None."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
impact = detector._calculate_volume_impact(Decimal("1000"), None)
|
||||
assert impact == 0.0
|
||||
|
||||
def test_volume_impact_zero_volume(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test volume impact returns 0 when volume is zero."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
impact = detector._calculate_volume_impact(Decimal("1000"), Decimal("0"))
|
||||
assert impact == 0.0
|
||||
|
||||
def test_volume_impact_negative_volume(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test volume impact returns 0 when volume is negative."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
impact = detector._calculate_volume_impact(Decimal("1000"), Decimal("-1000"))
|
||||
assert impact == 0.0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Book Impact Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestBookImpactCalculation:
|
||||
"""Tests for order book impact calculation."""
|
||||
|
||||
def test_book_impact_calculation(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test correct book impact calculation."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Trade size $5000, book depth $50000 = 10% impact
|
||||
impact = detector._calculate_book_impact(Decimal("5000"), Decimal("50000"))
|
||||
assert impact == pytest.approx(0.10)
|
||||
|
||||
def test_book_impact_none_depth(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test book impact returns 0 when depth is None."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
impact = detector._calculate_book_impact(Decimal("5000"), None)
|
||||
assert impact == 0.0
|
||||
|
||||
def test_book_impact_zero_depth(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test book impact returns 0 when depth is zero."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
impact = detector._calculate_book_impact(Decimal("5000"), Decimal("0"))
|
||||
assert impact == 0.0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Niche Market Detection Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestNicheMarketDetection:
|
||||
"""Tests for niche market detection."""
|
||||
|
||||
def test_niche_market_low_volume(
|
||||
self, mock_metadata_sync: AsyncMock, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test market is niche when volume below threshold."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Volume $40k < $50k threshold
|
||||
is_niche = detector._is_niche_market(sample_metadata, Decimal("40000"))
|
||||
assert is_niche is True
|
||||
|
||||
def test_not_niche_high_volume(
|
||||
self, mock_metadata_sync: AsyncMock, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test market is not niche when volume above threshold."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Volume $100k > $50k threshold
|
||||
is_niche = detector._is_niche_market(sample_metadata, Decimal("100000"))
|
||||
assert is_niche is False
|
||||
|
||||
def test_niche_market_unknown_volume_niche_category(
|
||||
self, mock_metadata_sync: AsyncMock, sample_token: Token
|
||||
) -> None:
|
||||
"""Test market is niche when volume unknown and category is niche-prone."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
for category in NICHE_PRONE_CATEGORIES:
|
||||
metadata = MarketMetadata(
|
||||
condition_id="test",
|
||||
question="Test",
|
||||
description="",
|
||||
tokens=(sample_token,),
|
||||
category=category,
|
||||
)
|
||||
is_niche = detector._is_niche_market(metadata, None)
|
||||
assert is_niche is True, f"Category {category} should be niche"
|
||||
|
||||
def test_not_niche_unknown_volume_mainstream_category(
|
||||
self, mock_metadata_sync: AsyncMock, sample_token: Token
|
||||
) -> None:
|
||||
"""Test market is not niche when volume unknown but category is mainstream."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
mainstream_categories = ["politics", "sports", "crypto", "entertainment"]
|
||||
for category in mainstream_categories:
|
||||
metadata = MarketMetadata(
|
||||
condition_id="test",
|
||||
question="Test",
|
||||
description="",
|
||||
tokens=(sample_token,),
|
||||
category=category,
|
||||
)
|
||||
is_niche = detector._is_niche_market(metadata, None)
|
||||
assert is_niche is False, f"Category {category} should not be niche"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Confidence Scoring Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestConfidenceScoring:
|
||||
"""Tests for confidence score calculation."""
|
||||
|
||||
def test_confidence_volume_impact_only(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test confidence with only volume impact."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Volume impact 3x threshold = max score 0.5
|
||||
confidence, factors = detector.calculate_confidence(
|
||||
volume_impact=0.06, # 3x the 0.02 threshold
|
||||
book_impact=0.0,
|
||||
is_niche=False,
|
||||
)
|
||||
|
||||
assert confidence == pytest.approx(0.5)
|
||||
assert "volume_impact" in factors
|
||||
assert factors["volume_impact"] == pytest.approx(0.5)
|
||||
|
||||
def test_confidence_book_impact_only(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test confidence with only book impact."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Book impact 3x threshold = max score 0.3
|
||||
confidence, factors = detector.calculate_confidence(
|
||||
volume_impact=0.0,
|
||||
book_impact=0.15, # 3x the 0.05 threshold
|
||||
is_niche=False,
|
||||
)
|
||||
|
||||
assert confidence == pytest.approx(0.3)
|
||||
assert "book_impact" in factors
|
||||
assert factors["book_impact"] == pytest.approx(0.3)
|
||||
|
||||
def test_confidence_combined_impacts(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test confidence with both volume and book impact."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Both at 3x threshold = 0.5 + 0.3 = 0.8
|
||||
confidence, factors = detector.calculate_confidence(
|
||||
volume_impact=0.06,
|
||||
book_impact=0.15,
|
||||
is_niche=False,
|
||||
)
|
||||
|
||||
assert confidence == pytest.approx(0.8)
|
||||
assert "volume_impact" in factors
|
||||
assert "book_impact" in factors
|
||||
|
||||
def test_confidence_niche_multiplier(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test niche multiplier increases confidence."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Volume impact 2x threshold = 0.33, with 1.5x niche = 0.5
|
||||
confidence, factors = detector.calculate_confidence(
|
||||
volume_impact=0.04, # 2x threshold
|
||||
book_impact=0.0,
|
||||
is_niche=True,
|
||||
)
|
||||
|
||||
assert confidence == pytest.approx(0.5, rel=0.01)
|
||||
assert "niche_multiplier" in factors
|
||||
assert factors["niche_multiplier"] == 1.5
|
||||
|
||||
def test_confidence_niche_only_base(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test niche market with no other signals gives base confidence."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# No threshold exceeded, but is niche
|
||||
confidence, factors = detector.calculate_confidence(
|
||||
volume_impact=0.01, # Below 0.02 threshold
|
||||
book_impact=0.01, # Below 0.05 threshold
|
||||
is_niche=True,
|
||||
)
|
||||
|
||||
assert confidence == 0.2
|
||||
assert "niche_base" in factors
|
||||
assert factors["niche_base"] == 0.2
|
||||
|
||||
def test_confidence_clamped_to_max(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test confidence is clamped to 1.0."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# High impacts with niche multiplier would exceed 1.0
|
||||
confidence, factors = detector.calculate_confidence(
|
||||
volume_impact=0.10, # 5x threshold (capped at 3x)
|
||||
book_impact=0.20, # 4x threshold (capped at 3x)
|
||||
is_niche=True, # 1.5x multiplier
|
||||
)
|
||||
|
||||
assert confidence == 1.0
|
||||
|
||||
def test_confidence_zero_no_signals(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test confidence is zero with no signals."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
confidence, factors = detector.calculate_confidence(
|
||||
volume_impact=0.01, # Below threshold
|
||||
book_impact=0.01, # Below threshold
|
||||
is_niche=False,
|
||||
)
|
||||
|
||||
assert confidence == 0.0
|
||||
assert len(factors) == 0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Analyze Method Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestAnalyzeMethod:
|
||||
"""Tests for the analyze method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_high_volume_impact(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
sample_metadata: MarketMetadata,
|
||||
) -> None:
|
||||
"""Test analyze detects high volume impact trade."""
|
||||
mock_metadata_sync.get_market.return_value = sample_metadata
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Trade notional = 6500, volume = 65000, impact = 10% > 2% threshold
|
||||
signal = await detector.analyze(
|
||||
sample_trade,
|
||||
daily_volume=Decimal("65000"),
|
||||
)
|
||||
|
||||
assert signal is not None
|
||||
assert signal.volume_impact == pytest.approx(0.10)
|
||||
assert signal.confidence > 0.1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_high_book_impact(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
sample_metadata: MarketMetadata,
|
||||
) -> None:
|
||||
"""Test analyze detects high book impact trade."""
|
||||
mock_metadata_sync.get_market.return_value = sample_metadata
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Trade notional = 6500, book depth = 32500, impact = 20% > 5% threshold
|
||||
signal = await detector.analyze(
|
||||
sample_trade,
|
||||
book_depth=Decimal("32500"),
|
||||
)
|
||||
|
||||
assert signal is not None
|
||||
assert signal.book_impact == pytest.approx(0.20)
|
||||
assert signal.confidence > 0.1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_niche_market(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
sample_metadata: MarketMetadata,
|
||||
) -> None:
|
||||
"""Test analyze detects niche market trade."""
|
||||
mock_metadata_sync.get_market.return_value = sample_metadata
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Low volume market (science category with volume unknown)
|
||||
signal = await detector.analyze(sample_trade)
|
||||
|
||||
assert signal is not None
|
||||
assert signal.is_niche_market is True
|
||||
assert signal.confidence == 0.2 # niche_base
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_no_anomaly(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
sample_token: Token,
|
||||
) -> None:
|
||||
"""Test analyze returns None for normal trade."""
|
||||
# Politics category is not niche
|
||||
metadata = MarketMetadata(
|
||||
condition_id="market_politics",
|
||||
question="Will Biden win?",
|
||||
description="",
|
||||
tokens=(sample_token,),
|
||||
category="politics",
|
||||
)
|
||||
mock_metadata_sync.get_market.return_value = metadata
|
||||
|
||||
trade = TradeEvent(
|
||||
market_id="market_politics",
|
||||
trade_id="tx_normal",
|
||||
wallet_address="0xnormal",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.50"),
|
||||
size=Decimal("100"), # Small trade = $50 notional
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_pol",
|
||||
)
|
||||
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# High volume, large book depth = low impact
|
||||
signal = await detector.analyze(
|
||||
trade,
|
||||
daily_volume=Decimal("1000000"),
|
||||
book_depth=Decimal("500000"),
|
||||
)
|
||||
|
||||
assert signal is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_creates_minimal_metadata_on_missing(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
) -> None:
|
||||
"""Test analyze creates minimal metadata when market not found."""
|
||||
mock_metadata_sync.get_market.return_value = None
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Should still work with minimal metadata (category="other" which is niche)
|
||||
signal = await detector.analyze(sample_trade)
|
||||
|
||||
assert signal is not None
|
||||
assert signal.market_metadata.condition_id == sample_trade.market_id
|
||||
assert signal.market_metadata.category == "other"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_handles_metadata_exception(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
sample_trade: TradeEvent,
|
||||
) -> None:
|
||||
"""Test analyze handles exception when fetching metadata."""
|
||||
mock_metadata_sync.get_market.side_effect = Exception("Redis error")
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# Should still work with minimal metadata
|
||||
signal = await detector.analyze(sample_trade)
|
||||
|
||||
assert signal is not None
|
||||
assert signal.market_metadata.category == "other"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_low_confidence_filtered(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
sample_token: Token,
|
||||
) -> None:
|
||||
"""Test analyze returns None when confidence is below 0.1."""
|
||||
# Use a mainstream category with below-threshold impacts
|
||||
metadata = MarketMetadata(
|
||||
condition_id="market_sports",
|
||||
question="Super Bowl winner?",
|
||||
description="",
|
||||
tokens=(sample_token,),
|
||||
category="sports",
|
||||
)
|
||||
mock_metadata_sync.get_market.return_value = metadata
|
||||
|
||||
trade = TradeEvent(
|
||||
market_id="market_sports",
|
||||
trade_id="tx_small",
|
||||
wallet_address="0xsmall",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.50"),
|
||||
size=Decimal("10"), # Tiny trade
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_sports",
|
||||
)
|
||||
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# High volume but below threshold impacts
|
||||
signal = await detector.analyze(
|
||||
trade,
|
||||
daily_volume=Decimal("10000000"), # $10M volume
|
||||
book_depth=Decimal("1000000"), # $1M depth
|
||||
)
|
||||
|
||||
assert signal is None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Batch Analysis Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestBatchAnalysis:
|
||||
"""Tests for batch analysis."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_batch_returns_signals(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
sample_metadata: MarketMetadata,
|
||||
) -> None:
|
||||
"""Test batch analysis returns signals for anomalous trades."""
|
||||
mock_metadata_sync.get_market.return_value = sample_metadata
|
||||
|
||||
trades = [
|
||||
TradeEvent(
|
||||
market_id="market_abc123",
|
||||
trade_id=f"tx_{i}",
|
||||
wallet_address=f"0xwallet{i}",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.50"),
|
||||
size=Decimal("10000"), # Large trade
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_123",
|
||||
)
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
signals = await detector.analyze_batch(trades)
|
||||
|
||||
# All trades are in niche category with unknown volume
|
||||
assert len(signals) == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_batch_with_volume_data(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
sample_metadata: MarketMetadata,
|
||||
) -> None:
|
||||
"""Test batch analysis uses provided volume data."""
|
||||
mock_metadata_sync.get_market.return_value = sample_metadata
|
||||
|
||||
trades = [
|
||||
TradeEvent(
|
||||
market_id="market_abc123",
|
||||
trade_id="tx_1",
|
||||
wallet_address="0xwallet1",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.50"),
|
||||
size=Decimal("10000"), # $5000 notional
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_123",
|
||||
)
|
||||
]
|
||||
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
|
||||
# $5000 trade / $50000 volume = 10% impact
|
||||
signals = await detector.analyze_batch(
|
||||
trades,
|
||||
volume_data={"market_abc123": Decimal("50000")},
|
||||
)
|
||||
|
||||
assert len(signals) == 1
|
||||
assert signals[0].volume_impact == pytest.approx(0.10)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_batch_handles_errors(
|
||||
self,
|
||||
mock_metadata_sync: AsyncMock,
|
||||
) -> None:
|
||||
"""Test batch analysis handles individual trade errors."""
|
||||
# First call succeeds, second fails
|
||||
mock_metadata_sync.get_market.side_effect = [
|
||||
Exception("Error"),
|
||||
None,
|
||||
]
|
||||
|
||||
trades = [
|
||||
TradeEvent(
|
||||
market_id=f"market_{i}",
|
||||
trade_id=f"tx_{i}",
|
||||
wallet_address=f"0xwallet{i}",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.50"),
|
||||
size=Decimal("10000"),
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token_123",
|
||||
)
|
||||
for i in range(2)
|
||||
]
|
||||
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
signals = await detector.analyze_batch(trades)
|
||||
|
||||
# Both should still produce signals (with minimal metadata fallback)
|
||||
assert len(signals) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_batch_empty_list(self, mock_metadata_sync: AsyncMock) -> None:
|
||||
"""Test batch analysis with empty list."""
|
||||
detector = SizeAnomalyDetector(mock_metadata_sync)
|
||||
signals = await detector.analyze_batch([])
|
||||
|
||||
assert signals == []
|
||||
@@ -0,0 +1,599 @@
|
||||
"""Tests for the SniperDetector module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from polymarket_insider_tracker.detector.models import SniperClusterSignal
|
||||
from polymarket_insider_tracker.detector.sniper import (
|
||||
ClusterInfo,
|
||||
MarketEntry,
|
||||
SniperDetector,
|
||||
)
|
||||
|
||||
|
||||
def create_mock_trade(
|
||||
wallet_address: str,
|
||||
market_id: str,
|
||||
timestamp: datetime,
|
||||
notional_value: Decimal = Decimal("1000"),
|
||||
) -> MagicMock:
|
||||
"""Create a mock TradeEvent for testing."""
|
||||
trade = MagicMock()
|
||||
trade.wallet_address = wallet_address
|
||||
trade.market_id = market_id
|
||||
trade.timestamp = timestamp
|
||||
trade.notional_value = notional_value
|
||||
return trade
|
||||
|
||||
|
||||
class TestSniperDetectorInit:
|
||||
"""Tests for SniperDetector initialization."""
|
||||
|
||||
def test_default_parameters(self) -> None:
|
||||
"""Test initialization with default parameters."""
|
||||
detector = SniperDetector()
|
||||
|
||||
assert detector.entry_threshold_seconds == 300
|
||||
assert detector.min_cluster_size == 3
|
||||
assert detector.eps == 0.5
|
||||
assert detector.min_samples == 2
|
||||
assert detector.min_entries_per_wallet == 2
|
||||
|
||||
def test_custom_parameters(self) -> None:
|
||||
"""Test initialization with custom parameters."""
|
||||
detector = SniperDetector(
|
||||
entry_threshold_seconds=600,
|
||||
min_cluster_size=5,
|
||||
eps=0.3,
|
||||
min_samples=3,
|
||||
min_entries_per_wallet=4,
|
||||
)
|
||||
|
||||
assert detector.entry_threshold_seconds == 600
|
||||
assert detector.min_cluster_size == 5
|
||||
assert detector.eps == 0.3
|
||||
assert detector.min_samples == 3
|
||||
assert detector.min_entries_per_wallet == 4
|
||||
|
||||
def test_empty_initial_state(self) -> None:
|
||||
"""Test that detector starts with empty state."""
|
||||
detector = SniperDetector()
|
||||
|
||||
assert detector.get_entry_count() == 0
|
||||
assert detector.get_wallet_count() == 0
|
||||
assert detector.get_cluster_count() == 0
|
||||
|
||||
|
||||
class TestRecordEntry:
|
||||
"""Tests for record_entry method."""
|
||||
|
||||
def test_records_entry_within_threshold(self) -> None:
|
||||
"""Test that entries within threshold are recorded."""
|
||||
detector = SniperDetector(entry_threshold_seconds=300)
|
||||
market_created = datetime.now(UTC)
|
||||
trade_time = market_created + timedelta(seconds=60)
|
||||
|
||||
trade = create_mock_trade(
|
||||
wallet_address="0x1111111111111111111111111111111111111111",
|
||||
market_id="market_001",
|
||||
timestamp=trade_time,
|
||||
)
|
||||
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
assert detector.get_entry_count() == 1
|
||||
assert detector.get_wallet_count() == 1
|
||||
|
||||
def test_ignores_entry_after_threshold(self) -> None:
|
||||
"""Test that entries after threshold are ignored."""
|
||||
detector = SniperDetector(entry_threshold_seconds=300)
|
||||
market_created = datetime.now(UTC)
|
||||
trade_time = market_created + timedelta(seconds=400) # 400s > 300s threshold
|
||||
|
||||
trade = create_mock_trade(
|
||||
wallet_address="0x1111111111111111111111111111111111111111",
|
||||
market_id="market_001",
|
||||
timestamp=trade_time,
|
||||
)
|
||||
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
assert detector.get_entry_count() == 0
|
||||
|
||||
def test_ignores_entry_before_market_creation(self) -> None:
|
||||
"""Test that entries before market creation are ignored."""
|
||||
detector = SniperDetector()
|
||||
market_created = datetime.now(UTC)
|
||||
trade_time = market_created - timedelta(seconds=60) # Before creation
|
||||
|
||||
trade = create_mock_trade(
|
||||
wallet_address="0x1111111111111111111111111111111111111111",
|
||||
market_id="market_001",
|
||||
timestamp=trade_time,
|
||||
)
|
||||
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
assert detector.get_entry_count() == 0
|
||||
|
||||
def test_tracks_multiple_wallets(self) -> None:
|
||||
"""Test tracking entries from multiple wallets."""
|
||||
detector = SniperDetector()
|
||||
market_created = datetime.now(UTC)
|
||||
|
||||
for i in range(5):
|
||||
trade = create_mock_trade(
|
||||
wallet_address=f"0x{i:040x}",
|
||||
market_id="market_001",
|
||||
timestamp=market_created + timedelta(seconds=30 * i),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
assert detector.get_entry_count() == 5
|
||||
assert detector.get_wallet_count() == 5
|
||||
|
||||
def test_tracks_wallet_across_markets(self) -> None:
|
||||
"""Test tracking one wallet across multiple markets."""
|
||||
detector = SniperDetector()
|
||||
wallet = "0x1111111111111111111111111111111111111111"
|
||||
|
||||
for i in range(3):
|
||||
market_created = datetime.now(UTC)
|
||||
trade = create_mock_trade(
|
||||
wallet_address=wallet,
|
||||
market_id=f"market_{i:03d}",
|
||||
timestamp=market_created + timedelta(seconds=60),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
assert detector.get_entry_count() == 3
|
||||
assert detector.get_wallet_count() == 1
|
||||
|
||||
|
||||
class TestRunClustering:
|
||||
"""Tests for run_clustering method."""
|
||||
|
||||
def test_returns_empty_with_insufficient_wallets(self) -> None:
|
||||
"""Test that clustering returns empty with too few wallets."""
|
||||
detector = SniperDetector(min_cluster_size=3)
|
||||
|
||||
# Add entries for only 2 wallets
|
||||
for i in range(2):
|
||||
market_created = datetime.now(UTC)
|
||||
for j in range(3): # Multiple entries per wallet
|
||||
trade = create_mock_trade(
|
||||
wallet_address=f"0x{i:040x}",
|
||||
market_id=f"market_{j:03d}",
|
||||
timestamp=market_created + timedelta(seconds=30),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
signals = detector.run_clustering()
|
||||
assert signals == []
|
||||
|
||||
def test_returns_empty_with_insufficient_entries_per_wallet(self) -> None:
|
||||
"""Test that wallets with few entries are excluded."""
|
||||
detector = SniperDetector(
|
||||
min_cluster_size=2,
|
||||
min_entries_per_wallet=3,
|
||||
)
|
||||
|
||||
# Add only 2 entries per wallet
|
||||
for i in range(5):
|
||||
for j in range(2): # Only 2 entries
|
||||
market_created = datetime.now(UTC)
|
||||
trade = create_mock_trade(
|
||||
wallet_address=f"0x{i:040x}",
|
||||
market_id=f"market_{j:03d}",
|
||||
timestamp=market_created + timedelta(seconds=30),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
signals = detector.run_clustering()
|
||||
assert signals == []
|
||||
|
||||
def test_detects_sniper_cluster(self) -> None:
|
||||
"""Test detection of a cluster of snipers with similar patterns."""
|
||||
detector = SniperDetector(
|
||||
min_cluster_size=3,
|
||||
min_entries_per_wallet=2,
|
||||
eps=1.0, # Larger eps for easier clustering
|
||||
min_samples=2,
|
||||
)
|
||||
|
||||
# Create 5 wallets that all enter markets within 30 seconds
|
||||
wallets = [f"0x{i:040x}" for i in range(5)]
|
||||
markets = ["market_001", "market_002", "market_003"]
|
||||
|
||||
for market in markets:
|
||||
market_created = datetime.now(UTC)
|
||||
for wallet in wallets:
|
||||
trade = create_mock_trade(
|
||||
wallet_address=wallet,
|
||||
market_id=market,
|
||||
timestamp=market_created + timedelta(seconds=30),
|
||||
notional_value=Decimal("1000"),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
signals = detector.run_clustering()
|
||||
|
||||
# Should detect at least one cluster
|
||||
assert len(signals) > 0
|
||||
assert all(isinstance(s, SniperClusterSignal) for s in signals)
|
||||
|
||||
def test_does_not_duplicate_signals(self) -> None:
|
||||
"""Test that same wallet isn't signaled twice."""
|
||||
detector = SniperDetector(
|
||||
min_cluster_size=3,
|
||||
min_entries_per_wallet=2,
|
||||
eps=1.0,
|
||||
min_samples=2,
|
||||
)
|
||||
|
||||
# Create cluster
|
||||
wallets = [f"0x{i:040x}" for i in range(5)]
|
||||
markets = ["market_001", "market_002"]
|
||||
|
||||
for market in markets:
|
||||
market_created = datetime.now(UTC)
|
||||
for wallet in wallets:
|
||||
trade = create_mock_trade(
|
||||
wallet_address=wallet,
|
||||
market_id=market,
|
||||
timestamp=market_created + timedelta(seconds=30),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
# Run clustering twice
|
||||
signals1 = detector.run_clustering()
|
||||
signals2 = detector.run_clustering()
|
||||
|
||||
# Second run should return empty (no new signals)
|
||||
assert len(signals2) == 0
|
||||
# All wallets from signals1 should be unique
|
||||
signaled_wallets = [s.wallet_address for s in signals1]
|
||||
assert len(signaled_wallets) == len(set(signaled_wallets))
|
||||
|
||||
|
||||
class TestIsSniper:
|
||||
"""Tests for is_sniper method."""
|
||||
|
||||
def test_returns_false_for_unknown_wallet(self) -> None:
|
||||
"""Test that unknown wallets return False."""
|
||||
detector = SniperDetector()
|
||||
assert detector.is_sniper("0x1111111111111111111111111111111111111111") is False
|
||||
|
||||
def test_returns_true_for_cluster_member(self) -> None:
|
||||
"""Test that cluster members return True."""
|
||||
detector = SniperDetector(
|
||||
min_cluster_size=3,
|
||||
min_entries_per_wallet=2,
|
||||
eps=1.0,
|
||||
min_samples=2,
|
||||
)
|
||||
|
||||
wallets = [f"0x{i:040x}" for i in range(5)]
|
||||
for market in ["market_001", "market_002"]:
|
||||
market_created = datetime.now(UTC)
|
||||
for wallet in wallets:
|
||||
trade = create_mock_trade(
|
||||
wallet_address=wallet,
|
||||
market_id=market,
|
||||
timestamp=market_created + timedelta(seconds=30),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
signals = detector.run_clustering()
|
||||
|
||||
# All signaled wallets should be snipers
|
||||
for signal in signals:
|
||||
assert detector.is_sniper(signal.wallet_address) is True
|
||||
|
||||
|
||||
class TestGetClusterForWallet:
|
||||
"""Tests for get_cluster_for_wallet method."""
|
||||
|
||||
def test_returns_none_for_unknown_wallet(self) -> None:
|
||||
"""Test that unknown wallets return None."""
|
||||
detector = SniperDetector()
|
||||
result = detector.get_cluster_for_wallet("0x1111111111111111111111111111111111111111")
|
||||
assert result is None
|
||||
|
||||
def test_returns_cluster_info_for_member(self) -> None:
|
||||
"""Test that cluster members get ClusterInfo."""
|
||||
detector = SniperDetector(
|
||||
min_cluster_size=3,
|
||||
min_entries_per_wallet=2,
|
||||
eps=1.0,
|
||||
min_samples=2,
|
||||
)
|
||||
|
||||
wallets = [f"0x{i:040x}" for i in range(5)]
|
||||
for market in ["market_001", "market_002"]:
|
||||
market_created = datetime.now(UTC)
|
||||
for wallet in wallets:
|
||||
trade = create_mock_trade(
|
||||
wallet_address=wallet,
|
||||
market_id=market,
|
||||
timestamp=market_created + timedelta(seconds=30),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
signals = detector.run_clustering()
|
||||
|
||||
if signals:
|
||||
wallet = signals[0].wallet_address
|
||||
cluster = detector.get_cluster_for_wallet(wallet)
|
||||
assert cluster is not None
|
||||
assert isinstance(cluster, ClusterInfo)
|
||||
assert wallet in cluster.wallet_addresses
|
||||
|
||||
|
||||
class TestClearEntries:
|
||||
"""Tests for clear_entries method."""
|
||||
|
||||
def test_clears_all_entries(self) -> None:
|
||||
"""Test that clear removes all entries."""
|
||||
detector = SniperDetector()
|
||||
market_created = datetime.now(UTC)
|
||||
|
||||
for i in range(5):
|
||||
trade = create_mock_trade(
|
||||
wallet_address=f"0x{i:040x}",
|
||||
market_id="market_001",
|
||||
timestamp=market_created + timedelta(seconds=30),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
assert detector.get_entry_count() == 5
|
||||
|
||||
detector.clear_entries()
|
||||
|
||||
assert detector.get_entry_count() == 0
|
||||
assert detector.get_wallet_count() == 0
|
||||
|
||||
|
||||
class TestMarketEntryDataclass:
|
||||
"""Tests for MarketEntry dataclass."""
|
||||
|
||||
def test_creation(self) -> None:
|
||||
"""Test MarketEntry creation."""
|
||||
entry = MarketEntry(
|
||||
wallet_address="0x1111",
|
||||
market_id="market_001",
|
||||
entry_delta_seconds=30.5,
|
||||
position_size=Decimal("1000"),
|
||||
timestamp=datetime.now(UTC),
|
||||
)
|
||||
|
||||
assert entry.wallet_address == "0x1111"
|
||||
assert entry.market_id == "market_001"
|
||||
assert entry.entry_delta_seconds == 30.5
|
||||
assert entry.position_size == Decimal("1000")
|
||||
|
||||
|
||||
class TestClusterInfoDataclass:
|
||||
"""Tests for ClusterInfo dataclass."""
|
||||
|
||||
def test_creation(self) -> None:
|
||||
"""Test ClusterInfo creation."""
|
||||
wallets = {"0x1111", "0x2222", "0x3333"}
|
||||
cluster = ClusterInfo(
|
||||
cluster_id="cluster_001",
|
||||
wallet_addresses=wallets,
|
||||
avg_entry_delta=45.0,
|
||||
markets_in_common=3,
|
||||
)
|
||||
|
||||
assert cluster.cluster_id == "cluster_001"
|
||||
assert cluster.wallet_addresses == wallets
|
||||
assert cluster.avg_entry_delta == 45.0
|
||||
assert cluster.markets_in_common == 3
|
||||
assert cluster.created_at is not None
|
||||
|
||||
|
||||
class TestSniperClusterSignalModel:
|
||||
"""Tests for SniperClusterSignal dataclass."""
|
||||
|
||||
def test_creation(self) -> None:
|
||||
"""Test SniperClusterSignal creation."""
|
||||
signal = SniperClusterSignal(
|
||||
wallet_address="0x1111",
|
||||
cluster_id="cluster_001",
|
||||
cluster_size=5,
|
||||
avg_entry_delta_seconds=30.0,
|
||||
markets_in_common=3,
|
||||
confidence=0.85,
|
||||
)
|
||||
|
||||
assert signal.wallet_address == "0x1111"
|
||||
assert signal.cluster_id == "cluster_001"
|
||||
assert signal.cluster_size == 5
|
||||
assert signal.avg_entry_delta_seconds == 30.0
|
||||
assert signal.markets_in_common == 3
|
||||
assert signal.confidence == 0.85
|
||||
|
||||
def test_is_high_confidence(self) -> None:
|
||||
"""Test is_high_confidence property."""
|
||||
high = SniperClusterSignal(
|
||||
wallet_address="0x1111",
|
||||
cluster_id="c1",
|
||||
cluster_size=5,
|
||||
avg_entry_delta_seconds=30.0,
|
||||
markets_in_common=3,
|
||||
confidence=0.75,
|
||||
)
|
||||
low = SniperClusterSignal(
|
||||
wallet_address="0x2222",
|
||||
cluster_id="c2",
|
||||
cluster_size=5,
|
||||
avg_entry_delta_seconds=30.0,
|
||||
markets_in_common=3,
|
||||
confidence=0.65,
|
||||
)
|
||||
|
||||
assert high.is_high_confidence is True
|
||||
assert low.is_high_confidence is False
|
||||
|
||||
def test_is_very_high_confidence(self) -> None:
|
||||
"""Test is_very_high_confidence property."""
|
||||
very_high = SniperClusterSignal(
|
||||
wallet_address="0x1111",
|
||||
cluster_id="c1",
|
||||
cluster_size=5,
|
||||
avg_entry_delta_seconds=30.0,
|
||||
markets_in_common=3,
|
||||
confidence=0.90,
|
||||
)
|
||||
high = SniperClusterSignal(
|
||||
wallet_address="0x2222",
|
||||
cluster_id="c2",
|
||||
cluster_size=5,
|
||||
avg_entry_delta_seconds=30.0,
|
||||
markets_in_common=3,
|
||||
confidence=0.80,
|
||||
)
|
||||
|
||||
assert very_high.is_very_high_confidence is True
|
||||
assert high.is_very_high_confidence is False
|
||||
|
||||
def test_to_dict(self) -> None:
|
||||
"""Test to_dict serialization."""
|
||||
signal = SniperClusterSignal(
|
||||
wallet_address="0x1111",
|
||||
cluster_id="cluster_001",
|
||||
cluster_size=5,
|
||||
avg_entry_delta_seconds=30.0,
|
||||
markets_in_common=3,
|
||||
confidence=0.85,
|
||||
)
|
||||
|
||||
result = signal.to_dict()
|
||||
|
||||
assert result["wallet_address"] == "0x1111"
|
||||
assert result["cluster_id"] == "cluster_001"
|
||||
assert result["cluster_size"] == 5
|
||||
assert result["avg_entry_delta_seconds"] == 30.0
|
||||
assert result["markets_in_common"] == 3
|
||||
assert result["confidence"] == 0.85
|
||||
assert "timestamp" in result
|
||||
|
||||
|
||||
class TestConfidenceCalculation:
|
||||
"""Tests for confidence calculation logic."""
|
||||
|
||||
def test_higher_confidence_with_larger_cluster(self) -> None:
|
||||
"""Test that larger clusters get higher confidence."""
|
||||
detector = SniperDetector(
|
||||
min_cluster_size=2,
|
||||
min_entries_per_wallet=2,
|
||||
eps=1.0,
|
||||
min_samples=2,
|
||||
)
|
||||
|
||||
# Create a cluster
|
||||
stats = {
|
||||
"avg_delta": 30.0,
|
||||
"markets_in_common": 3,
|
||||
}
|
||||
|
||||
small_cluster = {"0x1", "0x2", "0x3"}
|
||||
large_cluster = {"0x1", "0x2", "0x3", "0x4", "0x5", "0x6", "0x7", "0x8"}
|
||||
|
||||
small_conf = detector._calculate_confidence(small_cluster, stats)
|
||||
large_conf = detector._calculate_confidence(large_cluster, stats)
|
||||
|
||||
assert large_conf > small_conf
|
||||
|
||||
def test_higher_confidence_with_faster_entries(self) -> None:
|
||||
"""Test that faster entries get higher confidence."""
|
||||
detector = SniperDetector()
|
||||
|
||||
cluster = {"0x1", "0x2", "0x3"}
|
||||
|
||||
fast_stats = {"avg_delta": 10.0, "markets_in_common": 3}
|
||||
slow_stats = {"avg_delta": 200.0, "markets_in_common": 3}
|
||||
|
||||
fast_conf = detector._calculate_confidence(cluster, fast_stats)
|
||||
slow_conf = detector._calculate_confidence(cluster, slow_stats)
|
||||
|
||||
assert fast_conf > slow_conf
|
||||
|
||||
def test_higher_confidence_with_more_overlap(self) -> None:
|
||||
"""Test that more market overlap gets higher confidence."""
|
||||
detector = SniperDetector()
|
||||
|
||||
cluster = {"0x1", "0x2", "0x3"}
|
||||
|
||||
high_overlap = {"avg_delta": 30.0, "markets_in_common": 5}
|
||||
low_overlap = {"avg_delta": 30.0, "markets_in_common": 1}
|
||||
|
||||
high_conf = detector._calculate_confidence(cluster, high_overlap)
|
||||
low_conf = detector._calculate_confidence(cluster, low_overlap)
|
||||
|
||||
assert high_conf > low_conf
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests for the sniper detection workflow."""
|
||||
|
||||
def test_end_to_end_sniper_detection(self) -> None:
|
||||
"""Test complete workflow from entry recording to signal generation."""
|
||||
detector = SniperDetector(
|
||||
entry_threshold_seconds=300,
|
||||
min_cluster_size=3,
|
||||
min_entries_per_wallet=2,
|
||||
eps=1.0,
|
||||
min_samples=2,
|
||||
)
|
||||
|
||||
# Simulate 5 snipers hitting 3 markets in rapid succession
|
||||
sniper_wallets = [f"0x{'a' * 38}{i:02d}" for i in range(5)]
|
||||
markets = ["market_001", "market_002", "market_003"]
|
||||
|
||||
for market in markets:
|
||||
market_created = datetime.now(UTC)
|
||||
|
||||
for i, wallet in enumerate(sniper_wallets):
|
||||
# Each sniper enters within 10-60 seconds
|
||||
trade = create_mock_trade(
|
||||
wallet_address=wallet,
|
||||
market_id=market,
|
||||
timestamp=market_created + timedelta(seconds=10 + i * 10),
|
||||
notional_value=Decimal("5000"),
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
# Also add some normal traders (enter after threshold)
|
||||
for market in markets:
|
||||
market_created = datetime.now(UTC)
|
||||
for i in range(3):
|
||||
trade = create_mock_trade(
|
||||
wallet_address=f"0x{'b' * 38}{i:02d}",
|
||||
market_id=market,
|
||||
timestamp=market_created + timedelta(seconds=400), # After threshold
|
||||
)
|
||||
detector.record_entry(trade, market_created)
|
||||
|
||||
# Run clustering
|
||||
signals = detector.run_clustering()
|
||||
|
||||
# Should detect the sniper cluster
|
||||
assert len(signals) > 0
|
||||
|
||||
# All signals should be from sniper wallets
|
||||
for signal in signals:
|
||||
assert signal.wallet_address in [w.lower() for w in sniper_wallets]
|
||||
assert signal.cluster_size >= 3
|
||||
assert signal.confidence > 0
|
||||
|
||||
# Verify snipers are marked
|
||||
for wallet in sniper_wallets:
|
||||
# May or may not be in cluster depending on clustering
|
||||
if detector.is_sniper(wallet.lower()):
|
||||
cluster = detector.get_cluster_for_wallet(wallet)
|
||||
assert cluster is not None
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for ingestor module."""
|
||||
@@ -0,0 +1,360 @@
|
||||
"""Tests for ClobClient wrapper."""
|
||||
|
||||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.ingestor.clob_client import (
|
||||
ClobClient,
|
||||
RateLimiter,
|
||||
RetryError,
|
||||
with_retry,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.models import Market, Orderbook
|
||||
|
||||
|
||||
class TestRateLimiter:
|
||||
"""Tests for RateLimiter."""
|
||||
|
||||
def test_acquire_sync_no_wait_first_call(self) -> None:
|
||||
"""First call should not wait."""
|
||||
limiter = RateLimiter(max_requests_per_second=10)
|
||||
start = time.monotonic()
|
||||
limiter.acquire_sync()
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
# Should be nearly instant
|
||||
assert elapsed < 0.05
|
||||
|
||||
def test_acquire_sync_enforces_rate(self) -> None:
|
||||
"""Subsequent calls should be rate limited."""
|
||||
limiter = RateLimiter(max_requests_per_second=10) # 100ms between calls
|
||||
|
||||
# First call
|
||||
limiter.acquire_sync()
|
||||
|
||||
# Second call should wait
|
||||
start = time.monotonic()
|
||||
limiter.acquire_sync()
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
# Should wait at least 90ms (allowing some tolerance)
|
||||
assert elapsed >= 0.08
|
||||
|
||||
|
||||
class TestWithRetry:
|
||||
"""Tests for retry decorator."""
|
||||
|
||||
def test_success_first_try(self) -> None:
|
||||
"""Function succeeds on first try."""
|
||||
call_count = 0
|
||||
|
||||
@with_retry(max_retries=3)
|
||||
def succeed() -> str:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return "success"
|
||||
|
||||
result = succeed()
|
||||
|
||||
assert result == "success"
|
||||
assert call_count == 1
|
||||
|
||||
def test_success_after_retries(self) -> None:
|
||||
"""Function succeeds after some retries."""
|
||||
call_count = 0
|
||||
|
||||
@with_retry(max_retries=3, base_delay=0.01)
|
||||
def succeed_eventually() -> str:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count < 3:
|
||||
raise ValueError("Not yet")
|
||||
return "success"
|
||||
|
||||
result = succeed_eventually()
|
||||
|
||||
assert result == "success"
|
||||
assert call_count == 3
|
||||
|
||||
def test_exhausted_retries(self) -> None:
|
||||
"""Raises RetryError after exhausting retries."""
|
||||
|
||||
@with_retry(max_retries=2, base_delay=0.01)
|
||||
def always_fails() -> str:
|
||||
raise ValueError("Always fails")
|
||||
|
||||
with pytest.raises(RetryError) as exc_info:
|
||||
always_fails()
|
||||
|
||||
assert "3 attempts failed" in str(exc_info.value)
|
||||
assert isinstance(exc_info.value.last_exception, ValueError)
|
||||
|
||||
def test_specific_exception_types(self) -> None:
|
||||
"""Only retries on specified exception types."""
|
||||
call_count = 0
|
||||
|
||||
@with_retry(max_retries=3, base_delay=0.01, retry_on=(ValueError,))
|
||||
def raise_type_error() -> str:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
raise TypeError("Not retried")
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
raise_type_error()
|
||||
|
||||
# Should only be called once since TypeError is not in retry_on
|
||||
assert call_count == 1
|
||||
|
||||
|
||||
class TestClobClient:
|
||||
"""Tests for ClobClient wrapper."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_base_client(self) -> MagicMock:
|
||||
"""Create a mock base CLOB client."""
|
||||
with patch("polymarket_insider_tracker.ingestor.clob_client.BaseClobClient") as mock:
|
||||
yield mock.return_value
|
||||
|
||||
def test_init_defaults(self, mock_base_client: MagicMock) -> None: # noqa: ARG002
|
||||
"""Test client initialization with defaults."""
|
||||
client = ClobClient()
|
||||
|
||||
assert client._host == "https://clob.polymarket.com"
|
||||
assert client._max_retries == 3
|
||||
|
||||
def test_init_with_env_api_key(self, mock_base_client: MagicMock) -> None: # noqa: ARG002
|
||||
"""Test client reads API key from environment."""
|
||||
with patch.dict("os.environ", {"POLYMARKET_API_KEY": "test-key"}):
|
||||
client = ClobClient()
|
||||
assert client._api_key == "test-key"
|
||||
|
||||
def test_init_with_explicit_api_key(self, mock_base_client: MagicMock) -> None: # noqa: ARG002
|
||||
"""Test client uses explicitly provided API key."""
|
||||
client = ClobClient(api_key="explicit-key")
|
||||
assert client._api_key == "explicit-key"
|
||||
|
||||
def test_health_check_success(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test health check returns True when API responds OK."""
|
||||
mock_base_client.get_ok.return_value = "OK"
|
||||
|
||||
client = ClobClient()
|
||||
result = client.health_check()
|
||||
|
||||
assert result is True
|
||||
mock_base_client.get_ok.assert_called_once()
|
||||
|
||||
def test_health_check_failure(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test health check returns False on error."""
|
||||
mock_base_client.get_ok.side_effect = Exception("Connection failed")
|
||||
|
||||
client = ClobClient()
|
||||
result = client.health_check()
|
||||
|
||||
assert result is False
|
||||
|
||||
def test_get_server_time(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test getting server time."""
|
||||
mock_base_client.get_server_time.return_value = 1704067200000
|
||||
|
||||
client = ClobClient()
|
||||
result = client.get_server_time()
|
||||
|
||||
assert result == 1704067200000
|
||||
|
||||
def test_get_markets(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test fetching markets."""
|
||||
mock_base_client.get_simplified_markets.return_value = {
|
||||
"data": [
|
||||
{
|
||||
"condition_id": "0x123",
|
||||
"question": "Test market?",
|
||||
"tokens": [],
|
||||
"closed": False,
|
||||
},
|
||||
],
|
||||
"next_cursor": "LTE=",
|
||||
}
|
||||
|
||||
client = ClobClient()
|
||||
markets = client.get_markets()
|
||||
|
||||
assert len(markets) == 1
|
||||
assert isinstance(markets[0], Market)
|
||||
assert markets[0].condition_id == "0x123"
|
||||
|
||||
def test_get_markets_filters_closed(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test that closed markets are filtered when active_only=True."""
|
||||
mock_base_client.get_simplified_markets.return_value = {
|
||||
"data": [
|
||||
{"condition_id": "0x1", "closed": False},
|
||||
{"condition_id": "0x2", "closed": True},
|
||||
],
|
||||
"next_cursor": "LTE=",
|
||||
}
|
||||
|
||||
client = ClobClient()
|
||||
markets = client.get_markets(active_only=True)
|
||||
|
||||
assert len(markets) == 1
|
||||
assert markets[0].condition_id == "0x1"
|
||||
|
||||
def test_get_markets_includes_closed(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test that closed markets are included when active_only=False."""
|
||||
mock_base_client.get_simplified_markets.return_value = {
|
||||
"data": [
|
||||
{"condition_id": "0x1", "closed": False},
|
||||
{"condition_id": "0x2", "closed": True},
|
||||
],
|
||||
"next_cursor": "LTE=",
|
||||
}
|
||||
|
||||
client = ClobClient()
|
||||
markets = client.get_markets(active_only=False)
|
||||
|
||||
assert len(markets) == 2
|
||||
|
||||
def test_get_markets_pagination(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test that pagination is handled correctly."""
|
||||
mock_base_client.get_simplified_markets.side_effect = [
|
||||
{
|
||||
"data": [{"condition_id": "0x1"}],
|
||||
"next_cursor": "cursor2",
|
||||
},
|
||||
{
|
||||
"data": [{"condition_id": "0x2"}],
|
||||
"next_cursor": "LTE=",
|
||||
},
|
||||
]
|
||||
|
||||
client = ClobClient()
|
||||
markets = client.get_markets()
|
||||
|
||||
assert len(markets) == 2
|
||||
assert mock_base_client.get_simplified_markets.call_count == 2
|
||||
|
||||
def test_get_market(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test fetching a single market."""
|
||||
mock_base_client.get_market.return_value = {
|
||||
"condition_id": "0xabc",
|
||||
"question": "Will it happen?",
|
||||
"tokens": [
|
||||
{"token_id": "t1", "outcome": "Yes"},
|
||||
{"token_id": "t2", "outcome": "No"},
|
||||
],
|
||||
}
|
||||
|
||||
client = ClobClient()
|
||||
market = client.get_market("0xabc")
|
||||
|
||||
assert isinstance(market, Market)
|
||||
assert market.condition_id == "0xabc"
|
||||
assert len(market.tokens) == 2
|
||||
|
||||
def test_get_market_not_found(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test error handling when market not found.
|
||||
|
||||
When the underlying API call fails, the @with_retry decorator will
|
||||
retry the operation. After all retries are exhausted, it raises
|
||||
RetryError wrapping the original exception.
|
||||
"""
|
||||
mock_base_client.get_market.side_effect = Exception("Not found")
|
||||
|
||||
client = ClobClient()
|
||||
|
||||
with pytest.raises(RetryError) as exc_info:
|
||||
client.get_market("0xnotfound")
|
||||
|
||||
# The RetryError wraps the original exception
|
||||
assert "get_market" in str(exc_info.value)
|
||||
assert exc_info.value.last_exception is not None
|
||||
|
||||
def test_get_orderbook(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test fetching an orderbook."""
|
||||
mock_bid = MagicMock()
|
||||
mock_bid.price = "0.50"
|
||||
mock_bid.size = "100"
|
||||
|
||||
mock_ask = MagicMock()
|
||||
mock_ask.price = "0.52"
|
||||
mock_ask.size = "150"
|
||||
|
||||
mock_orderbook = MagicMock()
|
||||
mock_orderbook.market = "0xmarket"
|
||||
mock_orderbook.asset_id = "token123"
|
||||
mock_orderbook.tick_size = "0.01"
|
||||
mock_orderbook.bids = [mock_bid]
|
||||
mock_orderbook.asks = [mock_ask]
|
||||
|
||||
mock_base_client.get_order_book.return_value = mock_orderbook
|
||||
|
||||
client = ClobClient()
|
||||
orderbook = client.get_orderbook("token123")
|
||||
|
||||
assert isinstance(orderbook, Orderbook)
|
||||
assert orderbook.asset_id == "token123"
|
||||
assert len(orderbook.bids) == 1
|
||||
assert len(orderbook.asks) == 1
|
||||
|
||||
def test_get_orderbooks(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test fetching multiple orderbooks."""
|
||||
mock_ob1 = MagicMock()
|
||||
mock_ob1.market = "m1"
|
||||
mock_ob1.asset_id = "t1"
|
||||
mock_ob1.tick_size = "0.01"
|
||||
mock_ob1.bids = []
|
||||
mock_ob1.asks = []
|
||||
|
||||
mock_ob2 = MagicMock()
|
||||
mock_ob2.market = "m2"
|
||||
mock_ob2.asset_id = "t2"
|
||||
mock_ob2.tick_size = "0.01"
|
||||
mock_ob2.bids = []
|
||||
mock_ob2.asks = []
|
||||
|
||||
mock_base_client.get_order_books.return_value = [mock_ob1, mock_ob2]
|
||||
|
||||
client = ClobClient()
|
||||
orderbooks = client.get_orderbooks(["t1", "t2"])
|
||||
|
||||
assert len(orderbooks) == 2
|
||||
assert all(isinstance(ob, Orderbook) for ob in orderbooks)
|
||||
|
||||
def test_get_midpoint(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test fetching midpoint price."""
|
||||
mock_base_client.get_midpoint.return_value = {"mid": "0.55"}
|
||||
|
||||
client = ClobClient()
|
||||
result = client.get_midpoint("token123")
|
||||
|
||||
assert result == "0.55"
|
||||
|
||||
def test_get_midpoint_error(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test midpoint returns None on error."""
|
||||
mock_base_client.get_midpoint.side_effect = Exception("API error")
|
||||
|
||||
client = ClobClient()
|
||||
result = client.get_midpoint("token123")
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_get_price_buy(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test fetching buy price."""
|
||||
mock_base_client.get_price.return_value = {"price": "0.53"}
|
||||
|
||||
client = ClobClient()
|
||||
result = client.get_price("token123", side="BUY")
|
||||
|
||||
assert result == "0.53"
|
||||
mock_base_client.get_price.assert_called_with("token123", side="BUY")
|
||||
|
||||
def test_get_price_sell(self, mock_base_client: MagicMock) -> None:
|
||||
"""Test fetching sell price."""
|
||||
mock_base_client.get_price.return_value = {"price": "0.51"}
|
||||
|
||||
client = ClobClient()
|
||||
result = client.get_price("token123", side="SELL")
|
||||
|
||||
assert result == "0.51"
|
||||
mock_base_client.get_price.assert_called_with("token123", side="SELL")
|
||||
@@ -0,0 +1,613 @@
|
||||
"""Tests for the connection health monitor."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from polymarket_insider_tracker.ingestor.health import (
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
DEFAULT_STALE_THRESHOLD_SECONDS,
|
||||
HealthMonitor,
|
||||
HealthReport,
|
||||
HealthStatus,
|
||||
StreamHealth,
|
||||
StreamStatus,
|
||||
)
|
||||
|
||||
|
||||
class TestStreamHealth:
|
||||
"""Tests for the StreamHealth dataclass."""
|
||||
|
||||
def test_stream_health_defaults(self) -> None:
|
||||
"""Test default values."""
|
||||
health = StreamHealth(name="test-stream")
|
||||
|
||||
assert health.name == "test-stream"
|
||||
assert health.status == StreamStatus.DISCONNECTED
|
||||
assert health.last_event_time is None
|
||||
assert health.events_received == 0
|
||||
assert health.events_per_second == 0.0
|
||||
assert health.connected_since is None
|
||||
assert health.last_error is None
|
||||
|
||||
def test_stream_health_custom_values(self) -> None:
|
||||
"""Test with custom values."""
|
||||
now = time.time()
|
||||
health = StreamHealth(
|
||||
name="trades",
|
||||
status=StreamStatus.ACTIVE,
|
||||
last_event_time=now,
|
||||
events_received=100,
|
||||
events_per_second=5.0,
|
||||
connected_since=now - 3600,
|
||||
last_error=None,
|
||||
)
|
||||
|
||||
assert health.name == "trades"
|
||||
assert health.status == StreamStatus.ACTIVE
|
||||
assert health.events_received == 100
|
||||
|
||||
|
||||
class TestHealthReport:
|
||||
"""Tests for the HealthReport dataclass."""
|
||||
|
||||
def test_health_report_defaults(self) -> None:
|
||||
"""Test default values."""
|
||||
report = HealthReport(status=HealthStatus.HEALTHY)
|
||||
|
||||
assert report.status == HealthStatus.HEALTHY
|
||||
assert report.streams == {}
|
||||
assert report.total_events_received == 0
|
||||
assert report.total_events_per_second == 0.0
|
||||
assert report.uptime_seconds == 0.0
|
||||
assert report.timestamp > 0
|
||||
|
||||
def test_health_report_with_streams(self) -> None:
|
||||
"""Test with stream data."""
|
||||
stream = StreamHealth(name="trades", events_received=100)
|
||||
report = HealthReport(
|
||||
status=HealthStatus.DEGRADED,
|
||||
streams={"trades": stream},
|
||||
total_events_received=100,
|
||||
total_events_per_second=5.0,
|
||||
uptime_seconds=3600.0,
|
||||
)
|
||||
|
||||
assert report.status == HealthStatus.DEGRADED
|
||||
assert "trades" in report.streams
|
||||
assert report.total_events_received == 100
|
||||
|
||||
|
||||
class TestHealthMonitor:
|
||||
"""Tests for the HealthMonitor class."""
|
||||
|
||||
def test_init(self) -> None:
|
||||
"""Test initialization."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
assert monitor._stale_threshold == DEFAULT_STALE_THRESHOLD_SECONDS
|
||||
assert monitor._health_check_interval == DEFAULT_HEALTH_CHECK_INTERVAL
|
||||
assert not monitor.is_running
|
||||
|
||||
def test_init_custom_config(self) -> None:
|
||||
"""Test initialization with custom config."""
|
||||
monitor = HealthMonitor(
|
||||
stale_threshold_seconds=30,
|
||||
health_check_interval=10,
|
||||
)
|
||||
|
||||
assert monitor._stale_threshold == 30
|
||||
assert monitor._health_check_interval == 10
|
||||
|
||||
def test_register_stream(self) -> None:
|
||||
"""Test registering a stream."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.register_stream("trades")
|
||||
|
||||
assert "trades" in monitor._streams
|
||||
assert monitor._streams["trades"].name == "trades"
|
||||
assert monitor._streams["trades"].status == StreamStatus.DISCONNECTED
|
||||
|
||||
def test_register_stream_idempotent(self) -> None:
|
||||
"""Test that registering the same stream twice is idempotent."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.register_stream("trades")
|
||||
monitor.record_event("trades") # Adds an event
|
||||
monitor.register_stream("trades") # Should not reset
|
||||
|
||||
assert monitor._streams["trades"].events_received == 1
|
||||
|
||||
def test_set_stream_connected(self) -> None:
|
||||
"""Test marking a stream as connected."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.set_stream_connected("trades")
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.ACTIVE
|
||||
assert monitor._streams["trades"].connected_since is not None
|
||||
assert monitor._streams["trades"].last_error is None
|
||||
|
||||
def test_set_stream_disconnected(self) -> None:
|
||||
"""Test marking a stream as disconnected."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.set_stream_connected("trades")
|
||||
monitor.set_stream_disconnected("trades", error="Connection reset")
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.DISCONNECTED
|
||||
assert monitor._streams["trades"].connected_since is None
|
||||
assert monitor._streams["trades"].last_error == "Connection reset"
|
||||
|
||||
def test_record_event(self) -> None:
|
||||
"""Test recording an event."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
|
||||
stream = monitor._streams["trades"]
|
||||
assert stream.events_received == 1
|
||||
assert stream.last_event_time is not None
|
||||
assert stream.status == StreamStatus.ACTIVE
|
||||
|
||||
def test_record_event_multiple(self) -> None:
|
||||
"""Test recording multiple events."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
for _ in range(10):
|
||||
monitor.record_event("trades")
|
||||
|
||||
assert monitor._streams["trades"].events_received == 10
|
||||
|
||||
def test_record_event_with_processing_time(self) -> None:
|
||||
"""Test recording event with processing time."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
# Should not raise
|
||||
monitor.record_event("trades", processing_time=0.001)
|
||||
|
||||
assert monitor._streams["trades"].events_received == 1
|
||||
|
||||
def test_calculate_throughput_empty(self) -> None:
|
||||
"""Test throughput calculation with no events."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
rate = monitor._calculate_throughput("nonexistent")
|
||||
|
||||
assert rate == 0.0
|
||||
|
||||
def test_calculate_throughput(self) -> None:
|
||||
"""Test throughput calculation."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
# Add events
|
||||
for _ in range(10):
|
||||
monitor.record_event("trades")
|
||||
|
||||
rate = monitor._calculate_throughput("trades")
|
||||
|
||||
# Should have ~10 events in the window
|
||||
assert rate > 0
|
||||
|
||||
def test_check_stream_staleness_active(self) -> None:
|
||||
"""Test that active stream is not marked stale."""
|
||||
monitor = HealthMonitor(stale_threshold_seconds=60)
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor._check_stream_staleness()
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.ACTIVE
|
||||
|
||||
def test_check_stream_staleness_stale(self) -> None:
|
||||
"""Test that stream becomes stale after threshold."""
|
||||
monitor = HealthMonitor(stale_threshold_seconds=1)
|
||||
|
||||
monitor.record_event("trades")
|
||||
# Simulate time passing
|
||||
monitor._streams["trades"].last_event_time = time.time() - 2
|
||||
|
||||
monitor._check_stream_staleness()
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.STALE
|
||||
|
||||
def test_check_stream_staleness_connected_no_events(self) -> None:
|
||||
"""Test staleness when connected but no events received."""
|
||||
monitor = HealthMonitor(stale_threshold_seconds=1)
|
||||
|
||||
monitor.set_stream_connected("trades")
|
||||
# Simulate time passing since connection
|
||||
monitor._streams["trades"].connected_since = time.time() - 2
|
||||
|
||||
monitor._check_stream_staleness()
|
||||
|
||||
assert monitor._streams["trades"].status == StreamStatus.STALE
|
||||
|
||||
def test_determine_overall_status_no_streams(self) -> None:
|
||||
"""Test overall status with no streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.HEALTHY
|
||||
|
||||
def test_determine_overall_status_all_active(self) -> None:
|
||||
"""Test overall status with all active streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.record_event("orderbook")
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.HEALTHY
|
||||
|
||||
def test_determine_overall_status_some_stale(self) -> None:
|
||||
"""Test overall status with some stale streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.register_stream("orderbook")
|
||||
monitor._streams["orderbook"].status = StreamStatus.STALE
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.DEGRADED
|
||||
|
||||
def test_determine_overall_status_some_disconnected(self) -> None:
|
||||
"""Test overall status with some disconnected streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.set_stream_disconnected("orderbook")
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.DEGRADED
|
||||
|
||||
def test_determine_overall_status_all_disconnected(self) -> None:
|
||||
"""Test overall status with all disconnected streams."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.set_stream_disconnected("trades")
|
||||
monitor.set_stream_disconnected("orderbook")
|
||||
|
||||
status = monitor._determine_overall_status()
|
||||
|
||||
assert status == HealthStatus.UNHEALTHY
|
||||
|
||||
def test_get_health_report(self) -> None:
|
||||
"""Test getting a health report."""
|
||||
monitor = HealthMonitor()
|
||||
monitor._start_time = time.time() - 100
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.record_event("trades")
|
||||
|
||||
report = monitor.get_health_report()
|
||||
|
||||
assert report.status == HealthStatus.HEALTHY
|
||||
assert "trades" in report.streams
|
||||
assert report.total_events_received == 2
|
||||
assert report.uptime_seconds >= 100
|
||||
assert report.timestamp > 0
|
||||
|
||||
def test_get_health_report_calculates_throughput(self) -> None:
|
||||
"""Test that health report calculates throughput."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
for _ in range(10):
|
||||
monitor.record_event("trades")
|
||||
|
||||
report = monitor.get_health_report()
|
||||
|
||||
assert report.streams["trades"].events_per_second > 0
|
||||
assert report.total_events_per_second > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_stop(self) -> None:
|
||||
"""Test starting and stopping the monitor."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
await monitor.start()
|
||||
assert monitor.is_running
|
||||
assert monitor._health_task is not None
|
||||
|
||||
await monitor.stop()
|
||||
assert not monitor.is_running
|
||||
assert monitor._health_task is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_idempotent(self) -> None:
|
||||
"""Test that starting twice is safe."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
await monitor.start()
|
||||
await monitor.start() # Should not raise
|
||||
|
||||
assert monitor.is_running
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_when_not_running(self) -> None:
|
||||
"""Test that stopping when not running is safe."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
await monitor.stop() # Should not raise
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_manager(self) -> None:
|
||||
"""Test async context manager."""
|
||||
async with HealthMonitor() as monitor:
|
||||
assert monitor.is_running
|
||||
|
||||
assert not monitor.is_running
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_loop_updates_report(self) -> None:
|
||||
"""Test that health check loop updates the report."""
|
||||
monitor = HealthMonitor(health_check_interval=0.1)
|
||||
|
||||
await monitor.start()
|
||||
monitor.record_event("trades")
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
# Health should have been checked
|
||||
report = monitor.get_health_report()
|
||||
assert report.status == HealthStatus.HEALTHY
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_change_callback(self) -> None:
|
||||
"""Test that health change callback is invoked."""
|
||||
callback = AsyncMock()
|
||||
monitor = HealthMonitor(
|
||||
health_check_interval=0.1,
|
||||
on_health_change=callback,
|
||||
)
|
||||
|
||||
await monitor.start()
|
||||
monitor.record_event("trades")
|
||||
|
||||
# Wait for health check
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
# Callback should have been called at least once
|
||||
assert callback.called
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_change_callback_error_handling(self) -> None:
|
||||
"""Test that callback errors don't crash the loop."""
|
||||
callback = AsyncMock(side_effect=ValueError("test error"))
|
||||
monitor = HealthMonitor(
|
||||
health_check_interval=0.1,
|
||||
on_health_change=callback,
|
||||
)
|
||||
|
||||
await monitor.start()
|
||||
monitor.record_event("trades")
|
||||
|
||||
# Should not crash
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
|
||||
class TestHealthMonitorHTTPEndpoints:
|
||||
"""Tests for HTTP endpoints."""
|
||||
|
||||
@pytest.fixture
|
||||
def monitor(self) -> HealthMonitor:
|
||||
"""Create a monitor instance."""
|
||||
return HealthMonitor()
|
||||
|
||||
@pytest.fixture
|
||||
def app(self, monitor: HealthMonitor) -> web.Application:
|
||||
"""Create the aiohttp application."""
|
||||
return monitor._create_app()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_healthy(
|
||||
self, monitor: HealthMonitor, app: web.Application
|
||||
) -> None:
|
||||
"""Test /health endpoint when healthy."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.record_event("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/health")
|
||||
assert resp.status == 200
|
||||
|
||||
data = await resp.json()
|
||||
assert data["status"] == "healthy"
|
||||
assert "trades" in data["streams"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_unhealthy(
|
||||
self, monitor: HealthMonitor, app: web.Application
|
||||
) -> None:
|
||||
"""Test /health endpoint when unhealthy."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.set_stream_disconnected("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/health")
|
||||
assert resp.status == 503
|
||||
|
||||
data = await resp.json()
|
||||
assert data["status"] == "unhealthy"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_endpoint(self, monitor: HealthMonitor, app: web.Application) -> None:
|
||||
"""Test /metrics endpoint returns Prometheus format."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.record_event("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/metrics")
|
||||
assert resp.status == 200
|
||||
|
||||
content_type = resp.headers.get("Content-Type", "")
|
||||
assert "text/plain" in content_type
|
||||
|
||||
text = await resp.text()
|
||||
assert "polymarket_events_total" in text
|
||||
assert "polymarket_health_status" in text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ready_endpoint_ready(self, monitor: HealthMonitor, app: web.Application) -> None:
|
||||
"""Test /ready endpoint when ready."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.record_event("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/ready")
|
||||
assert resp.status == 200
|
||||
|
||||
data = await resp.json()
|
||||
assert data["ready"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ready_endpoint_not_ready(
|
||||
self, monitor: HealthMonitor, app: web.Application
|
||||
) -> None:
|
||||
"""Test /ready endpoint when not ready."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
monitor.set_stream_disconnected("trades")
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/ready")
|
||||
assert resp.status == 503
|
||||
|
||||
data = await resp.json()
|
||||
assert data["ready"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_endpoint(self, app: web.Application) -> None:
|
||||
"""Test /live endpoint always returns 200."""
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
resp = await client.get("/live")
|
||||
assert resp.status == 200
|
||||
|
||||
data = await resp.json()
|
||||
assert data["live"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_stop_http_server(self, monitor: HealthMonitor) -> None:
|
||||
"""Test starting and stopping HTTP server."""
|
||||
await monitor.start_http_server(port=18080)
|
||||
assert monitor._runner is not None
|
||||
|
||||
await monitor.stop_http_server()
|
||||
assert monitor._runner is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_http_server_idempotent(self, monitor: HealthMonitor) -> None:
|
||||
"""Test that starting HTTP server twice is safe."""
|
||||
await monitor.start_http_server(port=18081)
|
||||
await monitor.start_http_server(port=18081) # Should not raise
|
||||
|
||||
await monitor.stop_http_server()
|
||||
|
||||
|
||||
class TestPrometheusMetrics:
|
||||
"""Tests for Prometheus metric updates."""
|
||||
|
||||
def test_events_total_incremented(self) -> None:
|
||||
"""Test that events_total counter is incremented."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("test-metrics")
|
||||
monitor.record_event("test-metrics")
|
||||
|
||||
# Counter should have been incremented
|
||||
# (We can't easily test prometheus metrics directly, but at least verify no errors)
|
||||
|
||||
def test_stream_status_updated(self) -> None:
|
||||
"""Test that stream_status gauge is updated."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.set_stream_connected("test-status")
|
||||
# Gauge should be 1.0
|
||||
|
||||
monitor.set_stream_disconnected("test-status")
|
||||
# Gauge should be 0.0
|
||||
|
||||
def test_health_status_updated(self) -> None:
|
||||
"""Test that health_status gauge is updated."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("test-health")
|
||||
report = monitor.get_health_report()
|
||||
|
||||
assert report.status == HealthStatus.HEALTHY
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Tests for edge cases and error handling."""
|
||||
|
||||
def test_throughput_with_old_events(self) -> None:
|
||||
"""Test throughput calculation ignores old events."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
# Manually add old event to window
|
||||
monitor._event_windows["trades"].append(time.time() - 100)
|
||||
|
||||
rate = monitor._calculate_throughput("trades")
|
||||
|
||||
# Old event should be filtered out
|
||||
# Rate should only count recent events
|
||||
assert rate >= 0
|
||||
|
||||
def test_multiple_streams_independent(self) -> None:
|
||||
"""Test that multiple streams are tracked independently."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
monitor.record_event("trades")
|
||||
monitor.record_event("trades")
|
||||
monitor.set_stream_disconnected("orderbook")
|
||||
|
||||
assert monitor._streams["trades"].events_received == 2
|
||||
assert monitor._streams["trades"].status == StreamStatus.ACTIVE
|
||||
assert monitor._streams["orderbook"].events_received == 0
|
||||
assert monitor._streams["orderbook"].status == StreamStatus.DISCONNECTED
|
||||
|
||||
def test_report_streams_are_copied(self) -> None:
|
||||
"""Test that report streams are a copy."""
|
||||
monitor = HealthMonitor()
|
||||
monitor.record_event("trades")
|
||||
|
||||
report = monitor.get_health_report()
|
||||
|
||||
# Modifying report should not affect monitor
|
||||
report.streams["trades"].events_received = 999
|
||||
assert monitor._streams["trades"].events_received == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_cleans_up_http_server(self) -> None:
|
||||
"""Test that stop() also stops HTTP server."""
|
||||
monitor = HealthMonitor()
|
||||
|
||||
await monitor.start()
|
||||
await monitor.start_http_server(port=18082)
|
||||
|
||||
await monitor.stop()
|
||||
|
||||
assert not monitor.is_running
|
||||
assert monitor._runner is None
|
||||
@@ -0,0 +1,443 @@
|
||||
"""Tests for the market metadata synchronizer."""
|
||||
|
||||
import json
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.ingestor.clob_client import ClobClient
|
||||
from polymarket_insider_tracker.ingestor.metadata_sync import (
|
||||
DEFAULT_CACHE_TTL_SECONDS,
|
||||
DEFAULT_REDIS_KEY_PREFIX,
|
||||
DEFAULT_SYNC_INTERVAL_SECONDS,
|
||||
MarketMetadataSync,
|
||||
MetadataSyncError,
|
||||
SyncState,
|
||||
SyncStats,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.models import (
|
||||
Market,
|
||||
MarketMetadata,
|
||||
Token,
|
||||
derive_category,
|
||||
)
|
||||
|
||||
|
||||
# Test fixtures
|
||||
@pytest.fixture
|
||||
def sample_token() -> Token:
|
||||
"""Create a sample token."""
|
||||
return Token(token_id="token123", outcome="Yes", price=Decimal("0.65"))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_market(sample_token: Token) -> Market:
|
||||
"""Create a sample market."""
|
||||
return Market(
|
||||
condition_id="cond123",
|
||||
question="Will Bitcoin exceed $100k in 2026?",
|
||||
description="Market on BTC price",
|
||||
tokens=(sample_token,),
|
||||
end_date=datetime(2026, 12, 31, tzinfo=UTC),
|
||||
active=True,
|
||||
closed=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_metadata(sample_market: Market) -> MarketMetadata:
|
||||
"""Create sample metadata from market."""
|
||||
return MarketMetadata.from_market(sample_market)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis() -> AsyncMock:
|
||||
"""Create a mock Redis client."""
|
||||
redis = AsyncMock()
|
||||
redis.get = AsyncMock(return_value=None)
|
||||
redis.setex = AsyncMock()
|
||||
redis.delete = AsyncMock(return_value=1)
|
||||
redis.scan = AsyncMock(return_value=(0, []))
|
||||
return redis
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_clob(sample_market: Market) -> MagicMock:
|
||||
"""Create a mock CLOB client."""
|
||||
clob = MagicMock(spec=ClobClient)
|
||||
clob.get_markets = MagicMock(return_value=[sample_market])
|
||||
clob.get_market = MagicMock(return_value=sample_market)
|
||||
return clob
|
||||
|
||||
|
||||
class TestDeriveCategory:
|
||||
"""Tests for the derive_category function."""
|
||||
|
||||
def test_politics_keywords(self) -> None:
|
||||
"""Test political category detection."""
|
||||
assert derive_category("Will Trump win the 2024 election?") == "politics"
|
||||
assert derive_category("Who will be the next president?") == "politics"
|
||||
assert derive_category("Senate majority party after midterms?") == "politics"
|
||||
|
||||
def test_crypto_keywords(self) -> None:
|
||||
"""Test crypto category detection."""
|
||||
assert derive_category("Will Bitcoin hit $100k?") == "crypto"
|
||||
assert derive_category("Ethereum price by end of year?") == "crypto"
|
||||
assert derive_category("Next altcoin to moon?") == "crypto"
|
||||
|
||||
def test_sports_keywords(self) -> None:
|
||||
"""Test sports category detection."""
|
||||
assert derive_category("Who will win the Super Bowl?") == "sports"
|
||||
assert derive_category("NBA Finals champion?") == "sports"
|
||||
assert derive_category("Next UFC heavyweight champion?") == "sports"
|
||||
|
||||
def test_entertainment_keywords(self) -> None:
|
||||
"""Test entertainment category detection."""
|
||||
assert derive_category("Best Picture Oscar winner?") == "entertainment"
|
||||
assert derive_category("Next Grammy Album of the Year?") == "entertainment"
|
||||
assert derive_category("Highest box office movie this summer?") == "entertainment"
|
||||
|
||||
def test_finance_keywords(self) -> None:
|
||||
"""Test finance category detection."""
|
||||
assert derive_category("Fed interest rate decision?") == "finance"
|
||||
assert derive_category("Will we enter a recession?") == "finance"
|
||||
assert derive_category("S&P 500 by year end?") == "finance"
|
||||
|
||||
def test_tech_keywords(self) -> None:
|
||||
"""Test tech category detection."""
|
||||
assert derive_category("Will Apple release a new iPhone?") == "tech"
|
||||
assert derive_category("Next major AI breakthrough?") == "tech"
|
||||
assert derive_category("Tesla vehicle deliveries?") == "tech"
|
||||
|
||||
def test_science_keywords(self) -> None:
|
||||
"""Test science category detection."""
|
||||
assert derive_category("NASA Mars mission timeline?") == "science"
|
||||
assert derive_category("FDA approval for new drug?") == "science"
|
||||
assert derive_category("Climate change targets met?") == "science"
|
||||
|
||||
def test_other_category(self) -> None:
|
||||
"""Test fallback to 'other' category."""
|
||||
assert derive_category("Random obscure question?") == "other"
|
||||
assert derive_category("Will it be sunny tomorrow?") == "other"
|
||||
|
||||
def test_case_insensitive(self) -> None:
|
||||
"""Test case insensitivity."""
|
||||
assert derive_category("BITCOIN PRICE") == "crypto"
|
||||
assert derive_category("bitcoin price") == "crypto"
|
||||
assert derive_category("Bitcoin Price") == "crypto"
|
||||
|
||||
|
||||
class TestMarketMetadata:
|
||||
"""Tests for the MarketMetadata dataclass."""
|
||||
|
||||
def test_from_market(self, sample_market: Market) -> None:
|
||||
"""Test creating metadata from a market."""
|
||||
metadata = MarketMetadata.from_market(sample_market)
|
||||
|
||||
assert metadata.condition_id == sample_market.condition_id
|
||||
assert metadata.question == sample_market.question
|
||||
assert metadata.description == sample_market.description
|
||||
assert metadata.tokens == sample_market.tokens
|
||||
assert metadata.end_date == sample_market.end_date
|
||||
assert metadata.active == sample_market.active
|
||||
assert metadata.closed == sample_market.closed
|
||||
assert metadata.category == "crypto" # "Bitcoin" in question
|
||||
assert metadata.last_updated is not None
|
||||
|
||||
def test_to_dict(self, sample_metadata: MarketMetadata) -> None:
|
||||
"""Test serialization to dict."""
|
||||
data = sample_metadata.to_dict()
|
||||
|
||||
assert data["condition_id"] == sample_metadata.condition_id
|
||||
assert data["question"] == sample_metadata.question
|
||||
assert data["category"] == "crypto"
|
||||
assert len(data["tokens"]) == 1
|
||||
assert data["tokens"][0]["token_id"] == "token123"
|
||||
|
||||
def test_from_dict(self, sample_metadata: MarketMetadata) -> None:
|
||||
"""Test deserialization from dict."""
|
||||
data = sample_metadata.to_dict()
|
||||
restored = MarketMetadata.from_dict(data)
|
||||
|
||||
assert restored.condition_id == sample_metadata.condition_id
|
||||
assert restored.question == sample_metadata.question
|
||||
assert restored.category == sample_metadata.category
|
||||
assert len(restored.tokens) == 1
|
||||
|
||||
def test_roundtrip(self, sample_metadata: MarketMetadata) -> None:
|
||||
"""Test serialization roundtrip."""
|
||||
data = sample_metadata.to_dict()
|
||||
json_str = json.dumps(data)
|
||||
parsed = json.loads(json_str)
|
||||
restored = MarketMetadata.from_dict(parsed)
|
||||
|
||||
assert restored.condition_id == sample_metadata.condition_id
|
||||
assert restored.question == sample_metadata.question
|
||||
|
||||
|
||||
class TestSyncStats:
|
||||
"""Tests for the SyncStats dataclass."""
|
||||
|
||||
def test_defaults(self) -> None:
|
||||
"""Test default values."""
|
||||
stats = SyncStats()
|
||||
|
||||
assert stats.total_syncs == 0
|
||||
assert stats.successful_syncs == 0
|
||||
assert stats.failed_syncs == 0
|
||||
assert stats.markets_cached == 0
|
||||
assert stats.last_sync_time is None
|
||||
assert stats.last_error is None
|
||||
|
||||
|
||||
class TestMarketMetadataSync:
|
||||
"""Tests for the MarketMetadataSync class."""
|
||||
|
||||
def test_init(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test initialization."""
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
|
||||
assert sync.state == SyncState.STOPPED
|
||||
assert sync.stats.total_syncs == 0
|
||||
assert sync._sync_interval == DEFAULT_SYNC_INTERVAL_SECONDS
|
||||
assert sync._cache_ttl == DEFAULT_CACHE_TTL_SECONDS
|
||||
assert sync._key_prefix == DEFAULT_REDIS_KEY_PREFIX
|
||||
|
||||
def test_init_custom_config(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test initialization with custom config."""
|
||||
sync = MarketMetadataSync(
|
||||
redis=mock_redis,
|
||||
clob_client=mock_clob,
|
||||
sync_interval_seconds=60,
|
||||
cache_ttl_seconds=120,
|
||||
key_prefix="custom:",
|
||||
)
|
||||
|
||||
assert sync._sync_interval == 60
|
||||
assert sync._cache_ttl == 120
|
||||
assert sync._key_prefix == "custom:"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_stop(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test starting and stopping the sync service."""
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
|
||||
# Start
|
||||
await sync.start()
|
||||
assert sync.state == SyncState.IDLE
|
||||
assert sync.stats.total_syncs == 1
|
||||
assert sync.stats.successful_syncs == 1
|
||||
|
||||
# Stop
|
||||
await sync.stop()
|
||||
assert sync.state == SyncState.STOPPED
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_performs_initial_sync(
|
||||
self, mock_redis: AsyncMock, mock_clob: MagicMock
|
||||
) -> None:
|
||||
"""Test that start performs an initial sync."""
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
|
||||
await sync.start()
|
||||
|
||||
# Should have called get_markets
|
||||
mock_clob.get_markets.assert_called_once_with(True)
|
||||
|
||||
# Should have cached the market
|
||||
mock_redis.setex.assert_called()
|
||||
|
||||
await sync.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_failure(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test start failure handling."""
|
||||
mock_clob.get_markets.side_effect = Exception("API error")
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
|
||||
with pytest.raises(MetadataSyncError, match="initial sync failed"):
|
||||
await sync.start()
|
||||
|
||||
assert sync.state == SyncState.ERROR
|
||||
assert sync.stats.last_error == "API error"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_market_cache_hit(
|
||||
self,
|
||||
mock_redis: AsyncMock,
|
||||
mock_clob: MagicMock,
|
||||
sample_metadata: MarketMetadata,
|
||||
) -> None:
|
||||
"""Test get_market with cache hit."""
|
||||
# Setup cache hit
|
||||
cached_data = json.dumps(sample_metadata.to_dict())
|
||||
mock_redis.get = AsyncMock(return_value=cached_data)
|
||||
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
await sync.start()
|
||||
|
||||
result = await sync.get_market("cond123")
|
||||
|
||||
assert result is not None
|
||||
assert result.condition_id == "cond123"
|
||||
# Should not have called API
|
||||
mock_clob.get_market.assert_not_called()
|
||||
|
||||
await sync.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_market_cache_miss(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test get_market with cache miss."""
|
||||
# Setup cache miss
|
||||
mock_redis.get = AsyncMock(return_value=None)
|
||||
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
await sync.start()
|
||||
|
||||
result = await sync.get_market("cond123")
|
||||
|
||||
assert result is not None
|
||||
assert result.condition_id == "cond123"
|
||||
# Should have called API
|
||||
mock_clob.get_market.assert_called_with("cond123")
|
||||
# Should have cached the result
|
||||
assert mock_redis.setex.call_count >= 2 # Initial sync + cache miss
|
||||
|
||||
await sync.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_market_not_found(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test get_market when market doesn't exist."""
|
||||
mock_redis.get = AsyncMock(return_value=None)
|
||||
mock_clob.get_market.return_value = None
|
||||
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
await sync.start()
|
||||
|
||||
result = await sync.get_market("nonexistent")
|
||||
|
||||
assert result is None
|
||||
|
||||
await sync.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalidate_market(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test cache invalidation."""
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
await sync.start()
|
||||
|
||||
result = await sync.invalidate_market("cond123")
|
||||
|
||||
assert result is True
|
||||
mock_redis.delete.assert_called_with(f"{DEFAULT_REDIS_KEY_PREFIX}cond123")
|
||||
|
||||
await sync.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_force_sync(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test forced sync."""
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
await sync.start()
|
||||
|
||||
# Initial sync
|
||||
assert sync.stats.total_syncs == 1
|
||||
|
||||
# Force sync
|
||||
await sync.force_sync()
|
||||
|
||||
assert sync.stats.total_syncs == 2
|
||||
assert sync.stats.successful_syncs == 2
|
||||
|
||||
await sync.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_state_change_callback(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test state change callback."""
|
||||
states: list[SyncState] = []
|
||||
|
||||
def on_state_change(state: SyncState) -> None:
|
||||
states.append(state)
|
||||
|
||||
sync = MarketMetadataSync(
|
||||
redis=mock_redis,
|
||||
clob_client=mock_clob,
|
||||
on_state_change=on_state_change,
|
||||
)
|
||||
|
||||
await sync.start()
|
||||
await sync.stop()
|
||||
|
||||
assert SyncState.STARTING in states
|
||||
assert SyncState.SYNCING in states
|
||||
assert SyncState.IDLE in states
|
||||
assert SyncState.STOPPING in states
|
||||
assert SyncState.STOPPED in states
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_complete_callback(
|
||||
self, mock_redis: AsyncMock, mock_clob: MagicMock
|
||||
) -> None:
|
||||
"""Test sync complete callback."""
|
||||
sync_stats: list[SyncStats] = []
|
||||
|
||||
def on_sync_complete(stats: SyncStats) -> None:
|
||||
sync_stats.append(stats)
|
||||
|
||||
sync = MarketMetadataSync(
|
||||
redis=mock_redis,
|
||||
clob_client=mock_clob,
|
||||
on_sync_complete=on_sync_complete,
|
||||
)
|
||||
|
||||
await sync.start()
|
||||
|
||||
assert len(sync_stats) == 1
|
||||
assert sync_stats[0].successful_syncs == 1
|
||||
assert sync_stats[0].markets_cached == 1
|
||||
|
||||
await sync.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_markets_by_category(
|
||||
self, mock_redis: AsyncMock, mock_clob: MagicMock, sample_metadata: MarketMetadata
|
||||
) -> None:
|
||||
"""Test getting markets by category."""
|
||||
# Setup scan to return keys
|
||||
key = f"{DEFAULT_REDIS_KEY_PREFIX}cond123"
|
||||
mock_redis.scan = AsyncMock(return_value=(0, [key]))
|
||||
|
||||
# Setup get to return cached data
|
||||
cached_data = json.dumps(sample_metadata.to_dict())
|
||||
mock_redis.get = AsyncMock(return_value=cached_data)
|
||||
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
# Don't start to avoid initial sync complexity
|
||||
sync._state = SyncState.IDLE
|
||||
|
||||
results = await sync.get_markets_by_category("crypto")
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].category == "crypto"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cannot_start_twice(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test that starting twice doesn't double-start."""
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
|
||||
await sync.start()
|
||||
await sync.start() # Should be a no-op
|
||||
|
||||
assert sync.stats.total_syncs == 1 # Only one initial sync
|
||||
|
||||
await sync.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_when_stopped(self, mock_redis: AsyncMock, mock_clob: MagicMock) -> None:
|
||||
"""Test stopping when already stopped."""
|
||||
sync = MarketMetadataSync(redis=mock_redis, clob_client=mock_clob)
|
||||
|
||||
await sync.stop() # Should be a no-op
|
||||
|
||||
assert sync.state == SyncState.STOPPED
|
||||
@@ -0,0 +1,420 @@
|
||||
"""Tests for ingestor data models."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.ingestor.models import (
|
||||
Market,
|
||||
Orderbook,
|
||||
OrderbookLevel,
|
||||
Token,
|
||||
TradeEvent,
|
||||
)
|
||||
|
||||
|
||||
class TestToken:
|
||||
"""Tests for Token model."""
|
||||
|
||||
def test_from_dict_with_price(self) -> None:
|
||||
"""Test creating Token from dict with price."""
|
||||
data = {
|
||||
"token_id": "123abc",
|
||||
"outcome": "Yes",
|
||||
"price": "0.65",
|
||||
}
|
||||
token = Token.from_dict(data)
|
||||
|
||||
assert token.token_id == "123abc"
|
||||
assert token.outcome == "Yes"
|
||||
assert token.price == Decimal("0.65")
|
||||
|
||||
def test_from_dict_without_price(self) -> None:
|
||||
"""Test creating Token from dict without price."""
|
||||
data = {
|
||||
"token_id": "456def",
|
||||
"outcome": "No",
|
||||
}
|
||||
token = Token.from_dict(data)
|
||||
|
||||
assert token.token_id == "456def"
|
||||
assert token.outcome == "No"
|
||||
assert token.price is None
|
||||
|
||||
def test_frozen(self) -> None:
|
||||
"""Test that Token is immutable."""
|
||||
token = Token(token_id="123", outcome="Yes", price=Decimal("0.5"))
|
||||
with pytest.raises(AttributeError):
|
||||
token.token_id = "456" # type: ignore[misc]
|
||||
|
||||
|
||||
class TestMarket:
|
||||
"""Tests for Market model."""
|
||||
|
||||
def test_from_dict_full(self) -> None:
|
||||
"""Test creating Market from complete dict."""
|
||||
data = {
|
||||
"condition_id": "0xabc123",
|
||||
"question": "Will it rain tomorrow?",
|
||||
"description": "Market for weather prediction",
|
||||
"tokens": [
|
||||
{"token_id": "t1", "outcome": "Yes", "price": "0.7"},
|
||||
{"token_id": "t2", "outcome": "No", "price": "0.3"},
|
||||
],
|
||||
"end_date_iso": "2024-12-31T23:59:59Z",
|
||||
"active": True,
|
||||
"closed": False,
|
||||
}
|
||||
market = Market.from_dict(data)
|
||||
|
||||
assert market.condition_id == "0xabc123"
|
||||
assert market.question == "Will it rain tomorrow?"
|
||||
assert market.description == "Market for weather prediction"
|
||||
assert len(market.tokens) == 2
|
||||
assert market.tokens[0].outcome == "Yes"
|
||||
assert market.tokens[1].outcome == "No"
|
||||
assert market.end_date == datetime(2024, 12, 31, 23, 59, 59, tzinfo=UTC)
|
||||
assert market.active is True
|
||||
assert market.closed is False
|
||||
|
||||
def test_from_dict_minimal(self) -> None:
|
||||
"""Test creating Market from minimal dict."""
|
||||
data = {
|
||||
"condition_id": "0xdef456",
|
||||
}
|
||||
market = Market.from_dict(data)
|
||||
|
||||
assert market.condition_id == "0xdef456"
|
||||
assert market.question == ""
|
||||
assert market.description == ""
|
||||
assert market.tokens == ()
|
||||
assert market.end_date is None
|
||||
assert market.active is True
|
||||
assert market.closed is False
|
||||
|
||||
def test_from_dict_invalid_date(self) -> None:
|
||||
"""Test that invalid date is handled gracefully."""
|
||||
data = {
|
||||
"condition_id": "0x123",
|
||||
"end_date_iso": "not-a-valid-date",
|
||||
}
|
||||
market = Market.from_dict(data)
|
||||
|
||||
assert market.end_date is None
|
||||
|
||||
def test_frozen(self) -> None:
|
||||
"""Test that Market is immutable."""
|
||||
market = Market(
|
||||
condition_id="0x123",
|
||||
question="Test?",
|
||||
description="",
|
||||
tokens=(),
|
||||
)
|
||||
with pytest.raises(AttributeError):
|
||||
market.condition_id = "0x456" # type: ignore[misc]
|
||||
|
||||
|
||||
class TestOrderbookLevel:
|
||||
"""Tests for OrderbookLevel model."""
|
||||
|
||||
def test_from_dict(self) -> None:
|
||||
"""Test creating OrderbookLevel from dict."""
|
||||
data = {"price": "0.55", "size": "100.5"}
|
||||
level = OrderbookLevel.from_dict(data)
|
||||
|
||||
assert level.price == Decimal("0.55")
|
||||
assert level.size == Decimal("100.5")
|
||||
|
||||
def test_frozen(self) -> None:
|
||||
"""Test that OrderbookLevel is immutable."""
|
||||
level = OrderbookLevel(price=Decimal("0.5"), size=Decimal("10"))
|
||||
with pytest.raises(AttributeError):
|
||||
level.price = Decimal("0.6") # type: ignore[misc]
|
||||
|
||||
|
||||
class TestOrderbook:
|
||||
"""Tests for Orderbook model."""
|
||||
|
||||
def test_from_clob_orderbook(self) -> None:
|
||||
"""Test creating Orderbook from py-clob-client response."""
|
||||
# Create mock bid/ask objects
|
||||
mock_bid = MagicMock()
|
||||
mock_bid.price = "0.50"
|
||||
mock_bid.size = "100"
|
||||
|
||||
mock_ask = MagicMock()
|
||||
mock_ask.price = "0.52"
|
||||
mock_ask.size = "150"
|
||||
|
||||
mock_orderbook = MagicMock()
|
||||
mock_orderbook.market = "0xmarket123"
|
||||
mock_orderbook.asset_id = "token123"
|
||||
mock_orderbook.tick_size = "0.01"
|
||||
mock_orderbook.bids = [mock_bid]
|
||||
mock_orderbook.asks = [mock_ask]
|
||||
|
||||
orderbook = Orderbook.from_clob_orderbook(mock_orderbook)
|
||||
|
||||
assert orderbook.market == "0xmarket123"
|
||||
assert orderbook.asset_id == "token123"
|
||||
assert orderbook.tick_size == Decimal("0.01")
|
||||
assert len(orderbook.bids) == 1
|
||||
assert len(orderbook.asks) == 1
|
||||
assert orderbook.bids[0].price == Decimal("0.50")
|
||||
assert orderbook.asks[0].price == Decimal("0.52")
|
||||
|
||||
def test_from_clob_orderbook_empty(self) -> None:
|
||||
"""Test creating Orderbook with empty bids/asks."""
|
||||
mock_orderbook = MagicMock()
|
||||
mock_orderbook.market = "0xmarket"
|
||||
mock_orderbook.asset_id = "token"
|
||||
mock_orderbook.tick_size = "0.01"
|
||||
mock_orderbook.bids = None
|
||||
mock_orderbook.asks = []
|
||||
|
||||
orderbook = Orderbook.from_clob_orderbook(mock_orderbook)
|
||||
|
||||
assert orderbook.bids == ()
|
||||
assert orderbook.asks == ()
|
||||
|
||||
def test_best_bid(self) -> None:
|
||||
"""Test best_bid property."""
|
||||
orderbook = Orderbook(
|
||||
market="0x",
|
||||
asset_id="t",
|
||||
bids=(
|
||||
OrderbookLevel(Decimal("0.50"), Decimal("100")),
|
||||
OrderbookLevel(Decimal("0.49"), Decimal("50")),
|
||||
),
|
||||
asks=(),
|
||||
tick_size=Decimal("0.01"),
|
||||
)
|
||||
|
||||
assert orderbook.best_bid == Decimal("0.50")
|
||||
|
||||
def test_best_bid_empty(self) -> None:
|
||||
"""Test best_bid with no bids."""
|
||||
orderbook = Orderbook(
|
||||
market="0x",
|
||||
asset_id="t",
|
||||
bids=(),
|
||||
asks=(),
|
||||
tick_size=Decimal("0.01"),
|
||||
)
|
||||
|
||||
assert orderbook.best_bid is None
|
||||
|
||||
def test_best_ask(self) -> None:
|
||||
"""Test best_ask property."""
|
||||
orderbook = Orderbook(
|
||||
market="0x",
|
||||
asset_id="t",
|
||||
bids=(),
|
||||
asks=(
|
||||
OrderbookLevel(Decimal("0.52"), Decimal("100")),
|
||||
OrderbookLevel(Decimal("0.53"), Decimal("50")),
|
||||
),
|
||||
tick_size=Decimal("0.01"),
|
||||
)
|
||||
|
||||
assert orderbook.best_ask == Decimal("0.52")
|
||||
|
||||
def test_spread(self) -> None:
|
||||
"""Test spread calculation."""
|
||||
orderbook = Orderbook(
|
||||
market="0x",
|
||||
asset_id="t",
|
||||
bids=(OrderbookLevel(Decimal("0.50"), Decimal("100")),),
|
||||
asks=(OrderbookLevel(Decimal("0.52"), Decimal("100")),),
|
||||
tick_size=Decimal("0.01"),
|
||||
)
|
||||
|
||||
assert orderbook.spread == Decimal("0.02")
|
||||
|
||||
def test_spread_missing_data(self) -> None:
|
||||
"""Test spread with missing bid or ask."""
|
||||
orderbook = Orderbook(
|
||||
market="0x",
|
||||
asset_id="t",
|
||||
bids=(OrderbookLevel(Decimal("0.50"), Decimal("100")),),
|
||||
asks=(),
|
||||
tick_size=Decimal("0.01"),
|
||||
)
|
||||
|
||||
assert orderbook.spread is None
|
||||
|
||||
def test_midpoint(self) -> None:
|
||||
"""Test midpoint calculation."""
|
||||
orderbook = Orderbook(
|
||||
market="0x",
|
||||
asset_id="t",
|
||||
bids=(OrderbookLevel(Decimal("0.50"), Decimal("100")),),
|
||||
asks=(OrderbookLevel(Decimal("0.52"), Decimal("100")),),
|
||||
tick_size=Decimal("0.01"),
|
||||
)
|
||||
|
||||
assert orderbook.midpoint == Decimal("0.51")
|
||||
|
||||
def test_midpoint_missing_data(self) -> None:
|
||||
"""Test midpoint with missing data."""
|
||||
orderbook = Orderbook(
|
||||
market="0x",
|
||||
asset_id="t",
|
||||
bids=(),
|
||||
asks=(OrderbookLevel(Decimal("0.52"), Decimal("100")),),
|
||||
tick_size=Decimal("0.01"),
|
||||
)
|
||||
|
||||
assert orderbook.midpoint is None
|
||||
|
||||
|
||||
class TestTradeEvent:
|
||||
"""Tests for TradeEvent model."""
|
||||
|
||||
def test_from_websocket_message_full(self) -> None:
|
||||
"""Test creating TradeEvent from a full WebSocket message."""
|
||||
data = {
|
||||
"conditionId": "0xmarket123",
|
||||
"transactionHash": "0xtx456",
|
||||
"proxyWallet": "0xwallet789",
|
||||
"side": "BUY",
|
||||
"outcome": "Yes",
|
||||
"outcomeIndex": 0,
|
||||
"price": 0.65,
|
||||
"size": 100,
|
||||
"timestamp": 1704067200, # 2024-01-01 00:00:00 UTC
|
||||
"asset": "token123",
|
||||
"slug": "will-it-rain",
|
||||
"eventSlug": "weather-markets",
|
||||
"title": "Weather Predictions",
|
||||
"name": "Alice",
|
||||
"pseudonym": "AliceTrader",
|
||||
}
|
||||
trade = TradeEvent.from_websocket_message(data)
|
||||
|
||||
assert trade.market_id == "0xmarket123"
|
||||
assert trade.trade_id == "0xtx456"
|
||||
assert trade.wallet_address == "0xwallet789"
|
||||
assert trade.side == "BUY"
|
||||
assert trade.outcome == "Yes"
|
||||
assert trade.outcome_index == 0
|
||||
assert trade.price == Decimal("0.65")
|
||||
assert trade.size == Decimal("100")
|
||||
assert trade.timestamp == datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC)
|
||||
assert trade.asset_id == "token123"
|
||||
assert trade.market_slug == "will-it-rain"
|
||||
assert trade.event_slug == "weather-markets"
|
||||
assert trade.event_title == "Weather Predictions"
|
||||
assert trade.trader_name == "Alice"
|
||||
assert trade.trader_pseudonym == "AliceTrader"
|
||||
|
||||
def test_from_websocket_message_minimal(self) -> None:
|
||||
"""Test creating TradeEvent with minimal data."""
|
||||
data = {
|
||||
"conditionId": "0x123",
|
||||
"transactionHash": "0xtx",
|
||||
"proxyWallet": "0xwallet",
|
||||
"side": "SELL",
|
||||
"outcome": "No",
|
||||
"price": 0.35,
|
||||
"size": 50,
|
||||
}
|
||||
trade = TradeEvent.from_websocket_message(data)
|
||||
|
||||
assert trade.market_id == "0x123"
|
||||
assert trade.side == "SELL"
|
||||
assert trade.outcome == "No"
|
||||
assert trade.price == Decimal("0.35")
|
||||
assert trade.size == Decimal("50")
|
||||
assert trade.market_slug == ""
|
||||
assert trade.trader_name == ""
|
||||
|
||||
def test_from_websocket_message_lowercase_side(self) -> None:
|
||||
"""Test that lowercase side is normalized."""
|
||||
data = {
|
||||
"side": "buy",
|
||||
"price": 0.5,
|
||||
"size": 10,
|
||||
}
|
||||
trade = TradeEvent.from_websocket_message(data)
|
||||
|
||||
assert trade.side == "BUY"
|
||||
|
||||
def test_from_websocket_message_sell_side(self) -> None:
|
||||
"""Test SELL side handling."""
|
||||
data = {
|
||||
"side": "sell",
|
||||
"price": 0.5,
|
||||
"size": 10,
|
||||
}
|
||||
trade = TradeEvent.from_websocket_message(data)
|
||||
|
||||
assert trade.side == "SELL"
|
||||
|
||||
def test_is_buy(self) -> None:
|
||||
"""Test is_buy property."""
|
||||
buy_trade = TradeEvent(
|
||||
market_id="",
|
||||
trade_id="",
|
||||
wallet_address="",
|
||||
side="BUY",
|
||||
outcome="",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("10"),
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="",
|
||||
)
|
||||
sell_trade = TradeEvent(
|
||||
market_id="",
|
||||
trade_id="",
|
||||
wallet_address="",
|
||||
side="SELL",
|
||||
outcome="",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("10"),
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="",
|
||||
)
|
||||
|
||||
assert buy_trade.is_buy is True
|
||||
assert buy_trade.is_sell is False
|
||||
assert sell_trade.is_buy is False
|
||||
assert sell_trade.is_sell is True
|
||||
|
||||
def test_notional_value(self) -> None:
|
||||
"""Test notional value calculation."""
|
||||
trade = TradeEvent(
|
||||
market_id="",
|
||||
trade_id="",
|
||||
wallet_address="",
|
||||
side="BUY",
|
||||
outcome="",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.65"),
|
||||
size=Decimal("100"),
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="",
|
||||
)
|
||||
|
||||
assert trade.notional_value == Decimal("65")
|
||||
|
||||
def test_frozen(self) -> None:
|
||||
"""Test that TradeEvent is immutable."""
|
||||
trade = TradeEvent(
|
||||
market_id="0x123",
|
||||
trade_id="0xtx",
|
||||
wallet_address="0xwallet",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.5"),
|
||||
size=Decimal("10"),
|
||||
timestamp=datetime.now(UTC),
|
||||
asset_id="token",
|
||||
)
|
||||
with pytest.raises(AttributeError):
|
||||
trade.market_id = "0x456" # type: ignore[misc]
|
||||
@@ -0,0 +1,470 @@
|
||||
"""Tests for the Redis Streams event publisher."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from redis.exceptions import ResponseError
|
||||
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
from polymarket_insider_tracker.ingestor.publisher import (
|
||||
DEFAULT_BLOCK_MS,
|
||||
DEFAULT_COUNT,
|
||||
DEFAULT_MAX_LEN,
|
||||
DEFAULT_STREAM_NAME,
|
||||
ConsumerGroupExistsError,
|
||||
EventPublisher,
|
||||
StreamEntry,
|
||||
_deserialize_trade_event,
|
||||
_serialize_trade_event,
|
||||
)
|
||||
|
||||
|
||||
# Test fixtures
|
||||
@pytest.fixture
|
||||
def sample_trade_event() -> TradeEvent:
|
||||
"""Create a sample trade event."""
|
||||
return TradeEvent(
|
||||
market_id="0xmarket123",
|
||||
trade_id="0xtx456",
|
||||
wallet_address="0xwallet789",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.65"),
|
||||
size=Decimal("1000"),
|
||||
timestamp=datetime(2026, 1, 4, 12, 0, 0, tzinfo=UTC),
|
||||
asset_id="token123",
|
||||
market_slug="will-it-rain",
|
||||
event_slug="weather-markets",
|
||||
event_title="Weather Predictions",
|
||||
trader_name="Alice",
|
||||
trader_pseudonym="AliceTrader",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis() -> AsyncMock:
|
||||
"""Create a mock Redis client."""
|
||||
redis = AsyncMock()
|
||||
redis.xadd = AsyncMock(return_value="1704369600000-0")
|
||||
redis.xreadgroup = AsyncMock(return_value=[])
|
||||
redis.xack = AsyncMock(return_value=1)
|
||||
redis.xlen = AsyncMock(return_value=100)
|
||||
redis.xtrim = AsyncMock(return_value=0)
|
||||
redis.xinfo_stream = AsyncMock(return_value={"length": 100})
|
||||
redis.xgroup_create = AsyncMock()
|
||||
redis.pipeline = MagicMock()
|
||||
return redis
|
||||
|
||||
|
||||
class TestSerializationFunctions:
|
||||
"""Tests for serialization helper functions."""
|
||||
|
||||
def test_serialize_trade_event(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test serializing a trade event."""
|
||||
data = _serialize_trade_event(sample_trade_event)
|
||||
|
||||
assert data["market_id"] == "0xmarket123"
|
||||
assert data["trade_id"] == "0xtx456"
|
||||
assert data["wallet_address"] == "0xwallet789"
|
||||
assert data["side"] == "BUY"
|
||||
assert data["outcome"] == "Yes"
|
||||
assert data["outcome_index"] == "0"
|
||||
assert data["price"] == "0.65"
|
||||
assert data["size"] == "1000"
|
||||
assert data["timestamp"] == "2026-01-04T12:00:00+00:00"
|
||||
assert data["asset_id"] == "token123"
|
||||
assert data["market_slug"] == "will-it-rain"
|
||||
assert data["trader_name"] == "Alice"
|
||||
|
||||
def test_serialize_all_values_are_strings(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test that all serialized values are strings."""
|
||||
data = _serialize_trade_event(sample_trade_event)
|
||||
|
||||
for key, value in data.items():
|
||||
assert isinstance(key, str), f"Key {key} is not a string"
|
||||
assert isinstance(value, str), f"Value for {key} is not a string"
|
||||
|
||||
def test_deserialize_trade_event(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test deserializing a trade event."""
|
||||
data = _serialize_trade_event(sample_trade_event)
|
||||
restored = _deserialize_trade_event(data)
|
||||
|
||||
assert restored.market_id == sample_trade_event.market_id
|
||||
assert restored.trade_id == sample_trade_event.trade_id
|
||||
assert restored.wallet_address == sample_trade_event.wallet_address
|
||||
assert restored.side == sample_trade_event.side
|
||||
assert restored.outcome == sample_trade_event.outcome
|
||||
assert restored.outcome_index == sample_trade_event.outcome_index
|
||||
assert restored.price == sample_trade_event.price
|
||||
assert restored.size == sample_trade_event.size
|
||||
assert restored.timestamp == sample_trade_event.timestamp
|
||||
assert restored.asset_id == sample_trade_event.asset_id
|
||||
|
||||
def test_deserialize_with_bytes_keys(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test deserializing with bytes keys/values (as returned by Redis)."""
|
||||
data = _serialize_trade_event(sample_trade_event)
|
||||
# Convert to bytes like Redis returns
|
||||
bytes_data = {k.encode(): v.encode() for k, v in data.items()}
|
||||
|
||||
restored = _deserialize_trade_event(bytes_data)
|
||||
|
||||
assert restored.market_id == sample_trade_event.market_id
|
||||
assert restored.side == sample_trade_event.side
|
||||
|
||||
def test_deserialize_with_invalid_timestamp(self) -> None:
|
||||
"""Test deserializing with invalid timestamp falls back to now."""
|
||||
data = {
|
||||
"market_id": "0x123",
|
||||
"timestamp": "not-a-timestamp",
|
||||
"side": "BUY",
|
||||
"price": "0.5",
|
||||
"size": "100",
|
||||
}
|
||||
|
||||
event = _deserialize_trade_event(data)
|
||||
|
||||
assert event.market_id == "0x123"
|
||||
# Timestamp should be recent (within last minute)
|
||||
assert (datetime.now(UTC) - event.timestamp).total_seconds() < 60
|
||||
|
||||
def test_deserialize_with_missing_fields(self) -> None:
|
||||
"""Test deserializing with missing fields uses defaults."""
|
||||
data = {
|
||||
"market_id": "0x123",
|
||||
"side": "SELL",
|
||||
}
|
||||
|
||||
event = _deserialize_trade_event(data)
|
||||
|
||||
assert event.market_id == "0x123"
|
||||
assert event.side == "SELL"
|
||||
assert event.price == Decimal("0")
|
||||
assert event.outcome == ""
|
||||
|
||||
|
||||
class TestEventPublisher:
|
||||
"""Tests for the EventPublisher class."""
|
||||
|
||||
def test_init(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test initialization."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
assert publisher.stream_name == DEFAULT_STREAM_NAME
|
||||
assert publisher._max_len == DEFAULT_MAX_LEN
|
||||
|
||||
def test_init_custom_config(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test initialization with custom config."""
|
||||
publisher = EventPublisher(
|
||||
mock_redis,
|
||||
stream_name="custom-stream",
|
||||
max_len=50_000,
|
||||
)
|
||||
|
||||
assert publisher.stream_name == "custom-stream"
|
||||
assert publisher._max_len == 50_000
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish(self, mock_redis: AsyncMock, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test publishing a single event."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entry_id = await publisher.publish(sample_trade_event)
|
||||
|
||||
assert entry_id == "1704369600000-0"
|
||||
mock_redis.xadd.assert_called_once()
|
||||
call_args = mock_redis.xadd.call_args
|
||||
assert call_args[0][0] == DEFAULT_STREAM_NAME
|
||||
assert call_args[1]["maxlen"] == DEFAULT_MAX_LEN
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_returns_decoded_bytes(
|
||||
self, mock_redis: AsyncMock, sample_trade_event: TradeEvent
|
||||
) -> None:
|
||||
"""Test that publish handles bytes entry IDs."""
|
||||
mock_redis.xadd = AsyncMock(return_value=b"1704369600000-0")
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entry_id = await publisher.publish(sample_trade_event)
|
||||
|
||||
assert entry_id == "1704369600000-0"
|
||||
assert isinstance(entry_id, str)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_batch(
|
||||
self, mock_redis: AsyncMock, sample_trade_event: TradeEvent
|
||||
) -> None:
|
||||
"""Test batch publishing."""
|
||||
mock_pipeline = AsyncMock()
|
||||
mock_pipeline.xadd = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(return_value=["1704369600000-0", "1704369600000-1"])
|
||||
mock_redis.pipeline.return_value = mock_pipeline
|
||||
|
||||
publisher = EventPublisher(mock_redis)
|
||||
events = [sample_trade_event, sample_trade_event]
|
||||
|
||||
entry_ids = await publisher.publish_batch(events)
|
||||
|
||||
assert len(entry_ids) == 2
|
||||
assert entry_ids[0] == "1704369600000-0"
|
||||
assert entry_ids[1] == "1704369600000-1"
|
||||
assert mock_pipeline.xadd.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_batch_empty(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test batch publishing with empty list."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entry_ids = await publisher.publish_batch([])
|
||||
|
||||
assert entry_ids == []
|
||||
mock_redis.pipeline.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_consumer_group(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test creating a consumer group."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
await publisher.create_consumer_group("test-group")
|
||||
|
||||
mock_redis.xgroup_create.assert_called_once_with(
|
||||
DEFAULT_STREAM_NAME,
|
||||
"test-group",
|
||||
id="0",
|
||||
mkstream=True,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_consumer_group_custom_start_id(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test creating a consumer group with custom start ID."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
await publisher.create_consumer_group("test-group", start_id="$")
|
||||
|
||||
call_args = mock_redis.xgroup_create.call_args
|
||||
assert call_args[1]["id"] == "$"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_consumer_group_already_exists(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test creating a consumer group that already exists."""
|
||||
mock_redis.xgroup_create.side_effect = ResponseError(
|
||||
"BUSYGROUP Consumer Group name already exists"
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
with pytest.raises(ConsumerGroupExistsError):
|
||||
await publisher.create_consumer_group("existing-group")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_consumer_group_creates(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test ensure_consumer_group creates if not exists."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
created = await publisher.ensure_consumer_group("new-group")
|
||||
|
||||
assert created is True
|
||||
mock_redis.xgroup_create.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_consumer_group_exists(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test ensure_consumer_group returns False if exists."""
|
||||
mock_redis.xgroup_create.side_effect = ResponseError("BUSYGROUP")
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
created = await publisher.ensure_consumer_group("existing-group")
|
||||
|
||||
assert created is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_events(self, mock_redis: AsyncMock, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test reading events from stream."""
|
||||
serialized = _serialize_trade_event(sample_trade_event)
|
||||
mock_redis.xreadgroup = AsyncMock(
|
||||
return_value=[
|
||||
(
|
||||
"trades",
|
||||
[
|
||||
("1704369600000-0", serialized),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_events("test-group", "worker-1")
|
||||
|
||||
assert len(entries) == 1
|
||||
assert entries[0].entry_id == "1704369600000-0"
|
||||
assert entries[0].event.market_id == sample_trade_event.market_id
|
||||
mock_redis.xreadgroup.assert_called_once_with(
|
||||
"test-group",
|
||||
"worker-1",
|
||||
{DEFAULT_STREAM_NAME: ">"},
|
||||
count=DEFAULT_COUNT,
|
||||
block=DEFAULT_BLOCK_MS,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_events_empty(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test reading when no events available."""
|
||||
mock_redis.xreadgroup = AsyncMock(return_value=None)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_events("test-group", "worker-1")
|
||||
|
||||
assert entries == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_events_with_bytes(
|
||||
self, mock_redis: AsyncMock, sample_trade_event: TradeEvent
|
||||
) -> None:
|
||||
"""Test reading events with bytes data (as from real Redis)."""
|
||||
serialized = _serialize_trade_event(sample_trade_event)
|
||||
bytes_data = {k.encode(): v.encode() for k, v in serialized.items()}
|
||||
mock_redis.xreadgroup = AsyncMock(
|
||||
return_value=[
|
||||
(
|
||||
b"trades",
|
||||
[
|
||||
(b"1704369600000-0", bytes_data),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_events("test-group", "worker-1")
|
||||
|
||||
assert len(entries) == 1
|
||||
assert entries[0].entry_id == "1704369600000-0"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_pending(
|
||||
self, mock_redis: AsyncMock, sample_trade_event: TradeEvent
|
||||
) -> None:
|
||||
"""Test reading pending events."""
|
||||
serialized = _serialize_trade_event(sample_trade_event)
|
||||
mock_redis.xreadgroup = AsyncMock(
|
||||
return_value=[
|
||||
(
|
||||
"trades",
|
||||
[
|
||||
("1704369600000-0", serialized),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_pending("test-group", "worker-1")
|
||||
|
||||
assert len(entries) == 1
|
||||
# Should read from "0" not ">"
|
||||
call_args = mock_redis.xreadgroup.call_args
|
||||
assert call_args[0][2] == {DEFAULT_STREAM_NAME: "0"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_pending_skips_empty_data(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test that read_pending skips entries with no data (already acked)."""
|
||||
mock_redis.xreadgroup = AsyncMock(
|
||||
return_value=[
|
||||
(
|
||||
"trades",
|
||||
[
|
||||
("1704369600000-0", {}), # Empty = already acked
|
||||
("1704369600000-1", None), # None = already acked
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_pending("test-group", "worker-1")
|
||||
|
||||
assert entries == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ack(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test acknowledging entries."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
count = await publisher.ack("test-group", "1704369600000-0", "1704369600000-1")
|
||||
|
||||
assert count == 1 # Mocked return value
|
||||
mock_redis.xack.assert_called_once_with(
|
||||
DEFAULT_STREAM_NAME,
|
||||
"test-group",
|
||||
"1704369600000-0",
|
||||
"1704369600000-1",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ack_empty(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test ack with no entry IDs."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
count = await publisher.ack("test-group")
|
||||
|
||||
assert count == 0
|
||||
mock_redis.xack.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_stream_info(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting stream info."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
info = await publisher.get_stream_info()
|
||||
|
||||
assert info["length"] == 100
|
||||
mock_redis.xinfo_stream.assert_called_once_with(DEFAULT_STREAM_NAME)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_stream_info_not_exists(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting stream info when stream doesn't exist."""
|
||||
mock_redis.xinfo_stream.side_effect = ResponseError("ERR no such key")
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
info = await publisher.get_stream_info()
|
||||
|
||||
assert info == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_stream_length(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting stream length."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
length = await publisher.get_stream_length()
|
||||
|
||||
assert length == 100
|
||||
mock_redis.xlen.assert_called_once_with(DEFAULT_STREAM_NAME)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_stream(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test trimming stream."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
await publisher.trim_stream(50_000)
|
||||
|
||||
mock_redis.xtrim.assert_called_once_with(DEFAULT_STREAM_NAME, maxlen=50_000)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_stream_default(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test trimming stream with default max_len."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
await publisher.trim_stream()
|
||||
|
||||
mock_redis.xtrim.assert_called_once_with(DEFAULT_STREAM_NAME, maxlen=DEFAULT_MAX_LEN)
|
||||
|
||||
|
||||
class TestStreamEntry:
|
||||
"""Tests for the StreamEntry dataclass."""
|
||||
|
||||
def test_stream_entry(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test creating a StreamEntry."""
|
||||
entry = StreamEntry(entry_id="1704369600000-0", event=sample_trade_event)
|
||||
|
||||
assert entry.entry_id == "1704369600000-0"
|
||||
assert entry.event == sample_trade_event
|
||||
@@ -0,0 +1,347 @@
|
||||
"""Tests for WebSocket trade stream handler."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
from polymarket_insider_tracker.ingestor.websocket import (
|
||||
ConnectionState,
|
||||
StreamStats,
|
||||
TradeStreamHandler,
|
||||
)
|
||||
|
||||
|
||||
class TestStreamStats:
|
||||
"""Tests for StreamStats."""
|
||||
|
||||
def test_defaults(self) -> None:
|
||||
"""Test default values."""
|
||||
stats = StreamStats()
|
||||
|
||||
assert stats.trades_received == 0
|
||||
assert stats.reconnect_count == 0
|
||||
assert stats.last_trade_time is None
|
||||
assert stats.connected_since is None
|
||||
assert stats.last_error is None
|
||||
|
||||
|
||||
class TestTradeStreamHandler:
|
||||
"""Tests for TradeStreamHandler."""
|
||||
|
||||
@pytest.fixture
|
||||
def on_trade_mock(self) -> AsyncMock:
|
||||
"""Create mock trade callback."""
|
||||
return AsyncMock()
|
||||
|
||||
@pytest.fixture
|
||||
def on_state_change_mock(self) -> AsyncMock:
|
||||
"""Create mock state change callback."""
|
||||
return AsyncMock()
|
||||
|
||||
@pytest.fixture
|
||||
def handler(
|
||||
self, on_trade_mock: AsyncMock, on_state_change_mock: AsyncMock
|
||||
) -> TradeStreamHandler:
|
||||
"""Create handler with mocks."""
|
||||
return TradeStreamHandler(
|
||||
on_trade=on_trade_mock,
|
||||
on_state_change=on_state_change_mock,
|
||||
initial_reconnect_delay=0.01, # Fast reconnect for tests
|
||||
max_reconnect_delay=0.1,
|
||||
)
|
||||
|
||||
def test_init_defaults(self, on_trade_mock: AsyncMock) -> None:
|
||||
"""Test handler initialization with defaults."""
|
||||
handler = TradeStreamHandler(on_trade=on_trade_mock)
|
||||
|
||||
assert handler.state == ConnectionState.DISCONNECTED
|
||||
assert handler.stats.trades_received == 0
|
||||
assert handler._host == "wss://ws-live-data.polymarket.com"
|
||||
|
||||
def test_init_custom_host(self, on_trade_mock: AsyncMock) -> None:
|
||||
"""Test handler with custom host."""
|
||||
handler = TradeStreamHandler(
|
||||
on_trade=on_trade_mock,
|
||||
host="wss://custom.example.com",
|
||||
)
|
||||
|
||||
assert handler._host == "wss://custom.example.com"
|
||||
|
||||
def test_init_with_event_filter(self, on_trade_mock: AsyncMock) -> None:
|
||||
"""Test handler with event filter."""
|
||||
handler = TradeStreamHandler(
|
||||
on_trade=on_trade_mock,
|
||||
event_filter="presidential-election-2024",
|
||||
)
|
||||
|
||||
assert handler._event_filter == "presidential-election-2024"
|
||||
|
||||
def test_build_subscription_message_no_filter(self, handler: TradeStreamHandler) -> None:
|
||||
"""Test building subscription message without filters."""
|
||||
msg = handler._build_subscription_message()
|
||||
|
||||
assert msg == {"subscriptions": [{"topic": "activity", "type": "trades"}]}
|
||||
|
||||
def test_build_subscription_message_with_event_filter(self, on_trade_mock: AsyncMock) -> None:
|
||||
"""Test building subscription message with event filter."""
|
||||
handler = TradeStreamHandler(
|
||||
on_trade=on_trade_mock,
|
||||
event_filter="test-event",
|
||||
)
|
||||
msg = handler._build_subscription_message()
|
||||
|
||||
assert msg["subscriptions"][0]["filters"] == json.dumps({"event_slug": "test-event"})
|
||||
|
||||
def test_build_subscription_message_with_market_filter(self, on_trade_mock: AsyncMock) -> None:
|
||||
"""Test building subscription message with market filter."""
|
||||
handler = TradeStreamHandler(
|
||||
on_trade=on_trade_mock,
|
||||
market_filter="test-market",
|
||||
)
|
||||
msg = handler._build_subscription_message()
|
||||
|
||||
assert msg["subscriptions"][0]["filters"] == json.dumps({"market_slug": "test-market"})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_message_trade(
|
||||
self, handler: TradeStreamHandler, on_trade_mock: AsyncMock
|
||||
) -> None:
|
||||
"""Test handling a valid trade message."""
|
||||
message = json.dumps(
|
||||
{
|
||||
"topic": "activity",
|
||||
"type": "trades",
|
||||
"payload": {
|
||||
"conditionId": "0xmarket",
|
||||
"transactionHash": "0xtx",
|
||||
"proxyWallet": "0xwallet",
|
||||
"side": "BUY",
|
||||
"outcome": "Yes",
|
||||
"price": 0.65,
|
||||
"size": 100,
|
||||
"timestamp": 1704067200,
|
||||
"asset": "token123",
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
await handler._handle_message(message)
|
||||
|
||||
on_trade_mock.assert_called_once()
|
||||
trade: TradeEvent = on_trade_mock.call_args[0][0]
|
||||
assert trade.market_id == "0xmarket"
|
||||
assert trade.side == "BUY"
|
||||
assert trade.price == Decimal("0.65")
|
||||
assert handler.stats.trades_received == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_message_non_trade(
|
||||
self, handler: TradeStreamHandler, on_trade_mock: AsyncMock
|
||||
) -> None:
|
||||
"""Test handling a non-trade message."""
|
||||
message = json.dumps(
|
||||
{
|
||||
"topic": "comments",
|
||||
"type": "comment_created",
|
||||
"payload": {"body": "Hello"},
|
||||
}
|
||||
)
|
||||
|
||||
await handler._handle_message(message)
|
||||
|
||||
on_trade_mock.assert_not_called()
|
||||
assert handler.stats.trades_received == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_message_invalid_json(
|
||||
self, handler: TradeStreamHandler, on_trade_mock: AsyncMock
|
||||
) -> None:
|
||||
"""Test handling invalid JSON message."""
|
||||
await handler._handle_message("not valid json")
|
||||
|
||||
on_trade_mock.assert_not_called()
|
||||
assert handler.stats.trades_received == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_message_callback_error(
|
||||
self, handler: TradeStreamHandler, on_trade_mock: AsyncMock
|
||||
) -> None:
|
||||
"""Test that callback errors don't crash the handler."""
|
||||
on_trade_mock.side_effect = ValueError("Callback error")
|
||||
|
||||
message = json.dumps(
|
||||
{
|
||||
"topic": "activity",
|
||||
"type": "trades",
|
||||
"payload": {
|
||||
"conditionId": "0x",
|
||||
"transactionHash": "0x",
|
||||
"proxyWallet": "0x",
|
||||
"side": "BUY",
|
||||
"price": 0.5,
|
||||
"size": 10,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
# Should not raise
|
||||
await handler._handle_message(message)
|
||||
|
||||
# Trade was still counted
|
||||
assert handler.stats.trades_received == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_state_calls_callback(
|
||||
self, handler: TradeStreamHandler, on_state_change_mock: AsyncMock
|
||||
) -> None:
|
||||
"""Test that state changes trigger callback."""
|
||||
await handler._set_state(ConnectionState.CONNECTING)
|
||||
|
||||
on_state_change_mock.assert_called_once_with(ConnectionState.CONNECTING)
|
||||
assert handler.state == ConnectionState.CONNECTING
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_state_same_state_no_callback(
|
||||
self, handler: TradeStreamHandler, on_state_change_mock: AsyncMock
|
||||
) -> None:
|
||||
"""Test that same state doesn't trigger callback."""
|
||||
handler._state = ConnectionState.CONNECTED
|
||||
|
||||
await handler._set_state(ConnectionState.CONNECTED)
|
||||
|
||||
on_state_change_mock.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_when_not_running(self, handler: TradeStreamHandler) -> None:
|
||||
"""Test stop when handler is not running."""
|
||||
# Should not raise
|
||||
await handler.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_sends_subscription(
|
||||
self,
|
||||
handler: TradeStreamHandler,
|
||||
on_state_change_mock: AsyncMock, # noqa: ARG002
|
||||
) -> None:
|
||||
"""Test that connection sends subscription message."""
|
||||
mock_ws = AsyncMock()
|
||||
mock_ws.send = AsyncMock()
|
||||
|
||||
with patch("websockets.connect", AsyncMock(return_value=mock_ws)):
|
||||
ws = await handler._connect()
|
||||
|
||||
assert ws is mock_ws
|
||||
mock_ws.send.assert_called_once()
|
||||
|
||||
# Verify subscription message
|
||||
sent_msg = json.loads(mock_ws.send.call_args[0][0])
|
||||
assert "subscriptions" in sent_msg
|
||||
assert sent_msg["subscriptions"][0]["topic"] == "activity"
|
||||
assert sent_msg["subscriptions"][0]["type"] == "trades"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_closes_websocket(self, handler: TradeStreamHandler) -> None:
|
||||
"""Test that cleanup closes the WebSocket."""
|
||||
mock_ws = AsyncMock()
|
||||
mock_ws.close = AsyncMock()
|
||||
handler._ws = mock_ws
|
||||
|
||||
await handler._cleanup()
|
||||
|
||||
mock_ws.close.assert_called_once()
|
||||
assert handler._ws is None
|
||||
assert handler.state == ConnectionState.DISCONNECTED
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_manager(self, on_trade_mock: AsyncMock) -> None:
|
||||
"""Test async context manager."""
|
||||
handler = TradeStreamHandler(on_trade=on_trade_mock)
|
||||
|
||||
async with handler:
|
||||
pass
|
||||
|
||||
# Should be stopped after exiting context
|
||||
assert handler._running is False
|
||||
|
||||
|
||||
class TestTradeStreamHandlerIntegration:
|
||||
"""Integration tests for TradeStreamHandler.
|
||||
|
||||
These tests verify the full message flow with mocked WebSocket.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_and_receive_trades(self) -> None:
|
||||
"""Test starting handler and receiving trades."""
|
||||
received_trades: list[TradeEvent] = []
|
||||
|
||||
async def on_trade(trade: TradeEvent) -> None:
|
||||
received_trades.append(trade)
|
||||
|
||||
handler = TradeStreamHandler(
|
||||
on_trade=on_trade,
|
||||
initial_reconnect_delay=0.01,
|
||||
)
|
||||
|
||||
trade_message = json.dumps(
|
||||
{
|
||||
"topic": "activity",
|
||||
"type": "trades",
|
||||
"payload": {
|
||||
"conditionId": "0xtest",
|
||||
"transactionHash": "0xtx",
|
||||
"proxyWallet": "0xwallet",
|
||||
"side": "BUY",
|
||||
"outcome": "Yes",
|
||||
"price": 0.75,
|
||||
"size": 50,
|
||||
"timestamp": 1704067200,
|
||||
"asset": "token",
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
# Create a proper async iterable mock WebSocket
|
||||
class MockWebSocket:
|
||||
"""Mock WebSocket that yields one message then stops."""
|
||||
|
||||
def __init__(self, handler: TradeStreamHandler, message: str):
|
||||
self.handler = handler
|
||||
self.message = message
|
||||
self.sent = False
|
||||
|
||||
async def send(self, _msg: str) -> None:
|
||||
pass
|
||||
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> str:
|
||||
if not self.sent:
|
||||
self.sent = True
|
||||
return self.message
|
||||
# Stop the handler and raise StopAsyncIteration
|
||||
await self.handler.stop()
|
||||
raise StopAsyncIteration
|
||||
|
||||
mock_ws = MockWebSocket(handler, trade_message)
|
||||
|
||||
with patch("websockets.connect", AsyncMock(return_value=mock_ws)):
|
||||
# Run with timeout to prevent hanging
|
||||
try:
|
||||
await asyncio.wait_for(handler.start(), timeout=1.0)
|
||||
except TimeoutError:
|
||||
await handler.stop()
|
||||
|
||||
# Verify trade was received
|
||||
assert len(received_trades) == 1
|
||||
assert received_trades[0].market_id == "0xtest"
|
||||
assert received_trades[0].side == "BUY"
|
||||
assert received_trades[0].price == Decimal("0.75")
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for profiler module."""
|
||||
@@ -0,0 +1,394 @@
|
||||
"""Tests for the wallet analyzer."""
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.profiler.analyzer import (
|
||||
DEFAULT_FRESH_THRESHOLD,
|
||||
USDC_POLYGON_ADDRESS,
|
||||
WalletAnalyzer,
|
||||
)
|
||||
from polymarket_insider_tracker.profiler.models import Transaction, WalletInfo
|
||||
|
||||
# Valid Ethereum addresses for testing
|
||||
VALID_ADDRESS = "0x742d35Cc6634C0532925a3b844Bc9e7595f5eaE2"
|
||||
VALID_ADDRESS_2 = "0x8ba1f109551bD432803012645Ac136ddd64DBA72"
|
||||
|
||||
|
||||
class TestWalletAnalyzerInit:
|
||||
"""Tests for WalletAnalyzer initialization."""
|
||||
|
||||
def test_init_default(self) -> None:
|
||||
"""Test initialization with defaults."""
|
||||
client = AsyncMock()
|
||||
analyzer = WalletAnalyzer(client)
|
||||
|
||||
assert analyzer._client is client
|
||||
assert analyzer._redis is None
|
||||
assert analyzer._fresh_threshold == DEFAULT_FRESH_THRESHOLD
|
||||
assert analyzer._usdc_address == USDC_POLYGON_ADDRESS
|
||||
|
||||
def test_init_with_redis(self) -> None:
|
||||
"""Test initialization with Redis."""
|
||||
client = AsyncMock()
|
||||
redis = AsyncMock()
|
||||
analyzer = WalletAnalyzer(client, redis=redis)
|
||||
|
||||
assert analyzer._redis is redis
|
||||
|
||||
def test_init_custom_threshold(self) -> None:
|
||||
"""Test initialization with custom threshold."""
|
||||
client = AsyncMock()
|
||||
analyzer = WalletAnalyzer(client, fresh_threshold=10)
|
||||
|
||||
assert analyzer._fresh_threshold == 10
|
||||
|
||||
|
||||
class TestWalletAnalyzerAnalyze:
|
||||
"""Tests for the analyze method."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> AsyncMock:
|
||||
"""Create a mock PolygonClient."""
|
||||
client = AsyncMock()
|
||||
client.get_wallet_info = AsyncMock()
|
||||
client.get_token_balance = AsyncMock(return_value=Decimal("1000000"))
|
||||
return client
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis(self) -> AsyncMock:
|
||||
"""Create a mock Redis client."""
|
||||
redis = AsyncMock()
|
||||
redis.get = AsyncMock(return_value=None)
|
||||
redis.set = AsyncMock()
|
||||
return redis
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_fresh_wallet(self, mock_client: AsyncMock) -> None:
|
||||
"""Test analyzing a fresh wallet."""
|
||||
first_tx = Transaction(
|
||||
hash="0xabc",
|
||||
block_number=1000,
|
||||
timestamp=datetime.now(UTC) - timedelta(hours=12),
|
||||
from_address="0xfaucet",
|
||||
to_address=VALID_ADDRESS.lower(),
|
||||
value=Decimal("1000000000000000000"),
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50000000000"),
|
||||
)
|
||||
mock_client.get_wallet_info.return_value = WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=3,
|
||||
balance_wei=Decimal("5000000000000000000"),
|
||||
first_transaction=first_tx,
|
||||
)
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client)
|
||||
profile = await analyzer.analyze(VALID_ADDRESS)
|
||||
|
||||
assert profile.address == VALID_ADDRESS.lower()
|
||||
assert profile.nonce == 3
|
||||
assert profile.is_fresh is True
|
||||
assert profile.first_seen is not None
|
||||
assert profile.age_hours is not None
|
||||
assert 11 < profile.age_hours < 13
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_old_wallet(self, mock_client: AsyncMock) -> None:
|
||||
"""Test analyzing an old wallet with many transactions."""
|
||||
first_tx = Transaction(
|
||||
hash="0xabc",
|
||||
block_number=1000,
|
||||
timestamp=datetime.now(UTC) - timedelta(days=365),
|
||||
from_address="0xfaucet",
|
||||
to_address=VALID_ADDRESS.lower(),
|
||||
value=Decimal("1000000000000000000"),
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50000000000"),
|
||||
)
|
||||
mock_client.get_wallet_info.return_value = WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=500,
|
||||
balance_wei=Decimal("100000000000000000000"),
|
||||
first_transaction=first_tx,
|
||||
)
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client)
|
||||
profile = await analyzer.analyze(VALID_ADDRESS)
|
||||
|
||||
assert profile.nonce == 500
|
||||
assert profile.is_fresh is False
|
||||
assert profile.age_hours is not None
|
||||
assert profile.age_hours > 8000 # Over 365 days in hours
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_brand_new_wallet(self, mock_client: AsyncMock) -> None:
|
||||
"""Test analyzing a wallet with no transactions."""
|
||||
mock_client.get_wallet_info.return_value = WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=0,
|
||||
balance_wei=Decimal("1000000000000000000"),
|
||||
first_transaction=None,
|
||||
)
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client)
|
||||
profile = await analyzer.analyze(VALID_ADDRESS)
|
||||
|
||||
assert profile.nonce == 0
|
||||
assert profile.is_fresh is True
|
||||
assert profile.is_brand_new is True
|
||||
assert profile.first_seen is None
|
||||
assert profile.age_hours is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_uses_cache(self, mock_client: AsyncMock, mock_redis: AsyncMock) -> None:
|
||||
"""Test that analyze uses cached data."""
|
||||
cached_data = {
|
||||
"address": VALID_ADDRESS.lower(),
|
||||
"nonce": 2,
|
||||
"first_seen": datetime.now(UTC).isoformat(),
|
||||
"age_hours": 6.0,
|
||||
"is_fresh": True,
|
||||
"total_tx_count": 2,
|
||||
"matic_balance": "1000000000000000000",
|
||||
"usdc_balance": "500000",
|
||||
"analyzed_at": datetime.now(UTC).isoformat(),
|
||||
"fresh_threshold": 5,
|
||||
}
|
||||
mock_redis.get = AsyncMock(return_value=str(cached_data).replace("'", '"').encode())
|
||||
|
||||
# Actually mock it properly with json
|
||||
import json
|
||||
|
||||
mock_redis.get = AsyncMock(return_value=json.dumps(cached_data).encode())
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client, redis=mock_redis)
|
||||
profile = await analyzer.analyze(VALID_ADDRESS)
|
||||
|
||||
assert profile.address == VALID_ADDRESS.lower()
|
||||
assert profile.nonce == 2
|
||||
mock_client.get_wallet_info.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_force_refresh(
|
||||
self, mock_client: AsyncMock, mock_redis: AsyncMock
|
||||
) -> None:
|
||||
"""Test that force_refresh bypasses cache."""
|
||||
import json
|
||||
|
||||
cached_data = {
|
||||
"address": VALID_ADDRESS.lower(),
|
||||
"nonce": 1,
|
||||
"first_seen": None,
|
||||
"age_hours": None,
|
||||
"is_fresh": True,
|
||||
"total_tx_count": 1,
|
||||
"matic_balance": "1000",
|
||||
"usdc_balance": "0",
|
||||
"analyzed_at": datetime.now(UTC).isoformat(),
|
||||
"fresh_threshold": 5,
|
||||
}
|
||||
mock_redis.get = AsyncMock(return_value=json.dumps(cached_data).encode())
|
||||
|
||||
mock_client.get_wallet_info.return_value = WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=10,
|
||||
balance_wei=Decimal("5000000000000000000"),
|
||||
first_transaction=None,
|
||||
)
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client, redis=mock_redis)
|
||||
profile = await analyzer.analyze(VALID_ADDRESS, force_refresh=True)
|
||||
|
||||
assert profile.nonce == 10 # From fresh query, not cache
|
||||
mock_client.get_wallet_info.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_handles_usdc_error(self, mock_client: AsyncMock) -> None:
|
||||
"""Test that USDC balance error is handled gracefully."""
|
||||
mock_client.get_wallet_info.return_value = WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=1,
|
||||
balance_wei=Decimal("1000000000000000000"),
|
||||
first_transaction=None,
|
||||
)
|
||||
mock_client.get_token_balance.side_effect = Exception("Token query failed")
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client)
|
||||
profile = await analyzer.analyze(VALID_ADDRESS)
|
||||
|
||||
assert profile.usdc_balance == Decimal(0)
|
||||
|
||||
|
||||
class TestWalletAnalyzerIsFresh:
|
||||
"""Tests for the is_fresh method."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> AsyncMock:
|
||||
"""Create a mock PolygonClient."""
|
||||
client = AsyncMock()
|
||||
client.get_wallet_info = AsyncMock()
|
||||
client.get_token_balance = AsyncMock(return_value=Decimal("0"))
|
||||
return client
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_fresh_true(self, mock_client: AsyncMock) -> None:
|
||||
"""Test is_fresh returns True for fresh wallet."""
|
||||
mock_client.get_wallet_info.return_value = WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=2,
|
||||
balance_wei=Decimal("1000000000000000000"),
|
||||
first_transaction=None,
|
||||
)
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client)
|
||||
result = await analyzer.is_fresh(VALID_ADDRESS)
|
||||
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_fresh_false(self, mock_client: AsyncMock) -> None:
|
||||
"""Test is_fresh returns False for old wallet."""
|
||||
mock_client.get_wallet_info.return_value = WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=100,
|
||||
balance_wei=Decimal("1000000000000000000"),
|
||||
first_transaction=None,
|
||||
)
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client)
|
||||
result = await analyzer.is_fresh(VALID_ADDRESS)
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestWalletAnalyzerFreshnessLogic:
|
||||
"""Tests for freshness determination logic."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> AsyncMock:
|
||||
"""Create a mock PolygonClient."""
|
||||
return AsyncMock()
|
||||
|
||||
def test_is_wallet_fresh_low_nonce_no_age(self, mock_client: AsyncMock) -> None:
|
||||
"""Test fresh wallet with low nonce and unknown age."""
|
||||
analyzer = WalletAnalyzer(mock_client, fresh_threshold=5)
|
||||
result = analyzer._is_wallet_fresh(nonce=2, age_hours=None)
|
||||
assert result is True
|
||||
|
||||
def test_is_wallet_fresh_low_nonce_young_age(self, mock_client: AsyncMock) -> None:
|
||||
"""Test fresh wallet with low nonce and young age."""
|
||||
analyzer = WalletAnalyzer(mock_client, fresh_threshold=5)
|
||||
result = analyzer._is_wallet_fresh(nonce=2, age_hours=12.0)
|
||||
assert result is True
|
||||
|
||||
def test_is_wallet_fresh_low_nonce_old_age(self, mock_client: AsyncMock) -> None:
|
||||
"""Test not fresh when nonce is low but age is old."""
|
||||
analyzer = WalletAnalyzer(mock_client, fresh_threshold=5)
|
||||
result = analyzer._is_wallet_fresh(nonce=2, age_hours=100.0)
|
||||
assert result is False
|
||||
|
||||
def test_is_wallet_fresh_high_nonce(self, mock_client: AsyncMock) -> None:
|
||||
"""Test not fresh when nonce is high."""
|
||||
analyzer = WalletAnalyzer(mock_client, fresh_threshold=5)
|
||||
result = analyzer._is_wallet_fresh(nonce=10, age_hours=12.0)
|
||||
assert result is False
|
||||
|
||||
def test_is_wallet_fresh_at_threshold(self, mock_client: AsyncMock) -> None:
|
||||
"""Test not fresh when nonce equals threshold."""
|
||||
analyzer = WalletAnalyzer(mock_client, fresh_threshold=5)
|
||||
result = analyzer._is_wallet_fresh(nonce=5, age_hours=12.0)
|
||||
assert result is False
|
||||
|
||||
def test_is_wallet_fresh_at_age_boundary(self, mock_client: AsyncMock) -> None:
|
||||
"""Test at 48 hour boundary."""
|
||||
analyzer = WalletAnalyzer(mock_client, fresh_threshold=5)
|
||||
|
||||
result_under = analyzer._is_wallet_fresh(nonce=2, age_hours=47.9)
|
||||
assert result_under is True
|
||||
|
||||
result_over = analyzer._is_wallet_fresh(nonce=2, age_hours=48.1)
|
||||
assert result_over is False
|
||||
|
||||
|
||||
class TestWalletAnalyzerBatch:
|
||||
"""Tests for batch analysis."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> AsyncMock:
|
||||
"""Create a mock PolygonClient."""
|
||||
client = AsyncMock()
|
||||
client.get_wallet_info = AsyncMock()
|
||||
client.get_token_balance = AsyncMock(return_value=Decimal("0"))
|
||||
return client
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_batch(self, mock_client: AsyncMock) -> None:
|
||||
"""Test batch analysis."""
|
||||
mock_client.get_wallet_info.side_effect = [
|
||||
WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=2,
|
||||
balance_wei=Decimal("1000"),
|
||||
first_transaction=None,
|
||||
),
|
||||
WalletInfo(
|
||||
address=VALID_ADDRESS_2.lower(),
|
||||
transaction_count=100,
|
||||
balance_wei=Decimal("2000"),
|
||||
first_transaction=None,
|
||||
),
|
||||
]
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client)
|
||||
profiles = await analyzer.analyze_batch([VALID_ADDRESS, VALID_ADDRESS_2])
|
||||
|
||||
assert len(profiles) == 2
|
||||
assert profiles[VALID_ADDRESS.lower()].is_fresh is True
|
||||
assert profiles[VALID_ADDRESS_2.lower()].is_fresh is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_batch_handles_errors(self, mock_client: AsyncMock) -> None:
|
||||
"""Test batch analysis handles individual failures."""
|
||||
mock_client.get_wallet_info.side_effect = [
|
||||
WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=2,
|
||||
balance_wei=Decimal("1000"),
|
||||
first_transaction=None,
|
||||
),
|
||||
Exception("RPC error"),
|
||||
]
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client)
|
||||
profiles = await analyzer.analyze_batch([VALID_ADDRESS, VALID_ADDRESS_2])
|
||||
|
||||
assert len(profiles) == 1
|
||||
assert VALID_ADDRESS.lower() in profiles
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_fresh_wallets(self, mock_client: AsyncMock) -> None:
|
||||
"""Test filtering to only fresh wallets."""
|
||||
mock_client.get_wallet_info.side_effect = [
|
||||
WalletInfo(
|
||||
address=VALID_ADDRESS.lower(),
|
||||
transaction_count=2,
|
||||
balance_wei=Decimal("1000"),
|
||||
first_transaction=None,
|
||||
),
|
||||
WalletInfo(
|
||||
address=VALID_ADDRESS_2.lower(),
|
||||
transaction_count=100,
|
||||
balance_wei=Decimal("2000"),
|
||||
first_transaction=None,
|
||||
),
|
||||
]
|
||||
|
||||
analyzer = WalletAnalyzer(mock_client)
|
||||
fresh = await analyzer.get_fresh_wallets([VALID_ADDRESS, VALID_ADDRESS_2])
|
||||
|
||||
assert len(fresh) == 1
|
||||
assert VALID_ADDRESS.lower() in fresh
|
||||
@@ -0,0 +1,494 @@
|
||||
"""Tests for the Polygon blockchain client."""
|
||||
|
||||
import asyncio
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from web3.exceptions import Web3Exception
|
||||
|
||||
from polymarket_insider_tracker.profiler.chain import (
|
||||
DEFAULT_CACHE_TTL_SECONDS,
|
||||
PolygonClient,
|
||||
RateLimiter,
|
||||
RPCError,
|
||||
)
|
||||
|
||||
# Valid Ethereum addresses for testing
|
||||
VALID_ADDRESS = "0x742d35Cc6634C0532925a3b844Bc9e7595f5eaE2"
|
||||
VALID_ADDRESS_2 = "0x8ba1f109551bD432803012645Ac136ddd64DBA72"
|
||||
VALID_ADDRESS_3 = "0x1234567890AbCdEf1234567890ABcDeF12345678"
|
||||
VALID_TOKEN = "0x7D1AfA7B718fb893dB30A3aBc0Cfc608AaCfeBB0" # MATIC token
|
||||
|
||||
|
||||
class TestRateLimiter:
|
||||
"""Tests for the RateLimiter class."""
|
||||
|
||||
def test_create(self) -> None:
|
||||
"""Test creating a rate limiter."""
|
||||
limiter = RateLimiter.create(10.0)
|
||||
|
||||
assert limiter.max_tokens == 10.0
|
||||
assert limiter.refill_rate == 10.0
|
||||
assert limiter.tokens == 10.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_available(self) -> None:
|
||||
"""Test acquiring when tokens are available."""
|
||||
limiter = RateLimiter.create(10.0)
|
||||
|
||||
await limiter.acquire(1.0)
|
||||
|
||||
assert limiter.tokens < 10.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_multiple(self) -> None:
|
||||
"""Test acquiring multiple tokens."""
|
||||
limiter = RateLimiter.create(10.0)
|
||||
|
||||
for _ in range(5):
|
||||
await limiter.acquire(1.0)
|
||||
|
||||
assert limiter.tokens < 6.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_waits_when_empty(self) -> None:
|
||||
"""Test that acquire waits when tokens are depleted."""
|
||||
limiter = RateLimiter.create(2.0)
|
||||
|
||||
# Deplete tokens
|
||||
await limiter.acquire(2.0)
|
||||
|
||||
# This should wait briefly for refill
|
||||
start = asyncio.get_event_loop().time()
|
||||
await limiter.acquire(0.5)
|
||||
elapsed = asyncio.get_event_loop().time() - start
|
||||
|
||||
# Should have waited some time
|
||||
assert elapsed >= 0.1
|
||||
|
||||
|
||||
class TestPolygonClient:
|
||||
"""Tests for the PolygonClient class."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis(self) -> AsyncMock:
|
||||
"""Create a mock Redis client."""
|
||||
redis = AsyncMock()
|
||||
redis.get = AsyncMock(return_value=None)
|
||||
redis.set = AsyncMock()
|
||||
return redis
|
||||
|
||||
@pytest.fixture
|
||||
def mock_w3(self) -> MagicMock:
|
||||
"""Create a mock Web3 instance."""
|
||||
w3 = MagicMock()
|
||||
w3.eth = MagicMock()
|
||||
w3.eth.get_transaction_count = AsyncMock(return_value=42)
|
||||
w3.eth.get_balance = AsyncMock(return_value=1000000000000000000)
|
||||
w3.eth.get_block = AsyncMock(return_value={"timestamp": 1704369600})
|
||||
w3.eth.block_number = AsyncMock(return_value=50000000)
|
||||
return w3
|
||||
|
||||
def test_init(self) -> None:
|
||||
"""Test initialization."""
|
||||
client = PolygonClient("https://polygon-rpc.com")
|
||||
|
||||
assert client._rpc_url == "https://polygon-rpc.com"
|
||||
assert client._fallback_rpc_url is None
|
||||
assert client._cache_ttl == DEFAULT_CACHE_TTL_SECONDS
|
||||
|
||||
def test_init_with_fallback(self) -> None:
|
||||
"""Test initialization with fallback RPC."""
|
||||
client = PolygonClient(
|
||||
"https://polygon-rpc.com",
|
||||
fallback_rpc_url="https://fallback.com",
|
||||
)
|
||||
|
||||
assert client._fallback_rpc_url == "https://fallback.com"
|
||||
assert client._w3_fallback is not None
|
||||
|
||||
def test_init_custom_config(self) -> None:
|
||||
"""Test initialization with custom config."""
|
||||
client = PolygonClient(
|
||||
"https://polygon-rpc.com",
|
||||
cache_ttl_seconds=600,
|
||||
max_requests_per_second=50,
|
||||
max_retries=5,
|
||||
)
|
||||
|
||||
assert client._cache_ttl == 600
|
||||
assert client._max_retries == 5
|
||||
assert client._rate_limiter.max_tokens == 50
|
||||
|
||||
def test_cache_key(self) -> None:
|
||||
"""Test cache key generation."""
|
||||
client = PolygonClient("https://polygon-rpc.com")
|
||||
|
||||
key = client._cache_key("nonce", "0xAbC123")
|
||||
|
||||
assert key == "polygon:nonce:0xabc123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_miss(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test cache miss."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
result = await client._get_cached("test:key")
|
||||
|
||||
assert result is None
|
||||
mock_redis.get.assert_called_once_with("test:key")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_hit(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test cache hit."""
|
||||
mock_redis.get = AsyncMock(return_value=b"cached_value")
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
result = await client._get_cached("test:key")
|
||||
|
||||
assert result == "cached_value"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_error_handling(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test that cache errors are handled gracefully."""
|
||||
mock_redis.get = AsyncMock(side_effect=Exception("Redis error"))
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
result = await client._get_cached("test:key")
|
||||
|
||||
assert result is None # Should not raise
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_cached(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test setting cache."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
await client._set_cached("test:key", "value")
|
||||
|
||||
mock_redis.set.assert_called_once_with("test:key", "value", ex=DEFAULT_CACHE_TTL_SECONDS)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_cached_custom_ttl(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test setting cache with custom TTL."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
await client._set_cached("test:key", "value", ttl=3600)
|
||||
|
||||
mock_redis.set.assert_called_once_with("test:key", "value", ex=3600)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_transaction_count_cached(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting transaction count from cache."""
|
||||
mock_redis.get = AsyncMock(return_value=b"42")
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
count = await client.get_transaction_count(VALID_ADDRESS)
|
||||
|
||||
assert count == 42
|
||||
mock_redis.get.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_transaction_count_uncached(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting transaction count from blockchain."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
with patch.object(client, "_execute_with_retry", new_callable=AsyncMock) as mock_exec:
|
||||
mock_exec.return_value = 42
|
||||
|
||||
count = await client.get_transaction_count(VALID_ADDRESS)
|
||||
|
||||
assert count == 42
|
||||
mock_exec.assert_called_once()
|
||||
mock_redis.set.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_transaction_counts_batch(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test batch getting transaction counts."""
|
||||
mock_redis.get = AsyncMock(side_effect=[b"10", None, b"30"])
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
with patch.object(client, "get_transaction_count", new_callable=AsyncMock) as mock_get:
|
||||
mock_get.return_value = 20
|
||||
|
||||
addresses = [VALID_ADDRESS, VALID_ADDRESS_2, VALID_ADDRESS_3]
|
||||
counts = await client.get_transaction_counts(addresses)
|
||||
|
||||
assert counts[VALID_ADDRESS.lower()] == 10 # From cache
|
||||
assert counts[VALID_ADDRESS_2.lower()] == 20 # From blockchain
|
||||
assert counts[VALID_ADDRESS_3.lower()] == 30 # From cache
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_transaction_counts_empty(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test batch with empty list."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
counts = await client.get_transaction_counts([])
|
||||
|
||||
assert counts == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_balance_cached(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting balance from cache."""
|
||||
mock_redis.get = AsyncMock(return_value=b"1000000000000000000")
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
balance = await client.get_balance(VALID_ADDRESS)
|
||||
|
||||
assert balance == Decimal("1000000000000000000")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_balance_uncached(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting balance from blockchain."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
with patch.object(client, "_execute_with_retry", new_callable=AsyncMock) as mock_exec:
|
||||
mock_exec.return_value = 2000000000000000000
|
||||
|
||||
balance = await client.get_balance(VALID_ADDRESS)
|
||||
|
||||
assert balance == Decimal("2000000000000000000")
|
||||
mock_redis.set.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_wallet_info(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting aggregated wallet info."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
with (
|
||||
patch.object(client, "get_transaction_count", new_callable=AsyncMock) as mock_nonce,
|
||||
patch.object(client, "get_balance", new_callable=AsyncMock) as mock_balance,
|
||||
patch.object(client, "get_first_transaction", new_callable=AsyncMock) as mock_tx,
|
||||
):
|
||||
mock_nonce.return_value = 42
|
||||
mock_balance.return_value = Decimal("1000000000000000000")
|
||||
mock_tx.return_value = None
|
||||
|
||||
info = await client.get_wallet_info(VALID_ADDRESS)
|
||||
|
||||
assert info.address == VALID_ADDRESS.lower()
|
||||
assert info.transaction_count == 42
|
||||
assert info.balance_wei == Decimal("1000000000000000000")
|
||||
assert info.first_transaction is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_first_transaction_no_transactions(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test get_first_transaction when wallet has no transactions."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
with patch.object(client, "get_transaction_count", new_callable=AsyncMock) as mock_nonce:
|
||||
mock_nonce.return_value = 0
|
||||
|
||||
tx = await client.get_first_transaction(VALID_ADDRESS)
|
||||
|
||||
assert tx is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_success(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test successful health check."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
with patch.object(client, "_execute_with_retry", new_callable=AsyncMock) as mock_exec:
|
||||
mock_exec.return_value = 50000000
|
||||
|
||||
healthy = await client.health_check()
|
||||
|
||||
assert healthy is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_failure(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test failed health check."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
with patch.object(client, "_execute_with_retry", new_callable=AsyncMock) as mock_exec:
|
||||
mock_exec.side_effect = RPCError("Connection failed")
|
||||
|
||||
healthy = await client.health_check()
|
||||
|
||||
assert healthy is False
|
||||
|
||||
|
||||
class TestPolygonClientRetryLogic:
|
||||
"""Tests for retry and failover logic."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis(self) -> AsyncMock:
|
||||
"""Create a mock Redis client."""
|
||||
redis = AsyncMock()
|
||||
redis.get = AsyncMock(return_value=None)
|
||||
redis.set = AsyncMock()
|
||||
return redis
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_on_failure(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test that client retries on RPC failure."""
|
||||
client = PolygonClient(
|
||||
"https://polygon-rpc.com",
|
||||
redis=mock_redis,
|
||||
max_retries=3,
|
||||
retry_delay_seconds=0.01,
|
||||
)
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_get_tx_count(*_args: object, **_kwargs: object) -> int:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count < 3:
|
||||
raise Web3Exception("Temporary error")
|
||||
return 42
|
||||
|
||||
client._w3.eth.get_transaction_count = mock_get_tx_count
|
||||
|
||||
count = await client.get_transaction_count(VALID_ADDRESS)
|
||||
|
||||
assert count == 42
|
||||
assert call_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failover_to_secondary(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test failover to secondary RPC."""
|
||||
client = PolygonClient(
|
||||
"https://polygon-rpc.com",
|
||||
fallback_rpc_url="https://fallback.com",
|
||||
redis=mock_redis,
|
||||
max_retries=1,
|
||||
retry_delay_seconds=0.01,
|
||||
)
|
||||
|
||||
# Primary always fails
|
||||
async def primary_fail(*_args: object, **_kwargs: object) -> int:
|
||||
raise Web3Exception("Primary down")
|
||||
|
||||
client._w3.eth.get_transaction_count = primary_fail
|
||||
|
||||
# Fallback works
|
||||
client._w3_fallback.eth.get_transaction_count = AsyncMock(return_value=42)
|
||||
|
||||
count = await client.get_transaction_count(VALID_ADDRESS)
|
||||
|
||||
assert count == 42
|
||||
assert not client._primary_healthy
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_retries_exhausted(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test error when all retries are exhausted."""
|
||||
client = PolygonClient(
|
||||
"https://polygon-rpc.com",
|
||||
redis=mock_redis,
|
||||
max_retries=2,
|
||||
retry_delay_seconds=0.01,
|
||||
)
|
||||
|
||||
async def always_fail(*_args: object, **_kwargs: object) -> int:
|
||||
raise Web3Exception("Always fails")
|
||||
|
||||
client._w3.eth.get_transaction_count = always_fail
|
||||
|
||||
with pytest.raises(RPCError):
|
||||
await client.get_transaction_count(VALID_ADDRESS)
|
||||
|
||||
|
||||
class TestPolygonClientRateLimiting:
|
||||
"""Tests for rate limiting."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limiting_enforced(self) -> None:
|
||||
"""Test that rate limiting delays requests."""
|
||||
redis = AsyncMock()
|
||||
redis.get = AsyncMock(return_value=None)
|
||||
redis.set = AsyncMock()
|
||||
|
||||
client = PolygonClient(
|
||||
"https://polygon-rpc.com",
|
||||
redis=redis,
|
||||
max_requests_per_second=5.0,
|
||||
)
|
||||
|
||||
# Deplete rate limit
|
||||
client._rate_limiter.tokens = 0
|
||||
|
||||
with patch.object(client._w3.eth, "get_transaction_count", new_callable=AsyncMock) as mock:
|
||||
mock.return_value = 42
|
||||
|
||||
start = asyncio.get_event_loop().time()
|
||||
await client.get_transaction_count(VALID_ADDRESS)
|
||||
elapsed = asyncio.get_event_loop().time() - start
|
||||
|
||||
# Should have waited for token refill
|
||||
assert elapsed >= 0.1
|
||||
|
||||
|
||||
class TestPolygonClientTokenBalance:
|
||||
"""Tests for ERC20 token balance queries."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis(self) -> AsyncMock:
|
||||
"""Create a mock Redis client."""
|
||||
redis = AsyncMock()
|
||||
redis.get = AsyncMock(return_value=None)
|
||||
redis.set = AsyncMock()
|
||||
return redis
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_token_balance_cached(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting token balance from cache."""
|
||||
mock_redis.get = AsyncMock(return_value=b"1000000")
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
balance = await client.get_token_balance(VALID_ADDRESS, VALID_TOKEN)
|
||||
|
||||
assert balance == Decimal("1000000")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_token_balance_uncached(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting token balance from blockchain."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
# Mock the contract call
|
||||
mock_contract = MagicMock()
|
||||
mock_contract.functions.balanceOf.return_value.call = AsyncMock(return_value=5000000)
|
||||
client._w3.eth.contract = MagicMock(return_value=mock_contract)
|
||||
|
||||
balance = await client.get_token_balance(VALID_ADDRESS, VALID_TOKEN)
|
||||
|
||||
assert balance == Decimal("5000000")
|
||||
mock_redis.set.assert_called_once()
|
||||
|
||||
|
||||
class TestPolygonClientBlock:
|
||||
"""Tests for block queries."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis(self) -> AsyncMock:
|
||||
"""Create a mock Redis client."""
|
||||
redis = AsyncMock()
|
||||
redis.get = AsyncMock(return_value=None)
|
||||
redis.set = AsyncMock()
|
||||
return redis
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_block_cached(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting block from cache."""
|
||||
mock_redis.get = AsyncMock(return_value=b'{"timestamp": 1704369600}')
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
block = await client.get_block(50000000)
|
||||
|
||||
assert block["timestamp"] == 1704369600
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_block_uncached(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting block from blockchain."""
|
||||
client = PolygonClient("https://polygon-rpc.com", redis=mock_redis)
|
||||
|
||||
with patch.object(client, "_execute_with_retry", new_callable=AsyncMock) as mock_exec:
|
||||
mock_exec.return_value = {"timestamp": 1704369600, "number": 50000000}
|
||||
|
||||
block = await client.get_block(50000000)
|
||||
|
||||
assert block["timestamp"] == 1704369600
|
||||
# Block cache uses 1 hour TTL
|
||||
mock_redis.set.assert_called_once()
|
||||
call_args = mock_redis.set.call_args
|
||||
assert call_args[1]["ex"] == 3600
|
||||
@@ -0,0 +1,365 @@
|
||||
"""Tests for known entity registry."""
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.profiler.entities import EntityRegistry
|
||||
from polymarket_insider_tracker.profiler.entity_data import (
|
||||
BRIDGE_ADDRESSES,
|
||||
CEX_ADDRESSES,
|
||||
DEFI_ADDRESSES,
|
||||
DEX_ADDRESSES,
|
||||
TOKEN_ADDRESSES,
|
||||
EntityType,
|
||||
get_all_known_entities,
|
||||
)
|
||||
|
||||
# ============================================================================
|
||||
# EntityType Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestEntityType:
|
||||
"""Tests for EntityType enum."""
|
||||
|
||||
def test_cex_types_exist(self) -> None:
|
||||
"""Test that CEX entity types are defined."""
|
||||
assert EntityType.CEX_BINANCE.value == "cex_binance"
|
||||
assert EntityType.CEX_COINBASE.value == "cex_coinbase"
|
||||
assert EntityType.CEX_OTHER.value == "cex_other"
|
||||
|
||||
def test_bridge_types_exist(self) -> None:
|
||||
"""Test that bridge entity types are defined."""
|
||||
assert EntityType.BRIDGE_POLYGON.value == "bridge_polygon"
|
||||
assert EntityType.BRIDGE_MULTICHAIN.value == "bridge_multichain"
|
||||
|
||||
def test_dex_types_exist(self) -> None:
|
||||
"""Test that DEX entity types are defined."""
|
||||
assert EntityType.DEX_UNISWAP.value == "dex_uniswap"
|
||||
assert EntityType.DEX_SUSHISWAP.value == "dex_sushiswap"
|
||||
|
||||
def test_token_types_exist(self) -> None:
|
||||
"""Test that token entity types are defined."""
|
||||
assert EntityType.TOKEN_USDC.value == "token_usdc"
|
||||
assert EntityType.TOKEN_WETH.value == "token_weth"
|
||||
|
||||
def test_unknown_type(self) -> None:
|
||||
"""Test unknown entity type."""
|
||||
assert EntityType.UNKNOWN.value == "unknown"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Entity Data Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestEntityData:
|
||||
"""Tests for entity data mappings."""
|
||||
|
||||
def test_cex_addresses_populated(self) -> None:
|
||||
"""Test that CEX addresses are populated."""
|
||||
assert len(CEX_ADDRESSES) > 0
|
||||
# Check Binance address is present
|
||||
binance_found = any(entity == EntityType.CEX_BINANCE for entity in CEX_ADDRESSES.values())
|
||||
assert binance_found
|
||||
|
||||
def test_bridge_addresses_populated(self) -> None:
|
||||
"""Test that bridge addresses are populated."""
|
||||
assert len(BRIDGE_ADDRESSES) > 0
|
||||
|
||||
def test_dex_addresses_populated(self) -> None:
|
||||
"""Test that DEX addresses are populated."""
|
||||
assert len(DEX_ADDRESSES) > 0
|
||||
# Check Uniswap is present
|
||||
uniswap_found = any(entity == EntityType.DEX_UNISWAP for entity in DEX_ADDRESSES.values())
|
||||
assert uniswap_found
|
||||
|
||||
def test_token_addresses_include_usdc(self) -> None:
|
||||
"""Test that USDC address is in token addresses."""
|
||||
usdc_address = "0x2791bca1f2de4661ed88a30c99a7a9449aa84174"
|
||||
assert usdc_address in TOKEN_ADDRESSES
|
||||
assert TOKEN_ADDRESSES[usdc_address] == EntityType.TOKEN_USDC
|
||||
|
||||
def test_get_all_known_entities(self) -> None:
|
||||
"""Test combining all entity mappings."""
|
||||
all_entities = get_all_known_entities()
|
||||
total_expected = (
|
||||
len(CEX_ADDRESSES)
|
||||
+ len(BRIDGE_ADDRESSES)
|
||||
+ len(DEX_ADDRESSES)
|
||||
+ len(TOKEN_ADDRESSES)
|
||||
+ len(DEFI_ADDRESSES)
|
||||
)
|
||||
assert len(all_entities) == total_expected
|
||||
|
||||
def test_addresses_are_lowercase(self) -> None:
|
||||
"""Test that all addresses in get_all_known_entities are lowercase."""
|
||||
all_entities = get_all_known_entities()
|
||||
for address in all_entities:
|
||||
assert address == address.lower()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# EntityRegistry Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestEntityRegistryInit:
|
||||
"""Tests for EntityRegistry initialization."""
|
||||
|
||||
def test_default_initialization(self) -> None:
|
||||
"""Test registry initializes with default entities."""
|
||||
registry = EntityRegistry()
|
||||
assert len(registry) > 0
|
||||
|
||||
def test_without_defaults(self) -> None:
|
||||
"""Test registry without default entities."""
|
||||
registry = EntityRegistry(include_defaults=False)
|
||||
assert len(registry) == 0
|
||||
|
||||
def test_with_custom_entities(self) -> None:
|
||||
"""Test registry with custom entities."""
|
||||
custom = {"0x1234": EntityType.CEX_OTHER}
|
||||
registry = EntityRegistry(custom_entities=custom, include_defaults=False)
|
||||
assert len(registry) == 1
|
||||
assert registry.classify("0x1234") == EntityType.CEX_OTHER
|
||||
|
||||
def test_custom_entities_override_defaults(self) -> None:
|
||||
"""Test that custom entities can override defaults."""
|
||||
# USDC address is in defaults as TOKEN_USDC
|
||||
usdc_address = "0x2791bca1f2de4661ed88a30c99a7a9449aa84174"
|
||||
custom = {usdc_address: EntityType.CONTRACT}
|
||||
registry = EntityRegistry(custom_entities=custom)
|
||||
assert registry.classify(usdc_address) == EntityType.CONTRACT
|
||||
|
||||
|
||||
class TestEntityRegistryClassify:
|
||||
"""Tests for EntityRegistry.classify method."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self) -> EntityRegistry:
|
||||
"""Create a registry for testing."""
|
||||
return EntityRegistry()
|
||||
|
||||
def test_classify_known_cex(self, registry: EntityRegistry) -> None:
|
||||
"""Test classifying a known CEX address."""
|
||||
# Binance address
|
||||
binance = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
assert registry.classify(binance) == EntityType.CEX_BINANCE
|
||||
|
||||
def test_classify_case_insensitive(self, registry: EntityRegistry) -> None:
|
||||
"""Test that classification is case-insensitive."""
|
||||
binance_lower = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
binance_mixed = "0x28C6c06298D514db089934071355E5743bf21d60"
|
||||
assert registry.classify(binance_lower) == registry.classify(binance_mixed)
|
||||
|
||||
def test_classify_unknown(self, registry: EntityRegistry) -> None:
|
||||
"""Test classifying an unknown address."""
|
||||
unknown = "0x0000000000000000000000000000000000000000"
|
||||
assert registry.classify(unknown) == EntityType.UNKNOWN
|
||||
|
||||
def test_classify_usdc(self, registry: EntityRegistry) -> None:
|
||||
"""Test classifying USDC token contract."""
|
||||
usdc = "0x2791bca1f2de4661ed88a30c99a7a9449aa84174"
|
||||
assert registry.classify(usdc) == EntityType.TOKEN_USDC
|
||||
|
||||
|
||||
class TestEntityRegistryChecks:
|
||||
"""Tests for EntityRegistry type check methods."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self) -> EntityRegistry:
|
||||
"""Create a registry for testing."""
|
||||
return EntityRegistry()
|
||||
|
||||
def test_is_known_entity_true(self, registry: EntityRegistry) -> None:
|
||||
"""Test is_known_entity returns True for known addresses."""
|
||||
binance = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
assert registry.is_known_entity(binance) is True
|
||||
|
||||
def test_is_known_entity_false(self, registry: EntityRegistry) -> None:
|
||||
"""Test is_known_entity returns False for unknown addresses."""
|
||||
unknown = "0x0000000000000000000000000000000000000000"
|
||||
assert registry.is_known_entity(unknown) is False
|
||||
|
||||
def test_is_cex_true(self, registry: EntityRegistry) -> None:
|
||||
"""Test is_cex returns True for CEX addresses."""
|
||||
binance = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
assert registry.is_cex(binance) is True
|
||||
|
||||
def test_is_cex_false(self, registry: EntityRegistry) -> None:
|
||||
"""Test is_cex returns False for non-CEX addresses."""
|
||||
usdc = "0x2791bca1f2de4661ed88a30c99a7a9449aa84174"
|
||||
assert registry.is_cex(usdc) is False
|
||||
|
||||
def test_is_bridge_true(self, registry: EntityRegistry) -> None:
|
||||
"""Test is_bridge returns True for bridge addresses."""
|
||||
polygon_bridge = "0xa0c68c638235ee32657e8f720a23cec1bfc77c77"
|
||||
assert registry.is_bridge(polygon_bridge) is True
|
||||
|
||||
def test_is_bridge_false(self, registry: EntityRegistry) -> None:
|
||||
"""Test is_bridge returns False for non-bridge addresses."""
|
||||
binance = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
assert registry.is_bridge(binance) is False
|
||||
|
||||
def test_is_dex_true(self, registry: EntityRegistry) -> None:
|
||||
"""Test is_dex returns True for DEX addresses."""
|
||||
uniswap = "0xe592427a0aece92de3edee1f18e0157c05861564"
|
||||
assert registry.is_dex(uniswap) is True
|
||||
|
||||
def test_is_dex_false(self, registry: EntityRegistry) -> None:
|
||||
"""Test is_dex returns False for non-DEX addresses."""
|
||||
binance = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
assert registry.is_dex(binance) is False
|
||||
|
||||
|
||||
class TestEntityRegistryTerminal:
|
||||
"""Tests for EntityRegistry.is_terminal method."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self) -> EntityRegistry:
|
||||
"""Create a registry for testing."""
|
||||
return EntityRegistry()
|
||||
|
||||
def test_cex_is_terminal(self, registry: EntityRegistry) -> None:
|
||||
"""Test that CEX addresses are terminal."""
|
||||
binance = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
assert registry.is_terminal(binance) is True
|
||||
|
||||
def test_bridge_is_terminal(self, registry: EntityRegistry) -> None:
|
||||
"""Test that bridge addresses are terminal."""
|
||||
polygon_bridge = "0xa0c68c638235ee32657e8f720a23cec1bfc77c77"
|
||||
assert registry.is_terminal(polygon_bridge) is True
|
||||
|
||||
def test_dex_is_not_terminal(self, registry: EntityRegistry) -> None:
|
||||
"""Test that DEX addresses are not terminal."""
|
||||
uniswap = "0xe592427a0aece92de3edee1f18e0157c05861564"
|
||||
assert registry.is_terminal(uniswap) is False
|
||||
|
||||
def test_token_is_not_terminal(self, registry: EntityRegistry) -> None:
|
||||
"""Test that token addresses are not terminal."""
|
||||
usdc = "0x2791bca1f2de4661ed88a30c99a7a9449aa84174"
|
||||
assert registry.is_terminal(usdc) is False
|
||||
|
||||
def test_unknown_is_not_terminal(self, registry: EntityRegistry) -> None:
|
||||
"""Test that unknown addresses are not terminal."""
|
||||
unknown = "0x0000000000000000000000000000000000000000"
|
||||
assert registry.is_terminal(unknown) is False
|
||||
|
||||
|
||||
class TestEntityRegistryCategory:
|
||||
"""Tests for EntityRegistry.get_entity_category method."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self) -> EntityRegistry:
|
||||
"""Create a registry for testing."""
|
||||
return EntityRegistry()
|
||||
|
||||
def test_category_cex(self, registry: EntityRegistry) -> None:
|
||||
"""Test CEX category."""
|
||||
binance = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
assert registry.get_entity_category(binance) == "cex"
|
||||
|
||||
def test_category_bridge(self, registry: EntityRegistry) -> None:
|
||||
"""Test bridge category."""
|
||||
polygon_bridge = "0xa0c68c638235ee32657e8f720a23cec1bfc77c77"
|
||||
assert registry.get_entity_category(polygon_bridge) == "bridge"
|
||||
|
||||
def test_category_dex(self, registry: EntityRegistry) -> None:
|
||||
"""Test DEX category."""
|
||||
uniswap = "0xe592427a0aece92de3edee1f18e0157c05861564"
|
||||
assert registry.get_entity_category(uniswap) == "dex"
|
||||
|
||||
def test_category_token(self, registry: EntityRegistry) -> None:
|
||||
"""Test token category."""
|
||||
usdc = "0x2791bca1f2de4661ed88a30c99a7a9449aa84174"
|
||||
assert registry.get_entity_category(usdc) == "token"
|
||||
|
||||
def test_category_defi(self, registry: EntityRegistry) -> None:
|
||||
"""Test DeFi category."""
|
||||
aave = "0x794a61358d6845594f94dc1db02a252b5b4814ad"
|
||||
assert registry.get_entity_category(aave) == "defi"
|
||||
|
||||
def test_category_unknown(self, registry: EntityRegistry) -> None:
|
||||
"""Test unknown category."""
|
||||
unknown = "0x0000000000000000000000000000000000000000"
|
||||
assert registry.get_entity_category(unknown) == "unknown"
|
||||
|
||||
|
||||
class TestEntityRegistryMutations:
|
||||
"""Tests for EntityRegistry mutation methods."""
|
||||
|
||||
def test_add_entity(self) -> None:
|
||||
"""Test adding an entity."""
|
||||
registry = EntityRegistry(include_defaults=False)
|
||||
registry.add_entity("0x1234", EntityType.CEX_OTHER)
|
||||
assert registry.classify("0x1234") == EntityType.CEX_OTHER
|
||||
|
||||
def test_add_entity_normalizes_address(self) -> None:
|
||||
"""Test that add_entity normalizes addresses to lowercase."""
|
||||
registry = EntityRegistry(include_defaults=False)
|
||||
registry.add_entity("0xABCD", EntityType.CEX_OTHER)
|
||||
assert registry.classify("0xabcd") == EntityType.CEX_OTHER
|
||||
|
||||
def test_remove_entity(self) -> None:
|
||||
"""Test removing an entity."""
|
||||
registry = EntityRegistry(include_defaults=False)
|
||||
registry.add_entity("0x1234", EntityType.CEX_OTHER)
|
||||
assert registry.remove_entity("0x1234") is True
|
||||
assert registry.classify("0x1234") == EntityType.UNKNOWN
|
||||
|
||||
def test_remove_nonexistent(self) -> None:
|
||||
"""Test removing a non-existent entity."""
|
||||
registry = EntityRegistry(include_defaults=False)
|
||||
assert registry.remove_entity("0x1234") is False
|
||||
|
||||
|
||||
class TestEntityRegistryDunder:
|
||||
"""Tests for EntityRegistry dunder methods."""
|
||||
|
||||
def test_len(self) -> None:
|
||||
"""Test __len__ returns count of entities."""
|
||||
registry = EntityRegistry(include_defaults=False)
|
||||
registry.add_entity("0x1234", EntityType.CEX_OTHER)
|
||||
registry.add_entity("0x5678", EntityType.DEX_OTHER)
|
||||
assert len(registry) == 2
|
||||
|
||||
def test_contains(self) -> None:
|
||||
"""Test __contains__ for membership testing."""
|
||||
registry = EntityRegistry()
|
||||
binance = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
assert binance in registry
|
||||
assert "0x0000000000000000000000000000000000000000" not in registry
|
||||
|
||||
|
||||
class TestEntityRegistryContract:
|
||||
"""Tests for EntityRegistry.is_contract method."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self) -> EntityRegistry:
|
||||
"""Create a registry for testing."""
|
||||
return EntityRegistry()
|
||||
|
||||
def test_dex_is_contract(self, registry: EntityRegistry) -> None:
|
||||
"""Test that DEX addresses are contracts."""
|
||||
uniswap = "0xe592427a0aece92de3edee1f18e0157c05861564"
|
||||
assert registry.is_contract(uniswap) is True
|
||||
|
||||
def test_token_is_contract(self, registry: EntityRegistry) -> None:
|
||||
"""Test that token addresses are contracts."""
|
||||
usdc = "0x2791bca1f2de4661ed88a30c99a7a9449aa84174"
|
||||
assert registry.is_contract(usdc) is True
|
||||
|
||||
def test_defi_is_contract(self, registry: EntityRegistry) -> None:
|
||||
"""Test that DeFi protocol addresses are contracts."""
|
||||
aave = "0x794a61358d6845594f94dc1db02a252b5b4814ad"
|
||||
assert registry.is_contract(aave) is True
|
||||
|
||||
def test_cex_is_not_contract(self, registry: EntityRegistry) -> None:
|
||||
"""Test that CEX addresses are not contracts."""
|
||||
binance = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
assert registry.is_contract(binance) is False
|
||||
|
||||
def test_unknown_is_not_contract(self, registry: EntityRegistry) -> None:
|
||||
"""Test that unknown addresses are not contracts."""
|
||||
unknown = "0x0000000000000000000000000000000000000000"
|
||||
assert registry.is_contract(unknown) is False
|
||||
@@ -0,0 +1,738 @@
|
||||
"""Tests for the FundingTracer module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.profiler.entities import EntityRegistry
|
||||
from polymarket_insider_tracker.profiler.entity_data import EntityType
|
||||
from polymarket_insider_tracker.profiler.funding import (
|
||||
TRANSFER_EVENT_SIGNATURE,
|
||||
USDC_BRIDGED,
|
||||
USDC_NATIVE,
|
||||
FundingTracer,
|
||||
)
|
||||
from polymarket_insider_tracker.profiler.models import FundingChain, FundingTransfer
|
||||
|
||||
# Test addresses
|
||||
TEST_WALLET = "0x1234567890abcdef1234567890abcdef12345678"
|
||||
TEST_SOURCE = "0xabcdef1234567890abcdef1234567890abcdef12"
|
||||
BINANCE_HOT_WALLET = "0x28c6c06298d514db089934071355e5743bf21d60"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_polygon_client() -> MagicMock:
|
||||
"""Create a mock PolygonClient."""
|
||||
client = MagicMock()
|
||||
client._rate_limiter = MagicMock()
|
||||
client._rate_limiter.acquire = AsyncMock()
|
||||
client._primary_healthy = True
|
||||
client._w3 = MagicMock()
|
||||
client._w3_fallback = None
|
||||
client.get_block = AsyncMock(return_value={"timestamp": 1704067200})
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def entity_registry() -> EntityRegistry:
|
||||
"""Create an EntityRegistry with default entities."""
|
||||
return EntityRegistry()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def funding_tracer(
|
||||
mock_polygon_client: MagicMock,
|
||||
entity_registry: EntityRegistry,
|
||||
) -> FundingTracer:
|
||||
"""Create a FundingTracer with mocked dependencies."""
|
||||
return FundingTracer(
|
||||
polygon_client=mock_polygon_client,
|
||||
entity_registry=entity_registry,
|
||||
max_hops=3,
|
||||
)
|
||||
|
||||
|
||||
class TestFundingTracerInit:
|
||||
"""Tests for FundingTracer initialization."""
|
||||
|
||||
def test_init_with_defaults(self, mock_polygon_client: MagicMock) -> None:
|
||||
"""Test initialization with default parameters."""
|
||||
tracer = FundingTracer(mock_polygon_client)
|
||||
|
||||
assert tracer.polygon_client is mock_polygon_client
|
||||
assert tracer.max_hops == 3
|
||||
assert USDC_BRIDGED.lower() in tracer._usdc_addresses
|
||||
assert USDC_NATIVE.lower() in tracer._usdc_addresses
|
||||
|
||||
def test_init_with_custom_max_hops(self, mock_polygon_client: MagicMock) -> None:
|
||||
"""Test initialization with custom max_hops."""
|
||||
tracer = FundingTracer(mock_polygon_client, max_hops=5)
|
||||
assert tracer.max_hops == 5
|
||||
|
||||
def test_init_with_custom_usdc_addresses(self, mock_polygon_client: MagicMock) -> None:
|
||||
"""Test initialization with custom USDC addresses."""
|
||||
custom_addresses = ["0x1111111111111111111111111111111111111111"]
|
||||
tracer = FundingTracer(mock_polygon_client, usdc_addresses=custom_addresses)
|
||||
assert tracer._usdc_addresses == [custom_addresses[0].lower()]
|
||||
|
||||
def test_init_with_custom_entity_registry(self, mock_polygon_client: MagicMock) -> None:
|
||||
"""Test initialization with custom entity registry."""
|
||||
registry = EntityRegistry()
|
||||
tracer = FundingTracer(mock_polygon_client, entity_registry=registry)
|
||||
assert tracer.entity_registry is registry
|
||||
|
||||
def test_init_creates_default_entity_registry(self, mock_polygon_client: MagicMock) -> None:
|
||||
"""Test initialization creates default EntityRegistry if None."""
|
||||
tracer = FundingTracer(mock_polygon_client, entity_registry=None)
|
||||
assert isinstance(tracer.entity_registry, EntityRegistry)
|
||||
|
||||
|
||||
class TestFundingTracerTrace:
|
||||
"""Tests for the trace method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_terminates_at_known_cex(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test trace terminates when starting at a CEX address."""
|
||||
result = await funding_tracer.trace(BINANCE_HOT_WALLET)
|
||||
|
||||
assert result.target_address == BINANCE_HOT_WALLET.lower()
|
||||
assert result.origin_address == BINANCE_HOT_WALLET.lower()
|
||||
assert result.origin_type == EntityType.CEX_BINANCE.value
|
||||
assert result.hop_count == 0
|
||||
assert len(result.chain) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_no_transfers_found(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test trace when no USDC transfers are found."""
|
||||
funding_tracer._get_transfer_logs = AsyncMock(return_value=[])
|
||||
|
||||
result = await funding_tracer.trace(TEST_WALLET)
|
||||
|
||||
assert result.target_address == TEST_WALLET.lower()
|
||||
assert result.origin_address == TEST_WALLET.lower()
|
||||
assert result.origin_type == "unknown"
|
||||
assert result.hop_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_finds_cex_origin(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test trace finds CEX as funding origin."""
|
||||
# Mock a transfer from Binance to test wallet
|
||||
mock_log = _create_mock_log(
|
||||
from_address=BINANCE_HOT_WALLET,
|
||||
to_address=TEST_WALLET,
|
||||
amount=1000000, # 1 USDC
|
||||
tx_hash="0x" + "ab" * 32,
|
||||
block_number=50000000,
|
||||
)
|
||||
|
||||
funding_tracer._get_transfer_logs = AsyncMock(return_value=[mock_log])
|
||||
|
||||
result = await funding_tracer.trace(TEST_WALLET)
|
||||
|
||||
assert result.target_address == TEST_WALLET.lower()
|
||||
assert result.origin_address == BINANCE_HOT_WALLET.lower()
|
||||
assert result.origin_type == EntityType.CEX_BINANCE.value
|
||||
assert result.hop_count == 1
|
||||
assert len(result.chain) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_multiple_hops(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test trace follows multiple hops."""
|
||||
intermediate_wallet = "0x" + "11" * 20
|
||||
|
||||
# First call: TEST_WALLET received from intermediate
|
||||
# Second call: intermediate received from Binance
|
||||
mock_logs = [
|
||||
_create_mock_log(
|
||||
from_address=intermediate_wallet,
|
||||
to_address=TEST_WALLET,
|
||||
amount=1000000,
|
||||
tx_hash="0x" + "aa" * 32,
|
||||
block_number=50000001,
|
||||
),
|
||||
_create_mock_log(
|
||||
from_address=BINANCE_HOT_WALLET,
|
||||
to_address=intermediate_wallet,
|
||||
amount=1000000,
|
||||
tx_hash="0x" + "bb" * 32,
|
||||
block_number=50000000,
|
||||
),
|
||||
]
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_get_logs(*_args: Any, **_kwargs: Any) -> list[dict[str, Any]]:
|
||||
nonlocal call_count
|
||||
result = [mock_logs[call_count]] if call_count < len(mock_logs) else []
|
||||
call_count += 1
|
||||
return result
|
||||
|
||||
funding_tracer._get_transfer_logs = mock_get_logs
|
||||
|
||||
result = await funding_tracer.trace(TEST_WALLET)
|
||||
|
||||
assert result.hop_count == 2
|
||||
assert result.origin_address == BINANCE_HOT_WALLET.lower()
|
||||
assert result.origin_type == EntityType.CEX_BINANCE.value
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_respects_max_hops(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test trace stops at max_hops."""
|
||||
# Create a chain of unknown wallets
|
||||
wallets = [f"0x{i:040x}" for i in range(10)]
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_get_logs(*_args: Any, **_kwargs: Any) -> list[dict[str, Any]]:
|
||||
nonlocal call_count
|
||||
if call_count < len(wallets) - 1:
|
||||
log = _create_mock_log(
|
||||
from_address=wallets[call_count + 1],
|
||||
to_address=wallets[call_count],
|
||||
amount=1000000,
|
||||
tx_hash=f"0x{call_count:064x}",
|
||||
block_number=50000000 + call_count,
|
||||
)
|
||||
call_count += 1
|
||||
return [log]
|
||||
return []
|
||||
|
||||
funding_tracer._get_transfer_logs = mock_get_logs
|
||||
|
||||
result = await funding_tracer.trace(wallets[0], max_hops=3)
|
||||
|
||||
assert result.hop_count == 3
|
||||
assert result.origin_type == "unknown"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_override_max_hops(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test trace can override default max_hops."""
|
||||
funding_tracer._get_transfer_logs = AsyncMock(return_value=[])
|
||||
|
||||
# Override to 1 hop
|
||||
await funding_tracer.trace(TEST_WALLET, max_hops=1)
|
||||
|
||||
# Verify only 1 iteration (no hops since no transfers found)
|
||||
# The trace should have been called once for the target wallet
|
||||
|
||||
|
||||
class TestGetFirstUsdcTransfer:
|
||||
"""Tests for get_first_usdc_transfer method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_first_usdc_transfer_bridged(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test getting first USDC transfer from bridged contract."""
|
||||
mock_log = _create_mock_log(
|
||||
from_address=TEST_SOURCE,
|
||||
to_address=TEST_WALLET,
|
||||
amount=5000000,
|
||||
tx_hash="0x" + "cc" * 32,
|
||||
block_number=50000000,
|
||||
)
|
||||
|
||||
funding_tracer._get_transfer_logs = AsyncMock(return_value=[mock_log])
|
||||
|
||||
result = await funding_tracer.get_first_usdc_transfer(TEST_WALLET)
|
||||
|
||||
assert result is not None
|
||||
assert result.from_address == TEST_SOURCE.lower()
|
||||
assert result.to_address == TEST_WALLET.lower()
|
||||
assert result.amount == Decimal(5000000)
|
||||
assert result.token == "USDC"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_first_usdc_transfer_native(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test fallback to native USDC contract."""
|
||||
call_count = 0
|
||||
|
||||
async def mock_get_logs(*_args: Any, **_kwargs: Any) -> list[dict[str, Any]]:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1: # First call (bridged) returns nothing
|
||||
return []
|
||||
# Second call (native) returns a transfer
|
||||
return [
|
||||
_create_mock_log(
|
||||
from_address=TEST_SOURCE,
|
||||
to_address=TEST_WALLET,
|
||||
amount=1000000,
|
||||
tx_hash="0x" + "dd" * 32,
|
||||
block_number=50000000,
|
||||
)
|
||||
]
|
||||
|
||||
funding_tracer._get_transfer_logs = mock_get_logs
|
||||
|
||||
result = await funding_tracer.get_first_usdc_transfer(TEST_WALLET)
|
||||
|
||||
assert result is not None
|
||||
assert call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_first_usdc_transfer_none_found(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test returns None when no USDC transfers found."""
|
||||
funding_tracer._get_transfer_logs = AsyncMock(return_value=[])
|
||||
|
||||
result = await funding_tracer.get_first_usdc_transfer(TEST_WALLET)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestGetTransferLogs:
|
||||
"""Tests for _get_transfer_logs method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_transfer_logs_formats_topics_correctly(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
mock_polygon_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test that transfer logs query is formatted correctly."""
|
||||
mock_w3 = MagicMock()
|
||||
mock_w3.eth.get_logs = AsyncMock(return_value=[])
|
||||
mock_polygon_client._w3 = mock_w3
|
||||
|
||||
await funding_tracer._get_transfer_logs(
|
||||
to_address=TEST_WALLET,
|
||||
token_address=USDC_BRIDGED,
|
||||
)
|
||||
|
||||
mock_w3.eth.get_logs.assert_called_once()
|
||||
call_args = mock_w3.eth.get_logs.call_args[0][0]
|
||||
|
||||
# Verify topics structure
|
||||
assert len(call_args["topics"]) == 3
|
||||
assert call_args["topics"][0] == TRANSFER_EVENT_SIGNATURE.hex()
|
||||
assert call_args["topics"][1] is None # from (any)
|
||||
# to address should be padded to 32 bytes
|
||||
assert call_args["topics"][2].endswith(TEST_WALLET.lower().replace("0x", ""))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_transfer_logs_respects_limit(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
mock_polygon_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test that limit parameter works correctly."""
|
||||
mock_logs = [MagicMock() for _ in range(10)]
|
||||
mock_w3 = MagicMock()
|
||||
mock_w3.eth.get_logs = AsyncMock(return_value=mock_logs)
|
||||
mock_polygon_client._w3 = mock_w3
|
||||
|
||||
result = await funding_tracer._get_transfer_logs(
|
||||
to_address=TEST_WALLET,
|
||||
token_address=USDC_BRIDGED,
|
||||
limit=3,
|
||||
)
|
||||
|
||||
assert len(result) == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_transfer_logs_uses_fallback_when_primary_unhealthy(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
mock_polygon_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test fallback RPC is used when primary is unhealthy."""
|
||||
mock_polygon_client._primary_healthy = False
|
||||
mock_fallback = MagicMock()
|
||||
mock_fallback.eth.get_logs = AsyncMock(return_value=[])
|
||||
mock_polygon_client._w3_fallback = mock_fallback
|
||||
|
||||
await funding_tracer._get_transfer_logs(
|
||||
to_address=TEST_WALLET,
|
||||
token_address=USDC_BRIDGED,
|
||||
)
|
||||
|
||||
mock_fallback.eth.get_logs.assert_called_once()
|
||||
|
||||
|
||||
class TestLogToFundingTransfer:
|
||||
"""Tests for _log_to_funding_transfer method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_to_funding_transfer_parses_correctly(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test correct parsing of log to FundingTransfer."""
|
||||
mock_log = _create_mock_log(
|
||||
from_address=TEST_SOURCE,
|
||||
to_address=TEST_WALLET,
|
||||
amount=1500000,
|
||||
tx_hash="0x" + "ee" * 32,
|
||||
block_number=50000000,
|
||||
)
|
||||
|
||||
result = await funding_tracer._log_to_funding_transfer(mock_log, USDC_BRIDGED)
|
||||
|
||||
assert result.from_address == TEST_SOURCE.lower()
|
||||
assert result.to_address == TEST_WALLET.lower()
|
||||
assert result.amount == Decimal(1500000)
|
||||
assert result.token == "USDC"
|
||||
assert result.tx_hash == "ee" * 32
|
||||
assert result.block_number == 50000000
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_to_funding_transfer_handles_block_error(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
mock_polygon_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test graceful handling when block fetch fails."""
|
||||
mock_polygon_client.get_block = AsyncMock(side_effect=Exception("Block error"))
|
||||
|
||||
mock_log = _create_mock_log(
|
||||
from_address=TEST_SOURCE,
|
||||
to_address=TEST_WALLET,
|
||||
amount=1000000,
|
||||
tx_hash="0x" + "ff" * 32,
|
||||
block_number=50000000,
|
||||
)
|
||||
|
||||
result = await funding_tracer._log_to_funding_transfer(mock_log, USDC_BRIDGED)
|
||||
|
||||
# Should still return a valid transfer with current timestamp
|
||||
assert result.from_address == TEST_SOURCE.lower()
|
||||
assert result.timestamp is not None
|
||||
|
||||
|
||||
class TestGetFundingChainsBatch:
|
||||
"""Tests for get_funding_chains_batch method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_traces_multiple_addresses(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test batch tracing multiple addresses."""
|
||||
addresses = [f"0x{i:040x}" for i in range(3)]
|
||||
|
||||
# Mock trace to return simple chains
|
||||
async def mock_trace(
|
||||
addr: str,
|
||||
*,
|
||||
max_hops: int | None = None, # noqa: ARG001
|
||||
) -> FundingChain:
|
||||
return FundingChain(
|
||||
target_address=addr.lower(),
|
||||
origin_type="unknown",
|
||||
)
|
||||
|
||||
funding_tracer.trace = mock_trace
|
||||
|
||||
results = await funding_tracer.get_funding_chains_batch(addresses)
|
||||
|
||||
assert len(results) == 3
|
||||
for addr in addresses:
|
||||
assert addr.lower() in results
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_handles_exceptions(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test batch handles exceptions gracefully."""
|
||||
addresses = ["0x" + "11" * 20, "0x" + "22" * 20]
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_trace(
|
||||
addr: str,
|
||||
*,
|
||||
max_hops: int | None = None, # noqa: ARG001
|
||||
) -> FundingChain:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise ValueError("Test error")
|
||||
return FundingChain(
|
||||
target_address=addr.lower(),
|
||||
origin_type="cex_binance",
|
||||
)
|
||||
|
||||
funding_tracer.trace = mock_trace
|
||||
|
||||
results = await funding_tracer.get_funding_chains_batch(addresses)
|
||||
|
||||
assert len(results) == 2
|
||||
# First address should have error origin type
|
||||
assert results[addresses[0].lower()].origin_type == "error"
|
||||
# Second address should succeed
|
||||
assert results[addresses[1].lower()].origin_type == "cex_binance"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_empty_list(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test batch with empty address list."""
|
||||
results = await funding_tracer.get_funding_chains_batch([])
|
||||
assert results == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_respects_max_hops_override(
|
||||
self,
|
||||
funding_tracer: FundingTracer,
|
||||
) -> None:
|
||||
"""Test batch passes max_hops to individual traces."""
|
||||
addresses = ["0x" + "11" * 20]
|
||||
captured_max_hops: list[int | None] = []
|
||||
|
||||
async def mock_trace(addr: str, max_hops: int | None = None) -> FundingChain:
|
||||
captured_max_hops.append(max_hops)
|
||||
return FundingChain(target_address=addr.lower())
|
||||
|
||||
funding_tracer.trace = mock_trace
|
||||
|
||||
await funding_tracer.get_funding_chains_batch(addresses, max_hops=5)
|
||||
|
||||
assert captured_max_hops == [5]
|
||||
|
||||
|
||||
class TestGetSuspiciousnessScore:
|
||||
"""Tests for get_suspiciousness_score method."""
|
||||
|
||||
def test_cex_origin_low_score(self, funding_tracer: FundingTracer) -> None:
|
||||
"""Test CEX origin results in low suspiciousness."""
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="cex_binance",
|
||||
hop_count=1,
|
||||
)
|
||||
|
||||
score = funding_tracer.get_suspiciousness_score(chain)
|
||||
|
||||
assert score == 0.1
|
||||
|
||||
def test_bridge_origin_low_score(self, funding_tracer: FundingTracer) -> None:
|
||||
"""Test bridge origin results in low-medium suspiciousness."""
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="bridge_polygon",
|
||||
hop_count=1,
|
||||
)
|
||||
|
||||
score = funding_tracer.get_suspiciousness_score(chain)
|
||||
|
||||
assert score == 0.3
|
||||
|
||||
def test_unknown_no_transfers_high_score(self, funding_tracer: FundingTracer) -> None:
|
||||
"""Test unknown origin with no transfers is most suspicious."""
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="unknown",
|
||||
hop_count=0,
|
||||
)
|
||||
|
||||
score = funding_tracer.get_suspiciousness_score(chain)
|
||||
|
||||
assert score == 1.0
|
||||
|
||||
def test_unknown_max_hops_high_score(self, funding_tracer: FundingTracer) -> None:
|
||||
"""Test unknown origin at max hops is suspicious."""
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="unknown",
|
||||
hop_count=3, # Same as max_hops
|
||||
)
|
||||
|
||||
score = funding_tracer.get_suspiciousness_score(chain)
|
||||
|
||||
assert score == 0.7
|
||||
|
||||
def test_unknown_partial_hops_medium_score(self, funding_tracer: FundingTracer) -> None:
|
||||
"""Test unknown origin with partial hops is moderately suspicious."""
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="unknown",
|
||||
hop_count=1,
|
||||
)
|
||||
|
||||
score = funding_tracer.get_suspiciousness_score(chain)
|
||||
|
||||
# 0.5 + (0.3 * (1 - 1/3)) = 0.5 + 0.2 = 0.7
|
||||
assert 0.5 < score < 0.8
|
||||
|
||||
|
||||
class TestFundingTransferModel:
|
||||
"""Tests for FundingTransfer dataclass."""
|
||||
|
||||
def test_amount_formatted_usdc(self) -> None:
|
||||
"""Test formatted amount for USDC (6 decimals)."""
|
||||
transfer = FundingTransfer(
|
||||
from_address=TEST_SOURCE,
|
||||
to_address=TEST_WALLET,
|
||||
amount=Decimal("1500000"), # 1.5 USDC
|
||||
token="USDC",
|
||||
tx_hash="0x" + "aa" * 32,
|
||||
block_number=50000000,
|
||||
timestamp=datetime.now(UTC),
|
||||
)
|
||||
|
||||
assert transfer.amount_formatted == Decimal("1.5")
|
||||
|
||||
def test_amount_formatted_other(self) -> None:
|
||||
"""Test formatted amount for other tokens (18 decimals)."""
|
||||
transfer = FundingTransfer(
|
||||
from_address=TEST_SOURCE,
|
||||
to_address=TEST_WALLET,
|
||||
amount=Decimal("1500000000000000000"), # 1.5 MATIC
|
||||
token="MATIC",
|
||||
tx_hash="0x" + "aa" * 32,
|
||||
block_number=50000000,
|
||||
timestamp=datetime.now(UTC),
|
||||
)
|
||||
|
||||
assert transfer.amount_formatted == Decimal("1.5")
|
||||
|
||||
|
||||
class TestFundingChainModel:
|
||||
"""Tests for FundingChain dataclass."""
|
||||
|
||||
def test_is_cex_origin(self) -> None:
|
||||
"""Test is_cex_origin property."""
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="cex_binance",
|
||||
)
|
||||
assert chain.is_cex_origin is True
|
||||
|
||||
chain2 = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="bridge_polygon",
|
||||
)
|
||||
assert chain2.is_cex_origin is False
|
||||
|
||||
def test_is_bridge_origin(self) -> None:
|
||||
"""Test is_bridge_origin property."""
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="bridge_polygon",
|
||||
)
|
||||
assert chain.is_bridge_origin is True
|
||||
|
||||
chain2 = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="cex_coinbase",
|
||||
)
|
||||
assert chain2.is_bridge_origin is False
|
||||
|
||||
def test_is_unknown_origin(self) -> None:
|
||||
"""Test is_unknown_origin property."""
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
origin_type="unknown",
|
||||
)
|
||||
assert chain.is_unknown_origin is True
|
||||
|
||||
def test_total_amount_empty_chain(self) -> None:
|
||||
"""Test total_amount with empty chain."""
|
||||
chain = FundingChain(target_address=TEST_WALLET)
|
||||
assert chain.total_amount == Decimal("0")
|
||||
|
||||
def test_total_amount_with_transfers(self) -> None:
|
||||
"""Test total_amount returns first transfer amount."""
|
||||
transfer = FundingTransfer(
|
||||
from_address=TEST_SOURCE,
|
||||
to_address=TEST_WALLET,
|
||||
amount=Decimal("5000000"),
|
||||
token="USDC",
|
||||
tx_hash="0x" + "aa" * 32,
|
||||
block_number=50000000,
|
||||
timestamp=datetime.now(UTC),
|
||||
)
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
chain=[transfer],
|
||||
)
|
||||
|
||||
assert chain.total_amount == Decimal("5000000")
|
||||
|
||||
def test_funding_depth(self) -> None:
|
||||
"""Test funding_depth property."""
|
||||
chain = FundingChain(
|
||||
target_address=TEST_WALLET,
|
||||
hop_count=3,
|
||||
)
|
||||
assert chain.funding_depth == 3
|
||||
|
||||
|
||||
class TestConstants:
|
||||
"""Tests for module constants."""
|
||||
|
||||
def test_usdc_bridged_address(self) -> None:
|
||||
"""Test USDC bridged contract address."""
|
||||
assert USDC_BRIDGED == "0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174"
|
||||
|
||||
def test_usdc_native_address(self) -> None:
|
||||
"""Test USDC native contract address."""
|
||||
assert USDC_NATIVE == "0x3c499c542cEF5E3811e1192ce70d8cC03d5c3359"
|
||||
|
||||
def test_transfer_event_signature(self) -> None:
|
||||
"""Test Transfer event signature is correct keccak hash."""
|
||||
# Transfer(address,address,uint256) hash
|
||||
assert TRANSFER_EVENT_SIGNATURE is not None
|
||||
assert len(TRANSFER_EVENT_SIGNATURE) == 32
|
||||
|
||||
|
||||
# Helper functions
|
||||
|
||||
|
||||
def _create_mock_log(
|
||||
from_address: str,
|
||||
to_address: str,
|
||||
amount: int,
|
||||
tx_hash: str,
|
||||
block_number: int,
|
||||
) -> dict[str, Any]:
|
||||
"""Create a mock log entry for testing."""
|
||||
# Pad addresses to 32 bytes (topics format)
|
||||
from_padded = bytes.fromhex(from_address.replace("0x", "").zfill(64))
|
||||
to_padded = bytes.fromhex(to_address.replace("0x", "").zfill(64))
|
||||
|
||||
# Amount as 32-byte hex data
|
||||
amount_hex = bytes.fromhex(f"{amount:064x}")
|
||||
|
||||
return {
|
||||
"topics": [
|
||||
TRANSFER_EVENT_SIGNATURE,
|
||||
from_padded,
|
||||
to_padded,
|
||||
],
|
||||
"data": amount_hex,
|
||||
"transactionHash": bytes.fromhex(tx_hash.replace("0x", "")),
|
||||
"blockNumber": block_number,
|
||||
}
|
||||
@@ -0,0 +1,425 @@
|
||||
"""Tests for the profiler data models."""
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.profiler.models import Transaction, WalletInfo, WalletProfile
|
||||
|
||||
|
||||
class TestTransaction:
|
||||
"""Tests for the Transaction dataclass."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_transaction(self) -> Transaction:
|
||||
"""Create a sample transaction."""
|
||||
return Transaction(
|
||||
hash="0xabc123",
|
||||
block_number=50000000,
|
||||
timestamp=datetime(2026, 1, 4, 12, 0, 0, tzinfo=UTC),
|
||||
from_address="0xsender",
|
||||
to_address="0xreceiver",
|
||||
value=Decimal("1000000000000000000"), # 1 MATIC
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50000000000"), # 50 Gwei
|
||||
)
|
||||
|
||||
def test_transaction_creation(self, sample_transaction: Transaction) -> None:
|
||||
"""Test creating a transaction."""
|
||||
assert sample_transaction.hash == "0xabc123"
|
||||
assert sample_transaction.block_number == 50000000
|
||||
assert sample_transaction.from_address == "0xsender"
|
||||
assert sample_transaction.to_address == "0xreceiver"
|
||||
|
||||
def test_value_matic(self, sample_transaction: Transaction) -> None:
|
||||
"""Test value in MATIC conversion."""
|
||||
assert sample_transaction.value_matic == Decimal("1")
|
||||
|
||||
def test_value_matic_fractional(self) -> None:
|
||||
"""Test fractional MATIC value."""
|
||||
tx = Transaction(
|
||||
hash="0x123",
|
||||
block_number=1,
|
||||
timestamp=datetime.now(UTC),
|
||||
from_address="0x1",
|
||||
to_address="0x2",
|
||||
value=Decimal("500000000000000000"), # 0.5 MATIC
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50000000000"),
|
||||
)
|
||||
|
||||
assert tx.value_matic == Decimal("0.5")
|
||||
|
||||
def test_gas_cost_wei(self, sample_transaction: Transaction) -> None:
|
||||
"""Test gas cost in Wei."""
|
||||
expected = 21000 * 50000000000
|
||||
assert sample_transaction.gas_cost_wei == Decimal(expected)
|
||||
|
||||
def test_gas_cost_matic(self, sample_transaction: Transaction) -> None:
|
||||
"""Test gas cost in MATIC."""
|
||||
# 21000 * 50 Gwei = 1050000 Gwei = 0.00105 MATIC
|
||||
expected = Decimal("21000") * Decimal("50000000000") / Decimal("1000000000000000000")
|
||||
assert sample_transaction.gas_cost_matic == expected
|
||||
|
||||
def test_transaction_frozen(self, sample_transaction: Transaction) -> None:
|
||||
"""Test that transaction is immutable."""
|
||||
with pytest.raises(AttributeError):
|
||||
sample_transaction.hash = "0xnew" # type: ignore[misc]
|
||||
|
||||
def test_transaction_no_recipient(self) -> None:
|
||||
"""Test transaction with no recipient (contract creation)."""
|
||||
tx = Transaction(
|
||||
hash="0x123",
|
||||
block_number=1,
|
||||
timestamp=datetime.now(UTC),
|
||||
from_address="0x1",
|
||||
to_address=None,
|
||||
value=Decimal("0"),
|
||||
gas_used=100000,
|
||||
gas_price=Decimal("50000000000"),
|
||||
)
|
||||
|
||||
assert tx.to_address is None
|
||||
|
||||
|
||||
class TestWalletInfo:
|
||||
"""Tests for the WalletInfo dataclass."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_wallet(self) -> WalletInfo:
|
||||
"""Create a sample wallet info."""
|
||||
return WalletInfo(
|
||||
address="0xwallet123",
|
||||
transaction_count=100,
|
||||
balance_wei=Decimal("5000000000000000000"), # 5 MATIC
|
||||
first_transaction=None,
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def wallet_with_transaction(self) -> WalletInfo:
|
||||
"""Create a wallet with first transaction."""
|
||||
first_tx = Transaction(
|
||||
hash="0xfirst",
|
||||
block_number=1000000,
|
||||
timestamp=datetime.now(UTC) - timedelta(days=365),
|
||||
from_address="0xfaucet",
|
||||
to_address="0xwallet123",
|
||||
value=Decimal("1000000000000000000"),
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50000000000"),
|
||||
)
|
||||
return WalletInfo(
|
||||
address="0xwallet123",
|
||||
transaction_count=100,
|
||||
balance_wei=Decimal("5000000000000000000"),
|
||||
first_transaction=first_tx,
|
||||
)
|
||||
|
||||
def test_wallet_creation(self, sample_wallet: WalletInfo) -> None:
|
||||
"""Test creating wallet info."""
|
||||
assert sample_wallet.address == "0xwallet123"
|
||||
assert sample_wallet.transaction_count == 100
|
||||
assert sample_wallet.balance_wei == Decimal("5000000000000000000")
|
||||
|
||||
def test_balance_matic(self, sample_wallet: WalletInfo) -> None:
|
||||
"""Test balance in MATIC conversion."""
|
||||
assert sample_wallet.balance_matic == Decimal("5")
|
||||
|
||||
def test_is_fresh_false(self, sample_wallet: WalletInfo) -> None:
|
||||
"""Test that wallet with many transactions is not fresh."""
|
||||
assert sample_wallet.is_fresh is False
|
||||
|
||||
def test_is_fresh_true(self) -> None:
|
||||
"""Test that wallet with few transactions is fresh."""
|
||||
wallet = WalletInfo(
|
||||
address="0xnewwallet",
|
||||
transaction_count=5,
|
||||
balance_wei=Decimal("1000000000000000000"),
|
||||
)
|
||||
|
||||
assert wallet.is_fresh is True
|
||||
|
||||
def test_is_fresh_boundary(self) -> None:
|
||||
"""Test fresh wallet boundary (10 transactions)."""
|
||||
wallet_9 = WalletInfo(
|
||||
address="0x1",
|
||||
transaction_count=9,
|
||||
balance_wei=Decimal("0"),
|
||||
)
|
||||
wallet_10 = WalletInfo(
|
||||
address="0x2",
|
||||
transaction_count=10,
|
||||
balance_wei=Decimal("0"),
|
||||
)
|
||||
|
||||
assert wallet_9.is_fresh is True
|
||||
assert wallet_10.is_fresh is False
|
||||
|
||||
def test_wallet_age_days_no_transaction(self, sample_wallet: WalletInfo) -> None:
|
||||
"""Test wallet age when no first transaction."""
|
||||
assert sample_wallet.wallet_age_days is None
|
||||
|
||||
def test_wallet_age_days_with_transaction(self, wallet_with_transaction: WalletInfo) -> None:
|
||||
"""Test wallet age calculation."""
|
||||
age = wallet_with_transaction.wallet_age_days
|
||||
|
||||
assert age is not None
|
||||
# Should be approximately 365 days
|
||||
assert 364 < age < 366
|
||||
|
||||
def test_wallet_frozen(self, sample_wallet: WalletInfo) -> None:
|
||||
"""Test that wallet info is immutable."""
|
||||
with pytest.raises(AttributeError):
|
||||
sample_wallet.address = "0xnew" # type: ignore[misc]
|
||||
|
||||
def test_wallet_zero_balance(self) -> None:
|
||||
"""Test wallet with zero balance."""
|
||||
wallet = WalletInfo(
|
||||
address="0xempty",
|
||||
transaction_count=0,
|
||||
balance_wei=Decimal("0"),
|
||||
)
|
||||
|
||||
assert wallet.balance_matic == Decimal("0")
|
||||
assert wallet.is_fresh is True
|
||||
|
||||
def test_wallet_very_small_balance(self) -> None:
|
||||
"""Test wallet with very small balance."""
|
||||
wallet = WalletInfo(
|
||||
address="0xdust",
|
||||
transaction_count=1,
|
||||
balance_wei=Decimal("1"), # 1 Wei
|
||||
)
|
||||
|
||||
# Should be a very small fraction
|
||||
expected = Decimal("1") / Decimal("1000000000000000000")
|
||||
assert wallet.balance_matic == expected
|
||||
|
||||
|
||||
class TestTransactionEquality:
|
||||
"""Tests for transaction equality and hashing."""
|
||||
|
||||
def test_equal_transactions(self) -> None:
|
||||
"""Test that identical transactions are equal."""
|
||||
tx1 = Transaction(
|
||||
hash="0xabc",
|
||||
block_number=1,
|
||||
timestamp=datetime(2026, 1, 1, tzinfo=UTC),
|
||||
from_address="0x1",
|
||||
to_address="0x2",
|
||||
value=Decimal("1000"),
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50"),
|
||||
)
|
||||
tx2 = Transaction(
|
||||
hash="0xabc",
|
||||
block_number=1,
|
||||
timestamp=datetime(2026, 1, 1, tzinfo=UTC),
|
||||
from_address="0x1",
|
||||
to_address="0x2",
|
||||
value=Decimal("1000"),
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50"),
|
||||
)
|
||||
|
||||
assert tx1 == tx2
|
||||
|
||||
def test_different_transactions(self) -> None:
|
||||
"""Test that different transactions are not equal."""
|
||||
tx1 = Transaction(
|
||||
hash="0xabc",
|
||||
block_number=1,
|
||||
timestamp=datetime(2026, 1, 1, tzinfo=UTC),
|
||||
from_address="0x1",
|
||||
to_address="0x2",
|
||||
value=Decimal("1000"),
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50"),
|
||||
)
|
||||
tx2 = Transaction(
|
||||
hash="0xdef", # Different hash
|
||||
block_number=1,
|
||||
timestamp=datetime(2026, 1, 1, tzinfo=UTC),
|
||||
from_address="0x1",
|
||||
to_address="0x2",
|
||||
value=Decimal("1000"),
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50"),
|
||||
)
|
||||
|
||||
assert tx1 != tx2
|
||||
|
||||
def test_transaction_hashable(self) -> None:
|
||||
"""Test that transactions can be used in sets."""
|
||||
tx = Transaction(
|
||||
hash="0xabc",
|
||||
block_number=1,
|
||||
timestamp=datetime(2026, 1, 1, tzinfo=UTC),
|
||||
from_address="0x1",
|
||||
to_address="0x2",
|
||||
value=Decimal("1000"),
|
||||
gas_used=21000,
|
||||
gas_price=Decimal("50"),
|
||||
)
|
||||
|
||||
tx_set = {tx}
|
||||
assert tx in tx_set
|
||||
|
||||
|
||||
class TestWalletProfile:
|
||||
"""Tests for the WalletProfile dataclass."""
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_profile(self) -> WalletProfile:
|
||||
"""Create a fresh wallet profile."""
|
||||
return WalletProfile(
|
||||
address="0xfresh",
|
||||
nonce=2,
|
||||
first_seen=datetime.now(UTC) - timedelta(hours=6),
|
||||
age_hours=6.0,
|
||||
is_fresh=True,
|
||||
total_tx_count=2,
|
||||
matic_balance=Decimal("1000000000000000000"), # 1 MATIC
|
||||
usdc_balance=Decimal("1000000"), # 1 USDC
|
||||
fresh_threshold=5,
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def old_profile(self) -> WalletProfile:
|
||||
"""Create an old wallet profile."""
|
||||
return WalletProfile(
|
||||
address="0xold",
|
||||
nonce=500,
|
||||
first_seen=datetime.now(UTC) - timedelta(days=365),
|
||||
age_hours=365 * 24,
|
||||
is_fresh=False,
|
||||
total_tx_count=500,
|
||||
matic_balance=Decimal("100000000000000000000"), # 100 MATIC
|
||||
usdc_balance=Decimal("10000000000"), # 10000 USDC
|
||||
fresh_threshold=5,
|
||||
)
|
||||
|
||||
def test_profile_creation(self, fresh_profile: WalletProfile) -> None:
|
||||
"""Test creating a wallet profile."""
|
||||
assert fresh_profile.address == "0xfresh"
|
||||
assert fresh_profile.nonce == 2
|
||||
assert fresh_profile.is_fresh is True
|
||||
assert fresh_profile.age_hours == 6.0
|
||||
|
||||
def test_age_days(self, fresh_profile: WalletProfile) -> None:
|
||||
"""Test age_days property."""
|
||||
assert fresh_profile.age_days == 0.25 # 6 hours = 0.25 days
|
||||
|
||||
def test_age_days_none(self) -> None:
|
||||
"""Test age_days when age_hours is None."""
|
||||
profile = WalletProfile(
|
||||
address="0xnew",
|
||||
nonce=0,
|
||||
first_seen=None,
|
||||
age_hours=None,
|
||||
is_fresh=True,
|
||||
total_tx_count=0,
|
||||
matic_balance=Decimal("0"),
|
||||
usdc_balance=Decimal("0"),
|
||||
)
|
||||
assert profile.age_days is None
|
||||
|
||||
def test_matic_balance_formatted(self, fresh_profile: WalletProfile) -> None:
|
||||
"""Test MATIC balance formatting."""
|
||||
assert fresh_profile.matic_balance_formatted == Decimal("1")
|
||||
|
||||
def test_usdc_balance_formatted(self, fresh_profile: WalletProfile) -> None:
|
||||
"""Test USDC balance formatting."""
|
||||
assert fresh_profile.usdc_balance_formatted == Decimal("1")
|
||||
|
||||
def test_is_brand_new(self) -> None:
|
||||
"""Test is_brand_new property."""
|
||||
brand_new = WalletProfile(
|
||||
address="0xnew",
|
||||
nonce=0,
|
||||
first_seen=None,
|
||||
age_hours=None,
|
||||
is_fresh=True,
|
||||
total_tx_count=0,
|
||||
matic_balance=Decimal("0"),
|
||||
usdc_balance=Decimal("0"),
|
||||
)
|
||||
assert brand_new.is_brand_new is True
|
||||
|
||||
not_brand_new = WalletProfile(
|
||||
address="0xold",
|
||||
nonce=1,
|
||||
first_seen=datetime.now(UTC),
|
||||
age_hours=1.0,
|
||||
is_fresh=True,
|
||||
total_tx_count=1,
|
||||
matic_balance=Decimal("0"),
|
||||
usdc_balance=Decimal("0"),
|
||||
)
|
||||
assert not_brand_new.is_brand_new is False
|
||||
|
||||
def test_freshness_score_brand_new(self) -> None:
|
||||
"""Test freshness score for brand new wallet."""
|
||||
profile = WalletProfile(
|
||||
address="0xnew",
|
||||
nonce=0,
|
||||
first_seen=None,
|
||||
age_hours=None,
|
||||
is_fresh=True,
|
||||
total_tx_count=0,
|
||||
matic_balance=Decimal("0"),
|
||||
usdc_balance=Decimal("0"),
|
||||
fresh_threshold=5,
|
||||
)
|
||||
# nonce_score = 1.0 (0/5 = 0, 1-0 = 1)
|
||||
# age_score = 1.0 (None = assumed new)
|
||||
# score = 0.6 * 1.0 + 0.4 * 1.0 = 1.0
|
||||
assert profile.freshness_score == 1.0
|
||||
|
||||
def test_freshness_score_old_wallet(self, old_profile: WalletProfile) -> None:
|
||||
"""Test freshness score for old wallet."""
|
||||
# nonce_score = max(0, 1 - 500/5) = 0
|
||||
# age_score = max(0, 1 - 8760/48) = 0
|
||||
# score = 0
|
||||
assert old_profile.freshness_score == 0.0
|
||||
|
||||
def test_freshness_score_moderate(self) -> None:
|
||||
"""Test freshness score for moderately fresh wallet."""
|
||||
profile = WalletProfile(
|
||||
address="0xmoderate",
|
||||
nonce=2,
|
||||
first_seen=datetime.now(UTC) - timedelta(hours=24),
|
||||
age_hours=24.0,
|
||||
is_fresh=True,
|
||||
total_tx_count=2,
|
||||
matic_balance=Decimal("0"),
|
||||
usdc_balance=Decimal("0"),
|
||||
fresh_threshold=5,
|
||||
)
|
||||
# nonce_score = 1 - 2/5 = 0.6
|
||||
# age_score = 1 - 24/48 = 0.5
|
||||
# score = 0.6 * 0.6 + 0.4 * 0.5 = 0.36 + 0.2 = 0.56
|
||||
assert profile.freshness_score == pytest.approx(0.56)
|
||||
|
||||
def test_profile_frozen(self, fresh_profile: WalletProfile) -> None:
|
||||
"""Test that wallet profile is immutable."""
|
||||
with pytest.raises(AttributeError):
|
||||
fresh_profile.nonce = 100 # type: ignore[misc]
|
||||
|
||||
def test_analyzed_at_default(self) -> None:
|
||||
"""Test that analyzed_at has a default."""
|
||||
before = datetime.now(UTC)
|
||||
profile = WalletProfile(
|
||||
address="0x1",
|
||||
nonce=0,
|
||||
first_seen=None,
|
||||
age_hours=None,
|
||||
is_fresh=True,
|
||||
total_tx_count=0,
|
||||
matic_balance=Decimal("0"),
|
||||
usdc_balance=Decimal("0"),
|
||||
)
|
||||
after = datetime.now(UTC)
|
||||
|
||||
assert before <= profile.analyzed_at <= after
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for the storage layer."""
|
||||
@@ -0,0 +1,496 @@
|
||||
"""Tests for storage repositories."""
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
|
||||
from polymarket_insider_tracker.storage.models import Base
|
||||
from polymarket_insider_tracker.storage.repos import (
|
||||
FundingRepository,
|
||||
FundingTransferDTO,
|
||||
RelationshipRepository,
|
||||
WalletProfileDTO,
|
||||
WalletRelationshipDTO,
|
||||
WalletRepository,
|
||||
)
|
||||
|
||||
# ============================================================================
|
||||
# Fixtures
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def async_engine():
|
||||
"""Create an async SQLite engine for testing."""
|
||||
engine = create_async_engine(
|
||||
"sqlite+aiosqlite:///:memory:",
|
||||
echo=False,
|
||||
)
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
yield engine
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def async_session(async_engine) -> AsyncSession:
|
||||
"""Create an async session for testing."""
|
||||
session_factory = async_sessionmaker(bind=async_engine, expire_on_commit=False)
|
||||
async with session_factory() as session:
|
||||
yield session
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_wallet_dto() -> WalletProfileDTO:
|
||||
"""Create a sample wallet profile DTO."""
|
||||
return WalletProfileDTO(
|
||||
address="0x1234567890abcdef1234567890abcdef12345678",
|
||||
nonce=5,
|
||||
first_seen_at=datetime.now(UTC) - timedelta(hours=24),
|
||||
is_fresh=True,
|
||||
matic_balance=Decimal("1000000000000000000"),
|
||||
usdc_balance=Decimal("1000.00"),
|
||||
analyzed_at=datetime.now(UTC),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_transfer_dto() -> FundingTransferDTO:
|
||||
"""Create a sample funding transfer DTO."""
|
||||
return FundingTransferDTO(
|
||||
from_address="0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
|
||||
to_address="0x1234567890abcdef1234567890abcdef12345678",
|
||||
amount=Decimal("5000.00"),
|
||||
token="USDC",
|
||||
tx_hash="0x" + "a" * 64,
|
||||
block_number=12345678,
|
||||
timestamp=datetime.now(UTC),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_relationship_dto() -> WalletRelationshipDTO:
|
||||
"""Create a sample wallet relationship DTO."""
|
||||
return WalletRelationshipDTO(
|
||||
wallet_a="0x1234567890abcdef1234567890abcdef12345678",
|
||||
wallet_b="0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
|
||||
relationship_type="funded_by",
|
||||
confidence=Decimal("0.95"),
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# WalletRepository Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestWalletRepository:
|
||||
"""Tests for WalletRepository."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_address_not_found(self, async_session: AsyncSession) -> None:
|
||||
"""Test getting a non-existent wallet returns None."""
|
||||
repo = WalletRepository(async_session)
|
||||
result = await repo.get_by_address("0xnonexistent")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_creates_new(
|
||||
self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO
|
||||
) -> None:
|
||||
"""Test upserting a new wallet profile."""
|
||||
repo = WalletRepository(async_session)
|
||||
await repo.upsert(sample_wallet_dto)
|
||||
await async_session.commit()
|
||||
|
||||
result = await repo.get_by_address(sample_wallet_dto.address)
|
||||
assert result is not None
|
||||
assert result.address == sample_wallet_dto.address.lower()
|
||||
assert result.nonce == sample_wallet_dto.nonce
|
||||
assert result.is_fresh == sample_wallet_dto.is_fresh
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_updates_existing(
|
||||
self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO
|
||||
) -> None:
|
||||
"""Test upserting updates existing profile."""
|
||||
repo = WalletRepository(async_session)
|
||||
await repo.upsert(sample_wallet_dto)
|
||||
await async_session.commit()
|
||||
|
||||
# Update the DTO
|
||||
updated_dto = WalletProfileDTO(
|
||||
address=sample_wallet_dto.address,
|
||||
nonce=10,
|
||||
first_seen_at=sample_wallet_dto.first_seen_at,
|
||||
is_fresh=False,
|
||||
matic_balance=sample_wallet_dto.matic_balance,
|
||||
usdc_balance=Decimal("2000.00"),
|
||||
analyzed_at=datetime.now(UTC),
|
||||
)
|
||||
await repo.upsert(updated_dto)
|
||||
await async_session.commit()
|
||||
|
||||
result = await repo.get_by_address(sample_wallet_dto.address)
|
||||
assert result is not None
|
||||
assert result.nonce == 10
|
||||
assert result.is_fresh is False
|
||||
assert result.usdc_balance == Decimal("2000.00")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_many(
|
||||
self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO
|
||||
) -> None:
|
||||
"""Test getting multiple wallets."""
|
||||
repo = WalletRepository(async_session)
|
||||
|
||||
# Insert two wallets
|
||||
dto2 = WalletProfileDTO(
|
||||
address="0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
|
||||
nonce=3,
|
||||
first_seen_at=datetime.now(UTC),
|
||||
is_fresh=True,
|
||||
matic_balance=None,
|
||||
usdc_balance=None,
|
||||
analyzed_at=datetime.now(UTC),
|
||||
)
|
||||
await repo.upsert(sample_wallet_dto)
|
||||
await repo.upsert(dto2)
|
||||
await async_session.commit()
|
||||
|
||||
results = await repo.get_many([sample_wallet_dto.address, dto2.address])
|
||||
assert len(results) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_fresh_wallets(
|
||||
self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO
|
||||
) -> None:
|
||||
"""Test getting fresh wallets."""
|
||||
repo = WalletRepository(async_session)
|
||||
|
||||
# Insert fresh and non-fresh wallets
|
||||
non_fresh = WalletProfileDTO(
|
||||
address="0xcccccccccccccccccccccccccccccccccccccccc",
|
||||
nonce=100,
|
||||
first_seen_at=datetime.now(UTC) - timedelta(days=30),
|
||||
is_fresh=False,
|
||||
matic_balance=None,
|
||||
usdc_balance=None,
|
||||
analyzed_at=datetime.now(UTC),
|
||||
)
|
||||
await repo.upsert(sample_wallet_dto)
|
||||
await repo.upsert(non_fresh)
|
||||
await async_session.commit()
|
||||
|
||||
results = await repo.get_fresh_wallets()
|
||||
assert len(results) == 1
|
||||
assert results[0].is_fresh is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete(
|
||||
self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO
|
||||
) -> None:
|
||||
"""Test deleting a wallet profile."""
|
||||
repo = WalletRepository(async_session)
|
||||
await repo.upsert(sample_wallet_dto)
|
||||
await async_session.commit()
|
||||
|
||||
deleted = await repo.delete(sample_wallet_dto.address)
|
||||
await async_session.commit()
|
||||
assert deleted is True
|
||||
|
||||
result = await repo.get_by_address(sample_wallet_dto.address)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_not_found(self, async_session: AsyncSession) -> None:
|
||||
"""Test deleting non-existent wallet returns False."""
|
||||
repo = WalletRepository(async_session)
|
||||
deleted = await repo.delete("0xnonexistent")
|
||||
assert deleted is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_stale(
|
||||
self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO
|
||||
) -> None:
|
||||
"""Test marking a wallet as stale."""
|
||||
repo = WalletRepository(async_session)
|
||||
await repo.upsert(sample_wallet_dto)
|
||||
await async_session.commit()
|
||||
|
||||
marked = await repo.mark_stale(sample_wallet_dto.address)
|
||||
await async_session.commit()
|
||||
assert marked is True
|
||||
|
||||
result = await repo.get_by_address(sample_wallet_dto.address)
|
||||
assert result is not None
|
||||
assert result.analyzed_at.year == 2000
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# FundingRepository Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestFundingRepository:
|
||||
"""Tests for FundingRepository."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert(
|
||||
self, async_session: AsyncSession, sample_transfer_dto: FundingTransferDTO
|
||||
) -> None:
|
||||
"""Test inserting a funding transfer."""
|
||||
repo = FundingRepository(async_session)
|
||||
await repo.insert(sample_transfer_dto)
|
||||
await async_session.commit()
|
||||
|
||||
result = await repo.get_by_tx_hash(sample_transfer_dto.tx_hash)
|
||||
assert result is not None
|
||||
assert result.amount == sample_transfer_dto.amount
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_transfers_to(
|
||||
self, async_session: AsyncSession, sample_transfer_dto: FundingTransferDTO
|
||||
) -> None:
|
||||
"""Test getting transfers to an address."""
|
||||
repo = FundingRepository(async_session)
|
||||
await repo.insert(sample_transfer_dto)
|
||||
await async_session.commit()
|
||||
|
||||
results = await repo.get_transfers_to(sample_transfer_dto.to_address)
|
||||
assert len(results) == 1
|
||||
assert results[0].from_address == sample_transfer_dto.from_address.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_transfers_from(
|
||||
self, async_session: AsyncSession, sample_transfer_dto: FundingTransferDTO
|
||||
) -> None:
|
||||
"""Test getting transfers from an address."""
|
||||
repo = FundingRepository(async_session)
|
||||
await repo.insert(sample_transfer_dto)
|
||||
await async_session.commit()
|
||||
|
||||
results = await repo.get_transfers_from(sample_transfer_dto.from_address)
|
||||
assert len(results) == 1
|
||||
assert results[0].to_address == sample_transfer_dto.to_address.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_first_transfer_to(
|
||||
self, async_session: AsyncSession, sample_transfer_dto: FundingTransferDTO
|
||||
) -> None:
|
||||
"""Test getting first transfer to an address."""
|
||||
repo = FundingRepository(async_session)
|
||||
|
||||
# Insert multiple transfers with different timestamps
|
||||
earlier = FundingTransferDTO(
|
||||
from_address="0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee",
|
||||
to_address=sample_transfer_dto.to_address,
|
||||
amount=Decimal("100.00"),
|
||||
token="USDC",
|
||||
tx_hash="0x" + "b" * 64,
|
||||
block_number=12345670,
|
||||
timestamp=datetime.now(UTC) - timedelta(hours=2),
|
||||
)
|
||||
await repo.insert(earlier)
|
||||
await repo.insert(sample_transfer_dto)
|
||||
await async_session.commit()
|
||||
|
||||
result = await repo.get_first_transfer_to(sample_transfer_dto.to_address)
|
||||
assert result is not None
|
||||
assert result.tx_hash == earlier.tx_hash.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_many(self, async_session: AsyncSession) -> None:
|
||||
"""Test inserting multiple transfers."""
|
||||
repo = FundingRepository(async_session)
|
||||
|
||||
transfers = [
|
||||
FundingTransferDTO(
|
||||
from_address=f"0x{'a' * 40}",
|
||||
to_address=f"0x{'b' * 40}",
|
||||
amount=Decimal(f"{i * 100}.00"),
|
||||
token="USDC",
|
||||
tx_hash=f"0x{str(i) * 64}"[:66],
|
||||
block_number=12345678 + i,
|
||||
timestamp=datetime.now(UTC),
|
||||
)
|
||||
for i in range(1, 4)
|
||||
]
|
||||
|
||||
count = await repo.insert_many(transfers)
|
||||
await async_session.commit()
|
||||
assert count == 3
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# RelationshipRepository Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestRelationshipRepository:
|
||||
"""Tests for RelationshipRepository."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert(
|
||||
self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO
|
||||
) -> None:
|
||||
"""Test upserting a relationship."""
|
||||
repo = RelationshipRepository(async_session)
|
||||
await repo.upsert(sample_relationship_dto)
|
||||
await async_session.commit()
|
||||
|
||||
results = await repo.get_relationships(sample_relationship_dto.wallet_a)
|
||||
assert len(results) == 1
|
||||
assert results[0].confidence == sample_relationship_dto.confidence
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_relationships_filter_type(
|
||||
self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO
|
||||
) -> None:
|
||||
"""Test getting relationships with type filter."""
|
||||
repo = RelationshipRepository(async_session)
|
||||
await repo.upsert(sample_relationship_dto)
|
||||
|
||||
same_entity = WalletRelationshipDTO(
|
||||
wallet_a=sample_relationship_dto.wallet_a,
|
||||
wallet_b="0xdddddddddddddddddddddddddddddddddddddddd",
|
||||
relationship_type="same_entity",
|
||||
confidence=Decimal("0.80"),
|
||||
)
|
||||
await repo.upsert(same_entity)
|
||||
await async_session.commit()
|
||||
|
||||
funded_results = await repo.get_relationships(
|
||||
sample_relationship_dto.wallet_a, relationship_type="funded_by"
|
||||
)
|
||||
assert len(funded_results) == 1
|
||||
assert funded_results[0].relationship_type == "funded_by"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_related_wallets(
|
||||
self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO
|
||||
) -> None:
|
||||
"""Test getting related wallet addresses."""
|
||||
repo = RelationshipRepository(async_session)
|
||||
await repo.upsert(sample_relationship_dto)
|
||||
await async_session.commit()
|
||||
|
||||
related = await repo.get_related_wallets(sample_relationship_dto.wallet_a)
|
||||
assert len(related) == 1
|
||||
assert sample_relationship_dto.wallet_b.lower() in related
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_relationship(
|
||||
self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO
|
||||
) -> None:
|
||||
"""Test deleting a relationship."""
|
||||
repo = RelationshipRepository(async_session)
|
||||
await repo.upsert(sample_relationship_dto)
|
||||
await async_session.commit()
|
||||
|
||||
deleted = await repo.delete(
|
||||
sample_relationship_dto.wallet_a,
|
||||
sample_relationship_dto.wallet_b,
|
||||
sample_relationship_dto.relationship_type,
|
||||
)
|
||||
await async_session.commit()
|
||||
assert deleted is True
|
||||
|
||||
results = await repo.get_relationships(sample_relationship_dto.wallet_a)
|
||||
assert len(results) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_updates_confidence(
|
||||
self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO
|
||||
) -> None:
|
||||
"""Test that upserting updates the confidence."""
|
||||
repo = RelationshipRepository(async_session)
|
||||
await repo.upsert(sample_relationship_dto)
|
||||
await async_session.commit()
|
||||
|
||||
updated = WalletRelationshipDTO(
|
||||
wallet_a=sample_relationship_dto.wallet_a,
|
||||
wallet_b=sample_relationship_dto.wallet_b,
|
||||
relationship_type=sample_relationship_dto.relationship_type,
|
||||
confidence=Decimal("0.99"),
|
||||
)
|
||||
await repo.upsert(updated)
|
||||
await async_session.commit()
|
||||
|
||||
results = await repo.get_relationships(sample_relationship_dto.wallet_a)
|
||||
assert len(results) == 1
|
||||
assert results[0].confidence == Decimal("0.99")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# DTO Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestDTOs:
|
||||
"""Tests for Data Transfer Objects."""
|
||||
|
||||
def test_wallet_profile_dto_from_model(self) -> None:
|
||||
"""Test WalletProfileDTO.from_model works correctly."""
|
||||
from polymarket_insider_tracker.storage.models import WalletProfileModel
|
||||
|
||||
now = datetime.now(UTC)
|
||||
model = WalletProfileModel(
|
||||
id=1,
|
||||
address="0x1234",
|
||||
nonce=5,
|
||||
first_seen_at=now,
|
||||
is_fresh=True,
|
||||
matic_balance=Decimal("100"),
|
||||
usdc_balance=Decimal("50.00"),
|
||||
analyzed_at=now,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
|
||||
dto = WalletProfileDTO.from_model(model)
|
||||
assert dto.address == "0x1234"
|
||||
assert dto.nonce == 5
|
||||
assert dto.is_fresh is True
|
||||
|
||||
def test_funding_transfer_dto_from_model(self) -> None:
|
||||
"""Test FundingTransferDTO.from_model works correctly."""
|
||||
from polymarket_insider_tracker.storage.models import FundingTransferModel
|
||||
|
||||
now = datetime.now(UTC)
|
||||
model = FundingTransferModel(
|
||||
id=1,
|
||||
from_address="0xaaa",
|
||||
to_address="0xbbb",
|
||||
amount=Decimal("100.00"),
|
||||
token="USDC",
|
||||
tx_hash="0x123",
|
||||
block_number=12345,
|
||||
timestamp=now,
|
||||
created_at=now,
|
||||
)
|
||||
|
||||
dto = FundingTransferDTO.from_model(model)
|
||||
assert dto.from_address == "0xaaa"
|
||||
assert dto.amount == Decimal("100.00")
|
||||
|
||||
def test_wallet_relationship_dto_from_model(self) -> None:
|
||||
"""Test WalletRelationshipDTO.from_model works correctly."""
|
||||
from polymarket_insider_tracker.storage.models import WalletRelationshipModel
|
||||
|
||||
now = datetime.now(UTC)
|
||||
model = WalletRelationshipModel(
|
||||
id=1,
|
||||
wallet_a="0xaaa",
|
||||
wallet_b="0xbbb",
|
||||
relationship_type="funded_by",
|
||||
confidence=Decimal("0.95"),
|
||||
created_at=now,
|
||||
)
|
||||
|
||||
dto = WalletRelationshipDTO.from_model(model)
|
||||
assert dto.wallet_a == "0xaaa"
|
||||
assert dto.relationship_type == "funded_by"
|
||||
assert dto.confidence == Decimal("0.95")
|
||||
@@ -0,0 +1,329 @@
|
||||
"""Tests for configuration management service."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from polymarket_insider_tracker.config import (
|
||||
DatabaseSettings,
|
||||
DiscordSettings,
|
||||
PolygonSettings,
|
||||
PolymarketSettings,
|
||||
RedisSettings,
|
||||
Settings,
|
||||
TelegramSettings,
|
||||
clear_settings_cache,
|
||||
get_settings,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_cache() -> Iterator[None]:
|
||||
"""Clear settings cache before and after each test."""
|
||||
clear_settings_cache()
|
||||
yield
|
||||
clear_settings_cache()
|
||||
|
||||
|
||||
class TestDatabaseSettings:
|
||||
"""Tests for DatabaseSettings."""
|
||||
|
||||
def test_valid_postgresql_url(self) -> None:
|
||||
"""Test valid PostgreSQL URL."""
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://user:pass@localhost/db"}):
|
||||
settings = DatabaseSettings()
|
||||
assert settings.url == "postgresql://user:pass@localhost/db"
|
||||
|
||||
def test_valid_asyncpg_url(self) -> None:
|
||||
"""Test valid asyncpg URL."""
|
||||
with patch.dict(
|
||||
os.environ, {"DATABASE_URL": "postgresql+asyncpg://user:pass@localhost/db"}
|
||||
):
|
||||
settings = DatabaseSettings()
|
||||
assert settings.url == "postgresql+asyncpg://user:pass@localhost/db"
|
||||
|
||||
def test_invalid_url_raises(self) -> None:
|
||||
"""Test that invalid database URL raises validation error."""
|
||||
with (
|
||||
patch.dict(os.environ, {"DATABASE_URL": "mysql://user:pass@localhost/db"}),
|
||||
pytest.raises(ValidationError, match="PostgreSQL connection string"),
|
||||
):
|
||||
DatabaseSettings()
|
||||
|
||||
|
||||
class TestRedisSettings:
|
||||
"""Tests for RedisSettings."""
|
||||
|
||||
def test_default_url(self) -> None:
|
||||
"""Test default Redis URL."""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
settings = RedisSettings()
|
||||
assert settings.url == "redis://localhost:6379"
|
||||
|
||||
def test_custom_url(self) -> None:
|
||||
"""Test custom Redis URL."""
|
||||
with patch.dict(os.environ, {"REDIS_URL": "redis://redis:6380"}):
|
||||
settings = RedisSettings()
|
||||
assert settings.url == "redis://redis:6380"
|
||||
|
||||
def test_invalid_url_raises(self) -> None:
|
||||
"""Test that invalid Redis URL raises validation error."""
|
||||
with (
|
||||
patch.dict(os.environ, {"REDIS_URL": "http://localhost:6379"}),
|
||||
pytest.raises(ValidationError, match="redis://"),
|
||||
):
|
||||
RedisSettings()
|
||||
|
||||
|
||||
class TestPolygonSettings:
|
||||
"""Tests for PolygonSettings."""
|
||||
|
||||
def test_default_rpc_url(self) -> None:
|
||||
"""Test default Polygon RPC URL."""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
settings = PolygonSettings()
|
||||
assert settings.rpc_url == "https://polygon-rpc.com"
|
||||
assert settings.fallback_rpc_url is None
|
||||
|
||||
def test_custom_urls(self) -> None:
|
||||
"""Test custom Polygon RPC URLs."""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"POLYGON_RPC_URL": "https://alchemy.io/polygon",
|
||||
"POLYGON_FALLBACK_RPC_URL": "https://backup.polygon.io",
|
||||
},
|
||||
):
|
||||
settings = PolygonSettings()
|
||||
assert settings.rpc_url == "https://alchemy.io/polygon"
|
||||
assert settings.fallback_rpc_url == "https://backup.polygon.io"
|
||||
|
||||
def test_invalid_url_raises(self) -> None:
|
||||
"""Test that invalid RPC URL raises validation error."""
|
||||
with (
|
||||
patch.dict(os.environ, {"POLYGON_RPC_URL": "ws://polygon.io"}),
|
||||
pytest.raises(ValidationError, match="HTTP"),
|
||||
):
|
||||
PolygonSettings()
|
||||
|
||||
|
||||
class TestPolymarketSettings:
|
||||
"""Tests for PolymarketSettings."""
|
||||
|
||||
def test_default_ws_url(self) -> None:
|
||||
"""Test default Polymarket WebSocket URL."""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
settings = PolymarketSettings()
|
||||
assert "polymarket.com" in settings.ws_url
|
||||
assert settings.api_key is None
|
||||
|
||||
def test_custom_api_key(self) -> None:
|
||||
"""Test custom API key (secret)."""
|
||||
with patch.dict(os.environ, {"POLYMARKET_API_KEY": "secret-key-123"}):
|
||||
settings = PolymarketSettings()
|
||||
assert settings.api_key is not None
|
||||
assert settings.api_key.get_secret_value() == "secret-key-123"
|
||||
|
||||
def test_invalid_ws_url_raises(self) -> None:
|
||||
"""Test that invalid WebSocket URL raises validation error."""
|
||||
with (
|
||||
patch.dict(os.environ, {"POLYMARKET_WS_URL": "http://polymarket.com"}),
|
||||
pytest.raises(ValidationError, match="ws://"),
|
||||
):
|
||||
PolymarketSettings()
|
||||
|
||||
|
||||
class TestDiscordSettings:
|
||||
"""Tests for DiscordSettings."""
|
||||
|
||||
def test_disabled_by_default(self) -> None:
|
||||
"""Test Discord is disabled when no webhook URL."""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
settings = DiscordSettings()
|
||||
assert not settings.enabled
|
||||
assert settings.webhook_url is None
|
||||
|
||||
def test_enabled_with_webhook(self) -> None:
|
||||
"""Test Discord is enabled with webhook URL."""
|
||||
with patch.dict(os.environ, {"DISCORD_WEBHOOK_URL": "https://discord.com/webhook/123"}):
|
||||
settings = DiscordSettings()
|
||||
assert settings.enabled
|
||||
assert settings.webhook_url is not None
|
||||
|
||||
|
||||
class TestTelegramSettings:
|
||||
"""Tests for TelegramSettings."""
|
||||
|
||||
def test_disabled_by_default(self) -> None:
|
||||
"""Test Telegram is disabled when no credentials."""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
settings = TelegramSettings()
|
||||
assert not settings.enabled
|
||||
|
||||
def test_disabled_with_partial_config(self) -> None:
|
||||
"""Test Telegram is disabled with only token or chat_id."""
|
||||
with patch.dict(os.environ, {"TELEGRAM_BOT_TOKEN": "token123"}):
|
||||
settings = TelegramSettings()
|
||||
assert not settings.enabled
|
||||
|
||||
with patch.dict(os.environ, {"TELEGRAM_CHAT_ID": "12345"}):
|
||||
settings = TelegramSettings()
|
||||
assert not settings.enabled
|
||||
|
||||
def test_enabled_with_full_config(self) -> None:
|
||||
"""Test Telegram is enabled with both token and chat_id."""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"TELEGRAM_BOT_TOKEN": "token123",
|
||||
"TELEGRAM_CHAT_ID": "12345",
|
||||
},
|
||||
):
|
||||
settings = TelegramSettings()
|
||||
assert settings.enabled
|
||||
|
||||
|
||||
class TestSettings:
|
||||
"""Tests for main Settings class."""
|
||||
|
||||
def test_loads_with_required_vars(self) -> None:
|
||||
"""Test settings load with required environment variables."""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABASE_URL": "postgresql://user:pass@localhost/db",
|
||||
"REDIS_URL": "redis://localhost:6379",
|
||||
},
|
||||
):
|
||||
settings = Settings()
|
||||
assert settings.database.url == "postgresql://user:pass@localhost/db"
|
||||
assert settings.redis.url == "redis://localhost:6379"
|
||||
|
||||
def test_default_log_level(self) -> None:
|
||||
"""Test default log level is INFO."""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"DATABASE_URL": "postgresql://user:pass@localhost/db"},
|
||||
):
|
||||
settings = Settings()
|
||||
assert settings.log_level == "INFO"
|
||||
|
||||
def test_custom_log_level(self) -> None:
|
||||
"""Test custom log level."""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABASE_URL": "postgresql://user:pass@localhost/db",
|
||||
"LOG_LEVEL": "DEBUG",
|
||||
},
|
||||
):
|
||||
settings = Settings()
|
||||
assert settings.log_level == "DEBUG"
|
||||
|
||||
def test_invalid_log_level_raises(self) -> None:
|
||||
"""Test invalid log level raises validation error."""
|
||||
with (
|
||||
patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABASE_URL": "postgresql://user:pass@localhost/db",
|
||||
"LOG_LEVEL": "TRACE",
|
||||
},
|
||||
),
|
||||
pytest.raises(ValidationError),
|
||||
):
|
||||
Settings()
|
||||
|
||||
def test_health_port_validation(self) -> None:
|
||||
"""Test health port must be valid port number."""
|
||||
with (
|
||||
patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABASE_URL": "postgresql://user:pass@localhost/db",
|
||||
"HEALTH_PORT": "99999",
|
||||
},
|
||||
),
|
||||
pytest.raises(ValidationError, match="65535"),
|
||||
):
|
||||
Settings()
|
||||
|
||||
def test_get_logging_level(self) -> None:
|
||||
"""Test get_logging_level returns numeric level."""
|
||||
import logging
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABASE_URL": "postgresql://user:pass@localhost/db",
|
||||
"LOG_LEVEL": "WARNING",
|
||||
},
|
||||
):
|
||||
settings = Settings()
|
||||
assert settings.get_logging_level() == logging.WARNING
|
||||
|
||||
def test_redacted_summary(self) -> None:
|
||||
"""Test redacted_summary masks sensitive data."""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABASE_URL": "postgresql://user:secretpass@localhost/db",
|
||||
"REDIS_URL": "redis://localhost:6379",
|
||||
},
|
||||
):
|
||||
settings = Settings()
|
||||
summary = settings.redacted_summary()
|
||||
|
||||
# Database password should be redacted
|
||||
db_url = summary["database_url"]
|
||||
assert isinstance(db_url, str)
|
||||
assert "secretpass" not in db_url
|
||||
assert "***" in db_url
|
||||
assert "user" in db_url
|
||||
|
||||
|
||||
class TestGetSettings:
|
||||
"""Tests for get_settings singleton."""
|
||||
|
||||
def test_returns_same_instance(self) -> None:
|
||||
"""Test get_settings returns cached instance."""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"DATABASE_URL": "postgresql://user:pass@localhost/db"},
|
||||
):
|
||||
settings1 = get_settings()
|
||||
settings2 = get_settings()
|
||||
assert settings1 is settings2
|
||||
|
||||
def test_clear_cache_allows_reload(self) -> None:
|
||||
"""Test clear_settings_cache allows reloading settings."""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABASE_URL": "postgresql://user:pass@localhost/db",
|
||||
"LOG_LEVEL": "INFO",
|
||||
},
|
||||
):
|
||||
settings1 = get_settings()
|
||||
assert settings1.log_level == "INFO"
|
||||
|
||||
clear_settings_cache()
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABASE_URL": "postgresql://user:pass@localhost/db",
|
||||
"LOG_LEVEL": "DEBUG",
|
||||
},
|
||||
):
|
||||
settings2 = get_settings()
|
||||
assert settings2.log_level == "DEBUG"
|
||||
assert settings1 is not settings2
|
||||
@@ -0,0 +1,214 @@
|
||||
"""Tests for the CLI entry point."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.__main__ import (
|
||||
EXIT_CONFIG_ERROR,
|
||||
EXIT_SUCCESS,
|
||||
configure_logging,
|
||||
create_parser,
|
||||
main,
|
||||
print_banner,
|
||||
run_config_check,
|
||||
validate_config,
|
||||
)
|
||||
|
||||
|
||||
class TestCreateParser:
|
||||
"""Tests for argument parser creation."""
|
||||
|
||||
def test_parser_has_version(self):
|
||||
"""Parser should have version flag."""
|
||||
parser = create_parser()
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
parser.parse_args(["--version"])
|
||||
assert exc_info.value.code == 0
|
||||
|
||||
def test_parser_config_check(self):
|
||||
"""Parser should accept --config-check flag."""
|
||||
parser = create_parser()
|
||||
args = parser.parse_args(["--config-check"])
|
||||
assert args.config_check is True
|
||||
|
||||
def test_parser_log_level(self):
|
||||
"""Parser should accept --log-level option."""
|
||||
parser = create_parser()
|
||||
args = parser.parse_args(["--log-level", "DEBUG"])
|
||||
assert args.log_level == "DEBUG"
|
||||
|
||||
def test_parser_dry_run(self):
|
||||
"""Parser should accept --dry-run flag."""
|
||||
parser = create_parser()
|
||||
args = parser.parse_args(["--dry-run"])
|
||||
assert args.dry_run is True
|
||||
|
||||
def test_parser_health_port(self):
|
||||
"""Parser should accept --health-port option."""
|
||||
parser = create_parser()
|
||||
args = parser.parse_args(["--health-port", "9090"])
|
||||
assert args.health_port == 9090
|
||||
|
||||
def test_parser_default_values(self):
|
||||
"""Parser should have correct defaults."""
|
||||
parser = create_parser()
|
||||
args = parser.parse_args([])
|
||||
assert args.config_check is False
|
||||
assert args.log_level is None
|
||||
assert args.dry_run is False
|
||||
assert args.health_port is None
|
||||
|
||||
|
||||
class TestConfigureLogging:
|
||||
"""Tests for logging configuration."""
|
||||
|
||||
def test_configure_logging_info(self):
|
||||
"""Should configure logging at INFO level."""
|
||||
configure_logging("INFO")
|
||||
import logging
|
||||
|
||||
assert logging.getLogger().level == logging.INFO
|
||||
|
||||
def test_configure_logging_debug(self):
|
||||
"""Should configure logging at DEBUG level."""
|
||||
configure_logging("DEBUG")
|
||||
import logging
|
||||
|
||||
assert logging.getLogger().level == logging.DEBUG
|
||||
|
||||
|
||||
class TestPrintBanner:
|
||||
"""Tests for banner printing."""
|
||||
|
||||
def test_banner_contains_app_name(self, capsys):
|
||||
"""Banner should contain application name."""
|
||||
print_banner()
|
||||
captured = capsys.readouterr()
|
||||
assert "Polymarket Insider Tracker" in captured.out
|
||||
|
||||
def test_banner_contains_version(self, capsys):
|
||||
"""Banner should contain version."""
|
||||
print_banner()
|
||||
captured = capsys.readouterr()
|
||||
assert "v0.1.0" in captured.out
|
||||
|
||||
|
||||
class TestValidateConfig:
|
||||
"""Tests for configuration validation."""
|
||||
|
||||
def test_validate_config_success(self, monkeypatch):
|
||||
"""Should return settings on valid config."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://localhost/test")
|
||||
|
||||
settings = validate_config()
|
||||
assert settings is not None
|
||||
|
||||
def test_validate_config_failure(self, monkeypatch, capsys):
|
||||
"""Should return None on invalid config."""
|
||||
# Clear any existing DATABASE_URL
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
|
||||
settings = validate_config()
|
||||
assert settings is None
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert "Configuration validation failed" in captured.err
|
||||
|
||||
|
||||
class TestRunConfigCheck:
|
||||
"""Tests for config check mode."""
|
||||
|
||||
def test_config_check_prints_summary(self, monkeypatch, capsys):
|
||||
"""Config check should print configuration summary."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://localhost/test")
|
||||
|
||||
settings = validate_config()
|
||||
assert settings is not None
|
||||
|
||||
result = run_config_check(settings)
|
||||
assert result == EXIT_SUCCESS
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert "Configuration is valid!" in captured.out
|
||||
assert "Configuration:" in captured.out
|
||||
|
||||
|
||||
class TestMain:
|
||||
"""Tests for main entry point."""
|
||||
|
||||
def test_main_with_config_check(self, monkeypatch):
|
||||
"""Main should exit successfully with --config-check."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://localhost/test")
|
||||
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
main(["--config-check"])
|
||||
|
||||
assert exc_info.value.code == EXIT_SUCCESS
|
||||
|
||||
def test_main_with_invalid_config(self, monkeypatch):
|
||||
"""Main should exit with config error on invalid config."""
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
main([])
|
||||
|
||||
assert exc_info.value.code == EXIT_CONFIG_ERROR
|
||||
|
||||
def test_main_with_dry_run_and_config_check(self, monkeypatch):
|
||||
"""Main should handle dry-run with config-check."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://localhost/test")
|
||||
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
main(["--config-check", "--dry-run"])
|
||||
|
||||
assert exc_info.value.code == EXIT_SUCCESS
|
||||
|
||||
@patch("polymarket_insider_tracker.__main__.run_pipeline")
|
||||
@patch("polymarket_insider_tracker.__main__.asyncio.run")
|
||||
def test_main_runs_pipeline(self, mock_asyncio_run, _mock_run_pipeline, monkeypatch):
|
||||
"""Main should run pipeline when not in config-check mode."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://localhost/test")
|
||||
mock_asyncio_run.return_value = EXIT_SUCCESS
|
||||
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
main([])
|
||||
|
||||
assert exc_info.value.code == EXIT_SUCCESS
|
||||
mock_asyncio_run.assert_called_once()
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests for CLI invocation."""
|
||||
|
||||
def test_cli_help_option(self, capsys):
|
||||
"""CLI should display help with -h option."""
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
main(["-h"])
|
||||
|
||||
assert exc_info.value.code == 0
|
||||
captured = capsys.readouterr()
|
||||
assert "polymarket-insider-tracker" in captured.out
|
||||
assert "--config-check" in captured.out
|
||||
assert "--dry-run" in captured.out
|
||||
assert "--log-level" in captured.out
|
||||
|
||||
def test_cli_version_option(self, capsys):
|
||||
"""CLI should display version with --version option."""
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
main(["--version"])
|
||||
|
||||
assert exc_info.value.code == 0
|
||||
captured = capsys.readouterr()
|
||||
assert "0.1.0" in captured.out
|
||||
|
||||
def test_cli_invalid_log_level(self, capsys):
|
||||
"""CLI should reject invalid log level."""
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
main(["--log-level", "INVALID"])
|
||||
|
||||
assert exc_info.value.code != 0
|
||||
captured = capsys.readouterr()
|
||||
assert "invalid choice" in captured.err
|
||||
@@ -0,0 +1,394 @@
|
||||
"""Tests for the main pipeline orchestrator."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.config import Settings
|
||||
from polymarket_insider_tracker.detector.models import FreshWalletSignal
|
||||
from polymarket_insider_tracker.detector.scorer import SignalBundle
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
from polymarket_insider_tracker.pipeline import Pipeline, PipelineState
|
||||
from polymarket_insider_tracker.profiler.models import WalletProfile
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_settings():
|
||||
"""Create mock settings for testing."""
|
||||
# Create nested mock objects
|
||||
redis = MagicMock()
|
||||
redis.url = "redis://localhost:6379"
|
||||
|
||||
database = MagicMock()
|
||||
database.url = "postgresql+asyncpg://user:pass@localhost/db"
|
||||
|
||||
polygon = MagicMock()
|
||||
polygon.rpc_url = "https://polygon-rpc.com"
|
||||
polygon.fallback_rpc_url = None
|
||||
|
||||
polymarket = MagicMock()
|
||||
polymarket.ws_url = "wss://ws-subscriptions-clob.polymarket.com/ws/market"
|
||||
polymarket.api_key = None
|
||||
|
||||
discord = MagicMock()
|
||||
discord.enabled = False
|
||||
discord.webhook_url = None
|
||||
|
||||
telegram = MagicMock()
|
||||
telegram.enabled = False
|
||||
telegram.bot_token = None
|
||||
telegram.chat_id = None
|
||||
|
||||
settings = MagicMock(spec=Settings)
|
||||
settings.redis = redis
|
||||
settings.database = database
|
||||
settings.polygon = polygon
|
||||
settings.polymarket = polymarket
|
||||
settings.discord = discord
|
||||
settings.telegram = telegram
|
||||
settings.dry_run = True
|
||||
return settings
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_trade_event():
|
||||
"""Create a sample trade event for testing."""
|
||||
return TradeEvent(
|
||||
trade_id="0x" + "a" * 64,
|
||||
wallet_address="0x" + "b" * 40,
|
||||
market_id="0x" + "c" * 64,
|
||||
asset_id="asset_123",
|
||||
side="BUY",
|
||||
price=Decimal("0.65"),
|
||||
size=Decimal("5000"),
|
||||
timestamp=datetime.now(UTC),
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
event_title="Test Market",
|
||||
market_slug="test-market",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_wallet_profile():
|
||||
"""Create a sample wallet profile for testing."""
|
||||
return WalletProfile(
|
||||
address="0x" + "b" * 40,
|
||||
nonce=2,
|
||||
first_seen=datetime.now(UTC),
|
||||
age_hours=1.5,
|
||||
is_fresh=True,
|
||||
total_tx_count=2,
|
||||
matic_balance=Decimal("100"),
|
||||
usdc_balance=Decimal("5000"),
|
||||
fresh_threshold=5,
|
||||
)
|
||||
|
||||
|
||||
class TestPipelineState:
|
||||
"""Tests for pipeline state management."""
|
||||
|
||||
def test_initial_state_is_stopped(self, mock_settings):
|
||||
"""Pipeline should start in stopped state."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
assert pipeline.state == PipelineState.STOPPED
|
||||
|
||||
def test_is_running_property(self, mock_settings):
|
||||
"""is_running property should reflect state."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
assert not pipeline.is_running
|
||||
|
||||
pipeline._state = PipelineState.RUNNING
|
||||
assert pipeline.is_running
|
||||
|
||||
|
||||
class TestPipelineStats:
|
||||
"""Tests for pipeline statistics."""
|
||||
|
||||
def test_initial_stats(self, mock_settings):
|
||||
"""Pipeline should have zero stats initially."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
stats = pipeline.stats
|
||||
|
||||
assert stats.started_at is None
|
||||
assert stats.trades_processed == 0
|
||||
assert stats.signals_generated == 0
|
||||
assert stats.alerts_sent == 0
|
||||
assert stats.errors == 0
|
||||
|
||||
|
||||
class TestPipelineInitialization:
|
||||
"""Tests for pipeline initialization."""
|
||||
|
||||
def test_dry_run_from_settings(self, mock_settings):
|
||||
"""Pipeline should use dry_run from settings by default."""
|
||||
mock_settings.dry_run = True
|
||||
pipeline = Pipeline(mock_settings)
|
||||
assert pipeline._dry_run is True
|
||||
|
||||
mock_settings.dry_run = False
|
||||
pipeline = Pipeline(mock_settings)
|
||||
assert pipeline._dry_run is False
|
||||
|
||||
def test_dry_run_override(self, mock_settings):
|
||||
"""Pipeline should allow overriding dry_run."""
|
||||
mock_settings.dry_run = False
|
||||
pipeline = Pipeline(mock_settings, dry_run=True)
|
||||
assert pipeline._dry_run is True
|
||||
|
||||
def test_uses_get_settings_when_none_provided(self):
|
||||
"""Pipeline should call get_settings if no settings provided."""
|
||||
with patch("polymarket_insider_tracker.pipeline.get_settings") as mock_get:
|
||||
mock_get.return_value = MagicMock(spec=Settings)
|
||||
mock_get.return_value.dry_run = False
|
||||
Pipeline()
|
||||
mock_get.assert_called_once()
|
||||
|
||||
|
||||
class TestBuildAlertChannels:
|
||||
"""Tests for alert channel building."""
|
||||
|
||||
def test_no_channels_when_none_enabled(self, mock_settings):
|
||||
"""Should return empty list when no channels enabled."""
|
||||
mock_settings.discord.enabled = False
|
||||
mock_settings.telegram.enabled = False
|
||||
|
||||
pipeline = Pipeline(mock_settings)
|
||||
channels = pipeline._build_alert_channels()
|
||||
|
||||
assert channels == []
|
||||
|
||||
def test_discord_channel_when_enabled(self, mock_settings):
|
||||
"""Should add Discord channel when enabled."""
|
||||
mock_settings.discord.enabled = True
|
||||
mock_settings.discord.webhook_url = MagicMock()
|
||||
mock_settings.discord.webhook_url.get_secret_value.return_value = (
|
||||
"https://discord.com/webhook"
|
||||
)
|
||||
|
||||
pipeline = Pipeline(mock_settings)
|
||||
channels = pipeline._build_alert_channels()
|
||||
|
||||
assert len(channels) == 1
|
||||
assert channels[0].name == "discord"
|
||||
|
||||
def test_telegram_channel_when_enabled(self, mock_settings):
|
||||
"""Should add Telegram channel when enabled."""
|
||||
mock_settings.telegram.enabled = True
|
||||
mock_settings.telegram.bot_token = MagicMock()
|
||||
mock_settings.telegram.bot_token.get_secret_value.return_value = "bot_token"
|
||||
mock_settings.telegram.chat_id = "chat_123"
|
||||
|
||||
pipeline = Pipeline(mock_settings)
|
||||
channels = pipeline._build_alert_channels()
|
||||
|
||||
assert len(channels) == 1
|
||||
assert channels[0].name == "telegram"
|
||||
|
||||
def test_both_channels_when_both_enabled(self, mock_settings):
|
||||
"""Should add both channels when both enabled."""
|
||||
mock_settings.discord.enabled = True
|
||||
mock_settings.discord.webhook_url = MagicMock()
|
||||
mock_settings.discord.webhook_url.get_secret_value.return_value = (
|
||||
"https://discord.com/webhook"
|
||||
)
|
||||
mock_settings.telegram.enabled = True
|
||||
mock_settings.telegram.bot_token = MagicMock()
|
||||
mock_settings.telegram.bot_token.get_secret_value.return_value = "bot_token"
|
||||
mock_settings.telegram.chat_id = "chat_123"
|
||||
|
||||
pipeline = Pipeline(mock_settings)
|
||||
channels = pipeline._build_alert_channels()
|
||||
|
||||
assert len(channels) == 2
|
||||
|
||||
|
||||
class TestOnTrade:
|
||||
"""Tests for trade event processing."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_trade_increments_stats(self, mock_settings, sample_trade_event):
|
||||
"""Processing a trade should increment stats."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
pipeline._fresh_wallet_detector = AsyncMock(return_value=None)
|
||||
pipeline._size_anomaly_detector = AsyncMock(return_value=None)
|
||||
|
||||
await pipeline._on_trade(sample_trade_event)
|
||||
|
||||
assert pipeline.stats.trades_processed == 1
|
||||
assert pipeline.stats.last_trade_time is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_trade_runs_detectors_in_parallel(self, mock_settings, sample_trade_event):
|
||||
"""Detectors should run in parallel."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
pipeline._fresh_wallet_detector = AsyncMock()
|
||||
pipeline._size_anomaly_detector = AsyncMock()
|
||||
|
||||
# Make detectors take some time
|
||||
async def slow_detect(*_args):
|
||||
await asyncio.sleep(0.1)
|
||||
return None
|
||||
|
||||
pipeline._fresh_wallet_detector.analyze = slow_detect
|
||||
pipeline._size_anomaly_detector.analyze = slow_detect
|
||||
|
||||
start = asyncio.get_event_loop().time()
|
||||
await pipeline._on_trade(sample_trade_event)
|
||||
elapsed = asyncio.get_event_loop().time() - start
|
||||
|
||||
# Should complete in ~0.1s not ~0.2s
|
||||
assert elapsed < 0.15
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_trade_handles_detector_errors(self, mock_settings, sample_trade_event):
|
||||
"""Should handle detector errors gracefully."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
pipeline._fresh_wallet_detector = MagicMock()
|
||||
pipeline._fresh_wallet_detector.analyze = AsyncMock(side_effect=Exception("Detector error"))
|
||||
pipeline._size_anomaly_detector = MagicMock()
|
||||
pipeline._size_anomaly_detector.analyze = AsyncMock(return_value=None)
|
||||
|
||||
# Should not raise
|
||||
await pipeline._on_trade(sample_trade_event)
|
||||
|
||||
# Should still increment trades processed
|
||||
assert pipeline.stats.trades_processed == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_trade_calls_score_and_alert_when_signals(
|
||||
self, mock_settings, sample_trade_event, sample_wallet_profile
|
||||
):
|
||||
"""Should call score_and_alert when signals are detected."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
|
||||
# Create a signal
|
||||
fresh_signal = FreshWalletSignal(
|
||||
trade_event=sample_trade_event,
|
||||
wallet_profile=sample_wallet_profile,
|
||||
confidence=0.8,
|
||||
factors={"base": 0.5, "brand_new": 0.2},
|
||||
)
|
||||
|
||||
pipeline._fresh_wallet_detector = MagicMock()
|
||||
pipeline._fresh_wallet_detector.analyze = AsyncMock(return_value=fresh_signal)
|
||||
pipeline._size_anomaly_detector = MagicMock()
|
||||
pipeline._size_anomaly_detector.analyze = AsyncMock(return_value=None)
|
||||
|
||||
# Mock score_and_alert
|
||||
pipeline._score_and_alert = AsyncMock()
|
||||
|
||||
await pipeline._on_trade(sample_trade_event)
|
||||
|
||||
# Should call score_and_alert with the bundle
|
||||
pipeline._score_and_alert.assert_called_once()
|
||||
bundle = pipeline._score_and_alert.call_args[0][0]
|
||||
assert bundle.fresh_wallet_signal == fresh_signal
|
||||
assert pipeline.stats.signals_generated == 1
|
||||
|
||||
|
||||
class TestScoreAndAlert:
|
||||
"""Tests for scoring and alerting."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dry_run_skips_dispatch(
|
||||
self, mock_settings, sample_trade_event, sample_wallet_profile
|
||||
):
|
||||
"""Dry run should skip actual alert dispatch."""
|
||||
mock_settings.dry_run = True
|
||||
pipeline = Pipeline(mock_settings)
|
||||
|
||||
# Create mock components
|
||||
pipeline._risk_scorer = MagicMock()
|
||||
pipeline._risk_scorer.assess = AsyncMock(
|
||||
return_value=MagicMock(
|
||||
should_alert=True,
|
||||
wallet_address="0x" + "b" * 40,
|
||||
weighted_score=0.85,
|
||||
)
|
||||
)
|
||||
pipeline._alert_formatter = MagicMock()
|
||||
pipeline._alert_dispatcher = MagicMock()
|
||||
pipeline._alert_dispatcher.dispatch = AsyncMock()
|
||||
|
||||
bundle = SignalBundle(
|
||||
trade_event=sample_trade_event,
|
||||
fresh_wallet_signal=FreshWalletSignal(
|
||||
trade_event=sample_trade_event,
|
||||
wallet_profile=sample_wallet_profile,
|
||||
confidence=0.8,
|
||||
factors={},
|
||||
),
|
||||
)
|
||||
|
||||
await pipeline._score_and_alert(bundle)
|
||||
|
||||
# Dispatcher should NOT be called in dry run
|
||||
pipeline._alert_dispatcher.dispatch.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_alert_when_below_threshold(self, mock_settings, sample_trade_event):
|
||||
"""Should not alert when below threshold."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
|
||||
# Create mock components
|
||||
pipeline._risk_scorer = MagicMock()
|
||||
pipeline._risk_scorer.assess = AsyncMock(
|
||||
return_value=MagicMock(
|
||||
should_alert=False,
|
||||
weighted_score=0.4,
|
||||
)
|
||||
)
|
||||
pipeline._alert_formatter = MagicMock()
|
||||
pipeline._alert_dispatcher = MagicMock()
|
||||
|
||||
bundle = SignalBundle(trade_event=sample_trade_event)
|
||||
|
||||
await pipeline._score_and_alert(bundle)
|
||||
|
||||
# Formatter should NOT be called
|
||||
pipeline._alert_formatter.format.assert_not_called()
|
||||
|
||||
|
||||
class TestPipelineLifecycle:
|
||||
"""Tests for pipeline lifecycle methods."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cannot_start_when_not_stopped(self, mock_settings):
|
||||
"""Should raise error when starting non-stopped pipeline."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
pipeline._state = PipelineState.RUNNING
|
||||
|
||||
with pytest.raises(RuntimeError, match="Cannot start pipeline"):
|
||||
await pipeline.start()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_when_already_stopped(self, mock_settings):
|
||||
"""Stop should be no-op when already stopped."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
assert pipeline.state == PipelineState.STOPPED
|
||||
|
||||
# Should not raise
|
||||
await pipeline.stop()
|
||||
assert pipeline.state == PipelineState.STOPPED
|
||||
|
||||
|
||||
class TestPipelineContextManager:
|
||||
"""Tests for async context manager."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_manager_calls_start_and_stop(self, mock_settings):
|
||||
"""Context manager should call start and stop."""
|
||||
pipeline = Pipeline(mock_settings)
|
||||
pipeline.start = AsyncMock()
|
||||
pipeline.stop = AsyncMock()
|
||||
|
||||
async with pipeline:
|
||||
pipeline.start.assert_called_once()
|
||||
|
||||
pipeline.stop.assert_called_once()
|
||||
@@ -0,0 +1,20 @@
|
||||
"""Test that the project setup is working correctly."""
|
||||
|
||||
import polymarket_insider_tracker
|
||||
|
||||
|
||||
def test_version() -> None:
|
||||
"""Test that version is defined."""
|
||||
assert polymarket_insider_tracker.__version__ == "0.1.0"
|
||||
|
||||
|
||||
def test_import_modules() -> None:
|
||||
"""Test that all submodules can be imported."""
|
||||
from polymarket_insider_tracker import alerter, detector, ingestor, profiler, storage
|
||||
|
||||
# Just verify imports work
|
||||
assert ingestor is not None
|
||||
assert profiler is not None
|
||||
assert detector is not None
|
||||
assert alerter is not None
|
||||
assert storage is not None
|
||||
@@ -0,0 +1,319 @@
|
||||
"""Tests for the graceful shutdown handler."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import signal
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from polymarket_insider_tracker.shutdown import (
|
||||
DEFAULT_SHUTDOWN_TIMEOUT,
|
||||
SHUTDOWN_SIGNALS,
|
||||
GracefulShutdown,
|
||||
run_with_graceful_shutdown,
|
||||
)
|
||||
|
||||
|
||||
class TestGracefulShutdownInit:
|
||||
"""Tests for GracefulShutdown initialization."""
|
||||
|
||||
def test_default_timeout(self) -> None:
|
||||
"""Should use default timeout when not specified."""
|
||||
shutdown = GracefulShutdown()
|
||||
assert shutdown.timeout == DEFAULT_SHUTDOWN_TIMEOUT
|
||||
|
||||
def test_custom_timeout(self) -> None:
|
||||
"""Should accept custom timeout."""
|
||||
shutdown = GracefulShutdown(timeout=60.0)
|
||||
assert shutdown.timeout == 60.0
|
||||
|
||||
def test_initial_state(self) -> None:
|
||||
"""Should start in non-shutdown state."""
|
||||
shutdown = GracefulShutdown()
|
||||
assert shutdown.is_shutdown_requested is False
|
||||
assert shutdown.is_force_exit_requested is False
|
||||
|
||||
|
||||
class TestRequestShutdown:
|
||||
"""Tests for programmatic shutdown requests."""
|
||||
|
||||
async def test_request_shutdown_sets_flag(self) -> None:
|
||||
"""Should set shutdown requested flag."""
|
||||
shutdown = GracefulShutdown()
|
||||
shutdown.request_shutdown()
|
||||
assert shutdown.is_shutdown_requested is True
|
||||
|
||||
async def test_request_shutdown_sets_event(self) -> None:
|
||||
"""Should set the shutdown event when called."""
|
||||
shutdown = GracefulShutdown()
|
||||
shutdown._shutdown_event = asyncio.Event()
|
||||
|
||||
shutdown.request_shutdown()
|
||||
|
||||
assert shutdown._shutdown_event.is_set()
|
||||
|
||||
async def test_request_shutdown_idempotent(self) -> None:
|
||||
"""Multiple requests should be idempotent."""
|
||||
shutdown = GracefulShutdown()
|
||||
shutdown._shutdown_event = asyncio.Event()
|
||||
|
||||
shutdown.request_shutdown()
|
||||
shutdown.request_shutdown()
|
||||
shutdown.request_shutdown()
|
||||
|
||||
assert shutdown.is_shutdown_requested is True
|
||||
|
||||
|
||||
class TestWait:
|
||||
"""Tests for waiting for shutdown."""
|
||||
|
||||
async def test_wait_blocks_until_shutdown(self) -> None:
|
||||
"""Wait should block until shutdown is requested."""
|
||||
shutdown = GracefulShutdown()
|
||||
|
||||
async def request_after_delay() -> None:
|
||||
await asyncio.sleep(0.1)
|
||||
shutdown.request_shutdown()
|
||||
|
||||
asyncio.create_task(request_after_delay())
|
||||
|
||||
# Should complete after the delay
|
||||
await asyncio.wait_for(shutdown.wait(), timeout=1.0)
|
||||
|
||||
assert shutdown.is_shutdown_requested is True
|
||||
|
||||
async def test_wait_with_timeout_returns_true_on_shutdown(self) -> None:
|
||||
"""wait_with_timeout should return True when shutdown is requested."""
|
||||
shutdown = GracefulShutdown(timeout=1.0)
|
||||
|
||||
async def request_after_delay() -> None:
|
||||
await asyncio.sleep(0.05)
|
||||
shutdown.request_shutdown()
|
||||
|
||||
asyncio.create_task(request_after_delay())
|
||||
|
||||
result = await shutdown.wait_with_timeout()
|
||||
assert result is True
|
||||
|
||||
async def test_wait_with_timeout_returns_false_on_timeout(self) -> None:
|
||||
"""wait_with_timeout should return False when timeout occurs."""
|
||||
shutdown = GracefulShutdown(timeout=0.05)
|
||||
|
||||
# Don't request shutdown, let it timeout
|
||||
result = await shutdown.wait_with_timeout()
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestSignalHandlers:
|
||||
"""Tests for signal handler installation and removal."""
|
||||
|
||||
async def test_install_signal_handlers_creates_event(self) -> None:
|
||||
"""Installing handlers should create the shutdown event."""
|
||||
shutdown = GracefulShutdown()
|
||||
shutdown.install_signal_handlers()
|
||||
|
||||
try:
|
||||
assert shutdown._shutdown_event is not None
|
||||
assert not shutdown._shutdown_event.is_set()
|
||||
finally:
|
||||
shutdown.remove_signal_handlers()
|
||||
|
||||
@pytest.mark.skipif(sys.platform == "win32", reason="Unix-specific test")
|
||||
async def test_unix_signal_handler_installed(self) -> None:
|
||||
"""On Unix, should install loop signal handlers."""
|
||||
shutdown = GracefulShutdown()
|
||||
|
||||
with patch.object(asyncio.get_running_loop(), "add_signal_handler") as mock_add:
|
||||
shutdown.install_signal_handlers()
|
||||
|
||||
# Should have been called for both SIGTERM and SIGINT
|
||||
assert mock_add.call_count >= 1
|
||||
|
||||
shutdown.remove_signal_handlers()
|
||||
|
||||
@pytest.mark.skipif(sys.platform != "win32", reason="Windows-specific test")
|
||||
async def test_windows_signal_handler_installed(self) -> None:
|
||||
"""On Windows, should install signal.signal handlers."""
|
||||
shutdown = GracefulShutdown()
|
||||
|
||||
with patch("signal.signal") as mock_signal:
|
||||
shutdown.install_signal_handlers()
|
||||
|
||||
# Should have been called for available signals
|
||||
assert mock_signal.call_count >= 1
|
||||
|
||||
shutdown.remove_signal_handlers()
|
||||
|
||||
|
||||
class TestHandleSignal:
|
||||
"""Tests for signal handling behavior."""
|
||||
|
||||
async def test_first_signal_sets_shutdown_event(self) -> None:
|
||||
"""First signal should set the shutdown event."""
|
||||
shutdown = GracefulShutdown()
|
||||
shutdown._shutdown_event = asyncio.Event()
|
||||
|
||||
shutdown._handle_signal(signal.SIGTERM)
|
||||
|
||||
assert shutdown.is_shutdown_requested is True
|
||||
assert shutdown._shutdown_event.is_set()
|
||||
assert shutdown.is_force_exit_requested is False
|
||||
|
||||
async def test_second_signal_force_exits(self) -> None:
|
||||
"""Second signal should trigger force exit."""
|
||||
shutdown = GracefulShutdown()
|
||||
shutdown._shutdown_event = asyncio.Event()
|
||||
shutdown._shutdown_requested = True # First signal already received
|
||||
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
shutdown._handle_signal(signal.SIGTERM)
|
||||
|
||||
assert exc_info.value.code == 128 + signal.SIGTERM.value
|
||||
assert shutdown.is_force_exit_requested is True
|
||||
|
||||
|
||||
class TestCleanupCallbacks:
|
||||
"""Tests for cleanup callback registration and execution."""
|
||||
|
||||
async def test_register_sync_callback(self) -> None:
|
||||
"""Should register sync cleanup callbacks."""
|
||||
shutdown = GracefulShutdown()
|
||||
callback = MagicMock()
|
||||
|
||||
shutdown.register_cleanup(callback)
|
||||
|
||||
assert callback in shutdown._cleanup_callbacks
|
||||
|
||||
async def test_run_sync_cleanup_callback(self) -> None:
|
||||
"""Should run sync cleanup callbacks."""
|
||||
shutdown = GracefulShutdown()
|
||||
callback = MagicMock()
|
||||
shutdown.register_cleanup(callback)
|
||||
|
||||
await shutdown.run_cleanup_callbacks()
|
||||
|
||||
callback.assert_called_once()
|
||||
|
||||
async def test_run_async_cleanup_callback(self) -> None:
|
||||
"""Should run async cleanup callbacks."""
|
||||
shutdown = GracefulShutdown()
|
||||
called = False
|
||||
|
||||
async def async_callback() -> None:
|
||||
nonlocal called
|
||||
called = True
|
||||
|
||||
shutdown.register_cleanup(async_callback)
|
||||
|
||||
await shutdown.run_cleanup_callbacks()
|
||||
|
||||
assert called is True
|
||||
|
||||
async def test_cleanup_callback_error_logged(self) -> None:
|
||||
"""Cleanup callback errors should be logged, not raised."""
|
||||
shutdown = GracefulShutdown()
|
||||
|
||||
def failing_callback() -> None:
|
||||
raise ValueError("Cleanup failed")
|
||||
|
||||
shutdown.register_cleanup(failing_callback)
|
||||
|
||||
# Should not raise
|
||||
await shutdown.run_cleanup_callbacks()
|
||||
|
||||
|
||||
class TestAsyncContextManager:
|
||||
"""Tests for async context manager protocol."""
|
||||
|
||||
async def test_context_manager_installs_handlers(self) -> None:
|
||||
"""Entering context should install signal handlers."""
|
||||
shutdown = GracefulShutdown()
|
||||
|
||||
async with shutdown:
|
||||
assert shutdown._shutdown_event is not None
|
||||
assert shutdown._loop is not None
|
||||
|
||||
async def test_context_manager_removes_handlers(self) -> None:
|
||||
"""Exiting context should remove signal handlers."""
|
||||
shutdown = GracefulShutdown()
|
||||
|
||||
with patch.object(shutdown, "remove_signal_handlers") as mock_remove:
|
||||
async with shutdown:
|
||||
pass
|
||||
|
||||
mock_remove.assert_called_once()
|
||||
|
||||
async def test_context_manager_runs_cleanup(self) -> None:
|
||||
"""Exiting context should run cleanup callbacks."""
|
||||
shutdown = GracefulShutdown()
|
||||
callback = MagicMock()
|
||||
shutdown.register_cleanup(callback)
|
||||
|
||||
async with shutdown:
|
||||
pass
|
||||
|
||||
callback.assert_called_once()
|
||||
|
||||
|
||||
class TestRunWithGracefulShutdown:
|
||||
"""Tests for the run_with_graceful_shutdown helper."""
|
||||
|
||||
async def test_runs_coroutine_to_completion(self) -> None:
|
||||
"""Should run the coroutine to completion if no shutdown."""
|
||||
result = []
|
||||
|
||||
async def my_coro() -> None:
|
||||
result.append("done")
|
||||
|
||||
await run_with_graceful_shutdown(my_coro(), timeout=1.0)
|
||||
|
||||
assert result == ["done"]
|
||||
|
||||
async def test_stops_on_shutdown_signal(self) -> None:
|
||||
"""Should stop when shutdown is signaled."""
|
||||
result = []
|
||||
|
||||
async def long_running() -> None:
|
||||
try:
|
||||
await asyncio.sleep(10.0)
|
||||
except asyncio.CancelledError:
|
||||
result.append("cancelled")
|
||||
raise
|
||||
|
||||
# Create a wrapper that will trigger shutdown
|
||||
async def run_with_interrupt() -> None:
|
||||
async def trigger_shutdown() -> None:
|
||||
await asyncio.sleep(0.1)
|
||||
# We can't easily trigger a signal in tests, so we test the task
|
||||
# completion path instead
|
||||
pass
|
||||
|
||||
# Run both tasks
|
||||
task = asyncio.create_task(long_running())
|
||||
asyncio.create_task(trigger_shutdown())
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
task.cancel()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
result.append("finished")
|
||||
|
||||
await run_with_interrupt()
|
||||
assert "finished" in result
|
||||
|
||||
|
||||
class TestShutdownSignals:
|
||||
"""Tests for shutdown signal configuration."""
|
||||
|
||||
def test_shutdown_signals_includes_sigterm(self) -> None:
|
||||
"""SHUTDOWN_SIGNALS should include SIGTERM."""
|
||||
assert signal.SIGTERM in SHUTDOWN_SIGNALS
|
||||
|
||||
def test_shutdown_signals_includes_sigint(self) -> None:
|
||||
"""SHUTDOWN_SIGNALS should include SIGINT."""
|
||||
assert signal.SIGINT in SHUTDOWN_SIGNALS
|
||||
Reference in New Issue
Block a user