Add comprehensive tests for Phase 2 correlation engine

New Tests (13 total, all passing):
- test_get_alerts_before_trade: Retrieve prior alerts
- test_get_alerts_before_trade_no_alerts: Handle no alerts
- test_calculate_timing_score_high_suspicion: High score logic
- test_calculate_timing_score_no_alerts: Zero score handling
- test_calculate_timing_score_factors: Multi-factor scoring
- test_analyze_trade_full: Complete trade analysis
- test_analyze_recent_disclosures: Batch processing
- test_get_official_timing_pattern: Historical patterns
- test_get_official_timing_pattern_no_trades: Edge case
- test_get_ticker_timing_analysis: Per-ticker analysis
- test_get_ticker_timing_analysis_no_trades: Edge case
- test_alerts_outside_lookback_window: Date filtering
- test_different_ticker_alerts_excluded: Ticker filtering

Test Coverage:
- Alert-to-trade correlation
- Timing score calculation (all factors)
- Pattern analysis (officials & tickers)
- Batch analysis
- Edge cases & filtering
- Date range handling

Total Test Suite: 82 tests passing 
This commit is contained in:
ilia 2025-12-15 15:20:40 -05:00
parent 6b62ae96f7
commit a52313145b

View File

@ -0,0 +1,455 @@
"""Tests for disclosure correlation module."""
import pytest
from datetime import date, datetime, timedelta, timezone
from decimal import Decimal
from pote.monitoring.disclosure_correlator import DisclosureCorrelator
from pote.db.models import Official, Security, Trade, MarketAlert
@pytest.fixture
def trade_with_alerts(test_db_session):
"""Create a trade with prior market alerts."""
session = test_db_session
# Create official and security
pelosi = Official(name="Nancy Pelosi", chamber="House", party="Democrat", state="CA")
nvda = Security(ticker="NVDA", name="NVIDIA Corporation", sector="Technology")
session.add_all([pelosi, nvda])
session.flush()
# Create trade on Jan 15
trade = Trade(
official_id=pelosi.id,
security_id=nvda.id,
source="test",
transaction_date=date(2024, 1, 15),
filing_date=date(2024, 2, 1),
side="buy",
value_min=Decimal("15001"),
value_max=Decimal("50000"),
)
session.add(trade)
session.flush()
# Create alerts BEFORE trade (suspicious)
alerts = [
MarketAlert(
ticker="NVDA",
alert_type="unusual_volume",
timestamp=datetime(2024, 1, 10, 10, 30, tzinfo=timezone.utc), # 5 days before
details={"multiplier": 3.5},
price=Decimal("490.00"),
volume=100000000,
change_pct=Decimal("2.0"),
severity=8,
),
MarketAlert(
ticker="NVDA",
alert_type="price_spike",
timestamp=datetime(2024, 1, 12, 14, 15, tzinfo=timezone.utc), # 3 days before
details={"change_pct": 5.5},
price=Decimal("505.00"),
volume=85000000,
change_pct=Decimal("5.5"),
severity=7,
),
MarketAlert(
ticker="NVDA",
alert_type="high_volatility",
timestamp=datetime(2024, 1, 14, 16, 20, tzinfo=timezone.utc), # 1 day before
details={"multiplier": 2.5},
price=Decimal("510.00"),
volume=90000000,
change_pct=Decimal("1.5"),
severity=6,
),
]
session.add_all(alerts)
session.commit()
return {
"trade": trade,
"official": pelosi,
"security": nvda,
"alerts": alerts,
}
@pytest.fixture
def trade_without_alerts(test_db_session):
"""Create a trade without prior alerts (clean)."""
session = test_db_session
official = Official(name="John Smith", chamber="House", party="Republican", state="TX")
security = Security(ticker="MSFT", name="Microsoft Corporation", sector="Technology")
session.add_all([official, security])
session.flush()
trade = Trade(
official_id=official.id,
security_id=security.id,
source="test",
transaction_date=date(2024, 2, 1),
side="buy",
value_min=Decimal("10000"),
value_max=Decimal("50000"),
)
session.add(trade)
session.commit()
return {
"trade": trade,
"official": official,
"security": security,
}
def test_get_alerts_before_trade(test_db_session, trade_with_alerts):
"""Test retrieving alerts before a trade."""
session = test_db_session
correlator = DisclosureCorrelator(session)
trade = trade_with_alerts["trade"]
# Get alerts before trade
prior_alerts = correlator.get_alerts_before_trade(trade, lookback_days=30)
assert len(prior_alerts) == 3
assert all(alert.ticker == "NVDA" for alert in prior_alerts)
assert all(alert.timestamp.date() < trade.transaction_date for alert in prior_alerts)
def test_get_alerts_before_trade_no_alerts(test_db_session, trade_without_alerts):
"""Test retrieving alerts when none exist."""
session = test_db_session
correlator = DisclosureCorrelator(session)
trade = trade_without_alerts["trade"]
prior_alerts = correlator.get_alerts_before_trade(trade, lookback_days=30)
assert len(prior_alerts) == 0
def test_calculate_timing_score_high_suspicion(test_db_session, trade_with_alerts):
"""Test timing score calculation for suspicious trade."""
session = test_db_session
correlator = DisclosureCorrelator(session)
trade = trade_with_alerts["trade"]
alerts = trade_with_alerts["alerts"]
timing_analysis = correlator.calculate_timing_score(trade, alerts)
assert timing_analysis["timing_score"] > 60, "Should be suspicious with 3 alerts"
assert timing_analysis["suspicious"] is True
assert timing_analysis["alert_count"] == 3
assert timing_analysis["recent_alert_count"] > 0
assert timing_analysis["high_severity_count"] >= 2 # 2 alerts with severity 7+
assert "reason" in timing_analysis
def test_calculate_timing_score_no_alerts(test_db_session):
"""Test timing score with no prior alerts."""
session = test_db_session
correlator = DisclosureCorrelator(session)
# Create minimal trade
official = Official(name="Test", chamber="House", party="Democrat", state="CA")
security = Security(ticker="TEST", name="Test Corp")
session.add_all([official, security])
session.flush()
trade = Trade(
official_id=official.id,
security_id=security.id,
source="test",
transaction_date=date(2024, 1, 1),
side="buy",
value_min=Decimal("10000"),
)
session.add(trade)
session.commit()
timing_analysis = correlator.calculate_timing_score(trade, [])
assert timing_analysis["timing_score"] == 0
assert timing_analysis["suspicious"] is False
assert timing_analysis["alert_count"] == 0
def test_calculate_timing_score_factors(test_db_session):
"""Test that timing score considers all factors correctly."""
session = test_db_session
correlator = DisclosureCorrelator(session)
# Create trade
official = Official(name="Test", chamber="House", party="Democrat", state="CA")
security = Security(ticker="TEST", name="Test Corp")
session.add_all([official, security])
session.flush()
trade_date = date(2024, 1, 15)
trade = Trade(
official_id=official.id,
security_id=security.id,
source="test",
transaction_date=trade_date,
side="buy",
value_min=Decimal("10000"),
)
session.add(trade)
session.flush()
# Test with low severity alerts (should have lower score)
low_sev_alerts = [
MarketAlert(
ticker="TEST",
alert_type="test",
timestamp=datetime(2024, 1, 10, 12, 0, tzinfo=timezone.utc),
severity=3,
),
MarketAlert(
ticker="TEST",
alert_type="test",
timestamp=datetime(2024, 1, 11, 12, 0, tzinfo=timezone.utc),
severity=4,
),
]
session.add_all(low_sev_alerts)
session.commit()
low_score = correlator.calculate_timing_score(trade, low_sev_alerts)
# Test with high severity alerts (should have higher score)
high_sev_alerts = [
MarketAlert(
ticker="TEST",
alert_type="test",
timestamp=datetime(2024, 1, 13, 12, 0, tzinfo=timezone.utc), # Recent
severity=9,
),
MarketAlert(
ticker="TEST",
alert_type="test",
timestamp=datetime(2024, 1, 14, 12, 0, tzinfo=timezone.utc), # Very recent
severity=8,
),
]
session.add_all(high_sev_alerts)
session.commit()
high_score = correlator.calculate_timing_score(trade, high_sev_alerts)
# High severity + recent should score higher
assert high_score["timing_score"] > low_score["timing_score"]
assert high_score["recent_alert_count"] > 0
assert high_score["high_severity_count"] > 0
def test_analyze_trade_full(test_db_session, trade_with_alerts):
"""Test complete trade analysis."""
session = test_db_session
correlator = DisclosureCorrelator(session)
trade = trade_with_alerts["trade"]
analysis = correlator.analyze_trade(trade)
# Check all required fields
assert analysis["trade_id"] == trade.id
assert analysis["official_name"] == "Nancy Pelosi"
assert analysis["ticker"] == "NVDA"
assert analysis["side"] == "buy"
assert analysis["transaction_date"] == "2024-01-15"
assert analysis["timing_score"] > 0
assert "prior_alerts" in analysis
assert len(analysis["prior_alerts"]) == 3
# Check alert details
for alert_detail in analysis["prior_alerts"]:
assert "timestamp" in alert_detail
assert "alert_type" in alert_detail
assert "severity" in alert_detail
assert "days_before_trade" in alert_detail
assert alert_detail["days_before_trade"] >= 0
def test_analyze_recent_disclosures(test_db_session, trade_with_alerts, trade_without_alerts):
"""Test batch analysis of recent disclosures."""
session = test_db_session
correlator = DisclosureCorrelator(session)
# Both trades were created "recently" (in fixture setup)
suspicious_trades = correlator.analyze_recent_disclosures(
days=365, # Wide window to catch test data
min_timing_score=50
)
# Should find at least the suspicious trade
assert len(suspicious_trades) >= 1
# Check sorting (highest score first)
if len(suspicious_trades) > 1:
for i in range(len(suspicious_trades) - 1):
assert suspicious_trades[i]["timing_score"] >= suspicious_trades[i + 1]["timing_score"]
def test_get_official_timing_pattern(test_db_session, trade_with_alerts):
"""Test official timing pattern analysis."""
session = test_db_session
correlator = DisclosureCorrelator(session)
official = trade_with_alerts["official"]
# Use wide lookback to catch test data (trade is 2024-01-15)
pattern = correlator.get_official_timing_pattern(official.id, lookback_days=3650)
assert pattern["official_id"] == official.id
assert pattern["trade_count"] >= 1
assert pattern["trades_with_prior_alerts"] >= 1
assert pattern["suspicious_trade_count"] >= 0
assert "pattern" in pattern
assert "analyses" in pattern
def test_get_official_timing_pattern_no_trades(test_db_session):
"""Test official with no trades."""
session = test_db_session
correlator = DisclosureCorrelator(session)
official = Official(name="No Trades", chamber="House", party="Democrat", state="CA")
session.add(official)
session.commit()
pattern = correlator.get_official_timing_pattern(official.id)
assert pattern["trade_count"] == 0
assert "No trades" in pattern["pattern"]
def test_get_ticker_timing_analysis(test_db_session, trade_with_alerts):
"""Test ticker timing analysis."""
session = test_db_session
correlator = DisclosureCorrelator(session)
# Use wide lookback to catch test data
analysis = correlator.get_ticker_timing_analysis("NVDA", lookback_days=3650)
assert analysis["ticker"] == "NVDA"
assert analysis["trade_count"] >= 1
assert analysis["trades_with_alerts"] >= 1
assert "avg_timing_score" in analysis
assert "analyses" in analysis
def test_get_ticker_timing_analysis_no_trades(test_db_session):
"""Test ticker with no trades."""
session = test_db_session
correlator = DisclosureCorrelator(session)
analysis = correlator.get_ticker_timing_analysis("ZZZZ")
assert analysis["ticker"] == "ZZZZ"
assert analysis["trade_count"] == 0
assert "No trades" in analysis["pattern"]
def test_alerts_outside_lookback_window(test_db_session):
"""Test that alerts outside lookback window are excluded."""
session = test_db_session
correlator = DisclosureCorrelator(session)
# Create trade and alerts
official = Official(name="Test", chamber="House", party="Democrat", state="CA")
security = Security(ticker="TEST", name="Test Corp")
session.add_all([official, security])
session.flush()
trade_date = date(2024, 1, 15)
trade = Trade(
official_id=official.id,
security_id=security.id,
source="test",
transaction_date=trade_date,
side="buy",
value_min=Decimal("10000"),
)
session.add(trade)
session.flush()
# Alert 2 days before (within window)
recent_alert = MarketAlert(
ticker="TEST",
alert_type="test",
timestamp=datetime(2024, 1, 13, 12, 0, tzinfo=timezone.utc),
severity=7,
)
# Alert 40 days before (outside 30-day window)
old_alert = MarketAlert(
ticker="TEST",
alert_type="test",
timestamp=datetime(2023, 12, 6, 12, 0, tzinfo=timezone.utc),
severity=8,
)
session.add_all([recent_alert, old_alert])
session.commit()
# Should only get recent alert
alerts = correlator.get_alerts_before_trade(trade, lookback_days=30)
assert len(alerts) == 1
assert alerts[0].timestamp.date() == date(2024, 1, 13)
def test_different_ticker_alerts_excluded(test_db_session):
"""Test that alerts for different tickers are excluded."""
session = test_db_session
correlator = DisclosureCorrelator(session)
# Create trade for NVDA
official = Official(name="Test", chamber="House", party="Democrat", state="CA")
nvda = Security(ticker="NVDA", name="NVIDIA")
msft = Security(ticker="MSFT", name="Microsoft")
session.add_all([official, nvda, msft])
session.flush()
trade = Trade(
official_id=official.id,
security_id=nvda.id,
source="test",
transaction_date=date(2024, 1, 15),
side="buy",
value_min=Decimal("10000"),
)
session.add(trade)
session.flush()
# Create alerts for both tickers
nvda_alert = MarketAlert(
ticker="NVDA",
alert_type="test",
timestamp=datetime(2024, 1, 10, 12, 0, tzinfo=timezone.utc),
severity=7,
)
msft_alert = MarketAlert(
ticker="MSFT",
alert_type="test",
timestamp=datetime(2024, 1, 10, 12, 0, tzinfo=timezone.utc),
severity=8,
)
session.add_all([nvda_alert, msft_alert])
session.commit()
# Should only get NVDA alert
alerts = correlator.get_alerts_before_trade(trade, lookback_days=30)
assert len(alerts) == 1
assert alerts[0].ticker == "NVDA"