92 lines
3.4 KiB
Python
92 lines
3.4 KiB
Python
"""Tests for the event graph model."""
|
|
|
|
import json
|
|
import unittest
|
|
|
|
from deltatrace.event_graph import Event, EventGraph, EventType
|
|
|
|
|
|
class TestEvent(unittest.TestCase):
|
|
def test_create_event(self):
|
|
ev = Event.create(EventType.MD_TICK, {"symbol": "AAPL", "price": 185.0})
|
|
self.assertEqual(ev.event_type, EventType.MD_TICK)
|
|
self.assertEqual(ev.payload["symbol"], "AAPL")
|
|
self.assertGreater(ev.timestamp_ns, 0)
|
|
self.assertEqual(len(ev.id), 32)
|
|
|
|
def test_event_types(self):
|
|
for et in EventType:
|
|
ev = Event.create(et, {"test": True})
|
|
self.assertEqual(ev.event_type, et)
|
|
|
|
|
|
class TestEventGraph(unittest.TestCase):
|
|
def setUp(self):
|
|
self.graph = EventGraph()
|
|
self.tick = Event.create(EventType.MD_TICK, {"symbol": "AAPL"})
|
|
self.signal = Event.create(EventType.SIGNAL, {"direction": "buy"})
|
|
self.order = Event.create(EventType.ORDER, {"qty": 100})
|
|
|
|
def test_add_events(self):
|
|
self.graph.add_event(self.tick)
|
|
self.graph.add_event(self.signal)
|
|
self.assertEqual(len(self.graph.events), 2)
|
|
|
|
def test_add_edge(self):
|
|
self.graph.add_event(self.tick)
|
|
self.graph.add_event(self.signal)
|
|
edge = self.graph.add_edge(self.tick.id, self.signal.id, "tick->signal")
|
|
self.assertEqual(edge.from_id, self.tick.id)
|
|
self.assertEqual(edge.to_id, self.signal.id)
|
|
self.assertEqual(edge.label, "tick->signal")
|
|
|
|
def test_edge_missing_event_raises(self):
|
|
self.graph.add_event(self.tick)
|
|
with self.assertRaises(KeyError):
|
|
self.graph.add_edge(self.tick.id, "nonexistent")
|
|
|
|
def test_roots(self):
|
|
self.graph.add_event(self.tick)
|
|
self.graph.add_event(self.signal)
|
|
self.graph.add_edge(self.tick.id, self.signal.id)
|
|
roots = self.graph.get_roots()
|
|
self.assertEqual(len(roots), 1)
|
|
self.assertEqual(roots[0].id, self.tick.id)
|
|
|
|
def test_children_parents(self):
|
|
self.graph.add_event(self.tick)
|
|
self.graph.add_event(self.signal)
|
|
self.graph.add_edge(self.tick.id, self.signal.id)
|
|
children = self.graph.get_children(self.tick.id)
|
|
parents = self.graph.get_parents(self.signal.id)
|
|
self.assertEqual(len(children), 1)
|
|
self.assertEqual(children[0].id, self.signal.id)
|
|
self.assertEqual(len(parents), 1)
|
|
self.assertEqual(parents[0].id, self.tick.id)
|
|
|
|
def test_topological_order(self):
|
|
self.graph.add_event(self.tick)
|
|
self.graph.add_event(self.signal)
|
|
self.graph.add_event(self.order)
|
|
self.graph.add_edge(self.tick.id, self.signal.id)
|
|
self.graph.add_edge(self.signal.id, self.order.id)
|
|
topo = self.graph.topological_order()
|
|
ids = [e.id for e in topo]
|
|
self.assertLess(ids.index(self.tick.id), ids.index(self.signal.id))
|
|
self.assertLess(ids.index(self.signal.id), ids.index(self.order.id))
|
|
|
|
def test_serialization_roundtrip(self):
|
|
self.graph.add_event(self.tick)
|
|
self.graph.add_event(self.signal)
|
|
self.graph.add_edge(self.tick.id, self.signal.id, "causal")
|
|
data = self.graph.to_dict()
|
|
json_str = json.dumps(data)
|
|
restored = EventGraph.from_dict(json.loads(json_str))
|
|
self.assertEqual(len(restored.events), 2)
|
|
self.assertEqual(len(restored.edges), 1)
|
|
self.assertEqual(restored.get_event(self.tick.id).event_type, EventType.MD_TICK)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|