From e004922090921c2badf92e2dcbad25f256b41243 Mon Sep 17 00:00:00 2001
From: cszsolnai <csaba.zsolnai@epfl.ch>
Date: Fri, 27 Sep 2024 14:55:46 +0200
Subject: [PATCH] Added changelog

---
 CHANGELOG.md                                  |  1 +
 .../test_supervisor_multi_agent.py            | 35 ++++++++++---------
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index eb751ad..69bdafe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Update readme
+- Extra multi agent unit tests 
 
 ### Removed
 - Github action to create the docs.
diff --git a/tests/multi_agents/test_supervisor_multi_agent.py b/tests/multi_agents/test_supervisor_multi_agent.py
index c756826..4545a43 100644
--- a/tests/multi_agents/test_supervisor_multi_agent.py
+++ b/tests/multi_agents/test_supervisor_multi_agent.py
@@ -1,11 +1,10 @@
-from unittest.mock import MagicMock, AsyncMock
+from unittest.mock import AsyncMock, MagicMock
 
+import pytest
 from langchain_core.language_models import GenericFakeChatModel
 from langchain_core.messages import HumanMessage, SystemMessage
-
-import pytest
-
 from neuroagent.multi_agents.supervisor_multi_agent import AgentState
+
 from src.neuroagent.multi_agents import SupervisorMultiAgent
 
 
@@ -14,11 +13,9 @@ def test_create_main_agent():
     bind_function_result = MagicMock()
     bind_function_result.__ror__.return_value = {}
     mock_llm.bind_functions.return_value = bind_function_result
-    data = {
-        "llm": mock_llm,
-        "agents": [("agent1",)]
-    }
+    data = {"llm": mock_llm, "agents": [("agent1",)]}
     from src.neuroagent.multi_agents import SupervisorMultiAgent
+
     result = SupervisorMultiAgent.create_main_agent(data)
     assert "main_agent" in result
     assert "summarizer" in result
@@ -32,17 +29,16 @@ async def test_agent_node():
     )
 
     async def mock_ainvoke(_):
-        return {
-            "messages": [
-                mock_message
-            ]
-        }
+        return {"messages": [mock_message]}
 
     agent_state = MagicMock()
     agent = MagicMock()
     agent.ainvoke = mock_ainvoke
     from src.neuroagent.multi_agents import SupervisorMultiAgent
-    agent_node_test = await SupervisorMultiAgent.agent_node(agent_state, agent, "test_agent")
+
+    agent_node_test = await SupervisorMultiAgent.agent_node(
+        agent_state, agent, "test_agent"
+    )
 
     assert isinstance(agent_node_test, dict)
     assert "messages" in agent_node_test
@@ -61,9 +57,14 @@ def bind_tools(self, functions: list):
         def bind_functions(self, **kwargs):
             return self
 
-
-    fake_state = AgentState(messages=[HumanMessage(content="What is the airspeed velocity of an unladen swallow?"),
-                                      SystemMessage(content="11 m/s")])
+    fake_state = AgentState(
+        messages=[
+            HumanMessage(
+                content="What is the airspeed velocity of an unladen swallow?"
+            ),
+            SystemMessage(content="11 m/s"),
+        ]
+    )
 
     mock_llm = FakeChatModel(messages=iter([]))
     agent = SupervisorMultiAgent(agents=[("agent1", [])], llm=mock_llm)