File size: 2,948 Bytes
e2ce547 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | 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()
|