gitadmin/dzentra_bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Build 060.27: implement Persistent Market Data Storage

Browse Source
This commit is contained in:
Sergey
2026-08-01 03:22:25 +03:00
parent cb8acfe5fe
commit 58e5a12a4d
54 changed files with 10243 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
app/tests/unit/bootstrap/test_app_factory.py
View File

@@ -46,6 +46,7 @@ def make_settings() -> SimpleNamespace:
exchange_name="dzengi",
default_symbol="BTC/USD_LEVERAGE",
trade_stream=SimpleNamespace(enabled=True),
market_data_storage=SimpleNamespace(enabled=True),
)
@@ -56,6 +57,12 @@ def test_create_app_builds_one_application_composition(
bot = object()
dispatcher = object()
runtime = object()
storage_lifecycle = object()
storage_sink = object()
storage = SimpleNamespace(
lifecycle=storage_lifecycle,
trade_observation_sink=storage_sink,
)
journal = RecordingJournal()
observed_runtime_settings: list[object] = []
registered_bots: list[object] = []
@@ -82,8 +89,23 @@ def test_create_app_builds_one_application_composition(
lambda: journal,
)
def build_runtime(received_settings: object) -> object:
monkeypatch.setattr(
app_factory,
"build_market_data_storage",
lambda received_settings: (
storage
if received_settings is settings
else None
),
)
def build_runtime(
received_settings: object,
*,
trade_observation_sink: object,
) -> object:
observed_runtime_settings.append(received_settings)
assert trade_observation_sink is storage_sink
return runtime
monkeypatch.setattr(
@@ -118,10 +140,18 @@ def test_create_app_builds_one_application_composition(
assert application.bot is bot
assert application.dispatcher is dispatcher
assert application.trade_stream_runtime is runtime
assert (
application.market_data_storage_lifecycle
is storage_lifecycle
)
assert observed_runtime_settings == [settings]
assert registered_bots == [bot]
assert routed_dispatchers == [dispatcher]
assert journal.info_calls[0][2]["trade_stream_enabled"] is True
assert (
journal.info_calls[0][2]["market_data_storage_enabled"]
is True
)
def test_runtime_build_error_is_fatal(
@@ -149,9 +179,18 @@ def test_runtime_build_error_is_fatal(
"JournalService",
RecordingJournal,
)
monkeypatch.setattr(
app_factory,
"build_market_data_storage",
lambda settings: None,
)
def fail_runtime_build(settings: object) -> None:
del settings
def fail_runtime_build(
settings: object,
*,
trade_observation_sink: object,
) -> None:
del settings, trade_observation_sink
raise expected
monkeypatch.setattr(

281
app/tests/unit/bootstrap/test_application.py
View File

@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import threading
import pytest
@@ -46,12 +47,14 @@ class FakeDispatcher:
*,
return_immediately: bool = False,
error: BaseException | None = None,
calls: list[str] | None = None,
) -> None:
self.started = asyncio.Event()
self.release = asyncio.Event()
self.cancelled = asyncio.Event()
self.return_immediately = return_immediately
self.error = error
self.calls = calls
self.close_bot_session_values: list[bool] = []
async def start_polling(
@@ -61,6 +64,8 @@ class FakeDispatcher:
close_bot_session: bool,
) -> None:
del bot
if self.calls is not None:
self.calls.append("polling.start")
self.close_bot_session_values.append(close_bot_session)
self.started.set()
@@ -82,6 +87,7 @@ class FakeRuntime:
return_immediately: bool = False,
error: BaseException | None = None,
stop_error: BaseException | None = None,
calls: list[str] | None = None,
) -> None:
self.started = asyncio.Event()
self.release = asyncio.Event()
@@ -89,6 +95,7 @@ class FakeRuntime:
self.return_immediately = return_immediately
self.error = error
self.stop_error = stop_error
self.calls = calls
self.stop_calls = 0
@property
@@ -104,6 +111,8 @@ class FakeRuntime:
return self.state is TradeStreamProductionRuntimeState.RUNNING
async def run(self) -> None:
if self.calls is not None:
self.calls.append("runtime.run")
self.started.set()
if not self.return_immediately:
@@ -114,6 +123,8 @@ class FakeRuntime:
async def stop(self) -> None:
self.stop_calls += 1
if self.calls is not None:
self.calls.append("runtime.stop")
self.release.set()
self.stopped.set()
@@ -136,19 +147,289 @@ class BlockingStopRuntime(FakeRuntime):
self.stopped.set()
class FakeStorageLifecycle:
def __init__(
self,
*,
calls: list[str] | None = None,
start_error: BaseException | None = None,
stop_error: BaseException | None = None,
start_release: threading.Event | None = None,
stop_release: threading.Event | None = None,
) -> None:
self.calls = calls
self.start_error = start_error
self.stop_error = stop_error
self.start_release = start_release
self.stop_release = stop_release
self.start_entered = threading.Event()
self.stop_entered = threading.Event()
self.start_calls = 0
self.stop_calls = 0
self.started = False
def start(self) -> None:
self.start_calls += 1
if self.calls is not None:
self.calls.append("storage.start")
self.start_entered.set()
if self.start_release is not None:
if not self.start_release.wait(timeout=2.0):
raise TimeoutError("storage startup was not released")
if self.start_error is not None:
raise self.start_error
self.started = True
def stop(self) -> None:
self.stop_calls += 1
if self.calls is not None:
self.calls.append("storage.stop")
self.started = False
self.stop_entered.set()
if self.stop_release is not None:
if not self.stop_release.wait(timeout=2.0):
raise TimeoutError("storage shutdown was not released")
if self.stop_error is not None:
raise self.stop_error
def make_application(
*,
dispatcher: FakeDispatcher,
runtime: FakeRuntime | None,
bot: FakeBot | None = None,
storage_lifecycle: FakeStorageLifecycle | None = None,
) -> ApplicationComposition:
return ApplicationComposition(
bot=bot or FakeBot(), # type: ignore[arg-type]
dispatcher=dispatcher, # type: ignore[arg-type]
trade_stream_runtime=runtime,
market_data_storage_lifecycle=storage_lifecycle,
)
def test_storage_lifecycle_wraps_root_runtime_tasks() -> None:
async def scenario() -> None:
calls: list[str] = []
storage = FakeStorageLifecycle(calls=calls)
dispatcher = FakeDispatcher(
return_immediately=True,
calls=calls,
)
runtime = FakeRuntime(calls=calls)
await run_application(
make_application(
dispatcher=dispatcher,
runtime=runtime,
storage_lifecycle=storage,
)
)
assert storage.start_calls == 1
assert storage.stop_calls == 1
assert calls.index("storage.start") < calls.index(
"polling.start"
)
assert calls.index("storage.start") < calls.index(
"runtime.run"
)
assert calls.index("runtime.stop") < calls.index(
"storage.stop"
)
asyncio.run(scenario())
def test_storage_startup_failure_prevents_root_task_start() -> None:
async def scenario() -> None:
expected = RuntimeError("storage startup failed")
storage = FakeStorageLifecycle(start_error=expected)
dispatcher = FakeDispatcher()
runtime = FakeRuntime()
bot = FakeBot()
with pytest.raises(RuntimeError) as error_info:
await run_application(
make_application(
dispatcher=dispatcher,
runtime=runtime,
bot=bot,
storage_lifecycle=storage,
)
)
assert error_info.value is expected
assert dispatcher.started.is_set() is False
assert runtime.started.is_set() is False
assert runtime.stop_calls == 1
assert storage.stop_calls == 1
assert bot.session.close_calls == 1
asyncio.run(scenario())
def test_cancellation_waits_for_storage_startup_before_cleanup() -> None:
async def scenario() -> None:
start_release = threading.Event()
storage = FakeStorageLifecycle(
start_release=start_release,
)
dispatcher = FakeDispatcher()
runtime = FakeRuntime()
bot = FakeBot()
task = asyncio.create_task(
run_application(
make_application(
dispatcher=dispatcher,
runtime=runtime,
bot=bot,
storage_lifecycle=storage,
)
)
)
entered = await asyncio.to_thread(
storage.start_entered.wait,
1.0,
)
assert entered is True
task.cancel()
try:
await asyncio.sleep(0)
await asyncio.sleep(0)
assert task.done() is False
finally:
start_release.set()
with pytest.raises(asyncio.CancelledError):
await task
assert dispatcher.started.is_set() is False
assert runtime.started.is_set() is False
assert storage.stop_calls == 1
assert bot.session.close_calls == 1
await asyncio.sleep(0)
assert not {
child.get_name()
for child in asyncio.all_tasks()
if child is not asyncio.current_task()
and not child.done()
and child.get_name()
in {
"application-shutdown",
"market-data-storage-shutdown",
"market-data-storage-startup",
"telegram-polling",
"trade-stream-runtime",
}
}
asyncio.run(scenario())
def test_storage_shutdown_error_is_reported_and_bot_still_closes() -> None:
async def scenario() -> None:
expected = RuntimeError("storage close failed")
storage = FakeStorageLifecycle(stop_error=expected)
dispatcher = FakeDispatcher(return_immediately=True)
bot = FakeBot()
with pytest.raises(RuntimeError) as error_info:
await run_application(
make_application(
dispatcher=dispatcher,
runtime=None,
bot=bot,
storage_lifecycle=storage,
)
)
assert error_info.value is expected
assert storage.stop_calls == 1
assert bot.session.close_calls == 1
asyncio.run(scenario())
def test_storage_shutdown_error_does_not_replace_runtime_error() -> None:
async def scenario() -> None:
runtime_error = RuntimeError("runtime failed")
storage_error = RuntimeError("storage close failed")
storage = FakeStorageLifecycle(stop_error=storage_error)
dispatcher = FakeDispatcher()
runtime = FakeRuntime(
return_immediately=True,
error=runtime_error,
)
with pytest.raises(RuntimeError) as error_info:
await run_application(
make_application(
dispatcher=dispatcher,
runtime=runtime,
storage_lifecycle=storage,
)
)
assert error_info.value is runtime_error
assert any(
"cleanup also failed" in note
for note in getattr(runtime_error, "__notes__", ())
)
asyncio.run(scenario())
def test_repeated_cancellation_does_not_interrupt_storage_shutdown(
) -> None:
async def scenario() -> None:
stop_release = threading.Event()
storage = FakeStorageLifecycle(
stop_release=stop_release,
)
dispatcher = FakeDispatcher(return_immediately=True)
bot = FakeBot()
task = asyncio.create_task(
run_application(
make_application(
dispatcher=dispatcher,
runtime=None,
bot=bot,
storage_lifecycle=storage,
)
)
)
entered = await asyncio.to_thread(
storage.stop_entered.wait,
1.0,
)
assert entered is True
task.cancel()
task.cancel()
try:
await asyncio.sleep(0)
assert task.done() is False
finally:
stop_release.set()
with pytest.raises(asyncio.CancelledError):
await task
assert storage.stop_calls == 1
assert bot.session.close_calls == 1
asyncio.run(scenario())
def test_disabled_runtime_runs_only_polling_and_closes_bot() -> None:
async def scenario() -> None:
dispatcher = FakeDispatcher(return_immediately=True)

232
app/tests/unit/bootstrap/test_market_data_storage_bootstrap.py Normal file
View File

@@ -0,0 +1,232 @@
from __future__ import annotations
from typing import Any, cast
import pytest
from psycopg.conninfo import conninfo_to_dict
from src.bootstrap.market_data_storage import (
MarketDataStorageLifecycle,
MarketDataStorageLifecycleProtocol,
build_market_data_storage,
)
from src.core.config import (
MarketDataStorageSettings,
Settings,
TradeStreamSettings,
)
def make_settings(
*,
storage_enabled: bool = True,
trade_stream_enabled: bool = True,
) -> Settings:
return Settings(
bot_token="test-token",
bot_parse_mode="HTML",
app_env="test",
log_level="INFO",
tz="UTC",
exchange_enabled=True,
exchange_name=" dzengi ",
exchange_base_url="https://rest.example.test",
exchange_ws_url="",
exchange_api_key="",
exchange_api_secret="",
exchange_timeout_sec=10,
exchange_testnet=True,
default_symbol="BTC/USD_LEVERAGE",
trade_stream=TradeStreamSettings(
enabled=trade_stream_enabled,
websocket_url="wss://stream.example.test",
symbols=("BTC/USD_LEVERAGE",),
open_timeout_seconds=10.0,
probe_timeout_seconds=20.0,
close_timeout_seconds=10.0,
heartbeat_timeout_seconds=30.0,
scheduler_interval_seconds=5.0,
recovery_window_ms=3_599_999,
),
db_host="db.example.test",
db_port=5544,
db_name="dzentra",
db_user="market-data",
db_password="p@ss word",
market_data_storage=MarketDataStorageSettings(
enabled=storage_enabled,
pool_min_size=2,
pool_max_size=6,
pool_timeout_seconds=7.5,
),
debug_enabled=False,
journal_debug_enabled=False,
)
class RecordingPool:
def __init__(
self,
*,
close_error: BaseException | None = None,
) -> None:
self.close_error = close_error
self.events: list[str] = []
def open(self) -> None:
self.events.append("pool.open")
def close(self) -> None:
self.events.append("pool.close")
if self.close_error is not None:
raise self.close_error
class RecordingMigrationRunner:
def __init__(
self,
*,
pool: RecordingPool,
error: BaseException | None = None,
) -> None:
self._pool = pool
self._error = error
self.calls = 0
def run(self) -> tuple[int, ...]:
self.calls += 1
self._pool.events.append("migrations.run")
if self._error is not None:
raise self._error
return (1,)
def make_lifecycle(
*,
migration_error: BaseException | None = None,
close_error: BaseException | None = None,
) -> tuple[
MarketDataStorageLifecycle,
RecordingPool,
RecordingMigrationRunner,
]:
pool = RecordingPool(close_error=close_error)
runner = RecordingMigrationRunner(
pool=pool,
error=migration_error,
)
lifecycle = MarketDataStorageLifecycle(
connection_pool=pool,
migration_runner=runner,
)
return lifecycle, pool, runner
def test_disabled_storage_builds_nothing() -> None:
assert build_market_data_storage(
make_settings(storage_enabled=False),
) is None
def test_enabled_storage_requires_enabled_trade_stream() -> None:
with pytest.raises(RuntimeError, match="Trade Stream"):
build_market_data_storage(
make_settings(trade_stream_enabled=False),
)
def test_builds_shared_graph_without_opening_pool() -> None:
composition = build_market_data_storage(make_settings())
assert composition is not None
assert composition.connection_pool.is_open is False
assert composition.lifecycle.started is False
migration_provider = cast(
Any,
composition.migration_runner._connection_provider,
)
repository_provider = cast(
Any,
composition.trade_repository._connection_provider,
)
assert migration_provider.__self__ is composition.connection_pool
assert repository_provider.__self__ is composition.connection_pool
assert (
composition.trade_observation_sink._trade_storage
is composition.trade_repository
)
assert composition.trade_observation_sink._venue == "dzengi"
conninfo = conninfo_to_dict(composition.connection_pool._conninfo)
assert conninfo == {
"host": "db.example.test",
"port": "5544",
"dbname": "dzentra",
"user": "market-data",
"password": "p@ss word",
}
def test_lifecycle_opens_pool_before_migrations_and_is_idempotent() -> None:
lifecycle, pool, runner = make_lifecycle()
lifecycle.start()
lifecycle.start()
assert isinstance(lifecycle, MarketDataStorageLifecycleProtocol)
assert lifecycle.started is True
assert pool.events == [
"pool.open",
"migrations.run",
]
assert runner.calls == 1
lifecycle.stop()
lifecycle.stop()
assert lifecycle.started is False
assert pool.events == [
"pool.open",
"migrations.run",
"pool.close",
"pool.close",
]
def test_migration_failure_closes_pool_and_preserves_error() -> None:
migration_error = RuntimeError("migration failed")
lifecycle, pool, _ = make_lifecycle(
migration_error=migration_error,
)
with pytest.raises(RuntimeError) as error_info:
lifecycle.start()
assert error_info.value is migration_error
assert lifecycle.started is False
assert pool.events == [
"pool.open",
"migrations.run",
"pool.close",
]
def test_migration_error_keeps_cleanup_failure_as_note() -> None:
migration_error = RuntimeError("migration failed")
lifecycle, _, _ = make_lifecycle(
migration_error=migration_error,
close_error=RuntimeError("close failed"),
)
with pytest.raises(RuntimeError) as error_info:
lifecycle.start()
assert error_info.value is migration_error
assert migration_error.__notes__ == [
"Market Data Storage startup cleanup also failed: RuntimeError."
]

38
app/tests/unit/bootstrap/test_trade_stream_production_integration.py
View File

@@ -5,9 +5,10 @@ import json
import threading
import time
from collections.abc import Awaitable, Callable
from typing import Any
from typing import Any, cast
import pytest
from aiogram import Bot, Dispatcher
from websockets.protocol import State
import src.bootstrap.trade_stream_runtime as production_factory
@@ -15,7 +16,11 @@ from src.bootstrap.application import (
ApplicationComposition,
run_application,
)
from src.core.config import Settings, TradeStreamSettings
from src.core.config import (
MarketDataStorageSettings,
Settings,
TradeStreamSettings,
)
from src.market_data.acquisition.adapters.dzengi.websocket_transport import (
DzengiWebSocketTransport,
)
@@ -37,6 +42,8 @@ SYMBOL = "BTC/USD_LEVERAGE"
_OWNED_TASK_NAMES = frozenset(
{
"application-shutdown",
"market-data-storage-shutdown",
"market-data-storage-startup",
"telegram-polling",
"trade-stream-receive",
"trade-stream-runtime",
@@ -295,6 +302,12 @@ def make_settings(
db_name="test",
db_user="test",
db_password="test",
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,
)
@@ -367,8 +380,8 @@ def make_application(
bot: FakeBot,
) -> ApplicationComposition:
return ApplicationComposition(
bot=bot, # type: ignore[arg-type]
dispatcher=dispatcher, # type: ignore[arg-type]
bot=cast(Bot, cast(object, bot)),
dispatcher=cast(Dispatcher, cast(object, dispatcher)),
trade_stream_runtime=runtime,
)
@@ -400,14 +413,23 @@ async def assert_no_owned_tasks() -> None:
def state_store_from(
runtime: TradeStreamProductionRuntime,
) -> TradeStreamStateStore:
runtime_graph: Any = runtime
return (
runtime
runtime_graph
._reconnect_recovery_coordinator
._recovery_coordinator
._state_store
)
def subscription_keys_from(
runtime: TradeStreamProductionRuntime,
) -> tuple[str, ...]:
runtime_graph: Any = runtime
return runtime_graph._subscription_manager.subscription_keys
def test_disabled_feature_runs_only_telegram_without_runtime_graph(
monkeypatch: pytest.MonkeyPatch,
) -> None:
@@ -529,7 +551,7 @@ def test_production_factory_processes_ack_trade_and_shutdown(
assert runtime.state is (
TradeStreamProductionRuntimeState.STOPPED
)
assert runtime._subscription_manager.subscription_keys == ()
assert subscription_keys_from(runtime) == ()
assert bot.session.close_calls == 1
assert len(transports) == 1
await assert_no_owned_tasks()
@@ -579,7 +601,7 @@ def test_connection_startup_failure_is_fatal_and_leaves_no_tasks(
TradeStreamProductionRuntimeState.FAILED
)
assert runtime._session.is_connected is False
assert runtime._subscription_manager.subscription_keys == ()
assert subscription_keys_from(runtime) == ()
assert bot.session.close_calls == 1
await assert_no_owned_tasks()
@@ -633,7 +655,7 @@ def test_subscription_startup_failure_rolls_back_concrete_graph(
TradeStreamProductionRuntimeState.FAILED
)
assert runtime._session.is_connected is False
assert runtime._subscription_manager.subscription_keys == ()
assert subscription_keys_from(runtime) == ()
assert bot.session.close_calls == 1
await assert_no_owned_tasks()

83
app/tests/unit/bootstrap/test_trade_stream_runtime.py
View File

@@ -1,13 +1,28 @@
from __future__ import annotations
from typing import Any
from src.bootstrap.trade_stream_runtime import (
build_trade_stream_production_runtime,
)
from src.core.config import Settings, TradeStreamSettings
from src.core.config import (
MarketDataStorageSettings,
Settings,
TradeStreamSettings,
)
from src.market_data.acquisition.runtime.trade_stream_production_runtime import (
TradeStreamProductionRuntime,
TradeStreamProductionRuntimeState,
)
from src.market_data.acquisition.models.trade import Trade
class RecordingTradeObservationSink:
def __init__(self) -> None:
self.observations: list[Trade] = []
def persist(self, trade: Trade) -> None:
self.observations.append(trade)
def make_settings(
@@ -49,6 +64,12 @@ def make_settings(
db_name="test",
db_user="test",
db_password="test",
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,
)
@@ -80,33 +101,60 @@ def test_uses_one_shared_stateful_dependency_graph() -> None:
assert isinstance(runtime, TradeStreamProductionRuntime)
transport = runtime._transport
service = runtime._trade_stream_service
reconnect_recovery = runtime._reconnect_recovery_coordinator
runtime_graph: Any = runtime
transport = runtime_graph._transport
service = runtime_graph._trade_stream_service
reconnect_recovery = runtime_graph._reconnect_recovery_coordinator
recovery = reconnect_recovery._recovery_coordinator
assert runtime._session._transport is transport
assert runtime._subscription_manager._transport is transport
assert runtime._runtime_scheduler.liveness_probe is transport
assert runtime_graph._session._transport is transport
assert runtime_graph._subscription_manager._transport is transport
assert runtime_graph._runtime_scheduler.liveness_probe is transport
assert (
service._consistency_controller
is recovery._recovery_controller._consistency_controller
)
assert runtime._live_processing_gate is (
assert runtime_graph._live_processing_gate is (
reconnect_recovery.live_processing_gate
)
assert runtime._runtime_scheduler.runtime_supervisor is (
runtime._runtime_supervisor
assert runtime_graph._runtime_scheduler.runtime_supervisor is (
runtime_graph._runtime_supervisor
)
def test_optional_storage_sink_is_shared_by_live_and_recovery() -> None:
sink = RecordingTradeObservationSink()
runtime = build_trade_stream_production_runtime(
make_settings(),
trade_observation_sink=sink,
)
assert isinstance(runtime, TradeStreamProductionRuntime)
runtime_graph: Any = runtime
live_controller = (
runtime_graph._trade_stream_service._consistency_controller
)
recovery_controller = (
runtime_graph
._reconnect_recovery_coordinator
._recovery_coordinator
._recovery_controller
._consistency_controller
)
assert live_controller is recovery_controller
assert live_controller._trade_observation_sink is sink
def test_applies_explicit_transport_and_runtime_settings() -> None:
settings = make_settings()
runtime = build_trade_stream_production_runtime(settings)
assert isinstance(runtime, TradeStreamProductionRuntime)
transport = runtime._transport
runtime_graph: Any = runtime
transport = runtime_graph._transport
assert transport._url == "wss://stream.example.test/root/connect"
assert transport._headers == {
@@ -119,17 +167,17 @@ def test_applies_explicit_transport_and_runtime_settings() -> None:
assert transport._close_timeout == 9.0
assert transport._ping_interval is None
assert transport._ping_timeout is None
assert runtime._symbols == (
assert runtime_graph._symbols == (
"BTC/USD_LEVERAGE",
"ETH/USD_LEVERAGE",
)
assert runtime._runtime_scheduler.interval_seconds == 6.0
assert runtime_graph._runtime_scheduler.interval_seconds == 6.0
assert (
runtime._runtime_supervisor._heartbeat_monitor.timeout_seconds
runtime_graph._runtime_supervisor._heartbeat_monitor.timeout_seconds
== 31.0
)
assert (
runtime._reconnect_recovery_coordinator
runtime_graph._reconnect_recovery_coordinator
._recovery_coordinator
._window_planner
.max_window_ms
@@ -143,8 +191,9 @@ def test_recovery_rest_client_reuses_settings_snapshot() -> None:
runtime = build_trade_stream_production_runtime(settings)
assert isinstance(runtime, TradeStreamProductionRuntime)
runtime_graph: Any = runtime
document_source = (
runtime._reconnect_recovery_coordinator
runtime_graph._reconnect_recovery_coordinator
._recovery_coordinator
._recovery_controller
._document_source
@@ -154,4 +203,4 @@ def test_recovery_rest_client_reuses_settings_snapshot() -> None:
assert rest_client.settings is settings
assert rest_client.base_url == "https://rest.example.test"
assert rest_client.timeout == 17
assert "X-MBX-APIKEY" not in runtime._transport._headers
assert "X-MBX-APIKEY" not in runtime_graph._transport._headers

149
app/tests/unit/core/test_config.py
View File

@@ -17,6 +17,13 @@ _TRADE_STREAM_VARIABLES = (
"TRADE_STREAM_RECOVERY_WINDOW_MS",
)
_MARKET_DATA_STORAGE_VARIABLES = (
"MARKET_DATA_STORAGE_ENABLED",
"MARKET_DATA_STORAGE_POOL_MIN_SIZE",
"MARKET_DATA_STORAGE_POOL_MAX_SIZE",
"MARKET_DATA_STORAGE_POOL_TIMEOUT_SECONDS",
)
def prepare_environment(
monkeypatch: pytest.MonkeyPatch,
@@ -28,6 +35,9 @@ def prepare_environment(
for variable in _TRADE_STREAM_VARIABLES:
monkeypatch.delenv(variable, raising=False)
for variable in _MARKET_DATA_STORAGE_VARIABLES:
monkeypatch.delenv(variable, raising=False)
def enable_trade_stream(
monkeypatch: pytest.MonkeyPatch,
@@ -200,3 +210,142 @@ def test_symbols_must_not_contain_empty_items(
match="empty symbols",
):
load_settings()
def test_market_data_storage_is_disabled_by_default(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
settings = load_settings()
assert settings.market_data_storage.enabled is False
assert settings.market_data_storage.pool_min_size == 1
assert settings.market_data_storage.pool_max_size == 4
assert settings.market_data_storage.pool_timeout_seconds == 10.0
def test_disabled_market_data_storage_ignores_dependent_values(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MIN_SIZE",
"invalid",
)
settings = load_settings()
assert settings.market_data_storage.enabled is False
def test_market_data_storage_flag_is_strict(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"sometimes",
)
with pytest.raises(
ValueError,
match="MARKET_DATA_STORAGE_ENABLED",
):
load_settings()
def test_enabled_market_data_storage_requires_trade_stream(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"true",
)
with pytest.raises(
RuntimeError,
match="TRADE_STREAM_ENABLED",
):
load_settings()
def test_enabled_market_data_storage_parses_pool_settings(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
enable_trade_stream(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"true",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MIN_SIZE",
"2",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MAX_SIZE",
"7",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_TIMEOUT_SECONDS",
"12.5",
)
storage = load_settings().market_data_storage
assert storage.enabled is True
assert storage.pool_min_size == 2
assert storage.pool_max_size == 7
assert storage.pool_timeout_seconds == 12.5
@pytest.mark.parametrize(
("variable", "value"),
(
("MARKET_DATA_STORAGE_POOL_MIN_SIZE", "0"),
("MARKET_DATA_STORAGE_POOL_MAX_SIZE", "invalid"),
("MARKET_DATA_STORAGE_POOL_TIMEOUT_SECONDS", "nan"),
),
)
def test_enabled_market_data_storage_rejects_invalid_pool_settings(
monkeypatch: pytest.MonkeyPatch,
variable: str,
value: str,
) -> None:
prepare_environment(monkeypatch)
enable_trade_stream(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"true",
)
monkeypatch.setenv(variable, value)
with pytest.raises(ValueError, match=variable):
load_settings()
def test_pool_max_size_must_not_be_smaller_than_min_size(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
enable_trade_stream(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"true",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MIN_SIZE",
"5",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MAX_SIZE",
"4",
)
with pytest.raises(
ValueError,
match="POOL_MAX_SIZE",
):
load_settings()

151
app/tests/unit/market_data/acquisition/consistency/test_trade_stream_consistency_controller.py
View File

@@ -10,6 +10,9 @@ import pytest
from src.market_data.acquisition.consistency.trade_stream_consistency_controller import (
TradeStreamConsistencyController,
)
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.consistency.trade_stream_exceptions import (
TradeConsistencyError,
TradeOrderingError,
@@ -47,6 +50,7 @@ def _trade(
price: Decimal = Decimal("50000.00"),
quantity: Decimal = Decimal("0.25"),
aggressor_side: TradeAggressorSide = TradeAggressorSide.BUY,
source: str = "dzengi",
) -> Trade:
return Trade(
symbol=symbol,
@@ -62,10 +66,29 @@ def _trade(
tzinfo=timezone.utc,
),
aggressor_side=aggressor_side,
source="dzengi",
source=source,
)
class RecordingTradeObservationSink:
def __init__(
self,
*,
error: Exception | None = None,
) -> None:
self.error = error
self.observations: list[Trade] = []
def persist(
self,
trade: Trade,
) -> None:
self.observations.append(trade)
if self.error is not None:
raise self.error
def test_creates_state_for_first_symbol(
controller: TradeStreamConsistencyController,
state_store: TradeStreamStateStore,
@@ -179,4 +202,128 @@ def test_propagates_consistency_error(
trade_id=100,
price=Decimal("50001.00"),
),
)
)
def test_persists_trade_before_advancing_checkpoint(
state_store: TradeStreamStateStore,
) -> None:
sink = RecordingTradeObservationSink()
controller = TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=sink,
)
trade = _trade()
result = controller.accept(trade)
state = state_store.get(trade.symbol)
assert result is trade
assert sink.observations == [trade]
assert state.last_trade is trade
def test_persistence_failure_leaves_checkpoint_unchanged(
state_store: TradeStreamStateStore,
) -> None:
storage_error = RuntimeError("storage failed")
sink = RecordingTradeObservationSink()
controller = TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=sink,
)
first_trade = _trade(trade_id=100)
failed_trade = _trade(trade_id=101)
controller.accept(first_trade)
sink.error = storage_error
with pytest.raises(
RuntimeError,
match="storage failed",
) as error_info:
controller.accept(failed_trade)
state = state_store.get(first_trade.symbol)
assert error_info.value is storage_error
assert state.last_trade is first_trade
assert state.last_trade_id == first_trade.trade_id
assert failed_trade.trade_id not in state._trades
sink.error = None
assert controller.accept(failed_trade) is failed_trade
def test_valid_duplicate_is_persisted_without_checkpoint_advance(
state_store: TradeStreamStateStore,
) -> None:
sink = RecordingTradeObservationSink()
controller = TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=sink,
)
websocket_trade = _trade(
source="dzengi_websocket_trade",
)
rest_duplicate = _trade(
source="dzengi",
)
controller.accept(websocket_trade)
result = controller.accept(rest_duplicate)
state = state_store.get(websocket_trade.symbol)
assert result is None
assert sink.observations == [
websocket_trade,
rest_duplicate,
]
assert state.last_trade is websocket_trade
assert state.last_trade_id == websocket_trade.trade_id
def test_invalid_trades_do_not_reach_persistence_sink(
state_store: TradeStreamStateStore,
) -> None:
sink = RecordingTradeObservationSink()
controller = TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=sink,
)
first_trade = _trade(trade_id=100)
controller.accept(first_trade)
with pytest.raises(TradeOrderingError):
controller.accept(_trade(trade_id=99))
with pytest.raises(TradeConsistencyError):
controller.accept(
_trade(
trade_id=100,
price=Decimal("50001.00"),
)
)
assert sink.observations == [first_trade]
def test_rejects_invalid_trade_observation_sink(
state_store: TradeStreamStateStore,
) -> None:
with pytest.raises(
TypeError,
match="TradeObservationSinkProtocol",
):
TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=object(), # type: ignore[arg-type]
)
def test_recording_sink_implements_public_protocol() -> None:
assert isinstance(
RecordingTradeObservationSink(),
TradeObservationSinkProtocol,
)

101
app/tests/unit/market_data/acquisition/runtime/test_trade_stream_production_runtime.py
View File

@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import json
import threading
from collections import deque
from collections.abc import Awaitable, Callable
from typing import cast
@@ -376,15 +377,20 @@ class FakeTradeStreamService:
subscribe_error: Exception | None = None,
handle_error: Exception | None = None,
subscribe_gate: asyncio.Event | None = None,
handle_entered: threading.Event | None = None,
handle_release: threading.Event | None = None,
) -> None:
self._calls = calls
self._subscribe_error = subscribe_error
self._handle_error = handle_error
self._subscribe_gate = subscribe_gate
self._handle_entered = handle_entered
self._handle_release = handle_release
self.subscribe_calls: list[tuple[str, ...]] = []
self.subscribe_correlation_ids: list[str | None] = []
self.documents: list[object] = []
self.handled = asyncio.Event()
self.handled = threading.Event()
self.handle_thread_ids: list[int] = []
self.subscribe_entered = asyncio.Event()
async def subscribe(
@@ -408,10 +414,17 @@ class FakeTradeStreamService:
self,
document: object,
) -> Trade | None:
self.handle_thread_ids.append(threading.get_ident())
self.documents.append(document)
self._calls.append("service.handle_message")
self.handled.set()
if self._handle_entered is not None:
self._handle_entered.set()
if self._handle_release is not None:
self._handle_release.wait()
if self._handle_error is not None:
raise self._handle_error
@@ -592,6 +605,8 @@ class RuntimeDependencies:
stop_error: Exception | None = None,
subscribe_error: Exception | None = None,
handle_error: Exception | None = None,
handle_entered: threading.Event | None = None,
handle_release: threading.Event | None = None,
clear_error: Exception | None = None,
start_gate: asyncio.Event | None = None,
connected_publish_gate: asyncio.Event | None = None,
@@ -633,6 +648,8 @@ class RuntimeDependencies:
subscribe_error=subscribe_error,
handle_error=handle_error,
subscribe_gate=subscribe_gate,
handle_entered=handle_entered,
handle_release=handle_release,
)
self.live_processing_gate = RuntimeLiveProcessingGate()
self.reconnect_recovery = FakeReconnectRecoveryCoordinator(
@@ -709,6 +726,16 @@ async def wait_until(
raise AssertionError("condition was not reached")
async def wait_for_thread_event(
event: threading.Event,
) -> None:
reached = await asyncio.to_thread(
event.wait,
1.0,
)
assert reached is True
def test_implements_public_protocol_and_uses_slots() -> None:
dependencies = RuntimeDependencies()
@@ -944,7 +971,7 @@ def test_decodes_and_forwards_market_messages(
dependencies.runtime.run(),
)
await dependencies.service.handled.wait()
await wait_for_thread_event(dependencies.service.handled)
await dependencies.runtime.stop()
await runtime_task
@@ -1295,7 +1322,10 @@ def test_transport_error_runs_recovery_and_resumes_receive_loop() -> None:
dependencies.runtime.run(),
)
await dependencies.service.handled.wait()
await wait_for_thread_event(dependencies.service.handled)
await wait_until(
lambda: dependencies.transport.receive_calls == 3
)
assert dependencies.runtime.running is True
await dependencies.runtime.stop()
@@ -1342,7 +1372,7 @@ def test_buffered_market_waits_for_reconnect_recovery() -> None:
assert dependencies.service.documents == []
reconnect_release.set()
await dependencies.service.handled.wait()
await wait_for_thread_event(dependencies.service.handled)
await dependencies.runtime.stop()
await runtime_task
@@ -1568,6 +1598,69 @@ def test_market_handler_error_is_terminal_and_not_wrapped() -> None:
assert dependencies.runtime.state is (
TradeStreamProductionRuntimeState.FAILED
)
assert dependencies.live_processing_gate.failed is True
def test_market_processing_runs_outside_event_loop_thread() -> None:
async def scenario() -> tuple[RuntimeDependencies, int]:
event_loop_thread_id = threading.get_ident()
dependencies = RuntimeDependencies(
incoming=(MARKET_MESSAGE,),
)
runtime_task = asyncio.create_task(
dependencies.runtime.run(),
)
await wait_for_thread_event(dependencies.service.handled)
await dependencies.runtime.stop()
await runtime_task
return dependencies, event_loop_thread_id
dependencies, event_loop_thread_id = asyncio.run(scenario())
assert dependencies.service.handle_thread_ids
assert dependencies.service.handle_thread_ids[0] != event_loop_thread_id
def test_stop_waits_for_inflight_market_processing() -> None:
async def scenario() -> RuntimeDependencies:
handle_entered = threading.Event()
handle_release = threading.Event()
dependencies = RuntimeDependencies(
incoming=(MARKET_MESSAGE,),
handle_entered=handle_entered,
handle_release=handle_release,
)
runtime_task = asyncio.create_task(
dependencies.runtime.run(),
)
await wait_for_thread_event(handle_entered)
stop_task = asyncio.create_task(
dependencies.runtime.stop(),
)
try:
await asyncio.sleep(0)
await asyncio.sleep(0)
assert stop_task.done() is False
assert dependencies.live_processing_gate.locked is True
assert dependencies.runtime._market_processing_task is not None
finally:
handle_release.set()
await stop_task
await runtime_task
return dependencies
dependencies = asyncio.run(scenario())
assert dependencies.runtime.state is (
TradeStreamProductionRuntimeState.STOPPED
)
assert dependencies.runtime._market_processing_task is None
def test_connect_failure_publishes_event_and_rolls_back() -> None:

138
app/tests/unit/market_data/acquisition/test_trade_stream_runtime_composition.py
View File

@@ -13,6 +13,9 @@ import pytest
from src.market_data.acquisition.adapters.dzengi.rest import (
DzengiTradesDocumentSource,
)
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
@@ -315,6 +318,25 @@ class RecordingSleep:
self.calls.append(seconds)
class RecordingTradeObservationSink:
def __init__(
self,
*,
fail_on_trade_id: int | None = None,
) -> None:
self._fail_on_trade_id = fail_on_trade_id
self.observations: list[Trade] = []
def persist(
self,
trade: Trade,
) -> None:
self.observations.append(trade)
if trade.trade_id == self._fail_on_trade_id:
raise RuntimeError("storage failed")
@dataclass(slots=True)
class CompositionDependencies:
session: FakeSession
@@ -336,6 +358,7 @@ def create_composition(
scheduler_interval_seconds: float = 1.0,
max_recovery_window_ms: int = 3_599_999,
probe_results: tuple[bool, ...] = (True,),
trade_observation_sink: TradeObservationSinkProtocol | None = None,
) -> tuple[
TradeStreamRuntimeComposition,
CompositionDependencies,
@@ -370,6 +393,7 @@ def create_composition(
symbols=(SYMBOL,),
heartbeat_timeout_seconds=heartbeat_timeout_seconds,
scheduler_interval_seconds=scheduler_interval_seconds,
trade_observation_sink=trade_observation_sink,
max_recovery_window_ms=max_recovery_window_ms,
heartbeat_clock=dependencies.heartbeat_clock,
recovery_end_time_clock=(
@@ -491,6 +515,120 @@ def test_live_stream_and_recovery_share_consistency_state() -> None:
)
def test_live_and_recovery_share_optional_persistence_sink() -> None:
sink = RecordingTradeObservationSink()
recovered_trade_id = 101
composition, _ = create_composition(
trade_observation_sink=sink,
recovery_document=[
make_raw_trade(
trade_id=recovered_trade_id,
timestamp=CHECKPOINT_TIME_MS + 1_000,
),
],
)
live_trade = (
composition.trade_stream_acquisition_service.handle_message(
{
"destination": "internal.trade",
}
)
)
recovery_result = composition.runtime_recovery_coordinator.recover(
symbol=SYMBOL,
recovery_end_time=RECOVERY_END_TIME_MS,
)
assert composition.trade_observation_sink is sink
assert isinstance(sink, TradeObservationSinkProtocol)
assert sink.observations[0] is live_trade
assert sink.observations[1] is recovery_result.last_trade
assert [
trade.trade_id
for trade in sink.observations
] == [100, recovered_trade_id]
def test_recovery_duplicate_updates_persistence_without_checkpoint_change(
) -> None:
live_trade = Trade(
symbol=SYMBOL,
trade_id=100,
price=Decimal("64556.00"),
quantity=Decimal("0.003"),
executed_at=CHECKPOINT_TIME,
aggressor_side=TradeAggressorSide.BUY,
source="dzengi_websocket_trade",
)
sink = RecordingTradeObservationSink()
composition, _ = create_composition(
trade=live_trade,
trade_observation_sink=sink,
recovery_document=[
make_raw_trade(
trade_id=live_trade.trade_id,
timestamp=CHECKPOINT_TIME_MS,
),
],
)
composition.trade_stream_acquisition_service.handle_message(
{
"destination": "internal.trade",
}
)
result = composition.runtime_recovery_coordinator.recover(
symbol=SYMBOL,
recovery_end_time=RECOVERY_END_TIME_MS,
)
state = composition.state_store.get(SYMBOL)
assert result.is_empty is True
assert len(sink.observations) == 2
assert sink.observations[0] is live_trade
assert sink.observations[1].source == "dzengi"
assert state.last_trade is live_trade
def test_recovery_persistence_failure_does_not_advance_checkpoint() -> None:
sink = RecordingTradeObservationSink(
fail_on_trade_id=101,
)
composition, _ = create_composition(
trade_observation_sink=sink,
recovery_document=[
make_raw_trade(
trade_id=101,
timestamp=CHECKPOINT_TIME_MS + 1_000,
),
],
)
live_trade = (
composition.trade_stream_acquisition_service.handle_message(
{
"destination": "internal.trade",
}
)
)
with pytest.raises(
RuntimeError,
match="storage failed",
):
composition.runtime_recovery_coordinator.recover(
symbol=SYMBOL,
recovery_end_time=RECOVERY_END_TIME_MS,
)
state = composition.state_store.get(SYMBOL)
assert state.last_trade is live_trade
assert state.last_trade_id == 100
assert [trade.trade_id for trade in sink.observations] == [100, 101]
def test_runtime_components_share_lifecycle_dependencies() -> None:
composition, dependencies = create_composition()

121
app/tests/unit/market_data/storage/test_contracts.py Normal file
View File

@@ -0,0 +1,121 @@
from __future__ import annotations
from datetime import datetime
from typing import Any
import pytest
from src.market_data.storage import (
CandleStorageProtocol,
MarketDataBatchWriteResult,
MarketDataWriteResult,
MarketDataWriteStatus,
QuoteStorageProtocol,
TradeStorageProtocol,
)
class RecordingMarketDataStorage:
def store_trade(
self,
*,
venue: str,
trade: Any,
observed_at: datetime,
) -> MarketDataWriteResult:
return MarketDataWriteResult(MarketDataWriteStatus.INSERTED)
def store_trades(
self,
*,
venue: str,
trades: tuple[Any, ...],
observed_at: datetime,
) -> MarketDataBatchWriteResult:
return MarketDataBatchWriteResult(
inserted_count=len(trades),
duplicate_count=0,
provenance_updated_count=0,
)
def store_quote(
self,
*,
venue: str,
quote: Any,
) -> MarketDataWriteResult:
return MarketDataWriteResult(MarketDataWriteStatus.INSERTED)
def store_candle_revision(
self,
*,
venue: str,
candle: Any,
observed_at: datetime,
is_final: bool,
) -> MarketDataWriteResult:
return MarketDataWriteResult(MarketDataWriteStatus.INSERTED)
def test_storage_protocols_are_runtime_checkable() -> None:
storage = RecordingMarketDataStorage()
assert isinstance(storage, TradeStorageProtocol)
assert isinstance(storage, QuoteStorageProtocol)
assert isinstance(storage, CandleStorageProtocol)
@pytest.mark.parametrize("status", tuple(MarketDataWriteStatus))
def test_write_result_accepts_every_defined_status(
status: MarketDataWriteStatus,
) -> None:
result = MarketDataWriteResult(status=status)
assert result.status is status
def test_write_result_rejects_unknown_status() -> None:
with pytest.raises(TypeError, match="MarketDataWriteStatus"):
MarketDataWriteResult(status="inserted") # type: ignore[arg-type]
def test_batch_result_reports_total_count() -> None:
result = MarketDataBatchWriteResult(
inserted_count=2,
duplicate_count=3,
provenance_updated_count=5,
)
assert result.total_count == 10
@pytest.mark.parametrize(
"field_name",
(
"inserted_count",
"duplicate_count",
"provenance_updated_count",
),
)
def test_batch_result_rejects_negative_count(field_name: str) -> None:
values = {
"inserted_count": 0,
"duplicate_count": 0,
"provenance_updated_count": 0,
}
values[field_name] = -1
with pytest.raises(ValueError, match=field_name):
MarketDataBatchWriteResult(**values)
@pytest.mark.parametrize("invalid_count", (True, 1.5, "1", None))
def test_batch_result_rejects_non_integer_count(
invalid_count: object,
) -> None:
with pytest.raises(TypeError, match="inserted_count"):
MarketDataBatchWriteResult(
inserted_count=invalid_count, # type: ignore[arg-type]
duplicate_count=0,
provenance_updated_count=0,
)

206
app/tests/unit/market_data/storage/test_market_data_storage.py Normal file
View File

@@ -0,0 +1,206 @@
from __future__ import annotations
from dataclasses import dataclass, field
from datetime import datetime, timezone
from decimal import Decimal
import pytest
from src.market_data.acquisition.models.candle import Candle
from src.market_data.acquisition.models.quote import Quote
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
)
from src.market_data.storage import (
CandleStorageProtocol,
MarketDataBatchWriteResult,
MarketDataStorage,
MarketDataWriteResult,
MarketDataWriteStatus,
QuoteStorageProtocol,
TradeStorageProtocol,
)
NOW = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc)
@dataclass
class RecordingTradeStorage:
calls: list[tuple[str, object, datetime]] = field(default_factory=list)
def store_trade(
self,
*,
venue: str,
trade: Trade,
observed_at: datetime,
) -> MarketDataWriteResult:
self.calls.append((venue, trade, observed_at))
return MarketDataWriteResult(MarketDataWriteStatus.INSERTED)
def store_trades(
self,
*,
venue: str,
trades: tuple[Trade, ...],
observed_at: datetime,
) -> MarketDataBatchWriteResult:
self.calls.append((venue, trades, observed_at))
return MarketDataBatchWriteResult(len(trades), 0, 0)
@dataclass
class RecordingQuoteStorage:
calls: list[tuple[str, Quote]] = field(default_factory=list)
def store_quote(
self,
*,
venue: str,
quote: Quote,
) -> MarketDataWriteResult:
self.calls.append((venue, quote))
return MarketDataWriteResult(MarketDataWriteStatus.DUPLICATE)
@dataclass
class RecordingCandleStorage:
calls: list[tuple[str, Candle, datetime, bool]] = field(
default_factory=list
)
def store_candle_revision(
self,
*,
venue: str,
candle: Candle,
observed_at: datetime,
is_final: bool,
) -> MarketDataWriteResult:
self.calls.append((venue, candle, observed_at, is_final))
return MarketDataWriteResult(
MarketDataWriteStatus.PROVENANCE_UPDATED
)
def _trade() -> Trade:
return Trade(
symbol="BTC/USD_LEVERAGE",
trade_id=1,
price=Decimal("100"),
quantity=Decimal("1"),
executed_at=NOW,
aggressor_side=TradeAggressorSide.BUY,
source="dzengi",
)
def _quote() -> Quote:
return Quote(
symbol="BTC/USD_LEVERAGE",
last_price=Decimal("100"),
bid_price=Decimal("99"),
ask_price=Decimal("101"),
exchange_timestamp=NOW,
received_at=NOW,
source="dzengi",
)
def _candle() -> Candle:
return Candle(
symbol="BTC/USD_LEVERAGE",
interval="1m",
open_time=NOW,
open_price=Decimal("100"),
high_price=Decimal("110"),
low_price=Decimal("90"),
close_price=Decimal("105"),
volume=Decimal("10"),
source="dzengi",
)
def _storage() -> tuple[
MarketDataStorage,
RecordingTradeStorage,
RecordingQuoteStorage,
RecordingCandleStorage,
]:
trades = RecordingTradeStorage()
quotes = RecordingQuoteStorage()
candles = RecordingCandleStorage()
return (
MarketDataStorage(
trade_storage=trades,
quote_storage=quotes,
candle_storage=candles,
),
trades,
quotes,
candles,
)
def test_facade_matches_all_write_protocols() -> None:
storage, *_ = _storage()
assert isinstance(storage, TradeStorageProtocol)
assert isinstance(storage, QuoteStorageProtocol)
assert isinstance(storage, CandleStorageProtocol)
def test_facade_delegates_without_replacing_models_or_results() -> None:
storage, trades, quotes, candles = _storage()
trade = _trade()
quote = _quote()
candle = _candle()
trade_result = storage.store_trade(
venue="venue",
trade=trade,
observed_at=NOW,
)
batch_result = storage.store_trades(
venue="venue",
trades=(trade,),
observed_at=NOW,
)
quote_result = storage.store_quote(venue="venue", quote=quote)
candle_result = storage.store_candle_revision(
venue="venue",
candle=candle,
observed_at=NOW,
is_final=True,
)
assert trades.calls == [
("venue", trade, NOW),
("venue", (trade,), NOW),
]
assert quotes.calls == [("venue", quote)]
assert candles.calls == [("venue", candle, NOW, True)]
assert trade_result.status is MarketDataWriteStatus.INSERTED
assert batch_result.inserted_count == 1
assert quote_result.status is MarketDataWriteStatus.DUPLICATE
assert candle_result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
@pytest.mark.parametrize(
"dependency_name",
("trade_storage", "quote_storage", "candle_storage"),
)
def test_facade_rejects_dependency_outside_protocol(
dependency_name: str,
) -> None:
dependencies: dict[str, object] = {
"trade_storage": RecordingTradeStorage(),
"quote_storage": RecordingQuoteStorage(),
"candle_storage": RecordingCandleStorage(),
}
dependencies[dependency_name] = object()
with pytest.raises(TypeError, match=dependency_name):
MarketDataStorage(**dependencies) # type: ignore[arg-type]

862
app/tests/unit/market_data/storage/test_postgres_partitions.py Normal file
View File

@@ -0,0 +1,862 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
import re
import threading
from typing import Any
import pytest
from psycopg.sql import Composable
from src.market_data.storage import (
MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,
MarketDataPartitionResult,
MarketDataPartitionType,
MarketDataRetentionPolicy,
MarketDataStorageConfigurationError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
PostgresMarketDataPartitionManager,
PostgresMarketDataRetentionService,
build_monthly_partition,
)
UTC = timezone.utc
JUNE = datetime(2026, 6, 1, tzinfo=UTC)
JULY = datetime(2026, 7, 1, tzinfo=UTC)
AUGUST = datetime(2026, 8, 1, tzinfo=UTC)
RegistryKey = tuple[str, datetime]
RegistryRow = tuple[str, datetime, datetime, str]
RelationState = tuple[bool, bool, str | None]
@dataclass
class PartitionDatabase:
registry: dict[RegistryKey, RegistryRow] = field(default_factory=dict)
relations: dict[str, RelationState] = field(default_factory=dict)
def _partition_bound(range_start: datetime, range_end: datetime) -> str:
return f"FOR VALUES FROM ({range_start.isoformat()}) TO ({range_end.isoformat()})"
class PartitionCursor:
def __init__(self, connection: PartitionConnection) -> None:
self._connection = connection
self._fetchone: object = None
self._fetchall: list[tuple[Any, ...]] = []
self.rowcount = -1
def __enter__(self) -> PartitionCursor:
return self
def __exit__(self, *args: object) -> None:
return None
def execute(
self,
statement: str | Composable,
parameters: tuple[Any, ...] | None = None,
) -> None:
text = (
statement
if isinstance(statement, str)
else statement.as_string()
)
normalized = " ".join(text.split())
self._connection.calls.append((normalized, parameters))
self._fetchone = None
self._fetchall = []
self.rowcount = -1
if self._connection.interrupt_on in normalized:
raise KeyboardInterrupt
if self._connection.fail_on in normalized:
raise RuntimeError("database failed")
if normalized.startswith("SELECT pg_advisory_xact_lock"):
self._connection.acquire_advisory_lock()
return
if normalized.startswith("SELECT partition_name"):
assert parameters is not None
data_type, boundary = parameters
if "range_end <=" in normalized:
self._fetchall = sorted(
(
row
for (kind, _), row
in self._connection.working.registry.items()
if kind == data_type and row[2] <= boundary
),
key=lambda row: row[1],
)
else:
self._fetchone = self._connection.working.registry.get(
(data_type, boundary)
)
return
if "FROM pg_catalog.pg_class AS child" in normalized:
assert parameters is not None
partition_name = parameters[1]
self._fetchone = self._connection.working.relations.get(
partition_name,
(False, False, None),
)
return
if normalized.startswith("LOCK TABLE"):
if "SHARE ROW EXCLUSIVE" in normalized:
self._connection.acquire_write_lock()
return
if normalized.startswith("CREATE TABLE"):
partition_name = _second_qualified_identifier(normalized)
self._connection.working.relations[partition_name] = (
True,
False,
None,
)
return
if normalized.startswith("ALTER TABLE") and "ADD CONSTRAINT" in normalized:
return
if normalized.startswith("WITH moved_rows AS"):
self.rowcount = self._connection.moved_row_count
return
if normalized.startswith("ALTER TABLE") and "ATTACH PARTITION" in normalized:
partition_name = _attached_partition_identifier(normalized)
assert parameters is None
range_start, range_end = _timestamp_literals(normalized)
self._connection.working.relations[partition_name] = (
True,
True,
_partition_bound(range_start, range_end),
)
return
if normalized.startswith(
"INSERT INTO market_data.partition_registry"
):
assert parameters is not None
data_type, name, range_start, range_end, partition_bound = parameters
self._connection.working.registry[(data_type, range_start)] = (
name,
range_start,
range_end,
partition_bound,
)
self.rowcount = 1
return
if normalized.startswith("SELECT COUNT(*) FROM"):
partition_name = _first_qualified_identifier(normalized)
self._fetchone = (
self._connection.partition_row_counts.get(
partition_name,
0,
),
)
return
if normalized.startswith("DROP TABLE"):
partition_name = _first_qualified_identifier(normalized)
self._connection.working.relations.pop(partition_name, None)
return
if normalized.startswith(
"DELETE FROM market_data.partition_registry"
):
assert parameters is not None
data_type, name, range_start, range_end = parameters
key = (data_type, range_start)
row = self._connection.working.registry.get(key)
if row is not None and row[:3] == (
name,
range_start,
range_end,
):
del self._connection.working.registry[key]
self.rowcount = 1
else:
self.rowcount = 0
return
if normalized.startswith("DELETE FROM"):
parent_name = _first_qualified_identifier(normalized)
self.rowcount = self._connection.partial_delete_counts.get(
parent_name,
0,
)
return
raise AssertionError(f"Unexpected SQL: {normalized}")
def fetchone(self) -> object:
return self._fetchone
def fetchall(self) -> list[tuple[Any, ...]]:
return list(self._fetchall)
def _quoted_identifiers(statement: str) -> list[str]:
return statement.split('"')[1::2]
def _first_qualified_identifier(statement: str) -> str:
identifiers = _quoted_identifiers(statement)
return identifiers[1]
def _second_qualified_identifier(statement: str) -> str:
identifiers = _quoted_identifiers(statement)
return identifiers[1]
def _attached_partition_identifier(statement: str) -> str:
identifiers = _quoted_identifiers(statement)
return identifiers[3]
def _timestamp_literals(statement: str) -> tuple[datetime, datetime]:
values = tuple(
datetime.fromisoformat(value)
for value in re.findall(
r"'([^']+)'::timestamptz",
statement,
)
)
if len(values) != 2:
raise AssertionError(
f"Expected two timestamptz literals: {statement}"
)
return values
class PartitionConnection:
def __init__(self, database: PartitionDatabase) -> None:
self._database = database
self.working = deepcopy(database)
self.calls: list[tuple[str, tuple[Any, ...] | None]] = []
self.exit_exception_types: list[type[BaseException] | None] = []
self.fail_on = "__never_fail__"
self.interrupt_on = "__never_interrupt__"
self.moved_row_count = 0
self.partition_row_counts: dict[str, int] = {}
self.partial_delete_counts: dict[str, int] = {}
self.advisory_lock: threading.Lock | None = None
self.advisory_attempted: threading.Event | None = None
self.advisory_acquired: threading.Event | None = None
self.advisory_continue: threading.Event | None = None
self.write_lock: threading.Lock | None = None
self.write_lock_acquired: threading.Event | None = None
self.write_lock_continue: threading.Event | None = None
self._holds_advisory_lock = False
self._holds_write_lock = False
def __enter__(self) -> PartitionConnection:
self.working = deepcopy(self._database)
return self
def __exit__(
self,
exception_type: type[BaseException] | None,
exception: BaseException | None,
traceback: object,
) -> None:
self.exit_exception_types.append(exception_type)
try:
if exception_type is None:
self._database.registry = self.working.registry
self._database.relations = self.working.relations
finally:
if self._holds_write_lock:
assert self.write_lock is not None
self._holds_write_lock = False
self.write_lock.release()
if self._holds_advisory_lock:
assert self.advisory_lock is not None
self._holds_advisory_lock = False
self.advisory_lock.release()
return None
def cursor(self) -> PartitionCursor:
return PartitionCursor(self)
def acquire_advisory_lock(self) -> None:
lock = self.advisory_lock
if lock is None or self._holds_advisory_lock:
return
if self.advisory_attempted is not None:
self.advisory_attempted.set()
lock.acquire()
self._holds_advisory_lock = True
self.working = deepcopy(self._database)
if self.advisory_acquired is not None:
self.advisory_acquired.set()
if self.advisory_continue is not None:
if not self.advisory_continue.wait(timeout=3):
raise TimeoutError("advisory test gate timed out")
def acquire_write_lock(self) -> None:
lock = self.write_lock
if lock is None or self._holds_write_lock:
return
lock.acquire()
self._holds_write_lock = True
if self.write_lock_acquired is not None:
self.write_lock_acquired.set()
if self.write_lock_continue is not None:
if not self.write_lock_continue.wait(timeout=3):
raise TimeoutError("retention test gate timed out")
@dataclass
class PartitionProvider:
connection: PartitionConnection
calls: int = 0
def __call__(self) -> PartitionConnection:
self.calls += 1
return self.connection
class ConcurrentPartitionProvider:
def __init__(self, database: PartitionDatabase) -> None:
self.database = database
self.advisory_lock = threading.Lock()
self.first_acquired = threading.Event()
self.first_continue = threading.Event()
self.second_attempted = threading.Event()
self.connections: list[PartitionConnection] = []
self._provider_lock = threading.Lock()
def __call__(self) -> PartitionConnection:
with self._provider_lock:
caller_index = len(self.connections)
connection = PartitionConnection(self.database)
connection.advisory_lock = self.advisory_lock
if caller_index == 0:
connection.advisory_acquired = self.first_acquired
connection.advisory_continue = self.first_continue
elif caller_index == 1:
connection.advisory_attempted = self.second_attempted
self.connections.append(connection)
return connection
def _dependencies() -> tuple[
PartitionDatabase,
PartitionConnection,
PartitionProvider,
]:
database = PartitionDatabase()
connection = PartitionConnection(database)
provider = PartitionProvider(connection)
return database, connection, provider
def _register(
database: PartitionDatabase,
data_type: MarketDataPartitionType,
month: datetime,
) -> str:
partition = build_monthly_partition(data_type=data_type, month=month)
database.registry[(data_type.value, partition.range_start)] = (
partition.partition_name,
partition.range_start,
partition.range_end,
_partition_bound(partition.range_start, partition.range_end),
)
database.relations[partition.partition_name] = (
True,
True,
_partition_bound(partition.range_start, partition.range_end),
)
return partition.partition_name
def test_monthly_descriptor_uses_utc_and_handles_year_rollover() -> None:
local_time = datetime(
2027,
1,
1,
1,
tzinfo=timezone(timedelta(hours=3)),
)
partition = build_monthly_partition(
data_type=MarketDataPartitionType.TRADES,
month=local_time,
)
assert partition.partition_name == "trades_2026_12"
assert partition.range_start == datetime(2026, 12, 1, tzinfo=UTC)
assert partition.range_end == datetime(2027, 1, 1, tzinfo=UTC)
def test_monthly_descriptor_rejects_naive_datetime() -> None:
with pytest.raises(MarketDataStorageValidationError, match="month"):
build_monthly_partition(
data_type=MarketDataPartitionType.QUOTES,
month=JULY.replace(tzinfo=None),
)
def test_manager_creates_partition_moves_default_rows_and_registers_it() -> None:
database, connection, provider = _dependencies()
connection.moved_row_count = 3
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
result = manager.ensure_month_partition(
data_type=MarketDataPartitionType.QUOTES,
month=JULY + timedelta(days=20),
)
assert result.created is True
assert result.moved_row_count == 3
assert result.partition.partition_name == "quotes_2026_07"
assert database.relations["quotes_2026_07"] == (
True,
True,
_partition_bound(JULY, AUGUST),
)
assert database.registry[("quotes", JULY)] == (
"quotes_2026_07",
JULY,
AUGUST,
_partition_bound(JULY, AUGUST),
)
statements = [statement for statement, _ in connection.calls]
assert statements[0] == "SELECT pg_advisory_xact_lock(%s)"
assert connection.calls[0][1] == (
MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,
)
assert statements.index(
'LOCK TABLE "market_data"."quotes_default" IN ACCESS EXCLUSIVE MODE'
) < next(
index
for index, statement in enumerate(statements)
if statement.startswith("WITH moved_rows AS")
)
assert any(
'CREATE TABLE "market_data"."quotes_2026_07"' in statement
for statement in statements
)
ddl_calls = tuple(
(statement, parameters)
for statement, parameters in connection.calls
if statement.startswith("ALTER TABLE")
)
assert ddl_calls
assert all(parameters is None for _, parameters in ddl_calls)
assert all("::timestamptz" in statement for statement, _ in ddl_calls)
def test_manager_is_idempotent_after_registered_partition_exists() -> None:
_, connection, provider = _dependencies()
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
first = manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
call_count = len(connection.calls)
second = manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
assert first.created is True
assert second.created is False
assert second.partition == first.partition
assert not any(
statement.startswith("CREATE TABLE")
for statement, _ in connection.calls[call_count:]
)
def test_two_manager_callers_participate_in_single_serialized_creation() -> None:
database = PartitionDatabase()
provider = ConcurrentPartitionProvider(database)
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
results: dict[str, MarketDataPartitionResult] = {}
errors: list[BaseException] = []
def run(caller: str) -> None:
try:
results[caller] = manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
except BaseException as error:
errors.append(error)
first = threading.Thread(target=run, args=("first",))
second = threading.Thread(target=run, args=("second",))
first.start()
try:
assert provider.first_acquired.wait(timeout=1)
second.start()
assert provider.second_attempted.wait(timeout=1)
assert second.is_alive()
finally:
provider.first_continue.set()
first.join(timeout=2)
second.join(timeout=2)
assert first.is_alive() is False
assert second.is_alive() is False
assert errors == []
assert {result.created for result in results.values()} == {
True,
False,
}
assert sum(
statement.startswith("CREATE TABLE")
for connection in provider.connections
for statement, _ in connection.calls
) == 1
assert ("trades", JULY) in database.registry
def test_unregistered_relation_is_not_silently_adopted() -> None:
database, _, provider = _dependencies()
database.relations["trades_2026_07"] = (True, False, None)
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
with pytest.raises(
MarketDataStorageConfigurationError,
match="Unregistered",
):
manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
def test_registered_but_detached_partition_is_rejected() -> None:
database, _, provider = _dependencies()
name = _register(database, MarketDataPartitionType.TRADES, JULY)
database.relations[name] = (True, False, None)
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
with pytest.raises(
MarketDataStorageConfigurationError,
match="missing or detached",
):
manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
def test_retention_rejects_actual_partition_bound_mismatch_before_drop() -> None:
database, connection, provider = _dependencies()
name = _register(database, MarketDataPartitionType.TRADES, JUNE)
database.relations[name] = (
True,
True,
_partition_bound(JUNE, AUGUST),
)
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
with pytest.raises(
MarketDataStorageConfigurationError,
match="bound differs",
):
service.apply(
policy=MarketDataRetentionPolicy(enabled=True, trade_days=31),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
assert ("trades", JUNE) in database.registry
assert name in database.relations
assert not any(
statement.startswith("DROP TABLE")
for statement, _ in connection.calls
)
def test_partition_setup_error_rolls_back_created_relation_and_registry() -> None:
database, connection, provider = _dependencies()
connection.fail_on = "ATTACH PARTITION"
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
with pytest.raises(MarketDataStorageOperationError) as error_info:
manager.ensure_month_partition(
data_type=MarketDataPartitionType.CANDLE_REVISIONS,
month=JULY,
)
assert isinstance(error_info.value.__cause__, RuntimeError)
assert database.registry == {}
assert database.relations == {}
assert connection.exit_exception_types[-1] is RuntimeError
def test_partition_setup_interrupt_is_not_wrapped_and_rolls_back() -> None:
database, connection, provider = _dependencies()
connection.interrupt_on = "ATTACH PARTITION"
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
with pytest.raises(KeyboardInterrupt):
manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
assert database.registry == {}
assert database.relations == {}
assert connection.exit_exception_types[-1] is KeyboardInterrupt
def test_retention_is_disabled_without_borrowing_connection() -> None:
_, _, provider = _dependencies()
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
result = service.apply(
policy=MarketDataRetentionPolicy(),
now=JULY.replace(tzinfo=None),
)
assert result.entries == ()
assert result.total_removed_count == 0
assert provider.calls == 0
@pytest.mark.parametrize("days", (0, -1))
def test_retention_rejects_non_positive_window(days: int) -> None:
with pytest.raises(MarketDataStorageConfigurationError):
MarketDataRetentionPolicy(enabled=True, trade_days=days)
def test_enabled_retention_requires_at_least_one_window() -> None:
with pytest.raises(MarketDataStorageConfigurationError, match="window"):
MarketDataRetentionPolicy(enabled=True)
def test_retention_drops_complete_month_and_deletes_partial_history() -> None:
database, connection, provider = _dependencies()
june_name = _register(database, MarketDataPartitionType.TRADES, JUNE)
july_name = _register(database, MarketDataPartitionType.TRADES, JULY)
connection.partition_row_counts[june_name] = 10
connection.partition_row_counts[july_name] = 20
connection.partial_delete_counts["trades"] = 4
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
result = service.apply(
policy=MarketDataRetentionPolicy(enabled=True, trade_days=31),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
assert len(result.entries) == 1
entry = result.entries[0]
assert entry.data_type is MarketDataPartitionType.TRADES
assert entry.cutoff == datetime(2026, 7, 15, tzinfo=UTC)
assert entry.dropped_partitions == (june_name,)
assert entry.dropped_row_count == 10
assert entry.deleted_row_count == 4
assert entry.total_removed_count == 14
assert june_name not in database.relations
assert july_name in database.relations
assert ("trades", JUNE) not in database.registry
assert ("trades", JULY) in database.registry
assert any(
statement == (
'DELETE FROM "market_data"."trades" '
'WHERE "executed_at" < %s'
)
for statement, _ in connection.calls
)
def test_unconfigured_data_type_is_unlimited_and_not_touched() -> None:
database, connection, provider = _dependencies()
_register(database, MarketDataPartitionType.QUOTES, JUNE)
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
result = service.apply(
policy=MarketDataRetentionPolicy(enabled=True, trade_days=30),
now=AUGUST,
)
assert tuple(entry.data_type for entry in result.entries) == (
MarketDataPartitionType.TRADES,
)
assert not any(
'"quotes"' in statement or parameters == ("quotes", JULY)
for statement, parameters in connection.calls
)
def test_parallel_writer_waits_until_retention_transaction_finishes() -> None:
_, connection, provider = _dependencies()
write_lock = threading.Lock()
connection.write_lock = write_lock
connection.write_lock_acquired = threading.Event()
connection.write_lock_continue = threading.Event()
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
retention_errors: list[BaseException] = []
writer_attempted = threading.Event()
writer_finished = threading.Event()
def retain() -> None:
try:
service.apply(
policy=MarketDataRetentionPolicy(
enabled=True,
trade_days=31,
),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
except BaseException as error:
retention_errors.append(error)
def write() -> None:
writer_attempted.set()
with write_lock:
writer_finished.set()
retention_thread = threading.Thread(target=retain)
writer_thread = threading.Thread(target=write)
retention_thread.start()
try:
assert connection.write_lock_acquired.wait(timeout=1)
writer_thread.start()
assert writer_attempted.wait(timeout=1)
assert writer_finished.wait(timeout=0.05) is False
finally:
assert connection.write_lock_continue is not None
connection.write_lock_continue.set()
retention_thread.join(timeout=2)
writer_thread.join(timeout=2)
assert retention_thread.is_alive() is False
assert writer_thread.is_alive() is False
assert retention_errors == []
assert writer_finished.is_set()
statements = [statement for statement, _ in connection.calls]
lock_index = statements.index(
'LOCK TABLE "market_data"."trades" '
"IN SHARE ROW EXCLUSIVE MODE"
)
registry_index = next(
index
for index, statement in enumerate(statements)
if statement.startswith("SELECT partition_name")
)
assert lock_index < registry_index
def test_retention_error_rolls_back_all_configured_data_types() -> None:
database, connection, provider = _dependencies()
trade_name = _register(
database,
MarketDataPartitionType.TRADES,
JUNE,
)
quote_name = _register(
database,
MarketDataPartitionType.QUOTES,
JUNE,
)
connection.partition_row_counts[trade_name] = 2
connection.partition_row_counts[quote_name] = 3
connection.fail_on = 'DROP TABLE "market_data"."quotes_2026_06"'
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
with pytest.raises(MarketDataStorageOperationError):
service.apply(
policy=MarketDataRetentionPolicy(
enabled=True,
trade_days=31,
quote_days=31,
),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
assert ("trades", JUNE) in database.registry
assert ("quotes", JUNE) in database.registry
assert trade_name in database.relations
assert quote_name in database.relations
assert connection.exit_exception_types[-1] is RuntimeError
def test_retention_interrupt_rolls_back_and_is_not_wrapped() -> None:
database, connection, provider = _dependencies()
trade_name = _register(
database,
MarketDataPartitionType.TRADES,
JUNE,
)
connection.interrupt_on = "DROP TABLE"
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
with pytest.raises(KeyboardInterrupt):
service.apply(
policy=MarketDataRetentionPolicy(enabled=True, trade_days=31),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
assert ("trades", JUNE) in database.registry
assert trade_name in database.relations
assert connection.exit_exception_types[-1] is KeyboardInterrupt

641
app/tests/unit/market_data/storage/test_postgres_quote_candle_repositories.py Normal file
View File

@@ -0,0 +1,641 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass, field, replace
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from typing import Any
import pytest
from src.market_data.acquisition.models.candle import Candle
from src.market_data.acquisition.models.quote import Quote
from src.market_data.storage import (
CandleStorageProtocol,
MarketDataStorageConflictError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
MarketDataWriteStatus,
PostgresCandleRepository,
PostgresQuoteRepository,
QuoteStorageProtocol,
)
VENUE = "dzengi"
SYMBOL = "BTC/USD_LEVERAGE"
OPEN_TIME = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc)
RECEIVED_AT = OPEN_TIME + timedelta(seconds=1)
OBSERVED_AT = OPEN_TIME + timedelta(minutes=1)
QuoteKey = tuple[str, str, datetime]
CandleKey = tuple[str, str, str, datetime, datetime]
StorageRow = dict[str, Any]
@dataclass
class TransactionalMarketDataDatabase:
quote_rows: dict[QuoteKey, StorageRow] = field(default_factory=dict)
candle_rows: dict[CandleKey, StorageRow] = field(default_factory=dict)
class TransactionalMarketDataCursor:
def __init__(
self,
connection: TransactionalMarketDataConnection,
) -> None:
self._connection = connection
self._fetchone_result: object = None
def __enter__(self) -> TransactionalMarketDataCursor:
return self
def __exit__(self, *args: object) -> None:
return None
def execute(
self,
statement: str,
parameters: tuple[Any, ...],
) -> None:
normalized = " ".join(statement.split())
self._connection.calls.append((normalized, parameters))
if (
self._connection.interrupt_on is not None
and self._connection.interrupt_on in normalized
):
raise KeyboardInterrupt
if (
self._connection.fail_on is not None
and self._connection.fail_on in normalized
):
raise RuntimeError("database failed")
if normalized.startswith("INSERT INTO market_data.quotes"):
self._insert_quote(parameters)
return
if normalized.startswith("SELECT exchange_timestamp"):
self._select_quote(parameters)
return
if normalized.startswith("UPDATE market_data.quotes"):
self._update_quote(parameters)
return
if normalized.startswith(
"INSERT INTO market_data.candle_revisions"
):
self._insert_candle(parameters)
return
if normalized.startswith("SELECT open_price"):
self._select_candle(parameters)
return
if normalized.startswith(
"UPDATE market_data.candle_revisions"
):
self._update_candle(parameters)
return
raise AssertionError(f"Unexpected SQL: {normalized}")
def fetchone(self) -> object:
return self._fetchone_result
def _insert_quote(self, parameters: tuple[Any, ...]) -> None:
(
venue,
symbol,
received_at,
exchange_timestamp,
last_price,
bid_price,
ask_price,
source,
observation_sources,
canonical_schema_version,
) = parameters
key = (venue, symbol, received_at)
if key in self._connection.working_quote_rows:
self._fetchone_result = None
return
self._connection.working_quote_rows[key] = {
"exchange_timestamp": exchange_timestamp,
"last_price": last_price,
"bid_price": bid_price,
"ask_price": ask_price,
"source": source,
"observation_sources": list(observation_sources),
"canonical_schema_version": canonical_schema_version,
}
self._fetchone_result = (1,)
def _select_quote(self, parameters: tuple[Any, ...]) -> None:
row = self._connection.working_quote_rows.get(parameters)
if row is None:
self._fetchone_result = None
return
self._fetchone_result = (
row["exchange_timestamp"],
row["last_price"],
row["bid_price"],
row["ask_price"],
list(row["observation_sources"]),
row["canonical_schema_version"],
)
def _update_quote(self, parameters: tuple[Any, ...]) -> None:
sources, *identity = parameters
row = self._connection.working_quote_rows[tuple(identity)]
row["observation_sources"] = list(sources)
self._fetchone_result = None
def _insert_candle(self, parameters: tuple[Any, ...]) -> None:
(
venue,
symbol,
interval,
open_time,
observed_at,
open_price,
high_price,
low_price,
close_price,
volume,
is_final,
source,
observation_sources,
canonical_schema_version,
) = parameters
key = (venue, symbol, interval, open_time, observed_at)
if key in self._connection.working_candle_rows:
self._fetchone_result = None
return
self._connection.working_candle_rows[key] = {
"open_price": open_price,
"high_price": high_price,
"low_price": low_price,
"close_price": close_price,
"volume": volume,
"is_final": is_final,
"source": source,
"observation_sources": list(observation_sources),
"canonical_schema_version": canonical_schema_version,
}
self._fetchone_result = (1,)
def _select_candle(self, parameters: tuple[Any, ...]) -> None:
row = self._connection.working_candle_rows.get(parameters)
if row is None:
self._fetchone_result = None
return
self._fetchone_result = (
row["open_price"],
row["high_price"],
row["low_price"],
row["close_price"],
row["volume"],
row["is_final"],
list(row["observation_sources"]),
row["canonical_schema_version"],
)
def _update_candle(self, parameters: tuple[Any, ...]) -> None:
sources, *identity = parameters
row = self._connection.working_candle_rows[tuple(identity)]
row["observation_sources"] = list(sources)
self._fetchone_result = None
class TransactionalMarketDataConnection:
def __init__(self, database: TransactionalMarketDataDatabase) -> None:
self._database = database
self.calls: list[tuple[str, tuple[Any, ...]]] = []
self.exit_exception_types: list[type[BaseException] | None] = []
self.fail_on: str | None = None
self.interrupt_on: str | None = None
self.working_quote_rows: dict[QuoteKey, StorageRow] = {}
self.working_candle_rows: dict[CandleKey, StorageRow] = {}
def __enter__(self) -> TransactionalMarketDataConnection:
self.working_quote_rows = deepcopy(self._database.quote_rows)
self.working_candle_rows = deepcopy(self._database.candle_rows)
return self
def __exit__(
self,
exception_type: type[BaseException] | None,
exception: BaseException | None,
traceback: object,
) -> None:
self.exit_exception_types.append(exception_type)
if exception_type is None:
self._database.quote_rows = self.working_quote_rows
self._database.candle_rows = self.working_candle_rows
return None
def cursor(self) -> TransactionalMarketDataCursor:
return TransactionalMarketDataCursor(self)
@dataclass
class RecordingConnectionProvider:
connection: TransactionalMarketDataConnection
calls: int = 0
def __call__(self) -> TransactionalMarketDataConnection:
self.calls += 1
return self.connection
def _dependencies() -> tuple[
TransactionalMarketDataDatabase,
TransactionalMarketDataConnection,
RecordingConnectionProvider,
]:
database = TransactionalMarketDataDatabase()
connection = TransactionalMarketDataConnection(database)
provider = RecordingConnectionProvider(connection)
return database, connection, provider
def _quote(
*,
received_at: datetime = RECEIVED_AT,
exchange_timestamp: datetime | None = OPEN_TIME,
last_price: Decimal = Decimal("100"),
bid_price: Decimal = Decimal("99"),
ask_price: Decimal = Decimal("101"),
source: str = "dzengi",
) -> Quote:
return Quote(
symbol=SYMBOL,
last_price=last_price,
bid_price=bid_price,
ask_price=ask_price,
exchange_timestamp=exchange_timestamp,
received_at=received_at,
source=source,
)
def _candle(
*,
open_time: datetime = OPEN_TIME,
open_price: Decimal = Decimal("100"),
high_price: Decimal = Decimal("110"),
low_price: Decimal = Decimal("90"),
close_price: Decimal = Decimal("105"),
volume: Decimal = Decimal("10"),
source: str = "rest_klines:bid",
) -> Candle:
return Candle(
symbol=SYMBOL,
interval="1m",
open_time=open_time,
open_price=open_price,
high_price=high_price,
low_price=low_price,
close_price=close_price,
volume=volume,
source=source,
)
def test_quote_repository_matches_protocol() -> None:
_, _, provider = _dependencies()
assert isinstance(
PostgresQuoteRepository(connection_provider=provider),
QuoteStorageProtocol,
)
def test_quote_insert_normalizes_identity_and_preserves_payload() -> None:
database, _, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
quote = replace(_quote(), symbol=" btc/usd_leverage ")
result = repository.store_quote(venue=" dzengi ", quote=quote)
assert result.status is MarketDataWriteStatus.INSERTED
assert tuple(database.quote_rows) == ((VENUE, SYMBOL, RECEIVED_AT),)
row = next(iter(database.quote_rows.values()))
assert row["last_price"] == Decimal("100")
assert row["source"] == "dzengi"
assert row["observation_sources"] == ["dzengi"]
def test_quote_exact_duplicate_is_not_updated() -> None:
database, connection, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
quote = _quote()
repository.store_quote(venue=VENUE, quote=quote)
call_count = len(connection.calls)
result = repository.store_quote(venue=VENUE, quote=quote)
assert result.status is MarketDataWriteStatus.DUPLICATE
assert len(database.quote_rows) == 1
assert not any(
statement.startswith("UPDATE market_data.quotes")
for statement, _ in connection.calls[call_count:]
)
def test_quote_new_source_updates_only_provenance() -> None:
database, _, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
quote = _quote(source="dzengi")
repository.store_quote(venue=VENUE, quote=quote)
result = repository.store_quote(
venue=VENUE,
quote=replace(quote, source="dzengi_websocket_quote"),
)
row = next(iter(database.quote_rows.values()))
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert row["source"] == "dzengi"
assert row["observation_sources"] == [
"dzengi",
"dzengi_websocket_quote",
]
@pytest.mark.parametrize(
"conflicting_quote",
(
_quote(last_price=Decimal("102")),
_quote(bid_price=Decimal("98")),
_quote(ask_price=Decimal("102")),
_quote(exchange_timestamp=OPEN_TIME + timedelta(seconds=1)),
),
)
def test_quote_timestamp_collision_with_different_payload_is_conflict(
conflicting_quote: Quote,
) -> None:
database, connection, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
repository.store_quote(venue=VENUE, quote=_quote())
with pytest.raises(MarketDataStorageConflictError):
repository.store_quote(venue=VENUE, quote=conflicting_quote)
assert len(database.quote_rows) == 1
assert connection.exit_exception_types[-1] is (
MarketDataStorageConflictError
)
@pytest.mark.parametrize(
"invalid_quote",
(
_quote(received_at=RECEIVED_AT.replace(tzinfo=None)),
_quote(exchange_timestamp=OPEN_TIME.replace(tzinfo=None)),
_quote(bid_price=Decimal("102"), ask_price=Decimal("101")),
_quote(last_price=Decimal("NaN")),
),
)
def test_quote_validation_happens_before_connection(
invalid_quote: Quote,
) -> None:
_, _, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
with pytest.raises(MarketDataStorageValidationError):
repository.store_quote(venue=VENUE, quote=invalid_quote)
assert provider.calls == 0
def test_quote_database_error_is_wrapped_and_rolled_back() -> None:
database, connection, provider = _dependencies()
connection.fail_on = "INSERT INTO market_data.quotes"
repository = PostgresQuoteRepository(connection_provider=provider)
with pytest.raises(MarketDataStorageOperationError) as error_info:
repository.store_quote(venue=VENUE, quote=_quote())
assert isinstance(error_info.value.__cause__, RuntimeError)
assert database.quote_rows == {}
assert connection.exit_exception_types[-1] is RuntimeError
def test_candle_repository_matches_protocol() -> None:
_, _, provider = _dependencies()
assert isinstance(
PostgresCandleRepository(connection_provider=provider),
CandleStorageProtocol,
)
def test_candle_revision_insert_normalizes_identity_and_payload() -> None:
database, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
candle = replace(
_candle(),
symbol=" btc/usd_leverage ",
interval=" 1m ",
)
result = repository.store_candle_revision(
venue=" dzengi ",
candle=candle,
observed_at=OBSERVED_AT,
is_final=False,
)
assert result.status is MarketDataWriteStatus.INSERTED
assert tuple(database.candle_rows) == (
(VENUE, SYMBOL, "1m", OPEN_TIME, OBSERVED_AT),
)
row = next(iter(database.candle_rows.values()))
assert row["is_final"] is False
assert row["observation_sources"] == ["rest_klines:bid"]
def test_candle_interval_case_is_preserved() -> None:
database, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
repository.store_candle_revision(
venue=VENUE,
candle=replace(_candle(), interval="1M"),
observed_at=OBSERVED_AT,
is_final=True,
)
assert next(iter(database.candle_rows))[2] == "1M"
def test_candle_exact_duplicate_is_not_updated() -> None:
database, connection, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
candle = _candle()
repository.store_candle_revision(
venue=VENUE,
candle=candle,
observed_at=OBSERVED_AT,
is_final=False,
)
call_count = len(connection.calls)
result = repository.store_candle_revision(
venue=VENUE,
candle=candle,
observed_at=OBSERVED_AT,
is_final=False,
)
assert result.status is MarketDataWriteStatus.DUPLICATE
assert len(database.candle_rows) == 1
assert not any(
statement.startswith("UPDATE market_data.candle_revisions")
for statement, _ in connection.calls[call_count:]
)
def test_candle_new_source_updates_only_provenance() -> None:
database, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
candle = _candle(source="rest_klines:bid")
repository.store_candle_revision(
venue=VENUE,
candle=candle,
observed_at=OBSERVED_AT,
is_final=True,
)
result = repository.store_candle_revision(
venue=VENUE,
candle=replace(candle, source="secondary"),
observed_at=OBSERVED_AT,
is_final=True,
)
row = next(iter(database.candle_rows.values()))
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert row["source"] == "rest_klines:bid"
assert row["observation_sources"] == [
"rest_klines:bid",
"secondary",
]
def test_candle_final_state_at_later_observation_is_new_revision() -> None:
database, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
candle = _candle()
repository.store_candle_revision(
venue=VENUE,
candle=candle,
observed_at=OBSERVED_AT,
is_final=False,
)
result = repository.store_candle_revision(
venue=VENUE,
candle=replace(candle, close_price=Decimal("106")),
observed_at=OBSERVED_AT + timedelta(seconds=1),
is_final=True,
)
assert result.status is MarketDataWriteStatus.INSERTED
assert len(database.candle_rows) == 2
@pytest.mark.parametrize(
("conflicting_candle", "is_final"),
(
(replace(_candle(), close_price=Decimal("106")), False),
(replace(_candle(), volume=Decimal("11")), False),
(_candle(), True),
),
)
def test_candle_same_revision_identity_with_different_payload_is_conflict(
conflicting_candle: Candle,
is_final: bool,
) -> None:
database, connection, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
repository.store_candle_revision(
venue=VENUE,
candle=_candle(),
observed_at=OBSERVED_AT,
is_final=False,
)
with pytest.raises(MarketDataStorageConflictError):
repository.store_candle_revision(
venue=VENUE,
candle=conflicting_candle,
observed_at=OBSERVED_AT,
is_final=is_final,
)
assert len(database.candle_rows) == 1
assert connection.exit_exception_types[-1] is (
MarketDataStorageConflictError
)
@pytest.mark.parametrize(
("invalid_candle", "observed_at", "is_final"),
(
(_candle(), OPEN_TIME - timedelta(seconds=1), False),
(_candle(open_time=OPEN_TIME.replace(tzinfo=None)), OBSERVED_AT, False),
(_candle(volume=Decimal("-1")), OBSERVED_AT, False),
(_candle(high_price=Decimal("99")), OBSERVED_AT, False),
(_candle(), OBSERVED_AT, 1),
),
)
def test_candle_validation_happens_before_connection(
invalid_candle: Candle,
observed_at: datetime,
is_final: object,
) -> None:
_, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
with pytest.raises(MarketDataStorageValidationError):
repository.store_candle_revision(
venue=VENUE,
candle=invalid_candle,
observed_at=observed_at,
is_final=is_final, # type: ignore[arg-type]
)
assert provider.calls == 0
def test_candle_keyboard_interrupt_is_not_wrapped_and_rolls_back() -> None:
database, connection, provider = _dependencies()
connection.interrupt_on = "INSERT INTO market_data.candle_revisions"
repository = PostgresCandleRepository(connection_provider=provider)
with pytest.raises(KeyboardInterrupt):
repository.store_candle_revision(
venue=VENUE,
candle=_candle(),
observed_at=OBSERVED_AT,
is_final=False,
)
assert database.candle_rows == {}
assert connection.exit_exception_types[-1] is KeyboardInterrupt

638
app/tests/unit/market_data/storage/test_postgres_trade_repository.py Normal file
View File

@@ -0,0 +1,638 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass, field, replace
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from typing import Any
import pytest
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
)
from src.market_data.acquisition.trade_id_sequence import (
SIGNED_TRADE_ID_MAX,
SIGNED_TRADE_ID_MIN,
)
from src.market_data.storage import (
MarketDataStorageConflictError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
MarketDataWriteStatus,
PostgresTradeRepository,
TradeStorageProtocol,
)
VENUE = "dzengi"
SYMBOL = "BTC/USD_LEVERAGE"
EXECUTED_AT = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc)
OBSERVED_AT = EXECUTED_AT + timedelta(seconds=1)
TradeKey = tuple[str, str, int, datetime]
TradeRow = dict[str, Any]
@dataclass
class TransactionalTradeDatabase:
rows: dict[TradeKey, TradeRow] = field(default_factory=dict)
class TransactionalCursor:
def __init__(
self,
connection: TransactionalConnection,
) -> None:
self._connection = connection
self._fetchone_result: object = None
def __enter__(self) -> TransactionalCursor:
return self
def __exit__(self, *args: object) -> None:
return None
def execute(
self,
statement: str,
parameters: tuple[Any, ...],
) -> None:
normalized = " ".join(statement.split())
self._connection.calls.append((normalized, parameters))
if (
self._connection.interrupt_on is not None
and self._connection.interrupt_on in normalized
):
raise KeyboardInterrupt
if (
self._connection.fail_on is not None
and self._connection.fail_on in normalized
):
raise RuntimeError("database failed")
if normalized.startswith("INSERT INTO market_data.trades"):
self._insert(parameters)
return
if normalized.startswith("SELECT price, quantity"):
self._select(parameters)
return
if normalized.startswith("UPDATE market_data.trades"):
self._update(parameters)
return
raise AssertionError(f"Unexpected SQL: {normalized}")
def fetchone(self) -> object:
return self._fetchone_result
def _insert(self, parameters: tuple[Any, ...]) -> None:
(
venue,
symbol,
trade_id,
executed_at,
price,
quantity,
aggressor_side,
source,
first_observed_at,
last_observed_at,
observation_sources,
canonical_schema_version,
) = parameters
key = (venue, symbol, trade_id, executed_at)
working_rows = self._connection.working_rows
if key in working_rows:
self._fetchone_result = None
return
working_rows[key] = {
"price": price,
"quantity": quantity,
"aggressor_side": aggressor_side,
"source": source,
"first_observed_at": first_observed_at,
"last_observed_at": last_observed_at,
"observation_sources": list(observation_sources),
"canonical_schema_version": canonical_schema_version,
}
self._fetchone_result = (1,)
def _select(self, parameters: tuple[Any, ...]) -> None:
key = parameters
row = self._connection.working_rows.get(key)
if row is None:
self._fetchone_result = None
return
self._fetchone_result = (
row["price"],
row["quantity"],
row["aggressor_side"],
row["first_observed_at"],
row["last_observed_at"],
list(row["observation_sources"]),
row["canonical_schema_version"],
)
def _update(self, parameters: tuple[Any, ...]) -> None:
first_observed_at, last_observed_at, sources, *identity = parameters
row = self._connection.working_rows[tuple(identity)]
row["first_observed_at"] = first_observed_at
row["last_observed_at"] = last_observed_at
row["observation_sources"] = list(sources)
self._fetchone_result = None
class TransactionalConnection:
def __init__(self, database: TransactionalTradeDatabase) -> None:
self._database = database
self.calls: list[tuple[str, tuple[Any, ...]]] = []
self.exit_exception_types: list[type[BaseException] | None] = []
self.fail_on: str | None = None
self.interrupt_on: str | None = None
self.working_rows: dict[TradeKey, TradeRow] = {}
def __enter__(self) -> TransactionalConnection:
self.working_rows = deepcopy(self._database.rows)
return self
def __exit__(
self,
exception_type: type[BaseException] | None,
exception: BaseException | None,
traceback: object,
) -> None:
self.exit_exception_types.append(exception_type)
if exception_type is None:
self._database.rows = self.working_rows
return None
def cursor(self) -> TransactionalCursor:
return TransactionalCursor(self)
@dataclass
class RecordingConnectionProvider:
connection: TransactionalConnection
calls: int = 0
def __call__(self) -> TransactionalConnection:
self.calls += 1
return self.connection
def _trade(
*,
symbol: str = SYMBOL,
trade_id: int = 100,
executed_at: datetime = EXECUTED_AT,
price: Decimal = Decimal("64159.45"),
quantity: Decimal = Decimal("0.125"),
aggressor_side: TradeAggressorSide = TradeAggressorSide.BUY,
source: str = "dzengi_websocket_trade",
) -> Trade:
return Trade(
symbol=symbol,
trade_id=trade_id,
price=price,
quantity=quantity,
executed_at=executed_at,
aggressor_side=aggressor_side,
source=source,
)
def _repository() -> tuple[
PostgresTradeRepository,
TransactionalTradeDatabase,
TransactionalConnection,
RecordingConnectionProvider,
]:
database = TransactionalTradeDatabase()
connection = TransactionalConnection(database)
provider = RecordingConnectionProvider(connection)
repository = PostgresTradeRepository(
connection_provider=provider,
)
return repository, database, connection, provider
def _only_row(database: TransactionalTradeDatabase) -> TradeRow:
assert len(database.rows) == 1
return next(iter(database.rows.values()))
def test_repository_matches_trade_storage_protocol() -> None:
repository, _, _, _ = _repository()
assert isinstance(repository, TradeStorageProtocol)
def test_store_trade_inserts_canonical_payload_and_provenance() -> None:
repository, database, _, provider = _repository()
result = repository.store_trade(
venue=" dzengi ",
trade=_trade(symbol=" btc/usd_leverage "),
observed_at=OBSERVED_AT,
)
assert result.status is MarketDataWriteStatus.INSERTED
assert provider.calls == 1
assert tuple(database.rows) == (
(VENUE, SYMBOL, 100, EXECUTED_AT),
)
assert _only_row(database) == {
"price": Decimal("64159.45"),
"quantity": Decimal("0.125"),
"aggressor_side": "buy",
"source": "dzengi_websocket_trade",
"first_observed_at": OBSERVED_AT,
"last_observed_at": OBSERVED_AT,
"observation_sources": ["dzengi_websocket_trade"],
"canonical_schema_version": 1,
}
def test_exact_duplicate_does_not_update_row() -> None:
repository, database, connection, _ = _repository()
trade = _trade()
repository.store_trade(
venue=VENUE,
trade=trade,
observed_at=OBSERVED_AT,
)
calls_before_duplicate = len(connection.calls)
result = repository.store_trade(
venue=VENUE,
trade=trade,
observed_at=OBSERVED_AT,
)
assert result.status is MarketDataWriteStatus.DUPLICATE
assert _only_row(database)["observation_sources"] == [
"dzengi_websocket_trade"
]
assert not any(
statement.startswith("UPDATE market_data.trades")
for statement, _ in connection.calls[calls_before_duplicate:]
)
def test_second_transport_source_updates_provenance_not_market_fact() -> None:
repository, database, _, _ = _repository()
websocket_trade = _trade(source="dzengi_websocket_trade")
rest_trade = replace(websocket_trade, source="dzengi")
repository.store_trade(
venue=VENUE,
trade=websocket_trade,
observed_at=OBSERVED_AT,
)
result = repository.store_trade(
venue=VENUE,
trade=rest_trade,
observed_at=OBSERVED_AT + timedelta(seconds=5),
)
row = _only_row(database)
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert row["source"] == "dzengi_websocket_trade"
assert row["observation_sources"] == [
"dzengi_websocket_trade",
"dzengi",
]
assert row["first_observed_at"] == OBSERVED_AT
assert row["last_observed_at"] == OBSERVED_AT + timedelta(seconds=5)
def test_second_transport_source_updates_provenance_at_same_time() -> None:
repository, database, _, _ = _repository()
websocket_trade = _trade(source="dzengi_websocket_trade")
repository.store_trade(
venue=VENUE,
trade=websocket_trade,
observed_at=OBSERVED_AT,
)
result = repository.store_trade(
venue=VENUE,
trade=replace(websocket_trade, source="dzengi"),
observed_at=OBSERVED_AT,
)
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert _only_row(database)["observation_sources"] == [
"dzengi_websocket_trade",
"dzengi",
]
def test_earlier_observation_moves_only_first_observed_at() -> None:
repository, database, _, _ = _repository()
trade = _trade()
repository.store_trade(
venue=VENUE,
trade=trade,
observed_at=OBSERVED_AT,
)
earlier = OBSERVED_AT - timedelta(seconds=5)
result = repository.store_trade(
venue=VENUE,
trade=trade,
observed_at=earlier,
)
row = _only_row(database)
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert row["first_observed_at"] == earlier
assert row["last_observed_at"] == OBSERVED_AT
@pytest.mark.parametrize(
"conflicting_trade",
(
_trade(price=Decimal("64160")),
_trade(quantity=Decimal("1")),
_trade(aggressor_side=TradeAggressorSide.SELL),
),
)
def test_same_identity_with_different_market_fact_is_conflict(
conflicting_trade: Trade,
) -> None:
repository, database, connection, _ = _repository()
original = _trade()
repository.store_trade(
venue=VENUE,
trade=original,
observed_at=OBSERVED_AT,
)
with pytest.raises(MarketDataStorageConflictError):
repository.store_trade(
venue=VENUE,
trade=conflicting_trade,
observed_at=OBSERVED_AT,
)
assert _only_row(database)["price"] == original.price
assert connection.exit_exception_types[-1] is (
MarketDataStorageConflictError
)
def test_empty_batch_returns_zero_without_borrowing_connection() -> None:
repository, _, _, provider = _repository()
result = repository.store_trades(
venue=VENUE,
trades=(),
observed_at=OBSERVED_AT,
)
assert result.total_count == 0
assert provider.calls == 0
def test_batch_returns_insert_duplicate_and_provenance_counts() -> None:
repository, _, _, _ = _repository()
duplicate = _trade(symbol="B", trade_id=2)
provenance = _trade(symbol="C", trade_id=3)
repository.store_trade(
venue=VENUE,
trade=duplicate,
observed_at=OBSERVED_AT + timedelta(seconds=1),
)
repository.store_trade(
venue=VENUE,
trade=provenance,
observed_at=OBSERVED_AT,
)
result = repository.store_trades(
venue=VENUE,
trades=(
replace(provenance, source="dzengi"),
duplicate,
_trade(symbol="A", trade_id=1),
),
observed_at=OBSERVED_AT + timedelta(seconds=1),
)
assert result.inserted_count == 1
assert result.duplicate_count == 1
assert result.provenance_updated_count == 1
assert result.total_count == 3
def test_batch_uses_stable_identity_order() -> None:
repository, _, connection, _ = _repository()
repository.store_trades(
venue=VENUE,
trades=(
_trade(symbol="C", trade_id=3),
_trade(symbol="A", trade_id=1),
_trade(symbol="B", trade_id=2),
),
observed_at=OBSERVED_AT,
)
inserted_symbols = tuple(
parameters[1]
for statement, parameters in connection.calls
if statement.startswith("INSERT INTO market_data.trades")
)
assert inserted_symbols == ("A", "B", "C")
def test_batch_conflict_rolls_back_preceding_insert() -> None:
repository, database, connection, _ = _repository()
existing = _trade(symbol="B", trade_id=2)
repository.store_trade(
venue=VENUE,
trade=existing,
observed_at=OBSERVED_AT,
)
with pytest.raises(MarketDataStorageConflictError):
repository.store_trades(
venue=VENUE,
trades=(
_trade(symbol="A", trade_id=1),
replace(existing, price=Decimal("999")),
),
observed_at=OBSERVED_AT,
)
assert tuple(database.rows) == (
(VENUE, "B", 2, EXECUTED_AT),
)
assert connection.exit_exception_types[-1] is (
MarketDataStorageConflictError
)
def test_batch_accepts_signed_rollover_boundary_as_distinct_trades() -> None:
repository, database, _, _ = _repository()
result = repository.store_trades(
venue=VENUE,
trades=(
_trade(
trade_id=SIGNED_TRADE_ID_MAX,
executed_at=EXECUTED_AT,
),
_trade(
trade_id=SIGNED_TRADE_ID_MIN,
executed_at=EXECUTED_AT + timedelta(milliseconds=1),
),
),
observed_at=OBSERVED_AT,
)
assert result.inserted_count == 2
assert len(database.rows) == 2
def test_database_error_is_wrapped_and_transaction_is_rolled_back() -> None:
repository, database, connection, _ = _repository()
connection.fail_on = "INSERT INTO market_data.trades"
with pytest.raises(MarketDataStorageOperationError) as error_info:
repository.store_trade(
venue=VENUE,
trade=_trade(),
observed_at=OBSERVED_AT,
)
assert isinstance(error_info.value.__cause__, RuntimeError)
assert database.rows == {}
assert connection.exit_exception_types[-1] is RuntimeError
def test_keyboard_interrupt_is_not_wrapped_and_rolls_back() -> None:
repository, database, connection, _ = _repository()
connection.interrupt_on = "INSERT INTO market_data.trades"
with pytest.raises(KeyboardInterrupt):
repository.store_trade(
venue=VENUE,
trade=_trade(),
observed_at=OBSERVED_AT,
)
assert database.rows == {}
assert connection.exit_exception_types[-1] is KeyboardInterrupt
@pytest.mark.parametrize("venue", ("", " ", None, 1))
def test_rejects_invalid_venue_without_connection(venue: object) -> None:
repository, _, _, provider = _repository()
with pytest.raises(MarketDataStorageValidationError, match="venue"):
repository.store_trade(
venue=venue, # type: ignore[arg-type]
trade=_trade(),
observed_at=OBSERVED_AT,
)
assert provider.calls == 0
def test_rejects_naive_observed_at_without_connection() -> None:
repository, _, _, provider = _repository()
with pytest.raises(
MarketDataStorageValidationError,
match="observed_at",
):
repository.store_trade(
venue=VENUE,
trade=_trade(),
observed_at=OBSERVED_AT.replace(tzinfo=None),
)
assert provider.calls == 0
def test_rejects_naive_executed_at_without_connection() -> None:
repository, _, _, provider = _repository()
with pytest.raises(
MarketDataStorageValidationError,
match="executed_at",
):
repository.store_trade(
venue=VENUE,
trade=_trade(executed_at=EXECUTED_AT.replace(tzinfo=None)),
observed_at=OBSERVED_AT,
)
assert provider.calls == 0
@pytest.mark.parametrize(
"invalid_trade",
(
_trade(trade_id=SIGNED_TRADE_ID_MAX + 1),
_trade(price=Decimal("0")),
_trade(quantity=Decimal("NaN")),
replace(_trade(), aggressor_side="buy"), # type: ignore[arg-type]
),
)
def test_rejects_invalid_canonical_trade_before_connection(
invalid_trade: Trade,
) -> None:
repository, _, _, provider = _repository()
with pytest.raises(MarketDataStorageValidationError):
repository.store_trade(
venue=VENUE,
trade=invalid_trade,
observed_at=OBSERVED_AT,
)
assert provider.calls == 0
def test_batch_rejects_non_tuple_without_connection() -> None:
repository, _, _, provider = _repository()
with pytest.raises(MarketDataStorageValidationError, match="tuple"):
repository.store_trades(
venue=VENUE,
trades=[_trade()], # type: ignore[arg-type]
observed_at=OBSERVED_AT,
)
assert provider.calls == 0
def test_batch_validates_every_trade_before_connection() -> None:
repository, _, _, provider = _repository()
with pytest.raises(MarketDataStorageValidationError):
repository.store_trades(
venue=VENUE,
trades=(
_trade(symbol="A"),
_trade(symbol=" "),
),
observed_at=OBSERVED_AT,
)
assert provider.calls == 0

170
app/tests/unit/market_data/storage/test_trade_storage_observation_sink.py Normal file
View File

@@ -0,0 +1,170 @@
from __future__ import annotations
from datetime import datetime, timezone
from decimal import Decimal
import pytest
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
)
from src.market_data.storage import (
MarketDataBatchWriteResult,
MarketDataStorageValidationError,
MarketDataWriteResult,
MarketDataWriteStatus,
TradeStorageObservationSink,
)
VENUE = "dzengi"
OBSERVED_AT = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc)
def make_trade() -> Trade:
return Trade(
symbol="BTC/USD_LEVERAGE",
trade_id=100,
price=Decimal("64555.55"),
quantity=Decimal("0.002"),
executed_at=OBSERVED_AT,
aggressor_side=TradeAggressorSide.BUY,
source="dzengi_websocket_trade",
)
class RecordingTradeStorage:
def __init__(
self,
*,
result: object | None = None,
error: Exception | None = None,
) -> None:
self.result = result or MarketDataWriteResult(
status=MarketDataWriteStatus.INSERTED,
)
self.error = error
self.calls: list[tuple[str, Trade, datetime]] = []
def store_trade(
self,
*,
venue: str,
trade: Trade,
observed_at: datetime,
) -> MarketDataWriteResult:
self.calls.append((venue, trade, observed_at))
if self.error is not None:
raise self.error
return self.result # type: ignore[return-value]
def store_trades(
self,
*,
venue: str,
trades: tuple[Trade, ...],
observed_at: datetime,
) -> MarketDataBatchWriteResult:
raise AssertionError("Runtime persists observations one by one")
@pytest.mark.parametrize(
"status",
tuple(MarketDataWriteStatus),
)
def test_forwards_observation_and_accepts_all_success_statuses(
status: MarketDataWriteStatus,
) -> None:
storage = RecordingTradeStorage(
result=MarketDataWriteResult(status=status),
)
sink = TradeStorageObservationSink(
trade_storage=storage,
venue=" dzengi ",
clock=lambda: OBSERVED_AT,
)
trade = make_trade()
result = sink.persist(trade)
assert result is None
assert storage.calls == [
(
VENUE,
trade,
OBSERVED_AT,
)
]
def test_implements_acquisition_side_sink_protocol() -> None:
sink = TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(),
venue=VENUE,
)
assert isinstance(sink, TradeObservationSinkProtocol)
assert not hasattr(sink, "__dict__")
def test_storage_error_is_not_wrapped() -> None:
storage_error = RuntimeError("storage failed")
sink = TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(error=storage_error),
venue=VENUE,
clock=lambda: OBSERVED_AT,
)
with pytest.raises(
RuntimeError,
match="storage failed",
) as error_info:
sink.persist(make_trade())
assert error_info.value is storage_error
def test_rejects_invalid_storage_result() -> None:
sink = TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(result=object()),
venue=VENUE,
clock=lambda: OBSERVED_AT,
)
with pytest.raises(
TypeError,
match="MarketDataWriteResult",
):
sink.persist(make_trade())
def test_rejects_invalid_dependencies() -> None:
with pytest.raises(TypeError, match="TradeStorageProtocol"):
TradeStorageObservationSink(
trade_storage=object(), # type: ignore[arg-type]
venue=VENUE,
)
with pytest.raises(TypeError, match="clock must be callable"):
TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(),
venue=VENUE,
clock=object(), # type: ignore[arg-type]
)
def test_rejects_empty_venue_at_construction() -> None:
with pytest.raises(
MarketDataStorageValidationError,
match="venue must not be empty",
):
TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(),
venue=" ",
)

295
app/tests/unit/storage/test_migrations.py Normal file
View File

@@ -0,0 +1,295 @@
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
import pytest
from src.storage.exceptions import StorageMigrationError
from src.storage.migrations import (
STORAGE_MIGRATION_ADVISORY_LOCK_ID,
STORAGE_MIGRATIONS,
StorageMigration,
StorageMigrationRunner,
)
@dataclass
class RecordingCursor:
applied_rows: list[tuple[int, str]] = field(default_factory=list)
fail_on: str | None = None
calls: list[tuple[str, object | None]] = field(default_factory=list)
def __enter__(self) -> RecordingCursor:
return self
def __exit__(self, *args: object) -> None:
return None
def execute(
self,
statement: str,
parameters: object | None = None,
) -> None:
normalized = " ".join(statement.split())
self.calls.append((normalized, parameters))
if self.fail_on is not None and self.fail_on in normalized:
raise RuntimeError("database failed")
def fetchall(self) -> list[tuple[int, str]]:
return list(self.applied_rows)
@dataclass
class RecordingConnection:
cursor_value: RecordingCursor
entered: int = 0
exited: int = 0
def __enter__(self) -> RecordingConnection:
self.entered += 1
return self
def __exit__(self, *args: object) -> None:
self.exited += 1
return None
def cursor(self) -> RecordingCursor:
return self.cursor_value
@dataclass
class RecordingProvider:
connection: RecordingConnection
calls: int = 0
def __call__(self) -> RecordingConnection:
self.calls += 1
return self.connection
def _runner(
*,
applied_rows: list[tuple[int, str]] | None = None,
fail_on: str | None = None,
migrations: tuple[StorageMigration, ...] = STORAGE_MIGRATIONS,
) -> tuple[
StorageMigrationRunner,
RecordingCursor,
RecordingConnection,
RecordingProvider,
]:
cursor = RecordingCursor(
applied_rows=applied_rows or [],
fail_on=fail_on,
)
connection = RecordingConnection(cursor)
provider = RecordingProvider(connection)
runner = StorageMigrationRunner(
connection_provider=provider,
migrations=migrations,
)
return runner, cursor, connection, provider
def test_default_migrations_have_stable_order_and_names() -> None:
assert tuple(
(migration.version, migration.name)
for migration in STORAGE_MIGRATIONS
) == (
(1, "create_market_data_schema"),
(2, "create_canonical_trades"),
(3, "create_canonical_quotes"),
(4, "create_canonical_candle_revisions"),
(5, "add_trade_observation_sources"),
(6, "add_quote_and_candle_observation_sources"),
(7, "create_market_data_partition_registry"),
)
def test_default_schema_defines_partitions_identities_and_constraints() -> None:
sql = "\n".join(
statement
for migration in STORAGE_MIGRATIONS
for statement in migration.statements
)
assert "CREATE SCHEMA IF NOT EXISTS market_data" in sql
assert "CREATE TABLE market_data.trades" in sql
assert "PRIMARY KEY (venue, symbol, trade_id, executed_at)" in sql
assert "trade_id BETWEEN -2147483648 AND 2147483647" in sql
assert sql.count("CHECK (BTRIM(venue) <> '')") == 3
assert sql.count("CHECK (BTRIM(symbol) <> '')") == 3
assert sql.count("CHECK (BTRIM(source) <> '')") == 3
assert "PARTITION BY RANGE (executed_at)" in sql
assert "CREATE TABLE market_data.quotes" in sql
assert "PRIMARY KEY (venue, symbol, received_at)" in sql
assert "PARTITION BY RANGE (received_at)" in sql
assert "CREATE TABLE market_data.candle_revisions" in sql
assert "open_time,\n observed_at" in sql
assert "PARTITION BY RANGE (open_time)" in sql
assert sql.count("PARTITION OF") == 3
assert "ADD COLUMN observation_sources TEXT[]" in sql
assert "SET observation_sources = ARRAY[source]" in sql
assert "CARDINALITY(observation_sources) > 0" in sql
assert "ALTER TABLE market_data.quotes" in sql
assert "ALTER TABLE market_data.candle_revisions" in sql
assert "quotes_observation_sources_not_empty" in sql
assert "candle_revisions_observation_sources_not_empty" in sql
assert "CREATE TABLE market_data.partition_registry" in sql
assert "PRIMARY KEY (data_type, range_start)" in sql
assert "UNIQUE (partition_name)" in sql
assert "partition_bound TEXT NOT NULL" in sql
assert "BTRIM(partition_bound) <> ''" in sql
assert "range_end > range_start" in sql
def test_run_locks_and_applies_every_pending_migration_in_order() -> None:
runner, cursor, connection, provider = _runner()
result = runner.run()
assert result == (1, 2, 3, 4, 5, 6, 7)
assert provider.calls == 1
assert connection.entered == 1
assert connection.exited == 1
assert cursor.calls[0] == (
"SELECT pg_advisory_xact_lock(%s)",
(STORAGE_MIGRATION_ADVISORY_LOCK_ID,),
)
assert "public.storage_schema_migrations" in cursor.calls[1][0]
inserted_versions = tuple(
parameters[0]
for statement, parameters in cursor.calls
if statement.startswith(
"INSERT INTO public.storage_schema_migrations"
)
and isinstance(parameters, tuple)
)
assert inserted_versions == (1, 2, 3, 4, 5, 6, 7)
def test_run_skips_already_applied_migrations() -> None:
applied = [
(migration.version, migration.name)
for migration in STORAGE_MIGRATIONS
]
runner, cursor, _, _ = _runner(applied_rows=applied)
result = runner.run()
assert result == ()
assert not any(
statement.startswith("CREATE SCHEMA")
or statement.startswith("CREATE TABLE market_data")
for statement, _ in cursor.calls
)
def test_run_applies_only_migrations_after_existing_prefix() -> None:
applied = [
(migration.version, migration.name)
for migration in STORAGE_MIGRATIONS[:2]
]
runner, cursor, _, _ = _runner(applied_rows=applied)
result = runner.run()
assert result == (3, 4, 5, 6, 7)
inserted_versions = tuple(
parameters[0]
for statement, parameters in cursor.calls
if statement.startswith(
"INSERT INTO public.storage_schema_migrations"
)
and isinstance(parameters, tuple)
)
assert inserted_versions == (3, 4, 5, 6, 7)
def test_run_rejects_unknown_applied_version() -> None:
runner, _, _, _ = _runner(applied_rows=[(99, "future")])
with pytest.raises(StorageMigrationError, match="unknown.*99"):
runner.run()
def test_run_rejects_changed_applied_name() -> None:
runner, _, _, _ = _runner(applied_rows=[(1, "renamed")])
with pytest.raises(StorageMigrationError, match="name mismatch"):
runner.run()
def test_run_rejects_non_prefix_applied_history() -> None:
second = STORAGE_MIGRATIONS[1]
runner, _, _, _ = _runner(
applied_rows=[(second.version, second.name)],
)
with pytest.raises(StorageMigrationError, match="ordered prefix"):
runner.run()
def test_database_error_is_wrapped() -> None:
runner, _, _, _ = _runner(fail_on="CREATE SCHEMA")
with pytest.raises(StorageMigrationError) as error_info:
runner.run()
assert isinstance(error_info.value.__cause__, RuntimeError)
def test_keyboard_interrupt_is_not_wrapped() -> None:
def interrupted_provider() -> Any:
raise KeyboardInterrupt
runner = StorageMigrationRunner(
connection_provider=interrupted_provider,
)
with pytest.raises(KeyboardInterrupt):
runner.run()
def test_runner_rejects_duplicate_versions() -> None:
migration = StorageMigration(
version=1,
name="one",
statements=("SELECT 1",),
)
with pytest.raises(ValueError, match="unique"):
StorageMigrationRunner(
connection_provider=lambda: None, # type: ignore[arg-type]
migrations=(migration, migration),
)
def test_runner_rejects_out_of_order_versions() -> None:
first = StorageMigration(1, "one", ("SELECT 1",))
second = StorageMigration(2, "two", ("SELECT 2",))
with pytest.raises(ValueError, match="ordered"):
StorageMigrationRunner(
connection_provider=lambda: None, # type: ignore[arg-type]
migrations=(second, first),
)
@pytest.mark.parametrize(
"arguments",
(
{"version": 0, "name": "zero", "statements": ("SELECT 0",)},
{"version": True, "name": "one", "statements": ("SELECT 1",)},
{"version": 1, "name": " ", "statements": ("SELECT 1",)},
{"version": 1, "name": "one", "statements": ()},
{"version": 1, "name": "one", "statements": (" ",)},
),
)
def test_migration_rejects_invalid_definition(
arguments: dict[str, object],
) -> None:
with pytest.raises(ValueError):
StorageMigration(**arguments) # type: ignore[arg-type]

236
app/tests/unit/storage/test_postgres_pool.py Normal file
View File

@@ -0,0 +1,236 @@
from __future__ import annotations
from contextlib import nullcontext
from typing import Any
import pytest
from src.storage.exceptions import PostgresConnectionPoolError
from src.storage.postgres_pool import PostgresConnectionPool
class RecordingPool:
def __init__(
self,
*,
open_error: BaseException | None = None,
close_error: Exception | None = None,
) -> None:
self.close_calls: list[float] = []
self.close_error = close_error
self.connection_calls: list[float] = []
self.open_calls: list[tuple[bool, float]] = []
self.open_error = open_error
self.connection_value = object()
def open(self, *, wait: bool, timeout: float) -> None:
self.open_calls.append((wait, timeout))
if self.open_error is not None:
raise self.open_error
def connection(self, *, timeout: float):
self.connection_calls.append(timeout)
return nullcontext(self.connection_value)
def close(self, *, timeout: float) -> None:
self.close_calls.append(timeout)
if self.close_error is not None:
raise self.close_error
class RecordingPoolFactory:
def __init__(self, pool: RecordingPool) -> None:
self.calls: list[dict[str, Any]] = []
self.pool = pool
def __call__(self, **kwargs: Any) -> RecordingPool:
self.calls.append(kwargs)
return self.pool
def _connection_pool(
*,
pool: RecordingPool | None = None,
) -> tuple[
PostgresConnectionPool,
RecordingPool,
RecordingPoolFactory,
]:
recording_pool = pool or RecordingPool()
factory = RecordingPoolFactory(recording_pool)
connection_pool = PostgresConnectionPool(
conninfo="postgresql://db/dzentra",
min_size=2,
max_size=6,
timeout_seconds=7.5,
name="market-data",
pool_factory=factory,
)
return connection_pool, recording_pool, factory
def test_constructor_does_not_create_or_open_pool() -> None:
connection_pool, pool, factory = _connection_pool()
assert connection_pool.is_open is False
assert factory.calls == []
assert pool.open_calls == []
def test_open_creates_pool_with_explicit_configuration_and_waits() -> None:
connection_pool, pool, factory = _connection_pool()
connection_pool.open()
assert connection_pool.is_open is True
assert factory.calls == [
{
"conninfo": "postgresql://db/dzentra",
"min_size": 2,
"max_size": 6,
"timeout": 7.5,
"kwargs": {"autocommit": False},
"name": "market-data",
"open": False,
}
]
assert pool.open_calls == [(True, 7.5)]
def test_repeated_open_is_no_op() -> None:
connection_pool, pool, factory = _connection_pool()
connection_pool.open()
connection_pool.open()
assert len(factory.calls) == 1
assert pool.open_calls == [(True, 7.5)]
def test_connection_requires_open_pool() -> None:
connection_pool, _, _ = _connection_pool()
with pytest.raises(PostgresConnectionPoolError, match="not open"):
connection_pool.connection()
def test_connection_borrows_from_pool_with_timeout() -> None:
connection_pool, pool, _ = _connection_pool()
connection_pool.open()
with connection_pool.connection() as connection:
assert connection is pool.connection_value
assert pool.connection_calls == [7.5]
def test_close_is_idempotent() -> None:
connection_pool, pool, _ = _connection_pool()
connection_pool.open()
connection_pool.close()
connection_pool.close()
assert connection_pool.is_open is False
assert pool.close_calls == [7.5]
def test_open_failure_closes_partial_pool_and_is_retryable() -> None:
pool = RecordingPool(open_error=RuntimeError("database unavailable"))
connection_pool, _, factory = _connection_pool(pool=pool)
with pytest.raises(PostgresConnectionPoolError) as error_info:
connection_pool.open()
assert isinstance(error_info.value.__cause__, RuntimeError)
assert connection_pool.is_open is False
assert pool.close_calls == [7.5]
pool.open_error = None
connection_pool.open()
assert len(factory.calls) == 2
assert connection_pool.is_open is True
def test_open_failure_preserves_cleanup_failure_as_note() -> None:
pool = RecordingPool(
open_error=RuntimeError("open failed"),
close_error=RuntimeError("close failed"),
)
connection_pool, _, _ = _connection_pool(pool=pool)
with pytest.raises(PostgresConnectionPoolError) as error_info:
connection_pool.open()
cause = error_info.value.__cause__
assert isinstance(cause, RuntimeError)
assert cause.__notes__ == [
"PostgreSQL pool cleanup also failed: RuntimeError."
]
def test_close_failure_leaves_wrapper_closed() -> None:
pool = RecordingPool(close_error=RuntimeError("close failed"))
connection_pool, _, _ = _connection_pool(pool=pool)
connection_pool.open()
with pytest.raises(PostgresConnectionPoolError) as error_info:
connection_pool.close()
assert isinstance(error_info.value.__cause__, RuntimeError)
assert connection_pool.is_open is False
connection_pool.close()
assert pool.close_calls == [7.5]
def test_keyboard_interrupt_is_not_wrapped() -> None:
def interrupted_factory(**kwargs: object) -> object:
raise KeyboardInterrupt
connection_pool = PostgresConnectionPool(
conninfo="postgresql://db/dzentra",
pool_factory=interrupted_factory,
)
with pytest.raises(KeyboardInterrupt):
connection_pool.open()
def test_keyboard_interrupt_during_pool_open_closes_partial_pool() -> None:
pool = RecordingPool(open_error=KeyboardInterrupt())
connection_pool, _, _ = _connection_pool(pool=pool)
with pytest.raises(KeyboardInterrupt):
connection_pool.open()
assert connection_pool.is_open is False
assert pool.open_calls == [(True, 7.5)]
assert pool.close_calls == [7.5]
@pytest.mark.parametrize(
("override", "message"),
(
({"conninfo": ""}, "conninfo"),
({"min_size": 0}, "min_size"),
({"min_size": True}, "min_size"),
({"min_size": 2, "max_size": 1}, "max_size"),
({"timeout_seconds": 0}, "timeout_seconds"),
({"timeout_seconds": float("inf")}, "timeout_seconds"),
({"name": " "}, "name"),
),
)
def test_constructor_rejects_invalid_configuration(
override: dict[str, object],
message: str,
) -> None:
arguments: dict[str, object] = {
"conninfo": "postgresql://db/dzentra",
}
arguments.update(override)
with pytest.raises(ValueError, match=message):
PostgresConnectionPool(**arguments) # type: ignore[arg-type]

190
app/tests/unit/test_postgres_market_data_support.py Normal file
View File

@@ -0,0 +1,190 @@
from __future__ import annotations
from typing import Any
import pytest
from psycopg.conninfo import conninfo_to_dict
from tests.support.postgres_market_data import (
POSTGRES_TEST_APPLICATION_NAME,
POSTGRES_TEST_CONTROL_APPLICATION_NAME,
load_postgres_test_settings,
reset_postgres_test_database,
)
class RecordingCursor:
def __init__(
self,
*,
identity: tuple[str, str],
statements: list[str],
) -> None:
self._identity = identity
self._statements = statements
def __enter__(self) -> RecordingCursor:
return self
def __exit__(self, *_: object) -> None:
return None
def execute(
self,
statement: object,
parameters: object = None,
) -> None:
del parameters
self._statements.append(" ".join(str(statement).split()))
def fetchone(self) -> tuple[str, str]:
return self._identity
class RecordingControlConnection:
def __init__(
self,
*,
identity: tuple[str, str],
autocommit: bool = True,
) -> None:
self.autocommit = autocommit
self.identity = identity
self.statements: list[str] = []
def cursor(self) -> RecordingCursor:
return RecordingCursor(
identity=self.identity,
statements=self.statements,
)
def test_postgres_harness_is_disabled_without_explicit_flag() -> None:
assert load_postgres_test_settings({}) is None
def test_postgres_harness_requires_exact_opt_in() -> None:
with pytest.raises(ValueError, match="exactly '1'"):
load_postgres_test_settings(
{"DZENTRA_RUN_POSTGRES_TESTS": "true"}
)
def test_postgres_harness_requires_explicit_dsn() -> None:
with pytest.raises(ValueError, match="DZENTRA_TEST_POSTGRES_DSN"):
load_postgres_test_settings(
{"DZENTRA_RUN_POSTGRES_TESTS": "1"}
)
@pytest.mark.parametrize(
"dsn",
(
"postgresql://localhost/production",
"postgresql://db.example.test/dzentra_test_060_27",
"hostaddr=192.0.2.10 dbname=dzentra_test_060_27",
"postgresql://localhost/dzentra_test_unsafe-name",
"service=production dbname=dzentra_test_shadow",
"dbname=dzentra_test_implicit_endpoint",
"host=/tmp,db.example.test dbname=dzentra_test_socket_fallback",
(
"hostaddr=127.0.0.1,192.0.2.10 "
"dbname=dzentra_test_address_fallback"
),
"host=localhost, dbname=dzentra_test_implicit_fallback",
),
)
def test_postgres_harness_rejects_unsafe_database_targets(
dsn: str,
) -> None:
with pytest.raises(ValueError):
load_postgres_test_settings(
{
"DZENTRA_RUN_POSTGRES_TESTS": "1",
"DZENTRA_TEST_POSTGRES_DSN": dsn,
}
)
@pytest.mark.parametrize(
"dsn",
(
"postgresql://localhost:5544/dzentra_test_060_27",
"host=/tmp dbname=dzentra_test_socket user=test",
"hostaddr=127.0.0.1 dbname=dzentra_test_loopback",
"hostaddr=::1 dbname=dzentra_test_ipv6",
),
)
def test_postgres_harness_accepts_only_explicit_local_test_database(
dsn: str,
) -> None:
settings = load_postgres_test_settings(
{
"DZENTRA_RUN_POSTGRES_TESTS": "1",
"DZENTRA_TEST_POSTGRES_DSN": dsn,
}
)
assert settings is not None
parsed = conninfo_to_dict(settings.dsn)
assert parsed["dbname"] == settings.database_name
assert parsed["application_name"] == POSTGRES_TEST_APPLICATION_NAME
assert parsed["connect_timeout"] == "5"
def test_postgres_harness_does_not_accept_environment_endpoint_fallback(
) -> None:
with pytest.raises(ValueError, match="explicit local host"):
load_postgres_test_settings(
{
"DZENTRA_RUN_POSTGRES_TESTS": "1",
"DZENTRA_TEST_POSTGRES_DSN": (
"dbname=dzentra_test_environment_fallback"
),
"PGHOST": "localhost",
}
)
def test_postgres_reset_revalidates_database_and_control_identity() -> None:
database_name = "dzentra_test_cleanup"
connection: Any = RecordingControlConnection(
identity=(
database_name,
POSTGRES_TEST_CONTROL_APPLICATION_NAME,
)
)
reset_postgres_test_database(
connection,
expected_database_name=database_name,
)
assert connection.statements == [
"SELECT current_database(), current_setting('application_name')",
"DROP SCHEMA IF EXISTS market_data CASCADE",
"DROP TABLE IF EXISTS public.storage_schema_migrations",
]
@pytest.mark.parametrize(
"identity",
(
("production", POSTGRES_TEST_CONTROL_APPLICATION_NAME),
("dzentra_test_cleanup", "unexpected-application"),
),
)
def test_postgres_reset_refuses_unvalidated_connection_before_drop(
identity: tuple[str, str],
) -> None:
connection: Any = RecordingControlConnection(identity=identity)
with pytest.raises(RuntimeError, match="Refusing destructive reset"):
reset_postgres_test_database(
connection,
expected_database_name="dzentra_test_cleanup",
)
assert connection.statements == [
"SELECT current_database(), current_setting('application_name')",
]