tests/test_search.py

import time
import unittest
from test_base import BaseLangGraphTest

class TestLangGraphSearch(BaseLangGraphTest):
    """Tests for search tools (wikipedia, arxiv, web)"""

#     system_prompt = """You are a searching assistant. Use the available tools to perform arithmetic. 
# Always use the correct parameter names! Never invent new parameter names.
# Only respond to the user query by selecting the correct tool with appropriate inputs.
# Follow the structured formatting when delivering the output.
# """
    def tearDown(self):
        time.sleep(2)

    def test_wikipedia_search(self):
        response = self.run_agent("search wikipedia for Alan Turing")
        response_lower = response.lower()
        self.assertIn("alan", response_lower)
        self.assertIn("turing", response_lower)
        self.assertIn("<document", response_lower)

    def test_arxiv_search(self):
        response = self.run_agent("Find arxiv papers on quantum computing")
        response_lower = response.lower()
        self.assertIn("quantum", response_lower)
        self.assertIn("<document", response_lower)

    def test_web_search(self):
        response = self.run_agent("Search the web for current news on AI safety")
        self.assertIn("<document", response.lower())
        self.assertGreater(len(response), 50)

if __name__ == "__main__":
    unittest.main()