tests/test_tools.py

import json
from unittest.mock import patch, MagicMock
from agent.tools import get_financial_metrics, search_filing, search_transcript, search_news


def _sample_metrics():
    return {
        "ticker": "AAPL", "company_name": "Apple Inc.", "filing_date": "2024-11-01",
        "period": "FY2024", "form_type": "10-K", "revenue": 391035000000.0,
        "revenue_yoy_pct": 2.0, "eps": 6.11, "gross_margin": 0.461,
        "operating_margin": 0.314, "free_cash_flow": 108807000000.0,
        "guidance_disclosed": 0, "guidance_text": None, "ingested_at": "2026-05-04T12:00:00",
        "period_basis": "annual", "report_date": "2024-09-28",
        "accession": "0000320193-24-000123", "source_url": "https://www.sec.gov/example",
        "metric_contexts": {
            "revenue": {"source": "sec_companyfacts", "selection": "exact_annual", "concept": "Revenues"},
        },
        "quality_warnings": [], "data_quality_status": "VERIFIED",
    }


@patch("agent.tools.metrics_db.get_all_metrics")
def test_get_metrics_found(mock_get):
    mock_get.return_value = [_sample_metrics()]
    result = get_financial_metrics.invoke({"ticker": "AAPL"})
    assert "Apple Inc." in result
    assert "$391.0B" in result
    assert "6.11" in result
    payload = json.loads(result)
    assert payload["schema"] == "evidence.v1"
    assert payload["status"] == "OK"
    assert payload["records"][0]["ref"]["source"] == "metrics"


@patch("agent.tools.metrics_db.get_all_metrics")
def test_get_metrics_not_ingested(mock_get):
    mock_get.return_value = []
    result = get_financial_metrics.invoke({"ticker": "ZZZZ"})
    assert "not ingested" in result.lower()


@patch("agent.tools.metrics_db.get_all_metrics")
def test_get_metrics_filters_unverified_rows_and_exposes_lineage(mock_get):
    verified = _sample_metrics()
    check_required = {
        **_sample_metrics(),
        "period": "Q32024",
        "filing_date": "2024-08-28",
        "accession": "0000320193-24-000099",
        "source_url": "https://www.sec.gov/check-required-example",
        "data_quality_status": "CHECK_REQUIRED",
        "quality_warnings": ["free_cash_flow:fallback_non_sec"],
    }
    mock_get.return_value = [verified, check_required]

    payload = json.loads(get_financial_metrics.invoke({"ticker": "AAPL"}))

    assert payload["status"] == "OK"
    assert len(payload["records"]) == 2
    verified_record, check_required_record = payload["records"]
    assert verified_record["metadata"]["data_quality_status"] == "VERIFIED"
    assert "sec_companyfacts/exact_annual/Revenues" in verified_record["content"]
    assert check_required_record["ref"]["source_url"] == check_required["source_url"]
    assert check_required_record["metadata"]["accession"] == check_required["accession"]
    assert check_required_record["metadata"]["data_quality_status"] == "CHECK_REQUIRED"
    assert check_required_record["metadata"]["quality_warnings"] == [
        "free_cash_flow:fallback_non_sec"
    ]
    assert "Data Quality: CHECK_REQUIRED" in check_required_record["content"]
    assert "free_cash_flow:fallback_non_sec" in check_required_record["content"]


@patch("agent.tools.metrics_db.get_all_metrics")
def test_get_metrics_all_unverified_is_empty(mock_get):
    mock_get.return_value = [{
        **_sample_metrics(),
        "data_quality_status": "LEGACY_UNVERIFIED",
    }]

    payload = json.loads(get_financial_metrics.invoke({"ticker": "AAPL"}))

    assert payload["status"] == "EMPTY"
    assert payload["records"] == []
    assert "none carries SEC period lineage" in payload["message"]


@patch("agent.tools.metrics_db.get_all_metrics")
def test_get_metrics_preserves_newest_first_order_with_check_required(mock_get):
    newest = {
        **_sample_metrics(),
        "period": "Q12027",
        "filing_date": "2026-05-20",
        "data_quality_status": "CHECK_REQUIRED",
        "quality_warnings": ["free_cash_flow:fallback_non_sec"],
    }
    older = {
        **_sample_metrics(),
        "period": "Q42026",
        "filing_date": "2026-02-25",
    }
    mock_get.return_value = [newest, older]

    payload = json.loads(get_financial_metrics.invoke({"ticker": "NVDA"}))

    assert payload["status"] == "OK"
    assert [record["metadata"]["period"] for record in payload["records"]] == [
        "Q12027", "Q42026"
    ]
    assert payload["records"][0]["metadata"]["data_quality_status"] == "CHECK_REQUIRED"


@patch("agent.tools.metrics_db.get_all_metrics")
def test_get_metrics_includes_check_required_but_excludes_legacy(mock_get):
    check_required = {
        **_sample_metrics(),
        "period": "Q12025",
        "data_quality_status": "CHECK_REQUIRED",
        "quality_warnings": ["eps:fallback_non_sec"],
    }
    legacy = {
        **_sample_metrics(),
        "period": "FY2023",
        "data_quality_status": "LEGACY_UNVERIFIED",
    }
    mock_get.return_value = [check_required, legacy]

    payload = json.loads(get_financial_metrics.invoke({"ticker": "AAPL"}))

    assert payload["status"] == "OK"
    assert len(payload["records"]) == 1
    assert payload["records"][0]["metadata"]["period"] == "Q12025"
    assert payload["records"][0]["metadata"]["data_quality_status"] == "CHECK_REQUIRED"


@patch("agent.tools.vector_store.search")
def test_search_filing_returns_chunks(mock_search):
    mock_search.return_value = [
        {"text": "Revenue grew 5%.", "metadata": {"source": "10-Q", "section": "MD&A", "filing_date": "2024-11-01"}}
    ]
    result = search_filing.invoke({"query": "revenue growth", "ticker": "AAPL"})
    assert "Revenue grew 5%" in result
    assert "10-Q" in result
    payload = json.loads(result)
    assert payload["status"] == "OK"
    assert payload["records"][0]["ref"]["source"] == "10-Q"


@patch("agent.tools.vector_store.search")
def test_search_filing_header_uses_source_prefix_when_no_chunk_context(mock_search):
    # chunk_context absent → fallback header must say "Source:" so the LLM can tag correctly
    mock_search.return_value = [
        {"text": "Margin expanded.", "metadata": {"source": "10-K", "section": "MD&A", "filing_date": "2024-11-01"}}
    ]
    result = search_filing.invoke({"query": "margins", "ticker": "AAPL"})
    assert "Source:" in result
    assert "10-K" in result


@patch("agent.tools.vector_store.search")
def test_search_transcript_header_uses_source_prefix_when_no_chunk_context(mock_search):
    mock_search.return_value = [
        {"text": "Tim Cook: strong demand.", "metadata": {"date": "2024-11-01"}}
    ]
    result = search_transcript.invoke({"query": "tone", "ticker": "AAPL"})
    assert "Source:" in result
    assert "transcript" in result


@patch("agent.tools.vector_store.search")
def test_search_filing_with_since_parameter(mock_search):
    mock_search.return_value = [
        {"text": "Recent revenue data.", "metadata": {"source": "10-Q", "section": "MD&A", "filing_date": "2025-01-15"}}
    ]
    result = search_filing.invoke({"query": "revenue", "ticker": "AAPL", "since": "2024-01-01"})
    assert "Recent revenue data" in result
    # Verify min_filing_date was passed to the search function
    mock_search.assert_called_once()
    call_kwargs = mock_search.call_args[1]
    assert call_kwargs.get("min_filing_date") == "2024-01-01"


@patch("agent.tools.vector_store.search")
def test_search_filing_no_data(mock_search):
    mock_search.return_value = []
    result = search_filing.invoke({"query": "revenue", "ticker": "ZZZZ"})
    assert "ingest" in result.lower()
    assert json.loads(result)["status"] == "EMPTY"


@patch("agent.tools.vector_store.search", side_effect=RuntimeError("index unavailable"))
def test_search_filing_error_is_structured(mock_search):
    result = search_filing.invoke({"query": "revenue", "ticker": "AAPL"})
    payload = json.loads(result)
    assert payload["status"] == "ERROR"
    assert payload["records"] == []
    assert payload["error"]["code"] == "TOOL_EXECUTION_ERROR"


@patch("agent.tools.vector_store.search")
def test_search_transcript_returns_chunks(mock_search):
    mock_search.return_value = [
        {"text": "Tim Cook: Record quarter.", "metadata": {"source": "transcript", "date": "2024-11-01", "speaker": "mixed"}}
    ]
    result = search_transcript.invoke({"query": "CEO commentary", "ticker": "AAPL"})
    assert "Tim Cook" in result


@patch("agent.tools.vector_store.search")
def test_search_transcript_with_since_parameter(mock_search):
    mock_search.return_value = [
        {"text": "Recent CEO commentary.", "metadata": {"source": "transcript", "date": "2025-01-15"}}
    ]
    result = search_transcript.invoke({"query": "CEO commentary", "ticker": "AAPL", "since": "2024-01-01"})
    assert "Recent CEO commentary" in result
    # Verify min_filing_date was passed to the search function
    mock_search.assert_called_once()
    call_kwargs = mock_search.call_args[1]
    assert call_kwargs.get("min_filing_date") == "2024-01-01"


@patch.dict("os.environ", {"TAVILY_API_KEY": "test-key"})
@patch("agent.tools.news_cache.put")
@patch("agent.tools.news_cache.get", return_value=None)
@patch("agent.tools.metrics_db.get_all_metrics")
@patch("agent.tools.TavilyClient")
def test_search_news(mock_cls, mock_metrics, _mock_cache_get, _mock_cache_put):
    import agent.tools
    agent.tools._tavily_client = None  # reset cached client so patch takes effect
    mock_metrics.return_value = [_sample_metrics()]
    mock_client = MagicMock()
    mock_cls.return_value = mock_client
    mock_client.search.return_value = {"results": [
        {"title": "Apple beats estimates", "content": "Strong Q4.", "url": "https://example.com", "published_date": "2024-11-01"}
    ]}
    result = search_news.invoke({"query": "earnings", "ticker": "AAPL", "days": 30})
    assert "Apple beats estimates" in result
    # Verify Tavily was called with ticker-grounded query and days/topic params
    mock_client.search.assert_called_once()
    call_kwargs = mock_client.search.call_args[1]
    assert "AAPL" in call_kwargs.get("query", "")
    assert call_kwargs.get("days") == 30
    assert call_kwargs.get("topic") == "news"


@patch("agent.tools.news_cache.get", return_value="legacy unstructured result")
def test_search_news_legacy_cache_is_not_repackaged_as_evidence(_mock_cache_get):
    payload = json.loads(search_news.invoke({
        "query": "earnings", "ticker": "AAPL", "days": 30,
    }))

    assert payload["status"] == "EMPTY"
    assert payload["records"] == []
    assert "provenance" in payload["message"].lower()