1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
|
"""Tests for storage module."""
import pytest
from decimal import Decimal
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock
from shared.models import Candle
from data_collector.storage import CandleStorage
def _make_candle(symbol: str = "BTCUSDT") -> Candle:
return Candle(
symbol=symbol,
timeframe="1m",
open_time=datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
open=Decimal("30000"),
high=Decimal("30100"),
low=Decimal("29900"),
close=Decimal("30050"),
volume=Decimal("1.5"),
)
@pytest.mark.asyncio
async def test_storage_saves_to_db_and_publishes():
"""Verify that store() calls insert_candle on db and publish on broker."""
mock_db = MagicMock()
mock_db.insert_candle = AsyncMock()
mock_broker = MagicMock()
mock_broker.publish = AsyncMock()
storage = CandleStorage(db=mock_db, broker=mock_broker)
candle = _make_candle()
await storage.store(candle)
mock_db.insert_candle.assert_called_once_with(candle)
mock_broker.publish.assert_called_once()
stream_arg = mock_broker.publish.call_args[0][0]
assert stream_arg == "candles.BTCUSDT"
data_arg = mock_broker.publish.call_args[0][1]
assert data_arg["type"] == "CANDLE"
assert data_arg["data"]["symbol"] == "BTCUSDT"
@pytest.mark.asyncio
async def test_storage_batch_store():
"""Verify that store_batch() calls store for each candle."""
mock_db = MagicMock()
mock_db.insert_candle = AsyncMock()
mock_broker = MagicMock()
mock_broker.publish = AsyncMock()
storage = CandleStorage(db=mock_db, broker=mock_broker)
candles = [_make_candle() for _ in range(3)]
await storage.store_batch(candles)
assert mock_db.insert_candle.call_count == 3
assert mock_broker.publish.call_count == 3
|
