test_ws_subscribe.py
5.89 KB
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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
"""U-024 tests: c_ws_subscribe's pure pieces (handshake key, frame
encode/decode, event polling) tested directly, no live socket needed for
these -- the live-socket end-to-end path is covered in
src/test_backend_server.py's WsSubscribeRouteTest.
Run: python3 -m unittest src.test_ws_subscribe -v
"""
import io
import json
import unittest
import uuid
from src import schema, ws_subscribe
class ComputeAcceptKeyTest(unittest.TestCase):
def test_matches_rfc6455_example(self):
# RFC 6455 §1.3's own worked example
self.assertEqual(
ws_subscribe.compute_accept_key("dGhlIHNhbXBsZSBub25jZQ=="),
"s3pPLMBiTxaQ9kYGzzhZRbK+xOo=",
)
class _FakeSocket:
"""Minimal recv()-only stand-in so decode_client_frame can be tested
without a real socket -- feeds back a bytes buffer in arbitrary chunks."""
def __init__(self, data):
self._buf = io.BytesIO(data)
def recv(self, n):
return self._buf.read(n)
class FrameEncodeDecodeTest(unittest.TestCase):
def test_text_frame_roundtrip_short_payload(self):
frame = ws_subscribe.encode_text_frame("hello")
# server frames aren't masked -- decode via a masked wrapper to
# reuse decode_client_frame's parsing logic for the length/opcode
# bits (masking bit is separately exercised in the masked test).
self.assertEqual(frame[0], 0x81) # FIN=1, opcode=text
self.assertEqual(frame[1], 5) # unmasked, length=5
self.assertEqual(frame[2:], b"hello")
def test_text_frame_medium_payload_uses_16bit_length(self):
text = "x" * 200
frame = ws_subscribe.encode_text_frame(text)
self.assertEqual(frame[1], 126)
self.assertEqual(int.from_bytes(frame[2:4], "big"), 200)
self.assertEqual(frame[4:].decode(), text)
def test_decode_masked_client_frame(self):
payload = b"ping-ish"
mask_key = bytes([0x12, 0x34, 0x56, 0x78])
masked_payload = bytes(b ^ mask_key[i % 4] for i, b in enumerate(payload))
raw = bytes([0x81, 0x80 | len(payload)]) + mask_key + masked_payload
opcode, decoded = ws_subscribe.decode_client_frame(_FakeSocket(raw))
self.assertEqual(opcode, ws_subscribe.OPCODE_TEXT)
self.assertEqual(decoded, payload)
def test_decode_close_frame_empty_payload(self):
raw = bytes([0x88, 0x80]) + bytes([0, 0, 0, 0]) # masked, zero-length
opcode, decoded = ws_subscribe.decode_client_frame(_FakeSocket(raw))
self.assertEqual(opcode, ws_subscribe.OPCODE_CLOSE)
self.assertEqual(decoded, b"")
def test_decode_raises_connection_error_on_truncated_socket(self):
with self.assertRaises(ConnectionError):
ws_subscribe.decode_client_frame(_FakeSocket(b"\x81"))
class PollNewEventsTest(unittest.TestCase):
def setUp(self):
self.conn = schema.connect(":memory:")
schema.create_schema(self.conn)
self.conn.execute(
"INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
"VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')"
)
self.conn.execute(
"INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
"VALUES ('proj-2', 'P2', '/tmp/proj-2', 1, 'active')"
)
self.conn.execute(
"INSERT INTO task (task_id, project_id, goal_text, task_type, status, assigned_org_node) "
"VALUES ('task-1', 'proj-1', 'goal', 'goal', 'queued', 'planner')"
)
self.conn.commit()
def tearDown(self):
self.conn.close()
def _insert_event(self, project_id, event_type="task_status_changed", task_id=None, payload=None):
self.conn.execute(
"INSERT INTO execution_event (event_id, project_id, task_id, event_type, payload) "
"VALUES (?, ?, ?, ?, ?)",
(str(uuid.uuid4()), project_id, task_id, event_type, json.dumps(payload or {})),
)
self.conn.commit()
def test_no_new_events_returns_empty_list_and_same_rowid(self):
after = ws_subscribe.latest_rowid(self.conn, "proj-1")
events, new_after = ws_subscribe.poll_new_events(self.conn, "proj-1", after)
self.assertEqual(events, [])
self.assertEqual(new_after, after)
def test_new_event_after_baseline_is_returned(self):
after = ws_subscribe.latest_rowid(self.conn, "proj-1")
self._insert_event("proj-1", event_type="budget_updated", payload={"x": 1})
events, new_after = ws_subscribe.poll_new_events(self.conn, "proj-1", after)
self.assertEqual(len(events), 1)
self.assertEqual(events[0]["event_type"], "budget_updated")
self.assertEqual(events[0]["payload"], {"x": 1})
self.assertGreater(new_after, after)
def test_events_from_other_project_are_not_returned(self):
after = ws_subscribe.latest_rowid(self.conn, "proj-1")
self._insert_event("proj-2")
events, _ = ws_subscribe.poll_new_events(self.conn, "proj-1", after)
self.assertEqual(events, [])
def test_task_id_filter_narrows_results(self):
after = ws_subscribe.latest_rowid(self.conn, "proj-1")
self._insert_event("proj-1", task_id="task-1")
self._insert_event("proj-1", task_id=None)
events, _ = ws_subscribe.poll_new_events(self.conn, "proj-1", after, task_id="task-1")
self.assertEqual(len(events), 1)
self.assertEqual(events[0]["task_id"], "task-1")
def test_already_seen_events_not_returned_again(self):
after = ws_subscribe.latest_rowid(self.conn, "proj-1")
self._insert_event("proj-1")
events, after2 = ws_subscribe.poll_new_events(self.conn, "proj-1", after)
self.assertEqual(len(events), 1)
events_again, after3 = ws_subscribe.poll_new_events(self.conn, "proj-1", after2)
self.assertEqual(events_again, [])
self.assertEqual(after2, after3)
if __name__ == "__main__":
unittest.main()