40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
from src.agent import SearchAgent
|
|
|
|
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.",
|
|
}
|
|
]
|
|
|
|
# Mock generator pipeline
|
|
def mock_generate(prompt, max_length, num_return_sequences):
|
|
return [
|
|
{
|
|
"generated_text": f"{prompt} Summary: This is a test summary."
|
|
}
|
|
]
|
|
|
|
mock_pipeline.return_value = mock_generate
|
|
|
|
agent = SearchAgent(api_key="dummy")
|
|
result = agent.process_query("test query")
|
|
self.assertIn("This is a test summary.", result)
|
|
|
|
@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() |