wehub-resource-sync
47 lines
1.9 KiB
Python
47 lines
1.9 KiB
Python
import unittest
|
|
from unittest.mock import MagicMock, patch, AsyncMock
|
|
import asyncio
|
|
from gpt_researcher.agent import GPTResearcher
|
|
import os
|
|
|
|
class TestQuickSearch(unittest.TestCase):
|
|
|
|
@patch('gpt_researcher.agent.get_search_results', new_callable=AsyncMock)
|
|
@patch('gpt_researcher.agent.create_chat_completion', new_callable=AsyncMock)
|
|
@patch('langchain_openai.OpenAIEmbeddings')
|
|
def test_quick_search_no_summary(self, mock_embeddings, mock_create_chat, mock_search):
|
|
# Setup mocks
|
|
mock_search.return_value = [{'title': 'Test Result', 'content': 'Content', 'url': 'http://test.com'}]
|
|
|
|
# Initialize researcher with dummy config to avoid API key issues
|
|
researcher = GPTResearcher(query="test query")
|
|
|
|
# Run quick_search without summary
|
|
results = asyncio.run(researcher.quick_search("test query", aggregated_summary=False))
|
|
|
|
# Verify
|
|
self.assertEqual(len(results), 1)
|
|
self.assertEqual(results[0]['title'], 'Test Result')
|
|
mock_create_chat.assert_not_called()
|
|
|
|
@patch('gpt_researcher.agent.get_search_results', new_callable=AsyncMock)
|
|
@patch('gpt_researcher.agent.create_chat_completion', new_callable=AsyncMock)
|
|
@patch('langchain_openai.OpenAIEmbeddings')
|
|
def test_quick_search_with_summary(self, mock_embeddings, mock_create_chat, mock_search):
|
|
# Setup mocks
|
|
mock_search.return_value = [{'title': 'Test Result', 'content': 'Content', 'url': 'http://test.com'}]
|
|
mock_create_chat.return_value = "This is a summary."
|
|
|
|
# Initialize researcher
|
|
researcher = GPTResearcher(query="test query")
|
|
|
|
# Run quick_search with summary
|
|
summary = asyncio.run(researcher.quick_search("test query", aggregated_summary=True))
|
|
|
|
# Verify
|
|
self.assertEqual(summary, "This is a summary.")
|
|
mock_create_chat.assert_called_once()
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|