forked from HAL9000/cleveragents-core
151 lines
5.0 KiB
Python
151 lines
5.0 KiB
Python
import asyncio
|
|
from typing import Any
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from behave import given, then, when
|
|
|
|
from cleveragents.agents.base import Agent, AgentWithMemory
|
|
from cleveragents.core.exceptions import ExecutionError
|
|
|
|
|
|
class DefaultAgent(Agent):
|
|
def __init__(
|
|
self, name: str, config: dict[str, Any] | None = None, raise_error: bool = False
|
|
):
|
|
self.raise_error = raise_error
|
|
super().__init__(name, config)
|
|
self.last_message = None
|
|
self.last_context = None
|
|
|
|
async def process_message(
|
|
self, message: Any, context: dict[str, Any] | None = None
|
|
) -> str:
|
|
self.last_message = message
|
|
self.last_context = context
|
|
if self.raise_error:
|
|
raise RuntimeError("boom")
|
|
return f"processed:{message}"
|
|
|
|
def get_capabilities(self) -> list[str]:
|
|
return ["default"]
|
|
|
|
|
|
class MemoryAgent(AgentWithMemory):
|
|
async def process_message(
|
|
self, message: Any, context: dict[str, Any] | None = None
|
|
) -> Any:
|
|
await self.remember("last", message)
|
|
return f"mem:{message}"
|
|
|
|
def get_capabilities(self) -> list[str]:
|
|
return ["memory"]
|
|
|
|
|
|
@given("a basic test agent setup")
|
|
def step_basic_agent(context):
|
|
context.loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(context.loop)
|
|
context.agent = DefaultAgent("base-test")
|
|
context.captured_output = []
|
|
|
|
|
|
@when("I patch asyncio create_task and send a message")
|
|
def step_patch_create_task(context):
|
|
with patch("asyncio.create_task") as mock_task:
|
|
context.agent.send_message("hello", {"foo": "bar"})
|
|
context.create_task_called_with = None
|
|
if mock_task.call_args:
|
|
context.create_task_called_with = mock_task.call_args[0][0]
|
|
|
|
|
|
@then("the processing pipeline should schedule the process wrapper task")
|
|
def step_verify_create_task_called(context):
|
|
assert context.create_task_called_with is not None
|
|
assert context.create_task_called_with.__qualname__.endswith("_process_wrapper")
|
|
|
|
|
|
@when("I process a tuple message through the wrapper")
|
|
def step_process_tuple(context):
|
|
context.captured_output = []
|
|
context.agent.output_stream.subscribe(
|
|
lambda value: context.captured_output.append(value)
|
|
)
|
|
asyncio.run(context.agent._process_wrapper(("tuple-msg", {"ctx": 1})))
|
|
|
|
|
|
@then("the tuple message and context should be passed and emitted")
|
|
def step_verify_tuple_processing(context):
|
|
assert context.agent.last_message == "tuple-msg"
|
|
assert context.agent.last_context == {"ctx": 1}
|
|
assert context.captured_output == ["processed:tuple-msg"]
|
|
|
|
|
|
@when("I process a non tuple message through the wrapper")
|
|
def step_process_non_tuple(context):
|
|
context.captured_output = []
|
|
context.agent.output_stream.subscribe(
|
|
lambda value: context.captured_output.append(value)
|
|
)
|
|
asyncio.run(context.agent._process_wrapper("plain-msg"))
|
|
|
|
|
|
@then("the non tuple message should be processed with empty context")
|
|
def step_verify_non_tuple_processing(context):
|
|
assert context.agent.last_message == "plain-msg"
|
|
assert context.agent.last_context == {}
|
|
assert context.captured_output == ["processed:plain-msg"]
|
|
|
|
|
|
@when("process_message raises an error inside the wrapper")
|
|
def step_wrapper_raises(context):
|
|
error_agent = DefaultAgent("err-agent", raise_error=True)
|
|
context.wrapper_error = None
|
|
try:
|
|
asyncio.run(error_agent._process_wrapper("fail"))
|
|
except Exception as exc:
|
|
context.wrapper_error = exc
|
|
|
|
|
|
@then("an ExecutionError should be raised with the agent name")
|
|
def step_verify_execution_error(context):
|
|
assert isinstance(context.wrapper_error, ExecutionError)
|
|
assert "err-agent" in str(context.wrapper_error)
|
|
|
|
|
|
@when("I call process_message_sync on the agent")
|
|
def step_call_process_message_sync(context):
|
|
loop = context.loop
|
|
asyncio.set_event_loop(loop)
|
|
context.sync_result = context.agent.process_message_sync("sync-msg", {"sync": True})
|
|
|
|
|
|
@then("the synchronous result should match the async process output")
|
|
def step_verify_sync_result(context):
|
|
assert context.sync_result == "processed:sync-msg"
|
|
assert context.agent.last_context == {"sync": True}
|
|
|
|
|
|
@when("I call dispose on an agent with disposable streams")
|
|
def step_call_dispose(context):
|
|
context.agent.input_stream = MagicMock(dispose=MagicMock())
|
|
context.agent.output_stream = MagicMock(dispose=MagicMock())
|
|
context.agent.dispose()
|
|
|
|
|
|
@then("both input and output streams should be disposed if available")
|
|
def step_verify_dispose(context):
|
|
context.agent.input_stream.dispose.assert_called_once()
|
|
context.agent.output_stream.dispose.assert_called_once()
|
|
|
|
|
|
@when("I remember a value and recall it later")
|
|
def step_memory_remember_recall(context):
|
|
context.memory_agent = MemoryAgent("mem-agent")
|
|
asyncio.run(context.memory_agent.remember("answer", 42))
|
|
context.recalled = asyncio.run(context.memory_agent.recall("answer"))
|
|
|
|
|
|
@then("the remembered value should be returned from memory")
|
|
def step_verify_memory_recall(context):
|
|
assert context.recalled == 42
|