From 7257de2bf4cba4cbdd81a5fa20e28d1ee047e7a8 Mon Sep 17 00:00:00 2001
From: Brendan Quinn <brendan-quinn@northwestern.edu>
Date: Wed, 18 Dec 2024 05:26:42 +0000
Subject: [PATCH] Full test coverage for metrics callback

---
 chat/test/agent/callbacks/test_metrics.py | 113 ++++++++++++++++++++++
 1 file changed, 113 insertions(+)

diff --git a/chat/test/agent/callbacks/test_metrics.py b/chat/test/agent/callbacks/test_metrics.py
index d421f7e..3da9f70 100644
--- a/chat/test/agent/callbacks/test_metrics.py
+++ b/chat/test/agent/callbacks/test_metrics.py
@@ -92,3 +92,116 @@ def __init__(self, name, artifact):
         output = MockToolMessage("discover_fields", {})
         self.handler.on_tool_end(output)
         self.assertEqual(self.handler.artifacts, [])
+    
+    def test_on_llm_end_with_none_response(self):
+        self.handler.on_llm_end(None)
+        self.assertEqual(self.handler.answers, [])
+        self.assertEqual(self.handler.accumulator, {})
+    
+    def test_on_llm_end_with_empty_generations(self):
+        class MockLLMResult:
+            generations = [[]]  # Empty list
+
+        response = MockLLMResult()
+        self.handler.on_llm_end(response)
+        self.assertEqual(self.handler.answers, [])
+        self.assertEqual(self.handler.accumulator, {})
+    
+    def test_on_llm_end_with_empty_text(self):
+        class MockMessage:
+            def __init__(self):
+                self.text = ""
+                self.message = self
+                self.usage_metadata = {"input_tokens": 5}
+
+        class MockLLMResult:
+            generations = [[MockMessage()]]
+
+        response = MockLLMResult()
+        self.handler.on_llm_end(response)
+        self.assertEqual(self.handler.answers, [])
+        self.assertEqual(self.handler.accumulator, {"input_tokens": 5})
+    
+    def test_on_llm_end_missing_message(self):
+        class MockMessage:
+            def __init__(self):
+                self.text = "No message attribute"
+
+        class MockLLMResult:
+            generations = [[MockMessage()]]
+
+        response = MockLLMResult()
+        self.handler.on_llm_end(response)
+        self.assertEqual(self.handler.answers, ["No message attribute"])
+        self.assertEqual(self.handler.accumulator, {})
+    
+    def test_on_tool_end_invalid_json(self):
+        import json
+
+        # Mocking ToolMessage and an invalid artifact that simulates JSONDecodeError
+        class MockDoc:
+            def __init__(self, metadata):
+                self.metadata = metadata
+
+        class InvalidArtifact:
+            def __init__(self):
+                self.metadata = self  # Simulate the `metadata` attribute
+
+            def __getitem__(self, item):
+                if item == "api_link":
+                    raise json.decoder.JSONDecodeError("Expecting value", "doc", 0)
+                return None
+
+        class MockToolMessage:
+            def __init__(self, name, artifact, content):
+                self.name = name
+                self.artifact = artifact
+                self.content = content  # Adding content attribute
+
+        # Invalid artifact with a mocked `metadata` attribute
+        invalid_artifact = [
+            InvalidArtifact(),  # This will raise JSONDecodeError on access
+        ]
+
+        output = MockToolMessage("search", invalid_artifact, "example_content")
+
+        with patch("builtins.print") as mock_print:
+            self.handler.on_tool_end(output)
+            mock_print.assert_called_once_with(
+                "Invalid json (Expecting value: line 1 column 1 (char 0)) returned from search tool: example_content"
+            )
+            self.assertEqual(self.handler.artifacts, [])
+
+    def test_on_tool_end_unrecognized_tool(self):
+        class MockToolMessage:
+            def __init__(self, name, artifact):
+                self.name = name
+                self.artifact = artifact
+
+        output = MockToolMessage("unknown_tool", {})
+        self.handler.on_tool_end(output)
+        self.assertEqual(self.handler.artifacts, [])
+
+    def test_on_llm_end_with_none_metadata(self):
+        """
+        Test the on_llm_end method when usage_metadata is None.
+        Ensures that answers are processed correctly and accumulator remains unchanged.
+        """
+        # Mocking LLMResult with usage_metadata as None
+        class MockMessage:
+            def __init__(self, text, usage_metadata=None):
+                self.text = text
+                self.message = self  # For simplicity, reuse same object for .message
+                self.usage_metadata = usage_metadata
+
+        class MockLLMResult:
+            def __init__(self, text):
+                message = MockMessage(text, usage_metadata=None)
+                self.generations = [[message]]
+
+        response = MockLLMResult("Answer without metadata")
+        with patch.object(self.handler, "on_llm_end", wraps=self.handler.on_llm_end) as mock:
+            self.handler.on_llm_end(response)
+            mock.assert_called_once_with(response)
+            self.assertEqual(self.handler.answers, ["Answer without metadata"])
+            self.assertEqual(self.handler.accumulator, {})
\ No newline at end of file