72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
"""
|
|
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 os
|
|
import shutil
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from src.agent import CustomSearchAgent
|
|
|
|
|
|
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()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |