tests/test_orchestrator.py

"""Orchestrator pipeline and the tool-calling ask-agent loop, fully mocked."""

from datetime import datetime, timedelta
from zoneinfo import ZoneInfo

from app.agents import orchestrator
from app.providers import llm

TZ = ZoneInfo("Asia/Kolkata")


def arrive_by_iso(hours_ahead=3):
    return (datetime.now(TZ) + timedelta(hours=hours_ahead)).replace(
        second=0, microsecond=0
    ).isoformat()


def mock_pipeline_tools(monkeypatch, sweep=None):
    place = {"input": "x", "name": "Mock Place", "address": "Lucknow", "lat": 26.85, "lon": 80.95}
    monkeypatch.setattr(orchestrator.geocode, "geocode_place", lambda p: dict(place, input=p))

    default_sweep = {
        "arrive_by": arrive_by_iso(),
        "curve": [
            {"depart": "08:00", "depart_iso": "2026-06-12T08:00:00+05:30", "eta": "08:35",
             "travel_min": 35.0, "delay_min": 5.0, "no_traffic_min": 30.0,
             "distance_km": 12.0, "on_time": True, "margin_min": 55.0},
        ],
        "recommended": {"depart": "08:00", "depart_iso": "2026-06-12T08:00:00+05:30",
                        "eta": "08:35", "travel_min": 35.0, "delay_min": 5.0,
                        "no_traffic_min": 30.0, "distance_km": 12.0,
                        "on_time": True, "margin_min": 55.0},
    }
    monkeypatch.setattr(
        orchestrator.traffic, "compare_departures", lambda *a, **k: sweep or default_sweep
    )
    monkeypatch.setattr(
        orchestrator.traffic, "get_route",
        lambda *a, **k: [{"duration_min": 35.0, "distance_km": 12.0, "delay_min": 5.0,
                          "no_traffic_min": 30.0, "points": [[26.85, 80.95], [26.84, 80.92]]}],
    )
    monkeypatch.setattr(
        orchestrator.weather, "get_weather",
        lambda *a, **k: {"rain_expected": False, "max_rain_mm": 0, "summary": "clear"},
    )
    monkeypatch.setattr(
        orchestrator.festivals, "get_festivals", lambda *a, **k: {"festivals": [], "impact": 0}
    )
    monkeypatch.setattr(
        orchestrator.events, "get_events", lambda *a, **k: {"events": [], "impact": 0}
    )
    monkeypatch.setattr(
        orchestrator.advisories, "get_police_advisories",
        lambda *a, **k: {"advisories": [], "count": 0},
    )


def test_run_plan_yields_verdict(monkeypatch, no_llm_keys):
    mock_pipeline_tools(monkeypatch)
    events_out = list(orchestrator.run_plan("Gomti Nagar", "Charbagh", arrive_by_iso()))

    types = [event["type"] for event in events_out]
    assert "error" not in types
    assert types[-1] == "verdict"

    verdict = events_out[-1]["data"]
    assert verdict["recommended_departure"] == "08:00"
    assert verdict["risk"]["level"] == "LOW"
    assert verdict["route"]["points"]
    assert verdict["origin"]["name"] == "Mock Place"
    assert verdict["llm_provider"] == "deterministic"


def test_run_plan_streams_tool_results(monkeypatch, no_llm_keys):
    mock_pipeline_tools(monkeypatch)
    events_out = list(orchestrator.run_plan("A", "B", arrive_by_iso()))
    tool_names = {event["tool"] for event in events_out if event["type"] == "tool_result"}
    assert {"geocode", "traffic", "weather", "festivals", "events", "advisories"} <= tool_names


def test_run_plan_with_coords_skips_geocoding(monkeypatch, no_llm_keys):
    mock_pipeline_tools(monkeypatch)

    def must_not_be_called(place):
        raise AssertionError("geocode_place should not run when coords are provided")

    monkeypatch.setattr(orchestrator.geocode, "geocode_place", must_not_be_called)

    events_out = list(orchestrator.run_plan(
        "IIIT Lucknow", "Hazratganj", arrive_by_iso(),
        origin_coords=(26.8003, 81.0242), dest_coords=(26.8488, 80.9436),
    ))
    assert events_out[-1]["type"] == "verdict"
    verdict = events_out[-1]["data"]
    assert verdict["origin"]["lat"] == 26.8003
    assert verdict["origin"]["name"] == "IIIT Lucknow"


def test_run_plan_geocode_failure_yields_error(monkeypatch, no_llm_keys):
    def boom(place):
        raise RuntimeError("place not found")

    monkeypatch.setattr(orchestrator.geocode, "geocode_place", boom)
    events_out = list(orchestrator.run_plan("???", "???", arrive_by_iso()))
    assert events_out[-1]["type"] == "error"
    assert "place not found" in events_out[-1]["message"]


def test_run_plan_sweep_failure_yields_error(monkeypatch, no_llm_keys):
    mock_pipeline_tools(monkeypatch)

    def boom(*args, **kwargs):
        raise RuntimeError("tomtom down")

    monkeypatch.setattr(orchestrator.traffic, "compare_departures", boom)
    events_out = list(orchestrator.run_plan("A", "B", arrive_by_iso()))
    assert events_out[-1]["type"] == "error"


def test_run_plan_survives_partial_tool_failures(monkeypatch, no_llm_keys):
    mock_pipeline_tools(monkeypatch)

    def boom(*args, **kwargs):
        raise RuntimeError("weather api down")

    monkeypatch.setattr(orchestrator.weather, "get_weather", boom)
    events_out = list(orchestrator.run_plan("A", "B", arrive_by_iso()))
    # Weather failed but the plan still completes with a verdict
    assert events_out[-1]["type"] == "verdict"


# ---- ask agent loop ----------------------------------------------------------

def test_agent_ask_executes_tools_then_answers(monkeypatch, fake_keys):
    responses = [
        {"text": "", "provider": "gemini",
         "tool_calls": [{"name": "get_festivals", "args": {"date_str": "2026-06-09"}}]},
        {"text": "Yes — Bada Mangal on Tuesday will slow Hazratganj.", "provider": "gemini",
         "tool_calls": []},
    ]
    call_log = {"count": 0, "messages": None}

    def fake_chat(messages, tools=None, **kwargs):
        call_log["messages"] = messages
        response = responses[call_log["count"]]
        call_log["count"] += 1
        return response

    monkeypatch.setattr(orchestrator.llm, "chat", fake_chat)
    monkeypatch.setattr(
        orchestrator, "dispatch",
        lambda name, args: '{"festivals": [{"name": "Bada Mangal"}], "impact": 2}',
    )

    result = orchestrator.agent_ask("Any festival near Hazratganj on Tuesday?")
    assert "Bada Mangal" in result["answer"]
    assert result["steps"] == ["get_festivals({'date_str': '2026-06-09'})"]
    # The tool result was fed back into the conversation
    tool_messages = [m for m in call_log["messages"] if m["role"] == "tool"]
    assert tool_messages and tool_messages[0]["name"] == "get_festivals"


def test_agent_ask_stream_yields_tool_then_answer(monkeypatch, fake_keys):
    responses = [
        {"text": "", "provider": "gemini",
         "tool_calls": [{"name": "get_weather", "args": {"lat": 26.8}}]},
        {"text": "Light rain at 6 PM — leave 15 min early.", "provider": "gemini",
         "tool_calls": []},
    ]
    counter = {"n": 0}

    def fake_chat(messages, tools=None, **kwargs):
        response = responses[counter["n"]]
        counter["n"] += 1
        return response

    monkeypatch.setattr(orchestrator.llm, "chat", fake_chat)
    monkeypatch.setattr(orchestrator, "dispatch", lambda name, args: '{"rain_expected": true}')

    events = list(orchestrator.agent_ask_stream("Will it rain this evening?"))
    types = [event["type"] for event in events]
    assert types == ["tool", "answer"]
    assert events[0]["name"] == "get_weather"
    assert "rain" in events[1]["text"].lower()
    assert events[1]["provider"] == "gemini"


def test_agent_ask_stream_includes_history(monkeypatch, fake_keys):
    captured = {}

    def fake_chat(messages, tools=None, **kwargs):
        captured["messages"] = messages
        return {"text": "answer", "provider": "gemini", "tool_calls": []}

    monkeypatch.setattr(orchestrator.llm, "chat", fake_chat)
    history = [
        {"role": "user", "content": "Any festival tomorrow?"},
        {"role": "assistant", "content": "Yes, Bada Mangal."},
        {"role": "hacker", "content": "ignored"},  # invalid role dropped
    ]
    list(orchestrator.agent_ask_stream("Where exactly?", history))

    roles = [message["role"] for message in captured["messages"]]
    assert roles == ["system", "user", "assistant", "user"]
    assert "Bada Mangal" in captured["messages"][2]["content"]


def test_agent_ask_stream_llm_outage_yields_error(monkeypatch, fake_keys):
    def fake_chat(*args, **kwargs):
        raise orchestrator.llm.LLMUnavailable("all providers down")

    monkeypatch.setattr(orchestrator.llm, "chat", fake_chat)
    events = list(orchestrator.agent_ask_stream("hello"))
    assert events[-1]["type"] == "error"


def test_agent_ask_stops_at_iteration_limit(monkeypatch, fake_keys):
    def always_calls_tools(messages, tools=None, **kwargs):
        return {"text": "", "provider": "gemini",
                "tool_calls": [{"name": "get_police_advisories", "args": {}}]}

    monkeypatch.setattr(orchestrator.llm, "chat", always_calls_tools)
    monkeypatch.setattr(
        orchestrator, "dispatch", lambda name, args: '{"advisories": [], "count": 0}'
    )

    result = orchestrator.agent_ask("loop forever")
    assert "ran out" in result["answer"]
    assert len(result["steps"]) == orchestrator.MAX_AGENT_ITERATIONS