tests/test_base.py

import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

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


class BaseLangGraphTest(unittest.TestCase):
    """Default to a weak but higher rate model for testing agentic tool usage."""
    system_prompt = None

    @classmethod
    def setUpClass(cls):
        cls.graph = build_graph(test_mode=True, system_prompt=cls.system_prompt)

    def run_agent(self, query: str) -> str:
        messages = [HumanMessage(content=query)]
        result = self.graph.invoke({"messages": messages})
        return result["messages"][-1].content.strip().lower()