test_web_search.py

"""Offline regression tests for the bounded web-search path."""

import os
import unittest
from unittest.mock import patch

import tools


class FakeTavilyClient:
    def __init__(self, response):
        self.response = response
        self.calls = []

    def search(self, **kwargs):
        self.calls.append(kwargs)
        return self.response


class WebSearchTests(unittest.TestCase):
    def tearDown(self):
        tools._tavily_client = None

    def test_query_is_normalized_and_capped(self):
        self.assertEqual(tools._compact_query("  a   short\nquery "), "a short query")

        long_query = "word " * 200
        compact = tools._compact_query(long_query)
        self.assertLessEqual(len(compact), tools.WEB_QUERY_MAX_CHARS)
        self.assertFalse(compact.endswith(" "))

    def test_result_payload_has_hard_limits(self):
        response = {
            "answer": "a" * 3_000,
            "results": [
                {
                    "title": f"Result {index}",
                    "url": f"https://example.com/{index}",
                    "content": "x" * 4_000,
                }
                for index in range(8)
            ],
        }

        output = tools._format_search_results(response)

        self.assertLessEqual(len(output), tools.WEB_OUTPUT_MAX_CHARS)
        self.assertIn("Source 3", output)
        self.assertNotIn("Source 4", output)
        self.assertIn("[truncated]", output)

    def test_web_research_is_one_direct_tavily_call(self):
        fake = FakeTavilyClient(
            {
                "answer": "Paris",
                "results": [
                    {
                        "title": "France",
                        "url": "https://example.com/france",
                        "content": "Paris is the capital of France.",
                    }
                ],
            }
        )
        tools._tavily_client = fake

        output = tools.web_research.invoke(
            {"question": "  capital   of France "}
        )

        self.assertIn("Paris", output)
        self.assertEqual(len(fake.calls), 1)
        self.assertEqual(fake.calls[0]["query"], "capital of France")
        self.assertEqual(fake.calls[0]["max_results"], tools.WEB_MAX_RESULTS)
        self.assertTrue(fake.calls[0]["include_answer"])
        self.assertFalse(fake.calls[0]["include_raw_content"])
        self.assertFalse(fake.calls[0]["include_images"])

    def test_missing_api_key_fails_only_when_search_is_used(self):
        tools._tavily_client = None
        with patch.dict(os.environ, {}, clear=True):
            with self.assertRaisesRegex(RuntimeError, "TAVILY_API_KEY"):
                tools._get_tavily_client()

    def test_manager_exposes_only_five_bounded_tools(self):
        self.assertEqual(
            [tool.name for tool in tools.ALL_TOOLS],
            [
                "web_research",
                "file_analysis",
                "media_analysis",
                "python_exec",
                "calculator",
            ],
        )


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