Merge pull request #32 from pselamy/fix/5
feat: add Redis Streams event publisher for trade events (#5)
This commit is contained in:
Patrick Selamy
@@ -20,6 +20,12 @@ from polymarket_insider_tracker.ingestor.models import (
|
||||
TradeEvent,
|
||||
derive_category,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.publisher import (
|
||||
ConsumerGroupExistsError,
|
||||
EventPublisher,
|
||||
PublisherError,
|
||||
StreamEntry,
|
||||
)
|
||||
from polymarket_insider_tracker.ingestor.websocket import (
|
||||
ConnectionState,
|
||||
StreamStats as WebSocketStreamStats,
|
||||
@@ -45,6 +51,11 @@ __all__ = [
|
||||
"Token",
|
||||
"TradeEvent",
|
||||
"derive_category",
|
||||
# Publisher
|
||||
"ConsumerGroupExistsError",
|
||||
"EventPublisher",
|
||||
"PublisherError",
|
||||
"StreamEntry",
|
||||
# WebSocket
|
||||
"ConnectionState",
|
||||
"WebSocketStreamStats",
|
||||
|
||||
@@ -0,0 +1,419 @@
|
||||
"""Redis Streams event publisher for trade events.
|
||||
|
||||
This module provides an event publisher that writes normalized trade events
|
||||
to Redis Streams, enabling downstream consumers to process events independently.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Any, Literal
|
||||
|
||||
from redis.asyncio import Redis
|
||||
from redis.exceptions import ResponseError
|
||||
|
||||
from .models import TradeEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_STREAM_NAME = "trades"
|
||||
DEFAULT_MAX_LEN = 100_000 # 100k events
|
||||
DEFAULT_BLOCK_MS = 1000
|
||||
DEFAULT_COUNT = 10
|
||||
|
||||
|
||||
class PublisherError(Exception):
|
||||
"""Base exception for publisher errors."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ConsumerGroupExistsError(PublisherError):
|
||||
"""Raised when trying to create a consumer group that already exists."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamEntry:
|
||||
"""Represents an entry read from a Redis Stream."""
|
||||
|
||||
entry_id: str
|
||||
event: TradeEvent
|
||||
|
||||
|
||||
def _serialize_trade_event(event: TradeEvent) -> dict[str, str]:
|
||||
"""Serialize a TradeEvent to a dict suitable for Redis Streams.
|
||||
|
||||
Redis Streams require string key-value pairs, so we convert all
|
||||
values to strings.
|
||||
|
||||
Args:
|
||||
event: The TradeEvent to serialize.
|
||||
|
||||
Returns:
|
||||
Dictionary with string keys and values.
|
||||
"""
|
||||
return {
|
||||
"market_id": event.market_id,
|
||||
"trade_id": event.trade_id,
|
||||
"wallet_address": event.wallet_address,
|
||||
"side": event.side,
|
||||
"outcome": event.outcome,
|
||||
"outcome_index": str(event.outcome_index),
|
||||
"price": str(event.price),
|
||||
"size": str(event.size),
|
||||
"timestamp": event.timestamp.isoformat(),
|
||||
"asset_id": event.asset_id,
|
||||
"market_slug": event.market_slug,
|
||||
"event_slug": event.event_slug,
|
||||
"event_title": event.event_title,
|
||||
"trader_name": event.trader_name,
|
||||
"trader_pseudonym": event.trader_pseudonym,
|
||||
}
|
||||
|
||||
|
||||
def _deserialize_trade_event(data: dict[bytes | str, bytes | str]) -> TradeEvent:
|
||||
"""Deserialize a TradeEvent from Redis Stream data.
|
||||
|
||||
Args:
|
||||
data: The raw data from Redis Stream (may have bytes keys/values).
|
||||
|
||||
Returns:
|
||||
TradeEvent instance.
|
||||
"""
|
||||
# Convert bytes to strings if needed
|
||||
decoded: dict[str, str] = {}
|
||||
for k, v in data.items():
|
||||
key = k.decode() if isinstance(k, bytes) else k
|
||||
value = v.decode() if isinstance(v, bytes) else v
|
||||
decoded[key] = value
|
||||
|
||||
# Parse timestamp
|
||||
timestamp_str = decoded.get("timestamp", "")
|
||||
try:
|
||||
timestamp = datetime.fromisoformat(timestamp_str)
|
||||
except (ValueError, TypeError):
|
||||
timestamp = datetime.now(UTC)
|
||||
|
||||
# Parse side
|
||||
side_raw = decoded.get("side", "BUY").upper()
|
||||
side: Literal["BUY", "SELL"] = "BUY" if side_raw == "BUY" else "SELL"
|
||||
|
||||
return TradeEvent(
|
||||
market_id=decoded.get("market_id", ""),
|
||||
trade_id=decoded.get("trade_id", ""),
|
||||
wallet_address=decoded.get("wallet_address", ""),
|
||||
side=side,
|
||||
outcome=decoded.get("outcome", ""),
|
||||
outcome_index=int(decoded.get("outcome_index", "0")),
|
||||
price=Decimal(decoded.get("price", "0")),
|
||||
size=Decimal(decoded.get("size", "0")),
|
||||
timestamp=timestamp,
|
||||
asset_id=decoded.get("asset_id", ""),
|
||||
market_slug=decoded.get("market_slug", ""),
|
||||
event_slug=decoded.get("event_slug", ""),
|
||||
event_title=decoded.get("event_title", ""),
|
||||
trader_name=decoded.get("trader_name", ""),
|
||||
trader_pseudonym=decoded.get("trader_pseudonym", ""),
|
||||
)
|
||||
|
||||
|
||||
class EventPublisher:
|
||||
"""Event publisher using Redis Streams.
|
||||
|
||||
This class wraps Redis Streams to provide:
|
||||
- Publishing single or batch trade events
|
||||
- Consumer group management
|
||||
- Event reading for consumers
|
||||
|
||||
Example:
|
||||
```python
|
||||
redis = Redis.from_url("redis://localhost:6379")
|
||||
publisher = EventPublisher(redis)
|
||||
|
||||
# Publish events
|
||||
event_id = await publisher.publish(trade_event)
|
||||
|
||||
# Create consumer group for downstream processing
|
||||
await publisher.create_consumer_group("wallet-profiler")
|
||||
|
||||
# Read events as consumer
|
||||
entries = await publisher.read_events(
|
||||
group_name="wallet-profiler",
|
||||
consumer_name="worker-1"
|
||||
)
|
||||
for entry in entries:
|
||||
process(entry.event)
|
||||
await publisher.ack(entry.entry_id)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis: Redis,
|
||||
stream_name: str = DEFAULT_STREAM_NAME,
|
||||
*,
|
||||
max_len: int = DEFAULT_MAX_LEN,
|
||||
) -> None:
|
||||
"""Initialize the event publisher.
|
||||
|
||||
Args:
|
||||
redis: Redis async client.
|
||||
stream_name: Name of the Redis Stream.
|
||||
max_len: Maximum number of entries to keep in stream.
|
||||
"""
|
||||
self._redis = redis
|
||||
self._stream_name = stream_name
|
||||
self._max_len = max_len
|
||||
|
||||
@property
|
||||
def stream_name(self) -> str:
|
||||
"""Return the stream name."""
|
||||
return self._stream_name
|
||||
|
||||
async def publish(self, event: TradeEvent) -> str:
|
||||
"""Publish a single trade event to the stream.
|
||||
|
||||
Args:
|
||||
event: The TradeEvent to publish.
|
||||
|
||||
Returns:
|
||||
The entry ID assigned by Redis.
|
||||
"""
|
||||
data = _serialize_trade_event(event)
|
||||
entry_id = await self._redis.xadd(
|
||||
self._stream_name,
|
||||
data,
|
||||
maxlen=self._max_len,
|
||||
)
|
||||
# entry_id may be bytes or str
|
||||
if isinstance(entry_id, bytes):
|
||||
return entry_id.decode()
|
||||
return str(entry_id)
|
||||
|
||||
async def publish_batch(self, events: Sequence[TradeEvent]) -> list[str]:
|
||||
"""Publish multiple trade events atomically.
|
||||
|
||||
Uses a Redis pipeline for efficiency.
|
||||
|
||||
Args:
|
||||
events: Sequence of TradeEvents to publish.
|
||||
|
||||
Returns:
|
||||
List of entry IDs assigned by Redis.
|
||||
"""
|
||||
if not events:
|
||||
return []
|
||||
|
||||
pipe = self._redis.pipeline()
|
||||
for event in events:
|
||||
data = _serialize_trade_event(event)
|
||||
pipe.xadd(self._stream_name, data, maxlen=self._max_len)
|
||||
|
||||
results = await pipe.execute()
|
||||
|
||||
entry_ids: list[str] = []
|
||||
for entry_id in results:
|
||||
if isinstance(entry_id, bytes):
|
||||
entry_ids.append(entry_id.decode())
|
||||
else:
|
||||
entry_ids.append(str(entry_id))
|
||||
|
||||
return entry_ids
|
||||
|
||||
async def create_consumer_group(
|
||||
self,
|
||||
group_name: str,
|
||||
start_id: str = "0",
|
||||
*,
|
||||
mkstream: bool = True,
|
||||
) -> None:
|
||||
"""Create a consumer group for the stream.
|
||||
|
||||
Args:
|
||||
group_name: Name of the consumer group.
|
||||
start_id: ID to start reading from ("0" = beginning, "$" = new only).
|
||||
mkstream: Create the stream if it doesn't exist.
|
||||
|
||||
Raises:
|
||||
ConsumerGroupExistsError: If the group already exists.
|
||||
"""
|
||||
try:
|
||||
await self._redis.xgroup_create(
|
||||
self._stream_name,
|
||||
group_name,
|
||||
id=start_id,
|
||||
mkstream=mkstream,
|
||||
)
|
||||
logger.info(f"Created consumer group '{group_name}' on stream '{self._stream_name}'")
|
||||
except ResponseError as e:
|
||||
if "BUSYGROUP" in str(e):
|
||||
raise ConsumerGroupExistsError(
|
||||
f"Consumer group '{group_name}' already exists"
|
||||
) from e
|
||||
raise
|
||||
|
||||
async def ensure_consumer_group(
|
||||
self,
|
||||
group_name: str,
|
||||
start_id: str = "0",
|
||||
) -> bool:
|
||||
"""Ensure a consumer group exists, creating it if needed.
|
||||
|
||||
Args:
|
||||
group_name: Name of the consumer group.
|
||||
start_id: ID to start reading from if creating.
|
||||
|
||||
Returns:
|
||||
True if the group was created, False if it already existed.
|
||||
"""
|
||||
try:
|
||||
await self.create_consumer_group(group_name, start_id)
|
||||
return True
|
||||
except ConsumerGroupExistsError:
|
||||
return False
|
||||
|
||||
async def read_events(
|
||||
self,
|
||||
group_name: str,
|
||||
consumer_name: str,
|
||||
*,
|
||||
count: int = DEFAULT_COUNT,
|
||||
block_ms: int = DEFAULT_BLOCK_MS,
|
||||
) -> list[StreamEntry]:
|
||||
"""Read events from the stream as a consumer.
|
||||
|
||||
Args:
|
||||
group_name: Consumer group name.
|
||||
consumer_name: Name of this consumer within the group.
|
||||
count: Maximum number of entries to read.
|
||||
block_ms: Milliseconds to block waiting for new entries.
|
||||
|
||||
Returns:
|
||||
List of StreamEntry with entry ID and TradeEvent.
|
||||
"""
|
||||
# Read new entries (> means entries not delivered to this consumer)
|
||||
results = await self._redis.xreadgroup(
|
||||
group_name,
|
||||
consumer_name,
|
||||
{self._stream_name: ">"},
|
||||
count=count,
|
||||
block=block_ms,
|
||||
)
|
||||
|
||||
entries: list[StreamEntry] = []
|
||||
if not results:
|
||||
return entries
|
||||
|
||||
# Results format: [[stream_name, [(entry_id, data), ...]]]
|
||||
for _stream_name, stream_entries in results:
|
||||
for entry_id, data in stream_entries:
|
||||
# Decode entry_id
|
||||
entry_id_str = entry_id.decode() if isinstance(entry_id, bytes) else str(entry_id)
|
||||
|
||||
try:
|
||||
event = _deserialize_trade_event(data)
|
||||
entries.append(StreamEntry(entry_id=entry_id_str, event=event))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to deserialize entry {entry_id_str}: {e}")
|
||||
|
||||
return entries
|
||||
|
||||
async def read_pending(
|
||||
self,
|
||||
group_name: str,
|
||||
consumer_name: str,
|
||||
*,
|
||||
count: int = DEFAULT_COUNT,
|
||||
) -> list[StreamEntry]:
|
||||
"""Read pending (unacknowledged) entries for a consumer.
|
||||
|
||||
This is useful for recovering from crashes - entries that were
|
||||
delivered but not acknowledged will be re-read.
|
||||
|
||||
Args:
|
||||
group_name: Consumer group name.
|
||||
consumer_name: Name of this consumer.
|
||||
count: Maximum number of entries to read.
|
||||
|
||||
Returns:
|
||||
List of pending StreamEntry.
|
||||
"""
|
||||
# Read pending entries (0 means all pending entries)
|
||||
results = await self._redis.xreadgroup(
|
||||
group_name,
|
||||
consumer_name,
|
||||
{self._stream_name: "0"},
|
||||
count=count,
|
||||
)
|
||||
|
||||
entries: list[StreamEntry] = []
|
||||
if not results:
|
||||
return entries
|
||||
|
||||
for _stream_name, stream_entries in results:
|
||||
for entry_id, data in stream_entries:
|
||||
entry_id_str = entry_id.decode() if isinstance(entry_id, bytes) else str(entry_id)
|
||||
|
||||
# Skip entries with no data (already acked)
|
||||
if not data:
|
||||
continue
|
||||
|
||||
try:
|
||||
event = _deserialize_trade_event(data)
|
||||
entries.append(StreamEntry(entry_id=entry_id_str, event=event))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to deserialize pending entry {entry_id_str}: {e}")
|
||||
|
||||
return entries
|
||||
|
||||
async def ack(self, group_name: str, *entry_ids: str) -> int:
|
||||
"""Acknowledge that entries have been processed.
|
||||
|
||||
Args:
|
||||
group_name: Consumer group name.
|
||||
*entry_ids: Entry IDs to acknowledge.
|
||||
|
||||
Returns:
|
||||
Number of entries acknowledged.
|
||||
"""
|
||||
if not entry_ids:
|
||||
return 0
|
||||
return await self._redis.xack(self._stream_name, group_name, *entry_ids)
|
||||
|
||||
async def get_stream_info(self) -> dict[str, Any]:
|
||||
"""Get information about the stream.
|
||||
|
||||
Returns:
|
||||
Dictionary with stream info (length, groups, etc.).
|
||||
"""
|
||||
try:
|
||||
info = await self._redis.xinfo_stream(self._stream_name)
|
||||
return dict(info) if info else {}
|
||||
except ResponseError:
|
||||
return {}
|
||||
|
||||
async def get_stream_length(self) -> int:
|
||||
"""Get the current length of the stream.
|
||||
|
||||
Returns:
|
||||
Number of entries in the stream.
|
||||
"""
|
||||
return await self._redis.xlen(self._stream_name)
|
||||
|
||||
async def trim_stream(self, max_len: int | None = None) -> int:
|
||||
"""Trim the stream to a maximum length.
|
||||
|
||||
Args:
|
||||
max_len: Maximum entries to keep (uses default if not specified).
|
||||
|
||||
Returns:
|
||||
Number of entries removed.
|
||||
"""
|
||||
length = max_len or self._max_len
|
||||
return await self._redis.xtrim(self._stream_name, maxlen=length)
|
||||
@@ -0,0 +1,470 @@
|
||||
"""Tests for the Redis Streams event publisher."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from redis.exceptions import ResponseError
|
||||
|
||||
from polymarket_insider_tracker.ingestor.models import TradeEvent
|
||||
from polymarket_insider_tracker.ingestor.publisher import (
|
||||
DEFAULT_BLOCK_MS,
|
||||
DEFAULT_COUNT,
|
||||
DEFAULT_MAX_LEN,
|
||||
DEFAULT_STREAM_NAME,
|
||||
ConsumerGroupExistsError,
|
||||
EventPublisher,
|
||||
StreamEntry,
|
||||
_deserialize_trade_event,
|
||||
_serialize_trade_event,
|
||||
)
|
||||
|
||||
|
||||
# Test fixtures
|
||||
@pytest.fixture
|
||||
def sample_trade_event() -> TradeEvent:
|
||||
"""Create a sample trade event."""
|
||||
return TradeEvent(
|
||||
market_id="0xmarket123",
|
||||
trade_id="0xtx456",
|
||||
wallet_address="0xwallet789",
|
||||
side="BUY",
|
||||
outcome="Yes",
|
||||
outcome_index=0,
|
||||
price=Decimal("0.65"),
|
||||
size=Decimal("1000"),
|
||||
timestamp=datetime(2026, 1, 4, 12, 0, 0, tzinfo=UTC),
|
||||
asset_id="token123",
|
||||
market_slug="will-it-rain",
|
||||
event_slug="weather-markets",
|
||||
event_title="Weather Predictions",
|
||||
trader_name="Alice",
|
||||
trader_pseudonym="AliceTrader",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis() -> AsyncMock:
|
||||
"""Create a mock Redis client."""
|
||||
redis = AsyncMock()
|
||||
redis.xadd = AsyncMock(return_value="1704369600000-0")
|
||||
redis.xreadgroup = AsyncMock(return_value=[])
|
||||
redis.xack = AsyncMock(return_value=1)
|
||||
redis.xlen = AsyncMock(return_value=100)
|
||||
redis.xtrim = AsyncMock(return_value=0)
|
||||
redis.xinfo_stream = AsyncMock(return_value={"length": 100})
|
||||
redis.xgroup_create = AsyncMock()
|
||||
redis.pipeline = MagicMock()
|
||||
return redis
|
||||
|
||||
|
||||
class TestSerializationFunctions:
|
||||
"""Tests for serialization helper functions."""
|
||||
|
||||
def test_serialize_trade_event(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test serializing a trade event."""
|
||||
data = _serialize_trade_event(sample_trade_event)
|
||||
|
||||
assert data["market_id"] == "0xmarket123"
|
||||
assert data["trade_id"] == "0xtx456"
|
||||
assert data["wallet_address"] == "0xwallet789"
|
||||
assert data["side"] == "BUY"
|
||||
assert data["outcome"] == "Yes"
|
||||
assert data["outcome_index"] == "0"
|
||||
assert data["price"] == "0.65"
|
||||
assert data["size"] == "1000"
|
||||
assert data["timestamp"] == "2026-01-04T12:00:00+00:00"
|
||||
assert data["asset_id"] == "token123"
|
||||
assert data["market_slug"] == "will-it-rain"
|
||||
assert data["trader_name"] == "Alice"
|
||||
|
||||
def test_serialize_all_values_are_strings(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test that all serialized values are strings."""
|
||||
data = _serialize_trade_event(sample_trade_event)
|
||||
|
||||
for key, value in data.items():
|
||||
assert isinstance(key, str), f"Key {key} is not a string"
|
||||
assert isinstance(value, str), f"Value for {key} is not a string"
|
||||
|
||||
def test_deserialize_trade_event(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test deserializing a trade event."""
|
||||
data = _serialize_trade_event(sample_trade_event)
|
||||
restored = _deserialize_trade_event(data)
|
||||
|
||||
assert restored.market_id == sample_trade_event.market_id
|
||||
assert restored.trade_id == sample_trade_event.trade_id
|
||||
assert restored.wallet_address == sample_trade_event.wallet_address
|
||||
assert restored.side == sample_trade_event.side
|
||||
assert restored.outcome == sample_trade_event.outcome
|
||||
assert restored.outcome_index == sample_trade_event.outcome_index
|
||||
assert restored.price == sample_trade_event.price
|
||||
assert restored.size == sample_trade_event.size
|
||||
assert restored.timestamp == sample_trade_event.timestamp
|
||||
assert restored.asset_id == sample_trade_event.asset_id
|
||||
|
||||
def test_deserialize_with_bytes_keys(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test deserializing with bytes keys/values (as returned by Redis)."""
|
||||
data = _serialize_trade_event(sample_trade_event)
|
||||
# Convert to bytes like Redis returns
|
||||
bytes_data = {k.encode(): v.encode() for k, v in data.items()}
|
||||
|
||||
restored = _deserialize_trade_event(bytes_data)
|
||||
|
||||
assert restored.market_id == sample_trade_event.market_id
|
||||
assert restored.side == sample_trade_event.side
|
||||
|
||||
def test_deserialize_with_invalid_timestamp(self) -> None:
|
||||
"""Test deserializing with invalid timestamp falls back to now."""
|
||||
data = {
|
||||
"market_id": "0x123",
|
||||
"timestamp": "not-a-timestamp",
|
||||
"side": "BUY",
|
||||
"price": "0.5",
|
||||
"size": "100",
|
||||
}
|
||||
|
||||
event = _deserialize_trade_event(data)
|
||||
|
||||
assert event.market_id == "0x123"
|
||||
# Timestamp should be recent (within last minute)
|
||||
assert (datetime.now(UTC) - event.timestamp).total_seconds() < 60
|
||||
|
||||
def test_deserialize_with_missing_fields(self) -> None:
|
||||
"""Test deserializing with missing fields uses defaults."""
|
||||
data = {
|
||||
"market_id": "0x123",
|
||||
"side": "SELL",
|
||||
}
|
||||
|
||||
event = _deserialize_trade_event(data)
|
||||
|
||||
assert event.market_id == "0x123"
|
||||
assert event.side == "SELL"
|
||||
assert event.price == Decimal("0")
|
||||
assert event.outcome == ""
|
||||
|
||||
|
||||
class TestEventPublisher:
|
||||
"""Tests for the EventPublisher class."""
|
||||
|
||||
def test_init(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test initialization."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
assert publisher.stream_name == DEFAULT_STREAM_NAME
|
||||
assert publisher._max_len == DEFAULT_MAX_LEN
|
||||
|
||||
def test_init_custom_config(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test initialization with custom config."""
|
||||
publisher = EventPublisher(
|
||||
mock_redis,
|
||||
stream_name="custom-stream",
|
||||
max_len=50_000,
|
||||
)
|
||||
|
||||
assert publisher.stream_name == "custom-stream"
|
||||
assert publisher._max_len == 50_000
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish(self, mock_redis: AsyncMock, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test publishing a single event."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entry_id = await publisher.publish(sample_trade_event)
|
||||
|
||||
assert entry_id == "1704369600000-0"
|
||||
mock_redis.xadd.assert_called_once()
|
||||
call_args = mock_redis.xadd.call_args
|
||||
assert call_args[0][0] == DEFAULT_STREAM_NAME
|
||||
assert call_args[1]["maxlen"] == DEFAULT_MAX_LEN
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_returns_decoded_bytes(
|
||||
self, mock_redis: AsyncMock, sample_trade_event: TradeEvent
|
||||
) -> None:
|
||||
"""Test that publish handles bytes entry IDs."""
|
||||
mock_redis.xadd = AsyncMock(return_value=b"1704369600000-0")
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entry_id = await publisher.publish(sample_trade_event)
|
||||
|
||||
assert entry_id == "1704369600000-0"
|
||||
assert isinstance(entry_id, str)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_batch(
|
||||
self, mock_redis: AsyncMock, sample_trade_event: TradeEvent
|
||||
) -> None:
|
||||
"""Test batch publishing."""
|
||||
mock_pipeline = AsyncMock()
|
||||
mock_pipeline.xadd = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(return_value=["1704369600000-0", "1704369600000-1"])
|
||||
mock_redis.pipeline.return_value = mock_pipeline
|
||||
|
||||
publisher = EventPublisher(mock_redis)
|
||||
events = [sample_trade_event, sample_trade_event]
|
||||
|
||||
entry_ids = await publisher.publish_batch(events)
|
||||
|
||||
assert len(entry_ids) == 2
|
||||
assert entry_ids[0] == "1704369600000-0"
|
||||
assert entry_ids[1] == "1704369600000-1"
|
||||
assert mock_pipeline.xadd.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_batch_empty(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test batch publishing with empty list."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entry_ids = await publisher.publish_batch([])
|
||||
|
||||
assert entry_ids == []
|
||||
mock_redis.pipeline.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_consumer_group(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test creating a consumer group."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
await publisher.create_consumer_group("test-group")
|
||||
|
||||
mock_redis.xgroup_create.assert_called_once_with(
|
||||
DEFAULT_STREAM_NAME,
|
||||
"test-group",
|
||||
id="0",
|
||||
mkstream=True,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_consumer_group_custom_start_id(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test creating a consumer group with custom start ID."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
await publisher.create_consumer_group("test-group", start_id="$")
|
||||
|
||||
call_args = mock_redis.xgroup_create.call_args
|
||||
assert call_args[1]["id"] == "$"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_consumer_group_already_exists(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test creating a consumer group that already exists."""
|
||||
mock_redis.xgroup_create.side_effect = ResponseError(
|
||||
"BUSYGROUP Consumer Group name already exists"
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
with pytest.raises(ConsumerGroupExistsError):
|
||||
await publisher.create_consumer_group("existing-group")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_consumer_group_creates(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test ensure_consumer_group creates if not exists."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
created = await publisher.ensure_consumer_group("new-group")
|
||||
|
||||
assert created is True
|
||||
mock_redis.xgroup_create.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_consumer_group_exists(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test ensure_consumer_group returns False if exists."""
|
||||
mock_redis.xgroup_create.side_effect = ResponseError("BUSYGROUP")
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
created = await publisher.ensure_consumer_group("existing-group")
|
||||
|
||||
assert created is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_events(self, mock_redis: AsyncMock, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test reading events from stream."""
|
||||
serialized = _serialize_trade_event(sample_trade_event)
|
||||
mock_redis.xreadgroup = AsyncMock(
|
||||
return_value=[
|
||||
(
|
||||
"trades",
|
||||
[
|
||||
("1704369600000-0", serialized),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_events("test-group", "worker-1")
|
||||
|
||||
assert len(entries) == 1
|
||||
assert entries[0].entry_id == "1704369600000-0"
|
||||
assert entries[0].event.market_id == sample_trade_event.market_id
|
||||
mock_redis.xreadgroup.assert_called_once_with(
|
||||
"test-group",
|
||||
"worker-1",
|
||||
{DEFAULT_STREAM_NAME: ">"},
|
||||
count=DEFAULT_COUNT,
|
||||
block=DEFAULT_BLOCK_MS,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_events_empty(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test reading when no events available."""
|
||||
mock_redis.xreadgroup = AsyncMock(return_value=None)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_events("test-group", "worker-1")
|
||||
|
||||
assert entries == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_events_with_bytes(
|
||||
self, mock_redis: AsyncMock, sample_trade_event: TradeEvent
|
||||
) -> None:
|
||||
"""Test reading events with bytes data (as from real Redis)."""
|
||||
serialized = _serialize_trade_event(sample_trade_event)
|
||||
bytes_data = {k.encode(): v.encode() for k, v in serialized.items()}
|
||||
mock_redis.xreadgroup = AsyncMock(
|
||||
return_value=[
|
||||
(
|
||||
b"trades",
|
||||
[
|
||||
(b"1704369600000-0", bytes_data),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_events("test-group", "worker-1")
|
||||
|
||||
assert len(entries) == 1
|
||||
assert entries[0].entry_id == "1704369600000-0"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_pending(
|
||||
self, mock_redis: AsyncMock, sample_trade_event: TradeEvent
|
||||
) -> None:
|
||||
"""Test reading pending events."""
|
||||
serialized = _serialize_trade_event(sample_trade_event)
|
||||
mock_redis.xreadgroup = AsyncMock(
|
||||
return_value=[
|
||||
(
|
||||
"trades",
|
||||
[
|
||||
("1704369600000-0", serialized),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_pending("test-group", "worker-1")
|
||||
|
||||
assert len(entries) == 1
|
||||
# Should read from "0" not ">"
|
||||
call_args = mock_redis.xreadgroup.call_args
|
||||
assert call_args[0][2] == {DEFAULT_STREAM_NAME: "0"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_pending_skips_empty_data(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test that read_pending skips entries with no data (already acked)."""
|
||||
mock_redis.xreadgroup = AsyncMock(
|
||||
return_value=[
|
||||
(
|
||||
"trades",
|
||||
[
|
||||
("1704369600000-0", {}), # Empty = already acked
|
||||
("1704369600000-1", None), # None = already acked
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
entries = await publisher.read_pending("test-group", "worker-1")
|
||||
|
||||
assert entries == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ack(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test acknowledging entries."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
count = await publisher.ack("test-group", "1704369600000-0", "1704369600000-1")
|
||||
|
||||
assert count == 1 # Mocked return value
|
||||
mock_redis.xack.assert_called_once_with(
|
||||
DEFAULT_STREAM_NAME,
|
||||
"test-group",
|
||||
"1704369600000-0",
|
||||
"1704369600000-1",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ack_empty(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test ack with no entry IDs."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
count = await publisher.ack("test-group")
|
||||
|
||||
assert count == 0
|
||||
mock_redis.xack.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_stream_info(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting stream info."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
info = await publisher.get_stream_info()
|
||||
|
||||
assert info["length"] == 100
|
||||
mock_redis.xinfo_stream.assert_called_once_with(DEFAULT_STREAM_NAME)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_stream_info_not_exists(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting stream info when stream doesn't exist."""
|
||||
mock_redis.xinfo_stream.side_effect = ResponseError("ERR no such key")
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
info = await publisher.get_stream_info()
|
||||
|
||||
assert info == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_stream_length(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test getting stream length."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
length = await publisher.get_stream_length()
|
||||
|
||||
assert length == 100
|
||||
mock_redis.xlen.assert_called_once_with(DEFAULT_STREAM_NAME)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_stream(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test trimming stream."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
await publisher.trim_stream(50_000)
|
||||
|
||||
mock_redis.xtrim.assert_called_once_with(DEFAULT_STREAM_NAME, maxlen=50_000)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_stream_default(self, mock_redis: AsyncMock) -> None:
|
||||
"""Test trimming stream with default max_len."""
|
||||
publisher = EventPublisher(mock_redis)
|
||||
|
||||
await publisher.trim_stream()
|
||||
|
||||
mock_redis.xtrim.assert_called_once_with(DEFAULT_STREAM_NAME, maxlen=DEFAULT_MAX_LEN)
|
||||
|
||||
|
||||
class TestStreamEntry:
|
||||
"""Tests for the StreamEntry dataclass."""
|
||||
|
||||
def test_stream_entry(self, sample_trade_event: TradeEvent) -> None:
|
||||
"""Test creating a StreamEntry."""
|
||||
entry = StreamEntry(entry_id="1704369600000-0", event=sample_trade_event)
|
||||
|
||||
assert entry.entry_id == "1704369600000-0"
|
||||
assert entry.event == sample_trade_event
|
||||
Reference in New Issue
Block a user