test_agent.py

import unittest
from agent import build_graph
from langchain_core.messages import HumanMessage


class BaseLangGraphTest(unittest.TestCase):
    """Base test class for all agent unit-tests"""
    # Base prompt (will be given a default prompt from the agent file if none)
    system_prompt = None 

    @classmethod
    def setUpClass(cls):
        cls.graph = build_graph(test_mode=True, system_prompt=system_prompt_calculator)
    
    def run_agent(self, query: str) -> str:
        messages = [HumanMessage(content=query)]
        result = self.graph.invoke({"messages": messages})
        return result["messages"][-1].content.strip().lower()


class TestLangGraphCalculator(BaseLangGraphTest):
    """Tests the calculation capabilities of the llm
    by utilizing its pre-built tools"""

    system_prompt = """You are a calculator assistant. Use the available tools to perform arithmetic. 
Always use the correct parameter names: a and b. Never invent new parameter names.
Only respond to the user query by selecting the correct tool with appropriate inputs.
"""

    def test_combined_add_multiply(self):
        response = self.run_agent("What is 5 plus 3 then multiplied by 2?")
        self.assertIn("16", response)
    
    def test_combined_subtract_modulus(self):
        response = self.run_agent("What is (23 - 7) % 5")
        self.assertIn("1", response)
    
    def test_nested_expression(self):
        response = self.run_agent("What is (4 plus 6)/2 * 3 modulus 4?")
        self.assertIn("3", response)


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 test_wikipedia_search(self):
        response = self.run_agent("search wikipedia for Alan Turing")
        self.assertIn("alan turing", response)
        self.assertIn("<document", response)

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

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


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