FluxIP/tests/test_aws_static_ip.py

from __future__ import annotations

from types import SimpleNamespace
from typing import Any

import pytest
from botocore.exceptions import ClientError

import app.integrations.aws_client as aws_module
from app.core.errors import ExternalServiceError, ValidationAppError
from app.integrations.aws_client import LightsailClient


class FakeNotFoundException(Exception):
    pass


class FakeLightsail:
    exceptions = SimpleNamespace(NotFoundException=FakeNotFoundException)

    def __init__(self) -> None:
        self.calls: list[tuple[str, dict[str, Any]]] = []
        self.pages: dict[str | None, dict[str, Any]] = {}
        self.static_ip_response: dict[str, Any] | Exception = {}
        self.operation: dict[str, Any] = {
            "id": "operation-id",
            "isTerminal": True,
            "status": "Succeeded",
        }

    def get_instance(self, **kwargs: Any) -> dict[str, Any]:
        self.calls.append(("get_instance", kwargs))
        return {
            "instance": {
                "name": kwargs["instanceName"],
                "state": {"name": "RUNNING"},
                "publicIpAddress": "203.0.113.20",
                "isStaticIp": True,
            }
        }

    def get_static_ips(self, **kwargs: Any) -> dict[str, Any]:
        self.calls.append(("get_static_ips", kwargs))
        return self.pages[kwargs.get("pageToken")]

    def get_static_ip(self, **kwargs: Any) -> dict[str, Any]:
        self.calls.append(("get_static_ip", kwargs))
        if isinstance(self.static_ip_response, Exception):
            raise self.static_ip_response
        return self.static_ip_response

    def allocate_static_ip(self, **kwargs: Any) -> dict[str, Any]:
        return self._mutation("allocate_static_ip", kwargs)

    def detach_static_ip(self, **kwargs: Any) -> dict[str, Any]:
        return self._mutation("detach_static_ip", kwargs)

    def attach_static_ip(self, **kwargs: Any) -> dict[str, Any]:
        return self._mutation("attach_static_ip", kwargs)

    def release_static_ip(self, **kwargs: Any) -> dict[str, Any]:
        return self._mutation("release_static_ip", kwargs)

    def get_operation(self, **kwargs: Any) -> dict[str, Any]:
        self.calls.append(("get_operation", kwargs))
        return {"operation": dict(self.operation)}

    def _mutation(self, method: str, kwargs: dict[str, Any]) -> dict[str, Any]:
        self.calls.append((method, kwargs))
        return {"operations": [{"id": "operation-id"}]}


@pytest.fixture
def fake_client(monkeypatch: pytest.MonkeyPatch) -> tuple[FakeLightsail, LightsailClient]:
    fake = FakeLightsail()
    monkeypatch.setattr(aws_module.boto3, "client", lambda *_args, **_kwargs: fake)
    return fake, LightsailClient(region="us-east-1", instance_name="proxy-node")


def static_ip(
    name: str,
    ip_address: str,
    *,
    attached_to: str | None = None,
) -> dict[str, Any]:
    return {
        "name": name,
        "arn": f"arn:aws:lightsail:us-east-1:123456789012:StaticIp/{name}",
        "ipAddress": ip_address,
        "isAttached": attached_to is not None,
        "attachedTo": attached_to,
        "location": {"regionName": "us-east-1"},
    }


def test_lists_all_pages_and_reports_attached_static_ips(
    fake_client: tuple[FakeLightsail, LightsailClient],
) -> None:
    fake, client = fake_client
    fake.pages = {
        None: {
            "staticIps": [static_ip("proxy-old", "203.0.113.20", attached_to="proxy-node")],
            "nextPageToken": "page-2",
        },
        "page-2": {"staticIps": [static_ip("spare-ip", "203.0.113.21")]},
    }

    snapshots = client.list_static_ips()
    assert [snapshot.name for snapshot in snapshots] == ["proxy-old", "spare-ip"]
    assert snapshots[0].attached_to == "proxy-node"
    assert snapshots[0].region == "us-east-1"
    assert ("get_static_ips", {}) in fake.calls
    assert ("get_static_ips", {"pageToken": "page-2"}) in fake.calls

    status = client.test_connection()
    assert status["attached_static_ips"] == [
        {
            "name": "proxy-old",
            "ip_address": "203.0.113.20",
            "attached_to": "proxy-node",
        }
    ]


def test_get_static_ip_handles_success_not_found_and_other_errors(
    fake_client: tuple[FakeLightsail, LightsailClient],
) -> None:
    fake, client = fake_client
    fake.static_ip_response = {"staticIp": static_ip("proxy-old", "203.0.113.20")}
    snapshot = client.get_static_ip("proxy-old")
    assert snapshot is not None
    assert snapshot.ip_address == "203.0.113.20"
    assert fake.calls[-1] == ("get_static_ip", {"staticIpName": "proxy-old"})

    fake.static_ip_response = FakeNotFoundException()
    assert client.get_static_ip("proxy-old") is None

    fake.static_ip_response = ClientError(
        {"Error": {"Code": "ThrottlingException", "Message": "slow down"}},
        "GetStaticIp",
    )
    with pytest.raises(ExternalServiceError, match="slow down") as error:
        client.get_static_ip("proxy-old")
    assert error.value.code == "ThrottlingException"


def test_mutations_send_exact_parameters_and_wait_for_terminal_operation(
    fake_client: tuple[FakeLightsail, LightsailClient],
) -> None:
    fake, client = fake_client

    client.allocate_static_ip("proxy-new")
    client.detach_static_ip("proxy-old")
    client.attach_static_ip("proxy-new")
    client.attach_static_ip("proxy-new", "another-node")
    client.release_static_ip("proxy-old")

    assert ("allocate_static_ip", {"staticIpName": "proxy-new"}) in fake.calls
    assert ("detach_static_ip", {"staticIpName": "proxy-old"}) in fake.calls
    assert (
        "attach_static_ip",
        {"staticIpName": "proxy-new", "instanceName": "proxy-node"},
    ) in fake.calls
    assert (
        "attach_static_ip",
        {"staticIpName": "proxy-new", "instanceName": "another-node"},
    ) in fake.calls
    assert ("release_static_ip", {"staticIpName": "proxy-old"}) in fake.calls
    assert sum(method == "get_operation" for method, _kwargs in fake.calls) == 5


def test_mutation_maps_failed_and_timed_out_operations(
    fake_client: tuple[FakeLightsail, LightsailClient],
) -> None:
    fake, client = fake_client
    fake.operation = {
        "id": "operation-id",
        "isTerminal": True,
        "status": "Failed",
        "errorCode": "StaticIpLimitExceeded",
        "errorDetails": "quota exhausted",
    }
    with pytest.raises(ExternalServiceError, match="quota exhausted") as failed:
        client.allocate_static_ip("proxy-new")
    assert failed.value.code == "StaticIpLimitExceeded"

    fake.operation = {"id": "operation-id", "isTerminal": False, "status": "Started"}
    client.operation_timeout_seconds = 0
    with pytest.raises(ExternalServiceError, match="结果仍需重新查询确认") as timed_out:
        client.detach_static_ip("proxy-old")
    assert timed_out.value.code == "AWS_OPERATION_TIMEOUT"


@pytest.mark.parametrize(
    "invalid_name",
    ["a", "-proxy", "proxy-", "proxy_ip", "代理-ip", "a" * 256],
)
def test_rejects_non_conservative_static_ip_names(
    fake_client: tuple[FakeLightsail, LightsailClient],
    invalid_name: str,
) -> None:
    _fake, client = fake_client
    with pytest.raises(ValidationAppError) as error:
        client.allocate_static_ip(invalid_name)
    assert error.value.code == "INVALID_STATIC_IP_NAME"


def test_from_credentials_supports_explicit_and_default_credentials(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    constructor_calls: list[dict[str, Any]] = []

    def fake_boto_client(_service: str, **kwargs: Any) -> FakeLightsail:
        constructor_calls.append(kwargs)
        return FakeLightsail()

    monkeypatch.setattr(aws_module.boto3, "client", fake_boto_client)
    secrets = {
        "aws_access_key_id": "access-key",
        "aws_secret_access_key": "secret-key",
        "aws_session_token": "session-token",
    }
    LightsailClient.from_credentials(
        region="us-east-1",
        instance_name="proxy-node",
        use_default=False,
        secrets=secrets,
    )
    assert constructor_calls[-1]["aws_access_key_id"] == "access-key"
    assert constructor_calls[-1]["aws_session_token"] == "session-token"

    LightsailClient.from_credentials(
        region="us-east-1",
        instance_name="proxy-node",
        use_default=True,
        secrets=secrets,
    )
    assert "aws_access_key_id" not in constructor_calls[-1]

    with pytest.raises(ValidationAppError) as error:
        LightsailClient.from_credentials(
            region="us-east-1",
            instance_name="proxy-node",
            use_default=False,
            secrets={},
        )
    assert error.value.code == "AWS_CREDENTIALS_MISSING"