tracer-cloud--opensre/tests/tools/investigation/test_session_runner.py

"""Tests for surface-agnostic session investigation streaming."""

from __future__ import annotations

import threading
from collections.abc import Iterator
from typing import Any

import pytest

from core.domain.stream import StreamEvent
from tools.investigation import session_runner


def _collecting_renderer(
    final_state: dict[str, Any] | None = None,
) -> session_runner.StreamRendererFn:
    state = final_state if final_state is not None else {"root_cause": "done"}

    def _render(events: Iterator[StreamEvent]) -> dict[str, Any]:
        for _ in events:
            pass
        return dict(state)

    return _render


def test_run_investigation_for_session_invokes_renderer(monkeypatch: pytest.MonkeyPatch) -> None:
    captured: dict[str, Any] = {}

    def _fake_astream(**_kwargs: Any):
        async def _gen():
            yield StreamEvent(event_type="end", data={"output": {"root_cause": "ok"}})

        return _gen()

    monkeypatch.setattr(
        "tools.investigation.session_runner.check_llm_settings",
        lambda: None,
    )
    monkeypatch.setattr(
        "tools.investigation.capability.astream_investigation",
        _fake_astream,
    )

    def _render(events: Iterator[StreamEvent]) -> dict[str, Any]:
        list(events)
        captured["rendered"] = True
        return {"root_cause": "ok"}

    result = session_runner.run_investigation_for_session(
        alert_text="payments failing",
        render_stream=_render,
    )

    assert captured["rendered"] is True
    assert result["root_cause"] == "ok"


def test_run_session_alert_payload_cancel_requested_raises_keyboard_interrupt(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    cancel = threading.Event()

    def _slow_astream(**_kwargs: Any):
        async def _gen():
            cancel.set()
            while True:
                yield StreamEvent(event_type="events", data={})
                await __import__("asyncio").sleep(0.05)

        return _gen()

    monkeypatch.setattr(
        "tools.investigation.session_runner.check_llm_settings",
        lambda: None,
    )
    monkeypatch.setattr(
        "tools.investigation.capability.astream_investigation",
        _slow_astream,
    )

    with pytest.raises(KeyboardInterrupt):
        session_runner.run_session_alert_payload(
            raw_alert={"alert_name": "test"},
            cancel_requested=cancel,
            render_stream=_collecting_renderer(),
        )


def test_run_sample_alert_for_session_uses_template(monkeypatch: pytest.MonkeyPatch) -> None:
    captured_alert: dict[str, Any] = {}

    def _fake_astream(*, raw_alert: dict[str, Any], **_kwargs: Any):
        captured_alert.update(raw_alert)

        async def _gen():
            if False:
                yield StreamEvent(event_type="end")

        return _gen()

    monkeypatch.setattr(
        "tools.investigation.session_runner.check_llm_settings",
        lambda: None,
    )
    monkeypatch.setattr(
        "tools.investigation.capability.astream_investigation",
        _fake_astream,
    )

    session_runner.run_sample_alert_for_session(
        template_name="generic",
        render_stream=_collecting_renderer(),
    )

    assert captured_alert.get("pipeline_name") == "payments_etl"


def test_run_investigation_for_session_background_uses_renderer(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(
        "tools.investigation.session_runner.check_llm_settings",
        lambda: None,
    )

    async def _empty_gen():
        if False:
            yield StreamEvent(event_type="end")

    monkeypatch.setattr(
        "tools.investigation.capability.astream_investigation",
        lambda **_kwargs: _empty_gen(),
    )

    result = session_runner.run_investigation_for_session_background(
        alert_text="background test",
        render_stream=_collecting_renderer({"status": "silent"}),
    )

    assert result["status"] == "silent"


def test_run_session_alert_payload_does_not_mutate_raw_alert(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(
        "tools.investigation.session_runner.check_llm_settings",
        lambda: None,
    )

    async def _empty_gen():
        if False:
            yield StreamEvent(event_type="end")

    monkeypatch.setattr(
        "tools.investigation.capability.astream_investigation",
        lambda **_kwargs: _empty_gen(),
    )

    shared: dict[str, Any] = {"alert_name": "test", "annotations": {"keep": "yes"}}
    original = dict(shared)
    original["annotations"] = dict(shared["annotations"])

    session_runner.run_session_alert_payload(
        raw_alert=shared,
        context_overrides={"add": "value"},
        render_stream=_collecting_renderer(),
    )

    assert shared == original