from __future__ import annotations import asyncio from collections.abc import Awaitable, Callable from typing import Any, Protocol from src.bootstrap.trade_stream_runtime import ( build_trade_stream_production_runtime, ) from src.core.config import Settings, TradeStreamSettings from src.market_data.acquisition.consistency.trade_stream_state_store import ( TradeStreamStateStore, ) from src.market_data.acquisition.runtime.trade_stream_production_runtime import ( TradeStreamProductionRuntime, TradeStreamProductionRuntimeState, ) from tests.support.async_wait import wait_until SYMBOL = "BTC/USD_LEVERAGE" RUNTIME_CLEANUP_TIMEOUT_SECONDS = 5.0 OWNED_TASK_NAMES = frozenset( { "trade-stream-receive", "trade-stream-runtime", "trade-stream-runtime-recovery", "trade-stream-scheduler", "trade-stream-startup", } ) class ManagedRuntimeProtocol(Protocol): @property def state(self) -> TradeStreamProductionRuntimeState: ... async def run(self) -> None: ... async def stop(self) -> None: ... def run_scenario( scenario: Awaitable[None], *, timeout_seconds: float = 30.0, ) -> None: asyncio.run( asyncio.wait_for( scenario, timeout=timeout_seconds, ) ) def make_settings( *, websocket_url: str, rest_base_url: str, open_timeout_seconds: float = 0.5, probe_timeout_seconds: float = 0.2, close_timeout_seconds: float = 0.2, heartbeat_timeout_seconds: float = 60.0, scheduler_interval_seconds: float = 60.0, ) -> Settings: return Settings( bot_token="integration-test-token", bot_parse_mode="HTML", app_env="test", log_level="INFO", tz="UTC", exchange_enabled=True, exchange_name="dzengi", exchange_base_url=rest_base_url, exchange_ws_url="", exchange_api_key="", exchange_api_secret="", exchange_timeout_sec=2, exchange_testnet=True, default_symbol=SYMBOL, trade_stream=TradeStreamSettings( enabled=True, websocket_url=websocket_url, symbols=(SYMBOL,), open_timeout_seconds=open_timeout_seconds, probe_timeout_seconds=probe_timeout_seconds, close_timeout_seconds=close_timeout_seconds, heartbeat_timeout_seconds=heartbeat_timeout_seconds, scheduler_interval_seconds=scheduler_interval_seconds, recovery_window_ms=3_599_999, ), db_host="localhost", db_port=5432, db_name="integration", db_user="integration", db_password="", debug_enabled=False, journal_debug_enabled=False, ) def build_runtime( *, websocket_url: str, rest_base_url: str, probe_timeout_seconds: float = 0.2, close_timeout_seconds: float = 0.2, heartbeat_timeout_seconds: float = 60.0, scheduler_interval_seconds: float = 60.0, ) -> TradeStreamProductionRuntime: runtime = build_trade_stream_production_runtime( make_settings( websocket_url=websocket_url, rest_base_url=rest_base_url, probe_timeout_seconds=probe_timeout_seconds, close_timeout_seconds=close_timeout_seconds, heartbeat_timeout_seconds=heartbeat_timeout_seconds, scheduler_interval_seconds=scheduler_interval_seconds, ) ) assert runtime is not None return runtime def state_store_from( runtime: TradeStreamProductionRuntime, ) -> TradeStreamStateStore: runtime_graph: Any = runtime return ( runtime_graph ._reconnect_recovery_coordinator ._recovery_coordinator ._state_store ) def reconnect_coordinator_from( runtime: TradeStreamProductionRuntime, ) -> Any: runtime_graph: Any = runtime return runtime_graph._reconnect_recovery_coordinator def active_owned_task_names() -> tuple[str, ...]: current_task = asyncio.current_task() return tuple( sorted( task.get_name() for task in asyncio.all_tasks() if task is not current_task and not task.done() and task.get_name() in OWNED_TASK_NAMES ) ) async def start_runtime( runtime: ManagedRuntimeProtocol, *, timeout_seconds: float = 3.0, ) -> asyncio.Task[None]: task = asyncio.create_task( runtime.run(), name="trade-stream-runtime", ) try: await wait_until( lambda: runtime.state in { TradeStreamProductionRuntimeState.RUNNING, TradeStreamProductionRuntimeState.FAILED, }, timeout_seconds=timeout_seconds, ) if ( runtime.state is TradeStreamProductionRuntimeState.FAILED ): await task except BaseException as error: try: await stop_runtime(runtime, task) except BaseException as cleanup_error: error.add_note( "Runtime startup cleanup also failed: " f"{type(cleanup_error).__name__}." ) raise return task async def wait_until_or_runtime_exit( predicate: Callable[[], bool], *, runtime_task: asyncio.Task[None], timeout_seconds: float, ) -> None: """ Дождаться условия, немедленно распространяя завершение Runtime. Runtime task не отменяется: helper владеет только внутренней задачей ожидания условия. Если Runtime завершился без ошибки до выполнения условия, это считается ошибкой проверочного сценария. """ condition_task = asyncio.create_task( wait_until( predicate, timeout_seconds=timeout_seconds, ) ) try: await asyncio.wait( ( condition_task, runtime_task, ), return_when=asyncio.FIRST_COMPLETED, ) if runtime_task.done(): condition_task.cancel() try: await condition_task except asyncio.CancelledError: pass await runtime_task raise RuntimeError( "Trade Stream Runtime exited before the expected " "live condition was reached." ) await condition_task finally: if not condition_task.done(): condition_task.cancel() try: await condition_task except asyncio.CancelledError: pass async def stop_runtime( runtime: ManagedRuntimeProtocol, task: asyncio.Task[None], *, timeout_seconds: float = RUNTIME_CLEANUP_TIMEOUT_SECONDS, ) -> None: primary_error: BaseException | None = None if not task.done(): try: await asyncio.wait_for( runtime.stop(), timeout=timeout_seconds, ) except BaseException as error: primary_error = error if not task.done() and primary_error is not None: task.cancel() try: await asyncio.wait_for( asyncio.shield(task), timeout=timeout_seconds, ) except BaseException as error: if isinstance(error, TimeoutError): task.cancel() try: await asyncio.wait_for( task, timeout=timeout_seconds, ) except asyncio.CancelledError: pass except BaseException as cancellation_error: error.add_note( "Runtime task cancellation also failed: " f"{type(cancellation_error).__name__}." ) if primary_error is None: primary_error = error else: primary_error.add_note( "Runtime task cleanup also failed: " f"{type(error).__name__}." ) if primary_error is not None: raise primary_error async def assert_no_owned_tasks() -> None: await asyncio.sleep(0) assert active_owned_task_names() == ()