2024-06-28 10:22:44 -04:00
|
|
|
|
|
|
|
|
|
|
|
from dataclasses import dataclass
|
|
|
|
|
|
|
|
import pytest
|
|
|
|
from agnext.application import SingleThreadedAgentRuntime
|
|
|
|
|
|
|
|
from agnext.core import AgentRuntime, AgentId, CancellationToken
|
|
|
|
|
|
|
|
from agnext.components import ClosureAgent
|
|
|
|
|
|
|
|
|
|
|
|
import asyncio
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
class Message:
|
|
|
|
content: str
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
|
async def test_register_receives_publish() -> None:
|
|
|
|
runtime = SingleThreadedAgentRuntime()
|
|
|
|
|
|
|
|
queue = asyncio.Queue[tuple[str, str]]()
|
|
|
|
|
|
|
|
async def log_message(_runtime: AgentRuntime, id: AgentId, message: Message, cancellation_token: CancellationToken) -> None:
|
|
|
|
namespace = id.namespace
|
|
|
|
await queue.put((namespace, message.content))
|
|
|
|
|
2024-07-23 11:49:38 -07:00
|
|
|
await runtime.register("name", lambda: ClosureAgent("My agent", log_message))
|
2024-07-01 11:53:45 -04:00
|
|
|
run_context = runtime.start()
|
2024-06-28 10:22:44 -04:00
|
|
|
await runtime.publish_message(Message("first message"), namespace="default")
|
|
|
|
await runtime.publish_message(Message("second message"), namespace="default")
|
|
|
|
await runtime.publish_message(Message("third message"), namespace="default")
|
|
|
|
|
2024-07-01 11:53:45 -04:00
|
|
|
await run_context.stop_when_idle()
|
2024-06-28 10:22:44 -04:00
|
|
|
|
|
|
|
assert queue.qsize() == 3
|
|
|
|
assert queue.get_nowait() == ("default", "first message")
|
|
|
|
assert queue.get_nowait() == ("default", "second message")
|
|
|
|
assert queue.get_nowait() == ("default", "third message")
|
|
|
|
assert queue.empty()
|