add more tests

Author: mehrdadrad

tests/unit_tests/test_core.py

@@ -177,6 +177,39 @@ def test_llm_factory_openai(self, mock_settings, mock_chat_openai):
         llm_factory("default")
         mock_chat_openai.assert_called_once()
 
+    @patch("core.llm.AzureChatOpenAI")
+    @patch("core.llm.settings")
+    def test_llm_factory_azure(self, mock_settings, mock_azure):
+        mock_settings.llm.profiles = {
+            "azure": MagicMock(
+                provider="azure",
+                model="gpt-4",
+                api_key="key",
+                api_base="endpoint",
+                api_version="2023-05-15"
+            )
+        }
+        llm_factory("azure")
+        mock_azure.assert_called_once()
+
+    @patch("core.llm.ChatAnthropic")
+    @patch("core.llm.settings")
+    def test_llm_factory_anthropic(self, mock_settings, mock_anthropic):
+        mock_settings.llm.profiles = {
+            "claude": MagicMock(provider="anthropic", model="claude-3", api_key="key")
+        }
+        llm_factory("claude")
+        mock_anthropic.assert_called_once()
+
+    @patch("core.llm.ChatGoogleGenerativeAI")
+    @patch("core.llm.settings")
+    def test_llm_factory_google(self, mock_settings, mock_google):
+        mock_settings.llm.profiles = {
+            "gemini": MagicMock(provider="gemini", model="gemini-pro", credentials="creds")
+        }
+        llm_factory("gemini")
+        mock_google.assert_called_once()
+
     @patch("core.llm.settings")
     def test_llm_factory_unknown_provider(self, mock_settings):
         mock_settings.llm.profiles = {

tests/unit_tests/test_graph.py

@@ -173,6 +173,13 @@ def test_route_workflow(self):
             route_workflow({"next_worker": "common_agent", "messages": [dummy_msg]}), "common_agent"
         )
 
+    def test_route_workflow_unknown(self):
+        """Test routing with an unknown worker."""
+        dummy_msg = HumanMessage(content="test")
+        self.assertEqual(
+            route_workflow({"next_worker": "unknown", "messages": [dummy_msg]}), "unknown"
+        )
+
     def test_route_after_worker(self):
         """Test routing logic after a worker agent finishes."""
         # Tool calls present -> route to tools
@@ -241,6 +248,21 @@ def test_detect_tool_loop(self):
         }
         self.assertTrue(detect_tool_loop(state_alternating))
 
+    def test_detect_tool_loop_different_args(self):
+        """Test that same tool with different args is not a loop."""
+        call_a = {"name": "tool_a", "args": {"x": 1}, "id": "1"}
+        call_b = {"name": "tool_a", "args": {"x": 2}, "id": "2"}
+        state = {
+            "messages": [
+                AIMessage(content="", tool_calls=[call_a]),
+                ToolMessage(content="res", tool_call_id="1", name="tool_a"),
+                AIMessage(content="", tool_calls=[call_b]),
+                ToolMessage(content="res", tool_call_id="2", name="tool_a"),
+                AIMessage(content="", tool_calls=[call_a]),
+            ]
+        }
+        self.assertFalse(detect_tool_loop(state))
+
     def test_route_after_worker_loop_detection(self):
         """Test that route_after_worker catches loops."""
         tool_call = {"name": "tool_a", "args": {"x": 1}, "id": "1"}
@@ -578,6 +600,17 @@ def test_update_step_status(self):
         self.assertEqual(steps_status[0], "completed")
         self.assertEqual(steps_status[1], "skipped")
 
+    def test_update_step_status_invalid_id(self):
+        """Test update_step_status with an ID that doesn't exist."""
+        steps_status = ["pending"]
+        decision = MagicMock()
+        decision.current_step_id = 99
+        decision.current_step_status = "completed"
+        decision.skipped_step_ids = []
+
+        update_step_status(decision, steps_status)
+        self.assertEqual(steps_status[0], "pending")
+
     def test_check_completion(self):
         decision = MagicMock()
         decision.next_worker = "supervisor"

tests/unit_tests/test_utils.py

@@ -43,6 +43,11 @@ def test_context_formatting(self):
         res = StatusGenerator.parse_tool_call("test", {"items": [1, 2, 3]})
         self.assertIn(": 3 items", res)
 
+    def test_parse_tool_call_no_args(self):
+        res = StatusGenerator.parse_tool_call("simple_tool", {})
+        self.assertIn("Processing", res)
+        self.assertIn("Simple", res)
+
 class TestLoggerUtils(unittest.TestCase):
     def test_parse_size(self):
         self.assertEqual(_parse_size("10B"), 10)
@@ -59,6 +64,9 @@ def test_parse_size_case_insensitive(self):
     def test_parse_size_float(self):
         self.assertEqual(_parse_size("1.5 KB"), int(1.5 * 1024))
 
+    def test_parse_size_no_space(self):
+        self.assertEqual(_parse_size("1KB"), 1024)
+
     @patch("src.utils.logger.logging")
     @patch("src.utils.logger.RotatingFileHandler")
     @patch("src.utils.logger.settings")