| from __future__ import annotations |
|
|
| import asyncio |
| import gzip |
| import json |
| import tempfile |
| import time |
| import unittest |
| from pathlib import Path |
|
|
| from websockets.asyncio.server import serve |
|
|
| from polymarket_collector.state import CollectorState |
| from polymarket_collector.storage import RotatingJsonlWriter |
| from polymarket_collector.stream import MarketStream |
|
|
|
|
| class StreamTests(unittest.IsolatedAsyncioTestCase): |
| async def test_records_public_market_message(self) -> None: |
| subscription_received = asyncio.Event() |
|
|
| async def handler(websocket) -> None: |
| subscription = json.loads(await websocket.recv()) |
| self.assertEqual(subscription["assets_ids"], ["asset-1", "asset-2"]) |
| subscription_received.set() |
| await websocket.send( |
| json.dumps( |
| { |
| "event_type": "book", |
| "asset_id": "asset-1", |
| "market": "condition-1", |
| "bids": [], |
| "asks": [], |
| "timestamp": "1", |
| } |
| ) |
| ) |
| while True: |
| try: |
| message = await websocket.recv() |
| except Exception: |
| return |
| if message == "PING": |
| await websocket.send("PONG") |
|
|
| with tempfile.TemporaryDirectory() as temporary: |
| writer = RotatingJsonlWriter( |
| Path(temporary), |
| "stream", |
| rotation_seconds=3600, |
| max_file_bytes=1024 * 1024, |
| flush_interval_seconds=1, |
| ) |
| await writer.start() |
| state = CollectorState() |
| stop = asyncio.Event() |
| async with serve(handler, "127.0.0.1", 0) as websocket_server: |
| port = websocket_server.sockets[0].getsockname()[1] |
| stream = MarketStream( |
| f"ws://127.0.0.1:{port}", writer, state |
| ) |
| stream.set_assets(["asset-1", "asset-2"]) |
| task = asyncio.create_task(stream.run(stop)) |
| await asyncio.wait_for(subscription_received.wait(), 2) |
| deadline = time.monotonic() + 2 |
| while state.messages_total < 1 and time.monotonic() < deadline: |
| await asyncio.sleep(0.01) |
| stop.set() |
| await asyncio.wait_for(task, 3) |
| await writer.close() |
|
|
| self.assertEqual(state.messages_total, 1) |
| archive = next(Path(temporary).rglob("*.jsonl.gz")) |
| with gzip.open(archive, "rt", encoding="utf-8") as handle: |
| record = json.loads(handle.readline()) |
| self.assertEqual(record["payload"]["event_type"], "book") |
| self.assertIn("received_at_ns", record) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
|
|
