Codex
512 lines
18 KiB
Python
512 lines
18 KiB
Python
from datetime import datetime, timedelta, timezone
|
|
import json
|
|
import os
|
|
import unittest
|
|
|
|
os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://user:pass@localhost:5432/ai_ict_test")
|
|
|
|
from src.services.market_data.okx_client import OKXCandle
|
|
from src.services.market_state import (
|
|
DATA_GAP_CODE_DETECTED,
|
|
EconomicCalendarProvider,
|
|
MarketDataStreamService,
|
|
MarketStateService,
|
|
MarketTruthRuntimeState,
|
|
MarketTruthSupervisor,
|
|
NewsProviderChain,
|
|
STREAM_CODE_CONNECTED,
|
|
STREAM_CODE_DISCONNECTED,
|
|
TradingEconomicsCalendarProvider,
|
|
build_data_gap_state,
|
|
)
|
|
from src.services.market_state.news_event_provider import NEWS_CODE_PROVIDER_CACHED, NEWS_CODE_PROVIDER_ERROR, NEWS_CODE_PROVIDER_FRESH, NewsProviderSnapshot
|
|
|
|
|
|
class MarketDataStreamServiceTests(unittest.TestCase):
|
|
def test_stream_service_defaults_to_business_websocket(self) -> None:
|
|
service = MarketDataStreamService()
|
|
|
|
self.assertEqual(service.url, "wss://ws.okx.com:8443/ws/v5/business")
|
|
|
|
def test_stream_service_parses_ticker_and_candle_messages(self) -> None:
|
|
service = MarketDataStreamService(url="wss://example.test")
|
|
service.subscribe(
|
|
instrument_id=1,
|
|
symbol="BTC-USDT",
|
|
inst_id="BTC-USDT-SWAP",
|
|
timeframe="1m",
|
|
)
|
|
|
|
service._on_message(None, json.dumps({"event": "subscribe", "arg": {"channel": "candle1m", "instId": "BTC-USDT-SWAP"}}))
|
|
service._on_message(
|
|
None,
|
|
json.dumps(
|
|
{
|
|
"arg": {"channel": "candle1m", "instId": "BTC-USDT-SWAP"},
|
|
"data": [["1778032800000", "64990.1", "65010.0", "64980.0", "65000.5", "12.3", "0", "0", "1"]],
|
|
}
|
|
),
|
|
)
|
|
self.assertEqual(service.iter_subscriptions()[0].candle_channel, "candle1m")
|
|
self.assertEqual(service.iter_subscriptions()[0].active_channels, ["candle1m"])
|
|
|
|
snapshot = service.get_snapshot(instrument_id=1, timeframe="1m")
|
|
|
|
self.assertEqual(snapshot["code"], STREAM_CODE_CONNECTED)
|
|
self.assertTrue(snapshot["connected"])
|
|
self.assertTrue(snapshot["subscribed"])
|
|
self.assertEqual(snapshot["latest_price"], 65000.5)
|
|
self.assertEqual(snapshot["latest_candle_open"], 64990.1)
|
|
self.assertEqual(snapshot["latest_candle_close"], 65000.5)
|
|
self.assertIsNotNone(snapshot["last_confirmed_candle_close_ts"])
|
|
|
|
service_no_ticker = MarketDataStreamService(url="wss://example.test")
|
|
service_no_ticker.subscribe(
|
|
instrument_id=1,
|
|
symbol="BTC-USDT",
|
|
inst_id="BTC-USDT-SWAP",
|
|
timeframe="1m",
|
|
)
|
|
service_no_ticker._on_message(None, json.dumps({"event": "subscribe", "arg": {"channel": "candle1m", "instId": "BTC-USDT-SWAP"}}))
|
|
service_no_ticker._on_message(
|
|
None,
|
|
json.dumps(
|
|
{
|
|
"arg": {"channel": "candle1m", "instId": "BTC-USDT-SWAP"},
|
|
"data": [["1778032800000", "64990.1", "65010.0", "64980.0", "65000.5", "12.3", "0", "0", "1"]],
|
|
}
|
|
),
|
|
)
|
|
snapshot_no_ticker = service_no_ticker.get_snapshot(instrument_id=1, timeframe="1m")
|
|
self.assertEqual(snapshot_no_ticker["code"], STREAM_CODE_CONNECTED)
|
|
self.assertEqual(snapshot_no_ticker["latest_price"], 65000.5)
|
|
self.assertTrue(snapshot_no_ticker["subscribed"])
|
|
self.assertTrue(snapshot_no_ticker["connected"])
|
|
|
|
def test_build_data_gap_state_detects_missing_candle_interval(self) -> None:
|
|
now = datetime(2026, 5, 7, 12, 5, tzinfo=timezone.utc)
|
|
gap_state = build_data_gap_state(
|
|
stream_snapshot={"code": STREAM_CODE_CONNECTED, "gap_detected": False},
|
|
latest_candle_ts=now - timedelta(minutes=4),
|
|
timeframe="1m",
|
|
now=now,
|
|
gap_multiplier=2,
|
|
)
|
|
|
|
self.assertTrue(gap_state["is_gap"])
|
|
self.assertEqual(gap_state["code"], DATA_GAP_CODE_DETECTED)
|
|
|
|
|
|
class EconomicCalendarProviderTests(unittest.TestCase):
|
|
def test_provider_marks_active_high_impact_event(self) -> None:
|
|
now = datetime(2026, 5, 7, 12, 0, tzinfo=timezone.utc)
|
|
provider = EconomicCalendarProvider(
|
|
pre_window_minutes=15,
|
|
post_window_minutes=15,
|
|
high_impact_only=True,
|
|
)
|
|
|
|
snapshot = provider.fetch_snapshot(
|
|
now=now,
|
|
rows=[
|
|
{
|
|
"id": "cpi",
|
|
"title": "US CPI",
|
|
"country": "USD",
|
|
"impact": "High",
|
|
"dateUtc": now.isoformat(),
|
|
}
|
|
],
|
|
)
|
|
|
|
self.assertTrue(snapshot.is_active)
|
|
self.assertEqual(snapshot.code, "news_window_active")
|
|
self.assertEqual(snapshot.event_count, 1)
|
|
|
|
def test_tradingeconomics_provider_parses_high_impact_rows(self) -> None:
|
|
now = datetime(2026, 5, 7, 12, 0, tzinfo=timezone.utc)
|
|
provider = TradingEconomicsCalendarProvider(high_impact_only=True)
|
|
|
|
snapshot = provider.fetch_snapshot(
|
|
now=now,
|
|
rows=[
|
|
{
|
|
"CalendarId": "42",
|
|
"Event": "US CPI",
|
|
"Country": "United States",
|
|
"Importance": 3,
|
|
"Date": now.isoformat(),
|
|
}
|
|
],
|
|
)
|
|
|
|
self.assertEqual(snapshot.freshness_code, NEWS_CODE_PROVIDER_FRESH)
|
|
self.assertEqual(snapshot.code, "news_window_active")
|
|
self.assertEqual(snapshot.active_events[0]["title"], "US CPI")
|
|
|
|
def test_provider_chain_uses_first_successful_snapshot(self) -> None:
|
|
now = datetime(2026, 5, 7, 12, 0, tzinfo=timezone.utc)
|
|
chain = NewsProviderChain(providers=[_FailingNewsProvider(), _WorkingNewsProvider(now=now)])
|
|
snapshot = chain.fetch_snapshot(now=now)
|
|
|
|
self.assertEqual(snapshot.provider_name, "working_news_provider")
|
|
self.assertEqual(snapshot.freshness_code, NEWS_CODE_PROVIDER_FRESH)
|
|
|
|
|
|
class MarketTruthSupervisorTests(unittest.TestCase):
|
|
def test_supervisor_runs_rest_fallback_and_records_incident(self) -> None:
|
|
now = datetime(2026, 5, 7, 12, 0, tzinfo=timezone.utc)
|
|
repository = _SupervisorRepository(now=now)
|
|
candle_repository = _FallbackCandleRepository(repository)
|
|
supervisor = MarketTruthSupervisor(
|
|
repository=repository,
|
|
candle_repository=candle_repository,
|
|
stream_service=_DisconnectedStreamService(now=now),
|
|
rest_client=_FallbackOKXClient(now=now),
|
|
news_provider=_InactiveNewsProvider(now=now),
|
|
enabled=False,
|
|
)
|
|
|
|
state = supervisor.current_state(instrument_id=1, timeframe="1m", now=now)
|
|
|
|
self.assertEqual(state.stream_state["code"], STREAM_CODE_DISCONNECTED)
|
|
self.assertEqual(candle_repository.write_calls, 1)
|
|
self.assertTrue(any(item["incident_code"] == "rest_backfill_applied" for item in repository.incidents))
|
|
self.assertTrue(any(item["incident_code"] == STREAM_CODE_DISCONNECTED for item in repository.incidents))
|
|
|
|
def test_supervisor_keeps_running_when_rest_fallback_fails(self) -> None:
|
|
now = datetime(2026, 5, 7, 12, 0, tzinfo=timezone.utc)
|
|
repository = _SupervisorRepository(now=now)
|
|
supervisor = MarketTruthSupervisor(
|
|
repository=repository,
|
|
candle_repository=_FallbackCandleRepository(repository),
|
|
stream_service=_DisconnectedStreamService(now=now),
|
|
rest_client=_FailingOKXClient(),
|
|
news_provider=_InactiveNewsProvider(now=now),
|
|
enabled=False,
|
|
)
|
|
|
|
state = supervisor.current_state(instrument_id=1, timeframe="1m", now=now)
|
|
|
|
self.assertEqual(state.stream_state["code"], STREAM_CODE_DISCONNECTED)
|
|
self.assertTrue(any(item["incident_code"] == "rest_backfill_failed" for item in repository.incidents))
|
|
|
|
def test_supervisor_uses_cached_news_snapshot_when_provider_errors(self) -> None:
|
|
now = datetime(2026, 5, 7, 12, 0, tzinfo=timezone.utc)
|
|
repository = _SupervisorRepository(now=now)
|
|
repository.cached_news_snapshot = {
|
|
"provider_name": "tradingeconomics_snapshot",
|
|
"fetched_ts": now - timedelta(minutes=30),
|
|
"freshness_code": "news_provider_fresh",
|
|
"event_count": 1,
|
|
"is_active": False,
|
|
"code": "news_window_inactive",
|
|
"active_events": [],
|
|
"error": None,
|
|
"raw_payload": {"rows": []},
|
|
}
|
|
supervisor = MarketTruthSupervisor(
|
|
repository=repository,
|
|
candle_repository=_FallbackCandleRepository(repository),
|
|
stream_service=_DisconnectedStreamService(now=now),
|
|
rest_client=_FallbackOKXClient(now=now),
|
|
news_provider=_FailingChainProvider(),
|
|
enabled=False,
|
|
)
|
|
|
|
state = supervisor.current_state(instrument_id=1, timeframe="1m", now=now)
|
|
|
|
self.assertEqual(state.news_snapshot.freshness_code, NEWS_CODE_PROVIDER_CACHED)
|
|
self.assertEqual(state.news_snapshot.provider_name, "tradingeconomics_snapshot")
|
|
|
|
def test_market_state_service_uses_supervisor_runtime_state(self) -> None:
|
|
now = datetime(2026, 5, 7, 12, 0, tzinfo=timezone.utc)
|
|
repository = _SupervisorRepository(now=now)
|
|
supervisor = _StaticSupervisor(now=now)
|
|
service = MarketStateService(repository=repository, market_truth_supervisor=supervisor)
|
|
|
|
state = service.build_market_state(instrument_id=1, timeframe="1m", now=now)
|
|
|
|
payload = state.to_dict()
|
|
self.assertIn("stream_state", payload)
|
|
self.assertIn("data_gap_state", payload)
|
|
self.assertIn("news_state", payload)
|
|
self.assertEqual(payload["stream_state"]["code"], STREAM_CODE_CONNECTED)
|
|
self.assertEqual(payload["latest_candle"]["latest_price"], 65111.0)
|
|
|
|
|
|
class _SupervisorRepository:
|
|
def __init__(self, now: datetime) -> None:
|
|
self.now = now
|
|
self.latest_candle = {
|
|
"id": 1,
|
|
"instrument_id": 1,
|
|
"timeframe": "1m",
|
|
"ts_close": now - timedelta(minutes=10),
|
|
"open": 64000.0,
|
|
"close": 64010.0,
|
|
}
|
|
self.incidents: list[dict] = []
|
|
self.cached_news_snapshot = None
|
|
|
|
def fetch_active_instruments_by_symbols(self, symbols: list[str]) -> list[dict]:
|
|
return [{"id": 1, "symbol": "BTC-USDT", "active": True}] if "BTC-USDT" in symbols else []
|
|
|
|
def fetch_instrument_by_id(self, instrument_id: int):
|
|
return {"id": instrument_id, "symbol": "BTC-USDT", "active": True}
|
|
|
|
def fetch_latest_candle(self, instrument_id: int, timeframe: str):
|
|
return dict(self.latest_candle)
|
|
|
|
def fetch_session_at(self, instrument_id: int, ts):
|
|
return {
|
|
"id": 5,
|
|
"session_code": "LONDON",
|
|
"status": "open",
|
|
"start_ts": self.now - timedelta(minutes=20),
|
|
"end_ts": self.now + timedelta(hours=2),
|
|
}
|
|
|
|
def insert_market_truth_incident(self, **kwargs) -> int:
|
|
self.incidents.append(dict(kwargs))
|
|
return len(self.incidents)
|
|
|
|
def upsert_news_provider_snapshot(self, **kwargs) -> int:
|
|
self.cached_news_snapshot = dict(kwargs)
|
|
return 1
|
|
|
|
def fetch_latest_news_provider_snapshot(self):
|
|
return dict(self.cached_news_snapshot) if self.cached_news_snapshot else None
|
|
|
|
|
|
class _FallbackCandleRepository:
|
|
def __init__(self, repository: _SupervisorRepository) -> None:
|
|
self.repository = repository
|
|
self.write_calls = 0
|
|
|
|
def upsert_many(self, candles, instrument_id: int) -> int:
|
|
self.write_calls += 1
|
|
latest = candles[-1]
|
|
self.repository.latest_candle = {
|
|
"id": 99,
|
|
"instrument_id": instrument_id,
|
|
"timeframe": latest.timeframe,
|
|
"ts_close": datetime.fromtimestamp((latest.ts_open_ms + 60_000 - 1) / 1000, tz=timezone.utc),
|
|
"open": latest.open,
|
|
"close": latest.close,
|
|
}
|
|
return len(candles)
|
|
|
|
|
|
class _DisconnectedStreamService:
|
|
def __init__(self, now: datetime) -> None:
|
|
self.now = now
|
|
self.subscriptions = []
|
|
|
|
def subscribe(self, **kwargs) -> None:
|
|
self.subscriptions.append(dict(kwargs))
|
|
|
|
def start(self) -> None:
|
|
return
|
|
|
|
def stop(self) -> None:
|
|
return
|
|
|
|
def iter_subscriptions(self):
|
|
return []
|
|
|
|
def get_snapshot(self, *, instrument_id: int, timeframe: str, now=None) -> dict:
|
|
return {
|
|
"instrument_id": instrument_id,
|
|
"timeframe": timeframe,
|
|
"code": STREAM_CODE_DISCONNECTED,
|
|
"connected": False,
|
|
"subscribed": False,
|
|
"reconnect_count": 1,
|
|
"last_message_ts": None,
|
|
"heartbeat_age_seconds": None,
|
|
"lag_seconds": None,
|
|
"latest_price": None,
|
|
"latest_candle_open": None,
|
|
"latest_candle_close": None,
|
|
"latest_candle_open_ts": None,
|
|
"latest_candle_close_ts": None,
|
|
"last_confirmed_candle_close_ts": None,
|
|
"last_disconnect_reason": "socket_closed",
|
|
"last_error": "socket_closed",
|
|
"gap_detected": False,
|
|
"gap_seconds": None,
|
|
}
|
|
|
|
|
|
class _FallbackOKXClient:
|
|
def __init__(self, now: datetime) -> None:
|
|
self.now = now
|
|
|
|
def get_candles(self, inst_id: str, bar: str, limit: int):
|
|
return [
|
|
OKXCandle(
|
|
inst_id=inst_id,
|
|
bar=bar,
|
|
ts=str(int((self.now - timedelta(minutes=1)).timestamp() * 1000)),
|
|
open="65000.0",
|
|
high="65120.0",
|
|
low="64980.0",
|
|
close="65111.0",
|
|
volume="5.2",
|
|
confirm="1",
|
|
)
|
|
]
|
|
|
|
|
|
class _FailingOKXClient:
|
|
def get_candles(self, inst_id: str, bar: str, limit: int):
|
|
raise TimeoutError("rest_timeout")
|
|
|
|
|
|
class _InactiveNewsProvider:
|
|
def __init__(self, now: datetime) -> None:
|
|
self.now = now
|
|
|
|
def fetch_snapshot(self, *, now: datetime):
|
|
return NewsProviderSnapshot(
|
|
provider_name="test",
|
|
fetched_ts=now,
|
|
freshness_code="news_provider_fresh",
|
|
active_events=[],
|
|
event_count=0,
|
|
is_active=False,
|
|
code="news_window_inactive",
|
|
error=None,
|
|
raw_payload={},
|
|
)
|
|
|
|
def evaluate_cached_snapshot(self, *, snapshot, now: datetime):
|
|
return snapshot
|
|
|
|
|
|
class _FailingNewsProvider:
|
|
provider_name = "failing_news_provider"
|
|
|
|
def fetch_snapshot(self, *, now: datetime):
|
|
return NewsProviderSnapshot(
|
|
provider_name=self.provider_name,
|
|
fetched_ts=None,
|
|
freshness_code=NEWS_CODE_PROVIDER_ERROR,
|
|
active_events=[],
|
|
event_count=0,
|
|
is_active=False,
|
|
code=NEWS_CODE_PROVIDER_ERROR,
|
|
error="rate_limited",
|
|
raw_payload={},
|
|
)
|
|
|
|
|
|
class _WorkingNewsProvider:
|
|
provider_name = "working_news_provider"
|
|
|
|
def __init__(self, now: datetime) -> None:
|
|
self.now = now
|
|
|
|
def fetch_snapshot(self, *, now: datetime):
|
|
return NewsProviderSnapshot(
|
|
provider_name=self.provider_name,
|
|
fetched_ts=now,
|
|
freshness_code=NEWS_CODE_PROVIDER_FRESH,
|
|
active_events=[],
|
|
event_count=1,
|
|
is_active=False,
|
|
code="news_window_inactive",
|
|
error=None,
|
|
raw_payload={"rows": [{"Event": "US CPI"}]},
|
|
)
|
|
|
|
|
|
class _FailingChainProvider:
|
|
def fetch_snapshot(self, *, now: datetime):
|
|
return NewsProviderSnapshot(
|
|
provider_name="failing_chain",
|
|
fetched_ts=None,
|
|
freshness_code=NEWS_CODE_PROVIDER_ERROR,
|
|
active_events=[],
|
|
event_count=0,
|
|
is_active=False,
|
|
code=NEWS_CODE_PROVIDER_ERROR,
|
|
error="provider_unavailable",
|
|
raw_payload={},
|
|
)
|
|
|
|
def evaluate_cached_snapshot(self, *, snapshot, now: datetime):
|
|
if snapshot is None:
|
|
return NewsProviderSnapshot(
|
|
provider_name="failing_chain",
|
|
fetched_ts=None,
|
|
freshness_code="news_provider_stale",
|
|
active_events=[],
|
|
event_count=0,
|
|
is_active=False,
|
|
code="news_provider_stale",
|
|
error="snapshot_missing",
|
|
raw_payload={},
|
|
)
|
|
return NewsProviderSnapshot(
|
|
provider_name=snapshot.provider_name,
|
|
fetched_ts=snapshot.fetched_ts,
|
|
freshness_code=NEWS_CODE_PROVIDER_CACHED,
|
|
active_events=list(snapshot.active_events),
|
|
event_count=snapshot.event_count,
|
|
is_active=snapshot.is_active,
|
|
code=snapshot.code,
|
|
error=snapshot.error,
|
|
raw_payload=dict(snapshot.raw_payload or {}),
|
|
)
|
|
|
|
|
|
class _StaticSupervisor:
|
|
def __init__(self, now: datetime) -> None:
|
|
self.now = now
|
|
|
|
def current_state(self, *, instrument_id: int, timeframe: str, now: datetime):
|
|
return MarketTruthRuntimeState(
|
|
stream_state={
|
|
"code": STREAM_CODE_CONNECTED,
|
|
"connected": True,
|
|
"subscribed": True,
|
|
"reconnect_count": 0,
|
|
"last_message_ts": now.isoformat(),
|
|
"heartbeat_age_seconds": 1.0,
|
|
"lag_seconds": 1.0,
|
|
"latest_price": 65111.0,
|
|
"latest_candle_open": 65000.0,
|
|
"latest_candle_close": 65111.0,
|
|
"latest_candle_open_ts": (now - timedelta(minutes=1)).isoformat(),
|
|
"latest_candle_close_ts": now.isoformat(),
|
|
"last_confirmed_candle_close_ts": now.isoformat(),
|
|
"last_disconnect_reason": None,
|
|
"last_error": None,
|
|
},
|
|
data_gap_state={
|
|
"is_gap": False,
|
|
"code": "market_data_gap_clear",
|
|
"expected_interval_seconds": 60,
|
|
"gap_seconds": 0.0,
|
|
"last_candle_ts": now.isoformat(),
|
|
"source": "stream",
|
|
},
|
|
news_snapshot=NewsProviderSnapshot(
|
|
provider_name="test",
|
|
fetched_ts=now,
|
|
freshness_code="news_provider_fresh",
|
|
active_events=[],
|
|
event_count=0,
|
|
is_active=False,
|
|
code="news_window_inactive",
|
|
error=None,
|
|
raw_payload={},
|
|
),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|