Files
dzentra_bot/app/tests/support/trade_stream_runtime.py

340 lines
9.1 KiB
Python

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 (
MarketDataStorageSettings,
Settings,
TradeStreamSettings,
)
from src.market_data.acquisition.consistency.trade_stream_state_store import (
TradeStreamStateStore,
)
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
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(
{
"application-shutdown",
"market-data-storage-shutdown",
"market-data-storage-startup",
"telegram-polling",
"trade-stream-market-processing",
"trade-stream-receive",
"trade-stream-runtime",
"trade-stream-runtime-recovery",
"trade-stream-scheduler",
"trade-stream-startup-recovery",
"trade-stream-state-hydration",
"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="",
market_data_storage=MarketDataStorageSettings(
enabled=False,
pool_min_size=1,
pool_max_size=4,
pool_timeout_seconds=10.0,
),
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,
trade_observation_sink: TradeObservationSinkProtocol | None = None,
) -> 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,
),
trade_observation_sink=trade_observation_sink,
)
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
or task.get_name().startswith("persistent-application-")
)
)
)
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() == ()