feat: solution for '8. Самописный поисковый агент на основе deep agents from scratch'

tests/test_agent.py

@@ -1,40 +1,68 @@
+"""
+Unit tests for the SearchAgent implementation.
+"""
+
 import unittest
-from unittest.mock import patch, MagicMock
-from src.agent import SearchAgent
+
+from src.agent import SearchAgent, Result
+
 
 class TestSearchAgent(unittest.TestCase):
-    @patch("src.agent.bing_search")
-    @patch("src.agent.pipeline")
-    def test_process_query(self, mock_pipeline, mock_bing_search):
-        # Mock Bing search results
-        mock_bing_search.return_value = [
-            {
-                "name": "Test Page",
-                "url": "http://example.com",
-                "snippet": "This is a test snippet.",
-            }
+    def setUp(self):
+        self.corpus = [
+            "The quick brown fox jumps over the lazy dog.",
+            "Deep learning models can capture complex patterns in data.",
+            "PyTorch is a popular deep learning framework.",
+            "Natural language processing involves understanding text.",
         ]
+        self.agent = SearchAgent(self.corpus, embedding_dim=20)
 
-        # Mock generator pipeline
-        def mock_generate(prompt, max_length, num_return_sequences):
-            return [
-                {
-                    "generated_text": f"{prompt} Summary: This is a test summary."
-                }
-            ]
+    def test_vocab_size(self):
+        # Vocabulary should contain all unique words
+        vocab_size = self.agent.get_vocab_size()
+        # Count unique words manually
+        unique_words = set()
+        for doc in self.corpus:
+            unique_words.update(doc.lower().split())
+        self.assertEqual(vocab_size, len(unique_words))
 
-        mock_pipeline.return_value = mock_generate
+    def test_document_embeddings_shape(self):
+        # Document embeddings should have shape (n_docs, dim)
+        doc_emb = self.agent.doc_embeddings
+        self.assertEqual(doc_emb.shape, (len(self.corpus), self.agent.embedding_dim))
 
-        agent = SearchAgent(api_key="dummy")
-        result = agent.process_query("test query")
-        self.assertIn("This is a test summary.", result)
+    def test_query_encoding_shape(self):
+        query = "deep learning"
+        vec = self.agent._encode_text(query)
+        self.assertEqual(vec.shape, (self.agent.embedding_dim,))
+
+    def test_cosine_similarity(self):
+        # Compute similarity between two identical vectors
+        vec = self.agent._encode_text("deep learning")
+        sims = self.agent._cosine_similarity(vec.unsqueeze(0), vec.unsqueeze(0))
+        self.assertAlmostEqual(sims.item(), 1.0, places=5)
+
+    def test_search_ranking(self):
+        # Query that matches second document
+        results = self.agent.search("deep learning", top_k=2)
+        # The first result should be the second document (index 1)
+        self.assertEqual(results[0].doc_id, 1)
+        self.assertGreater(results[0].score, results[1].score)
+
+    def test_unknown_words(self):
+        # Query with unknown words should still return results
+        results = self.agent.search("xyz abc", top_k=1)
+        self.assertEqual(len(results), 1)
+        self.assertIsInstance(results[0], Result)
+
+    def test_empty_query(self):
+        # Empty query should return top documents based on zero vector
+        results = self.agent.search("", top_k=3)
+        self.assertEqual(len(results), 3)
+        # Scores should be non‑negative
+        for res in results:
+            self.assertGreaterEqual(res.score, 0.0)
 
-    @patch("src.agent.bing_search")
-    def test_no_results(self, mock_bing_search):
-        mock_bing_search.return_value = []
-        agent = SearchAgent(api_key="dummy")
-        result = agent.process_query("no results")
-        self.assertEqual(result, "No results found.")
 
 if __name__ == "__main__":
     unittest.main()