315 lines
8.3 KiB
Python
315 lines
8.3 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 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() == ()
|