mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-21 07:26:50 +00:00
rayhpeng
23eacf9533
Phase 2-A prerequisite for event storage: adds the unified run event stream interface (RunEventStore) with an in-memory implementation, RunEventsConfig, gateway integration, and comprehensive tests (27 cases). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
273 lines
11 KiB
Python
273 lines
11 KiB
Python
"""Tests for RunEventStore ABC + MemoryRunEventStore.
|
|
|
|
Covers:
|
|
- Basic write and query (put, seq assignment, cross-thread independence)
|
|
- list_messages (category filtering, pagination, cross-run ordering)
|
|
- list_events (run filtering, event_types filtering)
|
|
- list_messages_by_run
|
|
- count_messages
|
|
- put_batch
|
|
- delete_by_thread, delete_by_run
|
|
- Edge cases (empty thread/run)
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
|
|
|
|
|
@pytest.fixture
|
|
def store():
|
|
return MemoryRunEventStore()
|
|
|
|
|
|
# -- Basic write and query --
|
|
|
|
|
|
class TestPutAndSeq:
|
|
@pytest.mark.anyio
|
|
async def test_put_returns_dict_with_seq(self, store):
|
|
record = await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message", content="hello")
|
|
assert "seq" in record
|
|
assert record["seq"] == 1
|
|
assert record["thread_id"] == "t1"
|
|
assert record["run_id"] == "r1"
|
|
assert record["event_type"] == "human_message"
|
|
assert record["category"] == "message"
|
|
assert record["content"] == "hello"
|
|
assert record["metadata"] == {}
|
|
assert "created_at" in record
|
|
|
|
@pytest.mark.anyio
|
|
async def test_seq_strictly_increasing_same_thread(self, store):
|
|
r1 = await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
r2 = await store.put(thread_id="t1", run_id="r1", event_type="ai_message", category="message")
|
|
r3 = await store.put(thread_id="t1", run_id="r1", event_type="llm_end", category="trace")
|
|
assert r1["seq"] == 1
|
|
assert r2["seq"] == 2
|
|
assert r3["seq"] == 3
|
|
|
|
@pytest.mark.anyio
|
|
async def test_seq_independent_across_threads(self, store):
|
|
r1 = await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
r2 = await store.put(thread_id="t2", run_id="r2", event_type="human_message", category="message")
|
|
assert r1["seq"] == 1
|
|
assert r2["seq"] == 1
|
|
|
|
@pytest.mark.anyio
|
|
async def test_put_respects_provided_created_at(self, store):
|
|
ts = "2024-06-01T12:00:00+00:00"
|
|
record = await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message", created_at=ts)
|
|
assert record["created_at"] == ts
|
|
|
|
@pytest.mark.anyio
|
|
async def test_put_metadata_preserved(self, store):
|
|
meta = {"model": "gpt-4", "tokens": 100}
|
|
record = await store.put(thread_id="t1", run_id="r1", event_type="llm_end", category="trace", metadata=meta)
|
|
assert record["metadata"] == meta
|
|
|
|
|
|
# -- list_messages --
|
|
|
|
|
|
class TestListMessages:
|
|
@pytest.mark.anyio
|
|
async def test_only_returns_message_category(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="llm_end", category="trace")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="run_start", category="lifecycle")
|
|
messages = await store.list_messages("t1")
|
|
assert len(messages) == 1
|
|
assert messages[0]["category"] == "message"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ascending_seq_order(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message", content="first")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="ai_message", category="message", content="second")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message", content="third")
|
|
messages = await store.list_messages("t1")
|
|
seqs = [m["seq"] for m in messages]
|
|
assert seqs == sorted(seqs)
|
|
|
|
@pytest.mark.anyio
|
|
async def test_before_seq_pagination(self, store):
|
|
# Put 10 messages with seq 1..10
|
|
for i in range(10):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message", content=str(i))
|
|
messages = await store.list_messages("t1", before_seq=6, limit=3)
|
|
assert len(messages) == 3
|
|
assert [m["seq"] for m in messages] == [3, 4, 5]
|
|
|
|
@pytest.mark.anyio
|
|
async def test_after_seq_pagination(self, store):
|
|
for i in range(10):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message", content=str(i))
|
|
messages = await store.list_messages("t1", after_seq=7, limit=3)
|
|
assert len(messages) == 3
|
|
assert [m["seq"] for m in messages] == [8, 9, 10]
|
|
|
|
@pytest.mark.anyio
|
|
async def test_limit_restricts_count(self, store):
|
|
for _ in range(20):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
messages = await store.list_messages("t1", limit=5)
|
|
assert len(messages) == 5
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cross_run_unified_ordering(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="ai_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r2", event_type="human_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r2", event_type="ai_message", category="message")
|
|
messages = await store.list_messages("t1")
|
|
assert [m["seq"] for m in messages] == [1, 2, 3, 4]
|
|
assert messages[0]["run_id"] == "r1"
|
|
assert messages[2]["run_id"] == "r2"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_default_returns_latest(self, store):
|
|
for _ in range(10):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
messages = await store.list_messages("t1", limit=3)
|
|
assert [m["seq"] for m in messages] == [8, 9, 10]
|
|
|
|
|
|
# -- list_events --
|
|
|
|
|
|
class TestListEvents:
|
|
@pytest.mark.anyio
|
|
async def test_returns_all_categories_for_run(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="llm_end", category="trace")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="run_start", category="lifecycle")
|
|
events = await store.list_events("t1", "r1")
|
|
assert len(events) == 3
|
|
|
|
@pytest.mark.anyio
|
|
async def test_event_types_filter(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="llm_start", category="trace")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="llm_end", category="trace")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="tool_start", category="trace")
|
|
events = await store.list_events("t1", "r1", event_types=["llm_end"])
|
|
assert len(events) == 1
|
|
assert events[0]["event_type"] == "llm_end"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_only_returns_specified_run(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="llm_end", category="trace")
|
|
await store.put(thread_id="t1", run_id="r2", event_type="llm_end", category="trace")
|
|
events = await store.list_events("t1", "r1")
|
|
assert len(events) == 1
|
|
assert events[0]["run_id"] == "r1"
|
|
|
|
|
|
# -- list_messages_by_run --
|
|
|
|
|
|
class TestListMessagesByRun:
|
|
@pytest.mark.anyio
|
|
async def test_only_messages_for_specified_run(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="llm_end", category="trace")
|
|
await store.put(thread_id="t1", run_id="r2", event_type="human_message", category="message")
|
|
messages = await store.list_messages_by_run("t1", "r1")
|
|
assert len(messages) == 1
|
|
assert messages[0]["run_id"] == "r1"
|
|
assert messages[0]["category"] == "message"
|
|
|
|
|
|
# -- count_messages --
|
|
|
|
|
|
class TestCountMessages:
|
|
@pytest.mark.anyio
|
|
async def test_counts_only_message_category(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="ai_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="llm_end", category="trace")
|
|
assert await store.count_messages("t1") == 2
|
|
|
|
|
|
# -- put_batch --
|
|
|
|
|
|
class TestPutBatch:
|
|
@pytest.mark.anyio
|
|
async def test_batch_assigns_seq(self, store):
|
|
events = [
|
|
{"thread_id": "t1", "run_id": "r1", "event_type": "human_message", "category": "message", "content": "a"},
|
|
{"thread_id": "t1", "run_id": "r1", "event_type": "ai_message", "category": "message", "content": "b"},
|
|
{"thread_id": "t1", "run_id": "r1", "event_type": "llm_end", "category": "trace"},
|
|
]
|
|
results = await store.put_batch(events)
|
|
assert len(results) == 3
|
|
assert all("seq" in r for r in results)
|
|
|
|
@pytest.mark.anyio
|
|
async def test_batch_seq_strictly_increasing(self, store):
|
|
events = [
|
|
{"thread_id": "t1", "run_id": "r1", "event_type": "human_message", "category": "message"},
|
|
{"thread_id": "t1", "run_id": "r1", "event_type": "ai_message", "category": "message"},
|
|
]
|
|
results = await store.put_batch(events)
|
|
assert results[0]["seq"] == 1
|
|
assert results[1]["seq"] == 2
|
|
|
|
|
|
# -- delete --
|
|
|
|
|
|
class TestDelete:
|
|
@pytest.mark.anyio
|
|
async def test_delete_by_thread(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r1", event_type="ai_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r2", event_type="llm_end", category="trace")
|
|
count = await store.delete_by_thread("t1")
|
|
assert count == 3
|
|
assert await store.list_messages("t1") == []
|
|
assert await store.count_messages("t1") == 0
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_by_run(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r2", event_type="human_message", category="message")
|
|
await store.put(thread_id="t1", run_id="r2", event_type="llm_end", category="trace")
|
|
count = await store.delete_by_run("t1", "r2")
|
|
assert count == 2
|
|
# r1 events should still be there
|
|
messages = await store.list_messages("t1")
|
|
assert len(messages) == 1
|
|
assert messages[0]["run_id"] == "r1"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_nonexistent_thread_returns_zero(self, store):
|
|
assert await store.delete_by_thread("nope") == 0
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_nonexistent_run_returns_zero(self, store):
|
|
await store.put(thread_id="t1", run_id="r1", event_type="human_message", category="message")
|
|
assert await store.delete_by_run("t1", "nope") == 0
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_nonexistent_thread_for_run_returns_zero(self, store):
|
|
assert await store.delete_by_run("nope", "r1") == 0
|
|
|
|
|
|
# -- Edge cases --
|
|
|
|
|
|
class TestEdgeCases:
|
|
@pytest.mark.anyio
|
|
async def test_empty_thread_list_messages(self, store):
|
|
assert await store.list_messages("empty") == []
|
|
|
|
@pytest.mark.anyio
|
|
async def test_empty_run_list_events(self, store):
|
|
assert await store.list_events("empty", "r1") == []
|
|
|
|
@pytest.mark.anyio
|
|
async def test_empty_thread_count_messages(self, store):
|
|
assert await store.count_messages("empty") == 0
|