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

"""Tests for PrefectWorkerHealthTool (class-based, BaseTool subclass)."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from integrations.prefect.tools import PrefectWorkerHealthTool
from tests.tools.conftest import BaseToolContract, mock_agent_state


class TestPrefectWorkerHealthToolContract(BaseToolContract):
    def get_tool_under_test(self):
        return PrefectWorkerHealthTool()


@pytest.fixture
def tool() -> PrefectWorkerHealthTool:
    return PrefectWorkerHealthTool()


def test_is_available_when_connection_verified(tool: PrefectWorkerHealthTool) -> None:
    assert tool.is_available({"prefect": {"connection_verified": True}}) is True


def test_is_available_false_without_connection_verified(tool: PrefectWorkerHealthTool) -> None:
    assert tool.is_available({"prefect": {}}) is False
    assert tool.is_available({}) is False


def test_extract_params_maps_source_fields(tool: PrefectWorkerHealthTool) -> None:
    sources = mock_agent_state({"prefect": {"work_pool_name": "my-pool"}})

    params = tool.extract_params(sources)

    assert params["api_url"] == "http://localhost:4200/api"
    assert params["work_pool_name"] == "my-pool"
    assert params["pool_limit"] == 20
    assert params["worker_limit"] == 20


def test_run_returns_unavailable_without_api_url(tool: PrefectWorkerHealthTool) -> None:
    result = tool.run(api_url="")

    assert result["available"] is False
    assert result["work_pools"] == []
    assert result["workers"] == []


def test_run_returns_unavailable_for_whitespace_only_api_url(
    tool: PrefectWorkerHealthTool,
) -> None:
    result = tool.run(api_url="   ")

    assert result["available"] is False


def test_run_returns_unavailable_on_api_failure(tool: PrefectWorkerHealthTool) -> None:
    mock_client = MagicMock()
    mock_client.__enter__ = MagicMock(return_value=mock_client)
    mock_client.__exit__ = MagicMock(return_value=False)
    mock_client.get_work_pools.return_value = {
        "success": False,
        "error": "HTTP 403: forbidden",
    }

    with patch("integrations.prefect.tools.make_prefect_client", return_value=mock_client):
        result = tool.run(api_url="http://localhost:4200/api")

    assert result["available"] is False
    assert "403" in result["error"]
    assert result["work_pools"] == []


def test_run_returns_unavailable_when_client_none(tool: PrefectWorkerHealthTool) -> None:
    with patch("integrations.prefect.tools.make_prefect_client", return_value=None):
        result = tool.run(api_url="http://localhost:4200/api")

    assert result["available"] is False


def test_run_identifies_unhealthy_pools(tool: PrefectWorkerHealthTool) -> None:
    pools = [
        {"id": "pool_1", "name": "default", "status": "READY", "is_paused": False},
        {"id": "pool_2", "name": "batch", "status": "NOT_READY", "is_paused": False},
        {"id": "pool_3", "name": "nightly", "status": "READY", "is_paused": True},
    ]
    mock_client = MagicMock()
    mock_client.__enter__ = MagicMock(return_value=mock_client)
    mock_client.__exit__ = MagicMock(return_value=False)
    mock_client.get_work_pools.return_value = {
        "success": True,
        "work_pools": pools,
        "total": 3,
    }

    with patch("integrations.prefect.tools.make_prefect_client", return_value=mock_client):
        result = tool.run(api_url="http://localhost:4200/api")

    assert result["available"] is True
    assert result["total_pools"] == 3
    assert result["total_unhealthy_pools"] == 2
    names = {pool["name"] for pool in result["unhealthy_pools"]}
    assert "batch" in names
    assert "nightly" in names
    assert "default" not in names


def test_run_empty_work_pools(tool: PrefectWorkerHealthTool) -> None:
    mock_client = MagicMock()
    mock_client.__enter__ = MagicMock(return_value=mock_client)
    mock_client.__exit__ = MagicMock(return_value=False)
    mock_client.get_work_pools.return_value = {
        "success": True,
        "work_pools": [],
        "total": 0,
    }

    with patch("integrations.prefect.tools.make_prefect_client", return_value=mock_client):
        result = tool.run(api_url="http://localhost:4200/api")

    assert result["available"] is True
    assert result["total_pools"] == 0
    assert result["unhealthy_pools"] == []


def test_run_fetches_workers_when_pool_name_provided(tool: PrefectWorkerHealthTool) -> None:
    workers = [
        {
            "name": "worker-1",
            "status": "ONLINE",
            "last_heartbeat_time": "2026-04-05T10:00:00Z",
        },
        {
            "name": "worker-2",
            "status": "OFFLINE",
            "last_heartbeat_time": "2026-04-04T08:00:00Z",
        },
        {
            "name": "worker-3",
            "status": "UNHEALTHY",
            "last_heartbeat_time": "2026-04-05T09:00:00Z",
        },
    ]
    mock_client = MagicMock()
    mock_client.__enter__ = MagicMock(return_value=mock_client)
    mock_client.__exit__ = MagicMock(return_value=False)
    mock_client.get_work_pools.return_value = {"success": True, "work_pools": [], "total": 0}
    mock_client.get_workers.return_value = {"success": True, "workers": workers, "total": 3}

    with patch("integrations.prefect.tools.make_prefect_client", return_value=mock_client):
        result = tool.run(api_url="http://localhost:4200/api", work_pool_name="default")

    mock_client.get_workers.assert_called_once_with(work_pool_name="default", limit=20)
    assert result["total_workers"] == 3
    assert result["total_unhealthy_workers"] == 2
    names = {worker["name"] for worker in result["unhealthy_workers"]}
    assert "worker-2" in names
    assert "worker-3" in names
    assert "worker-1" not in names


def test_run_skips_worker_fetch_without_pool_name(tool: PrefectWorkerHealthTool) -> None:
    mock_client = MagicMock()
    mock_client.__enter__ = MagicMock(return_value=mock_client)
    mock_client.__exit__ = MagicMock(return_value=False)
    mock_client.get_work_pools.return_value = {"success": True, "work_pools": [], "total": 0}

    with patch("integrations.prefect.tools.make_prefect_client", return_value=mock_client):
        result = tool.run(api_url="http://localhost:4200/api")

    mock_client.get_workers.assert_not_called()
    assert result["workers"] == []
    assert result["unhealthy_workers"] == []


def test_metadata_is_valid(tool: PrefectWorkerHealthTool) -> None:
    meta = tool.metadata()

    assert meta.name == "prefect_worker_health"
    assert meta.source == "prefect"
    assert "required" in meta.input_schema
    assert "api_url" in meta.input_schema["required"]