90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
import json
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from control_backend.agents.bdi import (
|
|
BDIBeliefCollectorAgent,
|
|
)
|
|
from control_backend.core.agent_system import InternalMessage
|
|
from control_backend.core.config import settings
|
|
from control_backend.schemas.belief_message import Belief
|
|
|
|
|
|
@pytest.fixture
|
|
def agent():
|
|
agent = BDIBeliefCollectorAgent("belief_collector_agent")
|
|
return agent
|
|
|
|
|
|
def make_msg(body: dict, sender: str = "sender"):
|
|
return InternalMessage(to="collector", sender=sender, body=json.dumps(body))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_message_routes_belief_text(agent, mocker):
|
|
"""
|
|
Test that when a message is received, _handle_belief_text is called with that message.
|
|
"""
|
|
payload = {"type": "belief_extraction_text", "beliefs": {"user_said": [["hi"]]}}
|
|
spy = mocker.patch.object(agent, "_handle_belief_text", new_callable=AsyncMock)
|
|
|
|
await agent.handle_message(make_msg(payload))
|
|
|
|
spy.assert_awaited_once_with(payload, "sender")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_message_routes_emotion(agent, mocker):
|
|
payload = {"type": "emotion_extraction_text"}
|
|
spy = mocker.patch.object(agent, "_handle_emo_text", new_callable=AsyncMock)
|
|
|
|
await agent.handle_message(make_msg(payload))
|
|
|
|
spy.assert_awaited_once_with(payload, "sender")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_message_bad_json(agent, mocker):
|
|
agent._handle_belief_text = AsyncMock()
|
|
bad_msg = InternalMessage(to="collector", sender="sender", body="not json")
|
|
|
|
await agent.handle_message(bad_msg)
|
|
|
|
agent._handle_belief_text.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_belief_text_sends_when_beliefs_exist(agent, mocker):
|
|
payload = {"type": "belief_extraction_text", "beliefs": {"user_said": ["hello"]}}
|
|
spy = mocker.patch.object(agent, "_send_beliefs_to_bdi", new_callable=AsyncMock)
|
|
expected = [Belief(name="user_said", arguments=["hello"])]
|
|
|
|
await agent._handle_belief_text(payload, "origin")
|
|
|
|
spy.assert_awaited_once_with(expected, origin="origin")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_belief_text_no_send_when_empty(agent, mocker):
|
|
payload = {"type": "belief_extraction_text", "beliefs": {}}
|
|
spy = mocker.patch.object(agent, "_send_beliefs_to_bdi", new_callable=AsyncMock)
|
|
|
|
await agent._handle_belief_text(payload, "origin")
|
|
|
|
spy.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_beliefs_to_bdi(agent):
|
|
agent.send = AsyncMock()
|
|
beliefs = [Belief(name="user_said", arguments=["hello", "world"])]
|
|
|
|
await agent._send_beliefs_to_bdi(beliefs, origin="origin")
|
|
|
|
agent.send.assert_awaited_once()
|
|
sent: InternalMessage = agent.send.call_args.args[0]
|
|
assert sent.to == settings.agent_settings.bdi_core_name
|
|
assert sent.thread == "beliefs"
|
|
assert json.loads(sent.body)["beliefs"] == [belief.model_dump() for belief in beliefs]
|