mxaccessgw/clients/python/tests/test_alarms.py

"""Tests for the AcknowledgeAlarm + StreamAlarms client surface."""

from __future__ import annotations

from typing import Any

import grpc
import pytest

from mxgateway import ClientOptions, GatewayClient
from mxgateway.errors import MxGatewayAuthenticationError, MxGatewayAuthorizationError
from mxgateway.generated import mxaccess_gateway_pb2 as pb


@pytest.mark.asyncio
async def test_acknowledge_alarm_sends_request_and_returns_reply() -> None:
    stub = FakeGatewayStub()
    stub.acknowledge_alarm.replies = [
        pb.AcknowledgeAlarmReply(
            correlation_id="corr-7",
            protocol_status=pb.ProtocolStatus(code=pb.PROTOCOL_STATUS_CODE_OK),
            status=pb.MxStatusProxy(success=1, category=pb.MX_STATUS_CATEGORY_OK),
        ),
    ]
    client = await GatewayClient.connect(
        ClientOptions(endpoint="fake", api_key="mxgw_test_secret", plaintext=True),
        stub=stub,
    )

    reply = await client.acknowledge_alarm(
        pb.AcknowledgeAlarmRequest(
            client_correlation_id="corr-7",
            alarm_full_reference="Tank01.Level.HiHi",
            comment="investigating",
            operator_user="alice",
        ),
    )

    assert reply.protocol_status.code == pb.PROTOCOL_STATUS_CODE_OK
    assert reply.status.category == pb.MX_STATUS_CATEGORY_OK

    captured_request = stub.acknowledge_alarm.requests[0]
    assert captured_request.alarm_full_reference == "Tank01.Level.HiHi"
    assert captured_request.comment == "investigating"
    assert captured_request.operator_user == "alice"
    assert stub.acknowledge_alarm.metadata == (("authorization", "Bearer mxgw_test_secret"),)


@pytest.mark.asyncio
async def test_acknowledge_alarm_unauthenticated_raises_typed_error() -> None:
    stub = FakeGatewayStub()
    stub.acknowledge_alarm.exception = FakeRpcError(grpc.StatusCode.UNAUTHENTICATED, "expired key")
    client = await GatewayClient.connect(
        ClientOptions(endpoint="fake", api_key="mxgw_test_secret", plaintext=True),
        stub=stub,
    )

    with pytest.raises(MxGatewayAuthenticationError):
        await client.acknowledge_alarm(
            pb.AcknowledgeAlarmRequest(
                alarm_full_reference="Tank01.Level.HiHi",
                comment="",
                operator_user="alice",
            ),
        )


@pytest.mark.asyncio
async def test_acknowledge_alarm_permission_denied_raises_typed_error() -> None:
    stub = FakeGatewayStub()
    stub.acknowledge_alarm.exception = FakeRpcError(grpc.StatusCode.PERMISSION_DENIED, "missing scope")
    client = await GatewayClient.connect(
        ClientOptions(endpoint="fake", api_key="mxgw_test_secret", plaintext=True),
        stub=stub,
    )

    with pytest.raises(MxGatewayAuthorizationError):
        await client.acknowledge_alarm(
            pb.AcknowledgeAlarmRequest(
                alarm_full_reference="Tank01.Level.HiHi",
                comment="",
                operator_user="alice",
            ),
        )


@pytest.mark.asyncio
async def test_stream_alarms_streams_snapshot_then_snapshot_complete() -> None:
    messages = [
        pb.AlarmFeedMessage(
            active_alarm=pb.ActiveAlarmSnapshot(
                alarm_full_reference="Tank01.Level.HiHi",
                current_state=pb.ALARM_CONDITION_STATE_ACTIVE,
                severity=750,
            ),
        ),
        pb.AlarmFeedMessage(
            active_alarm=pb.ActiveAlarmSnapshot(
                alarm_full_reference="Tank02.Level.HiHi",
                current_state=pb.ALARM_CONDITION_STATE_ACTIVE_ACKED,
                severity=750,
            ),
        ),
        pb.AlarmFeedMessage(snapshot_complete=True),
    ]
    stream = FakeAlarmFeedStream(messages)
    stub = FakeGatewayStub(alarm_feed_stream=stream)
    client = await GatewayClient.connect(
        ClientOptions(endpoint="fake", api_key="mxgw_test_secret", plaintext=True),
        stub=stub,
    )

    received: list[pb.AlarmFeedMessage] = []
    async for message in client.stream_alarms(pb.StreamAlarmsRequest()):
        received.append(message)

    assert len(received) == 3
    assert received[0].active_alarm.alarm_full_reference == "Tank01.Level.HiHi"
    assert received[0].active_alarm.current_state == pb.ALARM_CONDITION_STATE_ACTIVE
    assert received[1].active_alarm.current_state == pb.ALARM_CONDITION_STATE_ACTIVE_ACKED
    assert received[2].snapshot_complete is True
    assert stub.stream_metadata == (("authorization", "Bearer mxgw_test_secret"),)


@pytest.mark.asyncio
async def test_stream_alarms_passes_filter_prefix() -> None:
    stream = FakeAlarmFeedStream([])
    stub = FakeGatewayStub(alarm_feed_stream=stream)
    client = await GatewayClient.connect(
        ClientOptions(endpoint="fake", api_key="mxgw_test_secret", plaintext=True),
        stub=stub,
    )

    iterator = client.stream_alarms(
        pb.StreamAlarmsRequest(alarm_filter_prefix="Tank01."),
    )
    # Drain to trigger the stub call.
    async for _ in iterator:
        pass

    assert stub.stream_request is not None
    assert stub.stream_request.alarm_filter_prefix == "Tank01."


@pytest.mark.asyncio
async def test_stream_alarms_cancels_underlying_stream_on_close() -> None:
    messages = [
        pb.AlarmFeedMessage(
            active_alarm=pb.ActiveAlarmSnapshot(
                alarm_full_reference="Tank01.Level.HiHi",
                current_state=pb.ALARM_CONDITION_STATE_ACTIVE,
            ),
        ),
    ]
    stream = FakeAlarmFeedStream(messages)
    stub = FakeGatewayStub(alarm_feed_stream=stream)
    client = await GatewayClient.connect(
        ClientOptions(endpoint="fake", api_key="mxgw_test_secret", plaintext=True),
        stub=stub,
    )

    iterator = client.stream_alarms(pb.StreamAlarmsRequest())
    first = await anext(iterator)
    await iterator.aclose()

    assert first.active_alarm.alarm_full_reference == "Tank01.Level.HiHi"
    assert stream.cancelled


class FakeGatewayStub:
    def __init__(self, alarm_feed_stream: "FakeAlarmFeedStream | None" = None) -> None:
        self.open_session = FakeUnary(
            [
                pb.OpenSessionReply(
                    session_id="session-1",
                    protocol_status=pb.ProtocolStatus(code=pb.PROTOCOL_STATUS_CODE_OK),
                ),
            ],
        )
        self.acknowledge_alarm = FakeUnary([])
        self.OpenSession = self.open_session
        self.AcknowledgeAlarm = self.acknowledge_alarm
        self._alarm_feed_stream = alarm_feed_stream or FakeAlarmFeedStream([])
        self.stream_request: pb.StreamAlarmsRequest | None = None
        self.stream_metadata: tuple[tuple[str, str], ...] | None = None

    def StreamAlarms(
        self,
        request: pb.StreamAlarmsRequest,
        *,
        metadata: tuple[tuple[str, str], ...],
    ) -> "FakeAlarmFeedStream":
        self.stream_request = request
        self.stream_metadata = metadata
        return self._alarm_feed_stream


class FakeUnary:
    def __init__(self, replies: list[Any]) -> None:
        self.replies = replies
        self.requests: list[Any] = []
        self.metadata: tuple[tuple[str, str], ...] | None = None
        self.exception: Exception | None = None

    async def __call__(
        self,
        request: Any,
        *,
        metadata: tuple[tuple[str, str], ...],
    ) -> Any:
        self.requests.append(request)
        self.metadata = metadata
        if self.exception is not None:
            raise self.exception
        return self.replies.pop(0)


class FakeAlarmFeedStream:
    def __init__(self, messages: list[pb.AlarmFeedMessage]) -> None:
        self._messages = list(messages)
        self.cancelled = False

    def __aiter__(self) -> "FakeAlarmFeedStream":
        return self

    async def __anext__(self) -> pb.AlarmFeedMessage:
        if not self._messages:
            raise StopAsyncIteration
        return self._messages.pop(0)

    def cancel(self) -> None:
        self.cancelled = True


class FakeRpcError(grpc.RpcError):
    def __init__(self, code: grpc.StatusCode, details: str) -> None:
        self._code = code
        self._details = details

    def code(self) -> grpc.StatusCode:  # noqa: D401
        return self._code

    def details(self) -> str:  # noqa: D401
        return self._details