From a65ad6a4e0c0497c7e8a42d15b86815f84c798b3 Mon Sep 17 00:00:00 2001
From: Sid <shoutsid@gmail.com>
Date: Tue, 17 Oct 2023 17:31:29 +0100
Subject: [PATCH] feat: Introduce PlannerService Class with Tests

- Implement PlannerService class to encapsulate planner-related logic.
- Add pytest tests for PlannerService to validate its functionality.
- Fix mocking issue in PlannerService tests.
---
 agents/services/planner_service.py            | 23 ++++++++++++++++++
 tests/agents/services/planner_service_test.py | 24 +++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 agents/services/planner_service.py
 create mode 100644 tests/agents/services/planner_service_test.py

diff --git a/agents/services/planner_service.py b/agents/services/planner_service.py
new file mode 100644
index 0000000..9d1de34
--- /dev/null
+++ b/agents/services/planner_service.py
@@ -0,0 +1,23 @@
+import sys
+import os
+
+sys.path.insert(
+    0,
+    os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")),
+)
+
+from autogen import AssistantAgent, UserProxyAgent
+
+
+class PlannerService:
+    def __init__(self, config_list):
+        self.planner = AssistantAgent(
+            name="planner", llm_config={"config_list": config_list}
+        )
+        self.planner_user = UserProxyAgent(
+            name="planner_user", max_consecutive_auto_reply=0, human_input_mode="NEVER"
+        )
+
+    def ask_planner(self, msg):
+        self.planner_user.initiate_chat(self.planner, message=msg)
+        return self.planner_user.last_message()["content"]
diff --git a/tests/agents/services/planner_service_test.py b/tests/agents/services/planner_service_test.py
new file mode 100644
index 0000000..6691c2e
--- /dev/null
+++ b/tests/agents/services/planner_service_test.py
@@ -0,0 +1,24 @@
+import pytest
+from unittest.mock import patch, MagicMock
+import sys
+import os
+
+sys.path.insert(
+    0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+)
+
+from agents.services.planner_service import PlannerService
+
+
+@pytest.mark.parametrize(
+    "msg, expected", [("Hello", "Hello, Planner!"), ("Test", "Test message!")]
+)
+def test_ask_planner(msg, expected):
+    with patch("agents.services.planner_service.AssistantAgent") as MockedAgent, patch(
+        "agents.services.planner_service.UserProxyAgent"
+    ) as MockedProxy:
+        mocked_last_message = MagicMock()
+        mocked_last_message.return_value = {"content": expected}
+        MockedProxy.return_value.last_message = mocked_last_message
+        service = PlannerService("some_config")
+        assert service.ask_planner(msg) == expected