Merge pull request #58 from pselamy/fix/54
feat: implement main pipeline orchestrator (#54)
This commit is contained in:
Patrick Selamy
+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,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,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()
|
||||
Reference in New Issue
Block a user