mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-11 07:19:03 +00:00
f956682f31
MemoryStreamBridge._resolve_start_offset scanned the retained event buffer
(up to queue_maxsize=256 entries) on every subscribe/reconnect carrying a
Last-Event-ID. Event ids are "{ts}-{seq}" where seq is a per-run monotonic
counter that equals the event's absolute offset, so the offset is computable
arithmetically. Parse seq, index into the buffer, and verify the id matches
exactly -- a stale/evicted/foreign/malformed id falls back to
replay-from-earliest, identical to the previous scan.
Co-authored-by: ly-wang19 <ly-wang19@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
407 lines
13 KiB
Python
407 lines
13 KiB
Python
"""Tests for the in-memory StreamBridge implementation."""
|
|
|
|
import asyncio
|
|
import re
|
|
|
|
import anyio
|
|
import pytest
|
|
|
|
from deerflow.runtime import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, make_stream_bridge
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Unit tests for MemoryStreamBridge
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def bridge() -> MemoryStreamBridge:
|
|
return MemoryStreamBridge(queue_maxsize=256)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_publish_subscribe(bridge: MemoryStreamBridge):
|
|
"""Three events followed by end should be received in order."""
|
|
run_id = "run-1"
|
|
|
|
await bridge.publish(run_id, "metadata", {"run_id": run_id})
|
|
await bridge.publish(run_id, "values", {"messages": []})
|
|
await bridge.publish(run_id, "updates", {"step": 1})
|
|
await bridge.publish_end(run_id)
|
|
|
|
received = []
|
|
async for entry in bridge.subscribe(run_id, heartbeat_interval=1.0):
|
|
received.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
assert len(received) == 4
|
|
assert received[0].event == "metadata"
|
|
assert received[1].event == "values"
|
|
assert received[2].event == "updates"
|
|
assert received[3] is END_SENTINEL
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_heartbeat(bridge: MemoryStreamBridge):
|
|
"""When no events arrive within the heartbeat interval, yield a heartbeat."""
|
|
run_id = "run-heartbeat"
|
|
bridge._get_or_create_stream(run_id) # ensure stream exists
|
|
|
|
received = []
|
|
|
|
async def consumer():
|
|
async for entry in bridge.subscribe(run_id, heartbeat_interval=0.1):
|
|
received.append(entry)
|
|
if entry is HEARTBEAT_SENTINEL:
|
|
break
|
|
|
|
await asyncio.wait_for(consumer(), timeout=2.0)
|
|
assert len(received) == 1
|
|
assert received[0] is HEARTBEAT_SENTINEL
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cleanup(bridge: MemoryStreamBridge):
|
|
"""After cleanup, the run's stream/event log is removed."""
|
|
run_id = "run-cleanup"
|
|
await bridge.publish(run_id, "test", {})
|
|
assert run_id in bridge._streams
|
|
|
|
await bridge.cleanup(run_id)
|
|
assert run_id not in bridge._streams
|
|
assert run_id not in bridge._counters
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_history_is_bounded():
|
|
"""Retained history should be bounded by queue_maxsize."""
|
|
bridge = MemoryStreamBridge(queue_maxsize=1)
|
|
run_id = "run-bp"
|
|
|
|
await bridge.publish(run_id, "first", {})
|
|
await bridge.publish(run_id, "second", {})
|
|
await bridge.publish_end(run_id)
|
|
|
|
received = []
|
|
async for entry in bridge.subscribe(run_id, heartbeat_interval=1.0):
|
|
received.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
assert len(received) == 2
|
|
assert received[0].event == "second"
|
|
assert received[1] is END_SENTINEL
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_multiple_runs(bridge: MemoryStreamBridge):
|
|
"""Two different run_ids should not interfere with each other."""
|
|
await bridge.publish("run-a", "event-a", {"a": 1})
|
|
await bridge.publish("run-b", "event-b", {"b": 2})
|
|
await bridge.publish_end("run-a")
|
|
await bridge.publish_end("run-b")
|
|
|
|
events_a = []
|
|
async for entry in bridge.subscribe("run-a", heartbeat_interval=1.0):
|
|
events_a.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
events_b = []
|
|
async for entry in bridge.subscribe("run-b", heartbeat_interval=1.0):
|
|
events_b.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
assert len(events_a) == 2
|
|
assert events_a[0].event == "event-a"
|
|
assert events_a[0].data == {"a": 1}
|
|
|
|
assert len(events_b) == 2
|
|
assert events_b[0].event == "event-b"
|
|
assert events_b[0].data == {"b": 2}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_event_id_format(bridge: MemoryStreamBridge):
|
|
"""Event IDs should use timestamp-sequence format."""
|
|
run_id = "run-id-format"
|
|
await bridge.publish(run_id, "test", {"key": "value"})
|
|
await bridge.publish_end(run_id)
|
|
|
|
received = []
|
|
async for entry in bridge.subscribe(run_id, heartbeat_interval=1.0):
|
|
received.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
event = received[0]
|
|
assert re.match(r"^\d+-\d+$", event.id), f"Expected timestamp-seq format, got {event.id}"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_subscribe_replays_after_last_event_id(bridge: MemoryStreamBridge):
|
|
"""Reconnect should replay buffered events after the provided Last-Event-ID."""
|
|
run_id = "run-replay"
|
|
await bridge.publish(run_id, "metadata", {"run_id": run_id})
|
|
await bridge.publish(run_id, "values", {"step": 1})
|
|
await bridge.publish(run_id, "updates", {"step": 2})
|
|
await bridge.publish_end(run_id)
|
|
|
|
first_pass = []
|
|
async for entry in bridge.subscribe(run_id, heartbeat_interval=1.0):
|
|
first_pass.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
received = []
|
|
async for entry in bridge.subscribe(
|
|
run_id,
|
|
last_event_id=first_pass[0].id,
|
|
heartbeat_interval=1.0,
|
|
):
|
|
received.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
assert [entry.event for entry in received[:-1]] == ["values", "updates"]
|
|
assert received[-1] is END_SENTINEL
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_slow_subscriber_does_not_skip_after_buffer_trim():
|
|
"""A slow subscriber should continue from the correct absolute offset."""
|
|
bridge = MemoryStreamBridge(queue_maxsize=2)
|
|
run_id = "run-slow-subscriber"
|
|
await bridge.publish(run_id, "e1", {"step": 1})
|
|
await bridge.publish(run_id, "e2", {"step": 2})
|
|
|
|
stream = bridge._streams[run_id]
|
|
e1_id = stream.events[0].id
|
|
assert stream.start_offset == 0
|
|
|
|
await bridge.publish(run_id, "e3", {"step": 3}) # trims e1
|
|
assert stream.start_offset == 1
|
|
assert [entry.event for entry in stream.events] == ["e2", "e3"]
|
|
|
|
resumed_after_e1 = []
|
|
async for entry in bridge.subscribe(
|
|
run_id,
|
|
last_event_id=e1_id,
|
|
heartbeat_interval=1.0,
|
|
):
|
|
resumed_after_e1.append(entry)
|
|
if len(resumed_after_e1) == 2:
|
|
break
|
|
|
|
assert [entry.event for entry in resumed_after_e1] == ["e2", "e3"]
|
|
e2_id = resumed_after_e1[0].id
|
|
|
|
await bridge.publish_end(run_id)
|
|
|
|
received = []
|
|
async for entry in bridge.subscribe(
|
|
run_id,
|
|
last_event_id=e2_id,
|
|
heartbeat_interval=1.0,
|
|
):
|
|
received.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
assert [entry.event for entry in received[:-1]] == ["e3"]
|
|
assert received[-1] is END_SENTINEL
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Stream termination tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_publish_end_terminates_even_when_history_is_full():
|
|
"""publish_end() should terminate subscribers without mutating retained history."""
|
|
bridge = MemoryStreamBridge(queue_maxsize=2)
|
|
run_id = "run-end-history-full"
|
|
|
|
await bridge.publish(run_id, "event-1", {"n": 1})
|
|
await bridge.publish(run_id, "event-2", {"n": 2})
|
|
stream = bridge._streams[run_id]
|
|
assert [entry.event for entry in stream.events] == ["event-1", "event-2"]
|
|
|
|
await bridge.publish_end(run_id)
|
|
assert [entry.event for entry in stream.events] == ["event-1", "event-2"]
|
|
|
|
events = []
|
|
async for entry in bridge.subscribe(run_id, heartbeat_interval=0.1):
|
|
events.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
assert [entry.event for entry in events[:-1]] == ["event-1", "event-2"]
|
|
assert events[-1] is END_SENTINEL
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_publish_end_without_history_yields_end_immediately():
|
|
"""Subscribers should still receive END when a run completes without events."""
|
|
bridge = MemoryStreamBridge(queue_maxsize=2)
|
|
run_id = "run-end-empty"
|
|
await bridge.publish_end(run_id)
|
|
|
|
events = []
|
|
async for entry in bridge.subscribe(run_id, heartbeat_interval=0.1):
|
|
events.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
assert len(events) == 1
|
|
assert events[0] is END_SENTINEL
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_publish_end_preserves_history_when_space_available():
|
|
"""When history has spare capacity, publish_end should preserve prior events."""
|
|
bridge = MemoryStreamBridge(queue_maxsize=10)
|
|
run_id = "run-no-evict"
|
|
|
|
await bridge.publish(run_id, "event-1", {"n": 1})
|
|
await bridge.publish(run_id, "event-2", {"n": 2})
|
|
await bridge.publish_end(run_id)
|
|
|
|
events = []
|
|
async for entry in bridge.subscribe(run_id, heartbeat_interval=0.1):
|
|
events.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
# All events plus END should be present
|
|
assert len(events) == 3
|
|
assert events[0].event == "event-1"
|
|
assert events[1].event == "event-2"
|
|
assert events[2] is END_SENTINEL
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_concurrent_tasks_end_sentinel():
|
|
"""Multiple concurrent producer/consumer pairs should all terminate properly.
|
|
|
|
Simulates the production scenario where multiple runs share a single
|
|
bridge instance — each must receive its own END sentinel.
|
|
"""
|
|
bridge = MemoryStreamBridge(queue_maxsize=4)
|
|
num_runs = 4
|
|
|
|
async def producer(run_id: str):
|
|
for i in range(10): # More events than queue capacity
|
|
await bridge.publish(run_id, f"event-{i}", {"i": i})
|
|
await bridge.publish_end(run_id)
|
|
|
|
async def consumer(run_id: str) -> list:
|
|
events = []
|
|
async for entry in bridge.subscribe(run_id, heartbeat_interval=0.1):
|
|
events.append(entry)
|
|
if entry is END_SENTINEL:
|
|
return events
|
|
return events # pragma: no cover
|
|
|
|
run_ids = [f"concurrent-{i}" for i in range(num_runs)]
|
|
results: dict[str, list] = {}
|
|
|
|
async def consume_into(run_id: str) -> None:
|
|
results[run_id] = await consumer(run_id)
|
|
|
|
with anyio.fail_after(10):
|
|
async with anyio.create_task_group() as task_group:
|
|
for run_id in run_ids:
|
|
task_group.start_soon(consume_into, run_id)
|
|
await anyio.sleep(0)
|
|
for run_id in run_ids:
|
|
task_group.start_soon(producer, run_id)
|
|
|
|
for run_id in run_ids:
|
|
events = results[run_id]
|
|
assert events[-1] is END_SENTINEL, f"Run {run_id} did not receive END sentinel"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Factory tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_make_stream_bridge_defaults():
|
|
"""make_stream_bridge() with no config yields a MemoryStreamBridge."""
|
|
async with make_stream_bridge() as bridge:
|
|
assert isinstance(bridge, MemoryStreamBridge)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _resolve_start_offset: O(1) seq-indexed resolution (parity with linear scan)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _linear_resolve(stream, last_event_id):
|
|
"""The original linear-scan resolver, kept as a parity reference."""
|
|
if last_event_id is None:
|
|
return stream.start_offset
|
|
for index, entry in enumerate(stream.events):
|
|
if entry.id == last_event_id:
|
|
return stream.start_offset + index + 1
|
|
return stream.start_offset
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"event_id,expected",
|
|
[
|
|
("1718000000000-0", 0),
|
|
("1718000000000-42", 42),
|
|
("garbage", None), # no separator
|
|
("1718000000000-x", None), # non-integer seq
|
|
("", None),
|
|
],
|
|
)
|
|
def test_parse_event_seq(event_id, expected):
|
|
assert MemoryStreamBridge._parse_event_seq(event_id) == expected
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_resolve_start_offset_matches_linear_scan():
|
|
"""The seq-indexed resolver must return exactly what the linear scan returned,
|
|
across retained, evicted, foreign (same seq / wrong ts), malformed, and None ids."""
|
|
bridge = MemoryStreamBridge(queue_maxsize=4)
|
|
run_id = "run-parity"
|
|
ids = []
|
|
for i in range(10):
|
|
await bridge.publish(run_id, f"e{i}", {"i": i})
|
|
ids.append(bridge._streams[run_id].events[-1].id) # includes ids that later evict
|
|
stream = bridge._streams[run_id]
|
|
assert stream.start_offset == 6 # 10 published, buffer of 4 retains seq 6..9
|
|
|
|
# A foreign id: a retained event's seq but a different timestamp -> must NOT match.
|
|
ts, _, seq_text = stream.events[0].id.rpartition("-")
|
|
foreign_id = f"{int(ts) + 1}-{seq_text}"
|
|
|
|
candidates = [None, "garbage", "1718000000000-x", "999999-999999", foreign_id, *ids]
|
|
for eid in candidates:
|
|
assert bridge._resolve_start_offset(stream, eid) == _linear_resolve(stream, eid), eid
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_subscribe_with_unknown_last_event_id_replays_from_earliest():
|
|
"""A foreign/garbage Last-Event-ID falls back to replaying retained events."""
|
|
bridge = MemoryStreamBridge(queue_maxsize=10)
|
|
run_id = "run-unknown-id"
|
|
await bridge.publish(run_id, "first", {})
|
|
await bridge.publish(run_id, "second", {})
|
|
await bridge.publish_end(run_id)
|
|
|
|
received = []
|
|
async for entry in bridge.subscribe(run_id, last_event_id="not-a-real-id", heartbeat_interval=1.0):
|
|
received.append(entry)
|
|
if entry is END_SENTINEL:
|
|
break
|
|
|
|
assert [entry.event for entry in received[:-1]] == ["first", "second"]
|
|
assert received[-1] is END_SENTINEL
|