103 lines
2.7 KiB
Python
103 lines
2.7 KiB
Python
# -*- coding: utf-8 -*-
|
|
""" Unit test for msghub."""
|
|
import unittest
|
|
from typing import Optional, Union, Sequence
|
|
|
|
from agentscope.agents import AgentBase
|
|
from agentscope import msghub
|
|
from agentscope.message import Msg
|
|
|
|
|
|
class TestAgent(AgentBase):
|
|
"""Test agent class for msghub."""
|
|
|
|
def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
|
|
"""Reply function for agent."""
|
|
if x is not None:
|
|
self.memory.add(x)
|
|
return x
|
|
else:
|
|
return {}
|
|
|
|
|
|
class MsgHubTest(unittest.TestCase):
|
|
"""
|
|
Test for MsgHub
|
|
"""
|
|
|
|
def setUp(self) -> None:
|
|
"""Init for ExampleTest."""
|
|
self.wisper = TestAgent("wisper")
|
|
self.agent1 = TestAgent("agent1")
|
|
self.agent2 = TestAgent("agent2")
|
|
self.agent3 = TestAgent("agent3")
|
|
|
|
def test_msghub_operation(self) -> None:
|
|
"""Test add, delete and broadcast operations"""
|
|
msg1 = Msg(name="a1", content="msg1")
|
|
msg2 = Msg(name="a2", content="msg2")
|
|
msg3 = Msg(name="a3", content="msg3")
|
|
msg4 = Msg(name="a4", content="msg4")
|
|
|
|
with msghub(participants=[self.agent1, self.agent2]) as hub:
|
|
self.agent1(msg1)
|
|
self.agent2(msg2)
|
|
|
|
hub.delete(self.agent1)
|
|
|
|
hub.add(self.agent3)
|
|
|
|
self.agent3(msg3)
|
|
|
|
hub.broadcast(msg4)
|
|
|
|
self.assertListEqual(
|
|
self.agent2.memory.get_memory(),
|
|
[
|
|
msg1,
|
|
msg2,
|
|
msg3,
|
|
msg4,
|
|
],
|
|
)
|
|
|
|
self.assertListEqual(self.agent1.memory.get_memory(), [msg1, msg2])
|
|
|
|
self.assertListEqual(self.agent3.memory.get_memory(), [msg3, msg4])
|
|
|
|
def test_msghub(self) -> None:
|
|
"""msghub test."""
|
|
|
|
ground_truth = [
|
|
Msg(
|
|
name="w1",
|
|
content="This secret that my password is 123456 can't be"
|
|
" leaked!",
|
|
role="wisper",
|
|
),
|
|
]
|
|
|
|
with msghub(participants=[self.wisper, self.agent1, self.agent2]):
|
|
self.wisper(ground_truth)
|
|
|
|
# agent1 and agent2 heard wisper's secret!
|
|
self.assertListEqual(
|
|
self.agent1.memory.get_memory(),
|
|
ground_truth,
|
|
)
|
|
|
|
self.assertListEqual(
|
|
self.agent2.memory.get_memory(),
|
|
ground_truth,
|
|
)
|
|
|
|
# agent3 didn't hear wisper's secret!
|
|
self.assertListEqual(
|
|
self.agent3.memory.get_memory(),
|
|
[],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|