tracer-cloud--opensre/tests/tools/test_investigation_registry_prioritization.py

from __future__ import annotations

from collections.abc import Callable
from typing import Any, cast

from core.tool_framework.registered_tool import RegisteredTool
from tools.investigation_registry import prioritization


def _tool(name: str, source: str, use_cases: list[str] | None = None) -> RegisteredTool:
    def _run(**_kwargs: Any) -> dict[str, Any]:
        return {"ok": True}

    return RegisteredTool(
        name=name,
        description=name,
        input_schema={"type": "object", "properties": {}, "additionalProperties": False},
        source=source,  # type: ignore[arg-type]
        run=cast(Callable[..., Any], _run),
        use_cases=use_cases or [],
    )


def test_prioritization_scores_source_and_keywords(monkeypatch: Any) -> None:
    actions = [
        _tool("query_datadog_logs", "datadog", ["inspect error logs"]),
        _tool("query_github_commits", "github", ["inspect commits"]),
    ]
    monkeypatch.setattr(prioritization, "get_available_actions", lambda: actions)

    prioritized, reasons = prioritization.get_prioritized_actions_with_reasons(
        sources=["datadog"],
        keywords=["error"],
    )

    assert [action.name for action in prioritized] == [
        "query_datadog_logs",
        "query_github_commits",
    ]
    assert reasons[0]["score"] > reasons[1]["score"]


def test_prioritization_marks_deterministic_fallback(monkeypatch: Any) -> None:
    actions = [
        _tool("query_github_commits", "github"),
        _tool("get_sre_guidance", "knowledge"),
    ]
    monkeypatch.setattr(prioritization, "get_available_actions", lambda: actions)

    _prioritized, reasons = prioritization.get_prioritized_actions_with_reasons(
        sources=["datadog"],
        keywords=["postgres"],
    )

    guidance = next(item for item in reasons if item["name"] == "get_sre_guidance")
    assert prioritization.DETERMINISTIC_FALLBACK_REASON in guidance["reasons"]