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

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=" ",
)