From ac97dc6d426c92efeefa018cea4471c0bbd6834f Mon Sep 17 00:00:00 2001 From: Andrew Barnes Date: Wed, 25 Mar 2026 12:20:50 -0400 Subject: [PATCH] test: add unit tests for TodoMiddleware (#1307) * test: add unit tests for TodoMiddleware Cover context-loss detection logic: - _todos_in_messages and _reminder_in_messages helpers - _format_todos formatting - Reminder injection when write_todos truncated - No-op when todos visible or reminder already present - abefore_model async delegation * test: fix event loop error in todo middleware async test Use asyncio.run() instead of get_event_loop().run_until_complete() to avoid RuntimeError on Python 3.12 where no default event loop exists in the main thread. --------- Co-authored-by: Willem Jiang --- backend/tests/test_todo_middleware.py | 156 ++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 backend/tests/test_todo_middleware.py diff --git a/backend/tests/test_todo_middleware.py b/backend/tests/test_todo_middleware.py new file mode 100644 index 0000000..8849384 --- /dev/null +++ b/backend/tests/test_todo_middleware.py @@ -0,0 +1,156 @@ +"""Tests for TodoMiddleware context-loss detection.""" + +import asyncio +from unittest.mock import MagicMock + +from langchain_core.messages import AIMessage, HumanMessage + +from deerflow.agents.middlewares.todo_middleware import ( + TodoMiddleware, + _format_todos, + _reminder_in_messages, + _todos_in_messages, +) + + +def _ai_with_write_todos(): + return AIMessage(content="", tool_calls=[{"name": "write_todos", "id": "tc_1", "args": {}}]) + + +def _reminder_msg(): + return HumanMessage(name="todo_reminder", content="reminder") + + +def _make_runtime(): + runtime = MagicMock() + runtime.context = {"thread_id": "test-thread"} + return runtime + + +def _sample_todos(): + return [ + {"status": "completed", "content": "Set up project"}, + {"status": "in_progress", "content": "Write tests"}, + {"status": "pending", "content": "Deploy"}, + ] + + +class TestTodosInMessages: + def test_true_when_write_todos_present(self): + msgs = [HumanMessage(content="hi"), _ai_with_write_todos()] + assert _todos_in_messages(msgs) is True + + def test_false_when_no_write_todos(self): + msgs = [ + HumanMessage(content="hi"), + AIMessage(content="hello", tool_calls=[{"name": "bash", "id": "tc_1", "args": {}}]), + ] + assert _todos_in_messages(msgs) is False + + def test_false_for_empty_list(self): + assert _todos_in_messages([]) is False + + def test_false_for_ai_without_tool_calls(self): + msgs = [AIMessage(content="hello")] + assert _todos_in_messages(msgs) is False + + +class TestReminderInMessages: + def test_true_when_reminder_present(self): + msgs = [HumanMessage(content="hi"), _reminder_msg()] + assert _reminder_in_messages(msgs) is True + + def test_false_when_no_reminder(self): + msgs = [HumanMessage(content="hi"), AIMessage(content="hello")] + assert _reminder_in_messages(msgs) is False + + def test_false_for_empty_list(self): + assert _reminder_in_messages([]) is False + + def test_false_for_human_without_name(self): + msgs = [HumanMessage(content="todo_reminder")] + assert _reminder_in_messages(msgs) is False + + +class TestFormatTodos: + def test_formats_multiple_items(self): + todos = _sample_todos() + result = _format_todos(todos) + assert "- [completed] Set up project" in result + assert "- [in_progress] Write tests" in result + assert "- [pending] Deploy" in result + + def test_empty_list(self): + assert _format_todos([]) == "" + + def test_missing_fields_use_defaults(self): + todos = [{"content": "No status"}, {"status": "done"}] + result = _format_todos(todos) + assert "- [pending] No status" in result + assert "- [done] " in result + + +class TestBeforeModel: + def test_returns_none_when_no_todos(self): + mw = TodoMiddleware() + state = {"messages": [HumanMessage(content="hi")], "todos": []} + assert mw.before_model(state, _make_runtime()) is None + + def test_returns_none_when_todos_is_none(self): + mw = TodoMiddleware() + state = {"messages": [HumanMessage(content="hi")], "todos": None} + assert mw.before_model(state, _make_runtime()) is None + + def test_returns_none_when_write_todos_still_visible(self): + mw = TodoMiddleware() + state = { + "messages": [_ai_with_write_todos()], + "todos": _sample_todos(), + } + assert mw.before_model(state, _make_runtime()) is None + + def test_returns_none_when_reminder_already_present(self): + mw = TodoMiddleware() + state = { + "messages": [HumanMessage(content="hi"), _reminder_msg()], + "todos": _sample_todos(), + } + assert mw.before_model(state, _make_runtime()) is None + + def test_injects_reminder_when_todos_exist_but_truncated(self): + mw = TodoMiddleware() + state = { + "messages": [HumanMessage(content="hi"), AIMessage(content="sure")], + "todos": _sample_todos(), + } + result = mw.before_model(state, _make_runtime()) + assert result is not None + msgs = result["messages"] + assert len(msgs) == 1 + assert isinstance(msgs[0], HumanMessage) + assert msgs[0].name == "todo_reminder" + + def test_reminder_contains_formatted_todos(self): + mw = TodoMiddleware() + state = { + "messages": [HumanMessage(content="hi")], + "todos": _sample_todos(), + } + result = mw.before_model(state, _make_runtime()) + content = result["messages"][0].content + assert "Set up project" in content + assert "Write tests" in content + assert "Deploy" in content + assert "system_reminder" in content + + +class TestAbeforeModel: + def test_delegates_to_sync(self): + mw = TodoMiddleware() + state = { + "messages": [HumanMessage(content="hi")], + "todos": _sample_todos(), + } + result = asyncio.run(mw.abefore_model(state, _make_runtime())) + assert result is not None + assert result["messages"][0].name == "todo_reminder"