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

tests/test_agent.py

@@ -1,17 +1,72 @@
 """
-Unit tests for the base Agent class.
+Unit tests for the custom search agent.
+
+These tests verify that the agent:
+1. Initializes correctly.
+2. Generates deterministic mock search results.
+3. Creates virtual files during search.
+4. Exports virtual files to disk.
 """
 
-import pytest
-from src.agent import Agent
+import os
+import shutil
+import tempfile
+import unittest
+from pathlib import Path
+
+from src.agent import CustomSearchAgent
 
 
-class DummyAgent(Agent):
-    def act(self, state):
-        return state
+class TestCustomSearchAgent(unittest.TestCase):
+    def setUp(self):
+        self.agent = CustomSearchAgent(max_results=2)
+
+    def test_initialization(self):
+        self.assertIsInstance(self.agent, CustomSearchAgent)
+        self.assertEqual(self.agent.max_results, 2)
+        self.assertEqual(self.agent.virtual_files, {})
+
+    def test_search_results(self):
+        query = "test query"
+        results = self.agent.search(query)
+        self.assertEqual(len(results), 2)
+        expected_titles = [
+            f"{query.title()} Result 1",
+            f"{query.title()} Result 2",
+        ]
+        actual_titles = [title for title, _ in results]
+        self.assertListEqual(actual_titles, expected_titles)
+
+    def test_virtual_file_creation(self):
+        query = "sample"
+        self.agent.search(query)
+        vfiles = self.agent.virtual_files
+        self.assertIn("result_1.txt", vfiles)
+        self.assertIn("result_2.txt", vfiles)
+        content = vfiles["result_1.txt"]
+        self.assertIn("Title: Sample Result 1", content)
+        self.assertIn("Snippet: This is a mock snippet for 'sample' (result 1).", content)
+
+    def test_export_virtual_files(self):
+        query = "export"
+        self.agent.search(query)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            out_dir = Path(tmpdir)
+            self.agent.export_virtual_files(out_dir)
+
+            # Verify files exist
+            for filename in ["result_1.txt", "result_2.txt"]:
+                file_path = out_dir / filename
+                self.assertTrue(file_path.is_file(), f"{filename} not found")
+                # Verify content matches
+                content = file_path.read_text(encoding="utf-8")
+                self.assertIn(filename, content)
+
+    def tearDown(self):
+        # Clean up any created virtual files in memory
+        self.agent._virtual_files.clear()
 
 
-def test_dummy_agent():
-    agent = DummyAgent()
-    assert agent.act(5) == 5
-    assert agent.act("hello") == "hello"
+if __name__ == "__main__":
+    unittest.main()