Prometheus

"""PrometheusClient for Cortex XSOAR - Unit Tests file"""

from CommonServerPython import Optional
import Packs.Prometheus.Integrations.Prometheus.Prometheus as mod
import pytest
import json


def util_load_json(path):
    with open(path, encoding="utf-8") as f:
        return json.loads(f.read())


class TestClient(mod.PrometheusClient):
    def __init__(self, response=None, raise_exc: Exception = None):
        self._response = response
        self._raise = raise_exc
        self.timeout = 30

    def instant_query(self, query: str, at_time: Optional[str] = None):
        if self._raise:
            raise self._raise
        self._last_query = query
        self._last_time = at_time
        return self._response


@pytest.fixture
def sample_vector_response():
    # Matches Prometheus /api/v1/query vector response shape
    return {
        "status": "success",
        "data": {
            "resultType": "vector",
            "result": [
                {"metric": {"__name__": "co2", "host": "wattle"}, "value": [str(1_700_000_000), "412.7"]},
                {"metric": {"__name__": "temperature", "room": "garage"}, "value": [str(1_700_000_005), "23.1"]},
            ],
        },
    }


@pytest.fixture
def sample_empty_vector_response():
    return {"status": "success", "data": {"resultType": "vector", "result": []}}


@pytest.fixture
def sample_error_response():
    return {"status": "error", "errorType": "bad_data", "error": "parse error"}


def test_build_name_regex_anchor_true():
    regex = mod.build_name_regex("co2| solar |load", anchor=True)
    assert regex == "^(co2|solar|load)$"


def test_build_name_regex_anchor_false():
    regex = mod.build_name_regex("co2|solar|load", anchor=False)
    assert regex == "co2|solar|load"


def test_build_name_regex_empty_raises():
    with pytest.raises(ValueError):
        mod.build_name_regex("   ||  |  ", anchor=True)


def test_format_result_rows_parses_vector(sample_vector_response):
    rows = mod.format_result_rows(sample_vector_response)
    assert len(rows) == 2

    # Row 1 checks
    r0 = rows[0]
    assert r0["name"] == "co2"
    assert isinstance(r0["value"], float)
    assert r0["value"] == pytest.approx(412.7)
    assert r0["labels"] == {"host": "wattle"}
    assert r0["ts_unix"] == 1_700_000_000
    # ISO 8601 with UTC tz
    assert r0["ts"].endswith("+00:00")

    # Row 2 checks
    r1 = rows[1]
    assert r1["name"] == "temperature"
    assert r1["labels"] == {"room": "garage"}
    assert r1["ts_unix"] == 1_700_000_005


def test_format_result_rows_non_vector_type_returns_empty():
    non_vector = {"status": "success", "data": {"resultType": "matrix", "result": []}}
    rows = mod.format_result_rows(non_vector)
    assert rows == []


def test_prometheus_query_builds_query_from_fields(sample_vector_response):
    client = TestClient(response=sample_vector_response)
    args = {"fields": "co2|temperature", "anchor": "true", "time": "2025-09-03T09:00:00Z"}

    response = mod.prometheus_query_command(client, args, default_fields=None)
    assert client._last_query == '{__name__=~"^(co2|temperature)$"}'
    assert client._last_time == "2025-09-03T09:00:00Z"
    assert response.outputs_prefix == "Prometheus.Metrics"
    assert isinstance(response.outputs, list)
    assert len(response.outputs) == 2
    assert {r["name"] for r in response.outputs} == {"co2", "temperature"}


def test_prometheus_query_uses_default_fields_when_not_in_args(sample_vector_response):
    client = TestClient(response=sample_vector_response)
    args = {"anchor": "false"}
    default_fields = "co2|temperature"
    response = mod.prometheus_query_command(client, args, default_fields=default_fields)

    assert client._last_query == '{__name__=~"co2|temperature"}'
    assert len(response.outputs) == 2


def test_prometheus_query_uses_raw_query_when_provided(sample_vector_response):
    client = TestClient(response=sample_vector_response)
    args = {"query": '{__name__=~"(co2|temperature)"}', "time": "1700000000"}
    response = mod.prometheus_query_command(client, args, default_fields=None)
    assert client._last_query == '{__name__=~"(co2|temperature)"}'
    assert client._last_time == "1700000000"
    assert len(response.outputs) == 2


def test_prometheus_query_raises_when_no_fields_and_no_default():
    client = TestClient(response={"status": "success", "data": {"resultType": "vector", "result": []}})
    with pytest.raises(mod.DemistoException) as exc:
        mod.prometheus_query_command(client, args={}, default_fields=None)
    assert 'You must provide "fields"' in str(exc.value)


def test_prometheus_query_raises_on_api_error(sample_error_response):
    client = TestClient(response=sample_error_response)
    with pytest.raises(mod.DemistoException) as exc:
        mod.prometheus_query_command(client, args={"fields": "co2"}, default_fields=None)
    assert "Prometheus API error" in str(exc.value)


def test_prometheus_raw_happy_path(sample_vector_response):
    client = TestClient(response=sample_vector_response)
    response = mod.prometheus_raw_command(client, args={"query": "vector(1)"})
    assert client._last_query == "vector(1)"
    assert len(response.outputs) == 2
    names = {r["name"] for r in response.outputs}
    assert "co2" in names
    assert "temperature" in names


def test_prometheus_raw_requires_query_arg():
    client = TestClient(response={"status": "success", "data": {"resultType": "vector", "result": []}})
    with pytest.raises(mod.DemistoException):
        mod.prometheus_raw_command(client, args={})


def test_prometheus_raw_raises_on_api_error(sample_error_response):
    client = TestClient(response=sample_error_response)
    with pytest.raises(mod.DemistoException):
        mod.prometheus_raw_command(client, args={"query": "up"})


def test_test_module_ok():
    client = TestClient(response={"status": "success", "data": {"resultType": "vector", "result": []}})
    assert mod.test_module(client) == "ok"


def test_test_module_failure_message():
    client = TestClient(response={"status": "error", "error": "boom"})
    assert mod.test_module(client).startswith("Failed:")


def test_test_module_exception_is_caught():
    client = TestClient(response=None, raise_exc=RuntimeError("network down"))
    out = mod.test_module(client)
    assert out.startswith("Failed:")
    assert "network down" in out