cleveragents-core/features/steps/cloud_aws_helpers.py

"""Shared helpers for AWS SDK step definitions.

Provides environment management, resource factories, and mock session
builders used across the cloud_aws_* step definition files.
"""

from __future__ import annotations

import contextlib
import os
from typing import Any
from unittest.mock import MagicMock

from cleveragents.domain.models.core.resource import (
    PhysVirt,
    Resource,
    ResourceCapabilities,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_CLOUD_ENV_VARS = [
    "AWS_ACCESS_KEY_ID",
    "AWS_SECRET_ACCESS_KEY",
    "AWS_SESSION_TOKEN",
    "AWS_REGION",
    "AWS_PROFILE",
    "GOOGLE_APPLICATION_CREDENTIALS",
    "GCLOUD_PROJECT",
    "GCP_REGION",
    "AZURE_SUBSCRIPTION_ID",
    "AZURE_TENANT_ID",
    "AZURE_CLIENT_ID",
    "AZURE_CLIENT_SECRET",
    "AZURE_REGION",
]


def _clear_cloud_env() -> dict[str, str]:
    saved: dict[str, str] = {}
    for var in _CLOUD_ENV_VARS:
        val = os.environ.pop(var, None)
        if val is not None:
            saved[var] = val
    return saved


def _restore_cloud_env(saved: dict[str, str]) -> None:
    for var in _CLOUD_ENV_VARS:
        if var in saved:
            os.environ[var] = saved[var]
        else:
            os.environ.pop(var, None)


def _make_resource(
    type_name: str,
    properties: dict[str, Any] | None = None,
    location: str | None = None,
) -> Resource:
    return Resource(
        resource_id="01KJ5C5TPMP8GGX3QC83E2MAQS",
        resource_type_name=type_name,
        classification=PhysVirt.PHYSICAL,
        location=location,
        properties=properties or {},
        capabilities=ResourceCapabilities(
            readable=True,
            writable=False,
            sandboxable=False,
            checkpointable=False,
        ),
    )


def _make_mock_session(
    account_id: str = "123456789012",
    sts_error: Exception | None = None,
    ec2_vpcs: list[dict[str, Any]] | None = None,
    ec2_instances: list[dict[str, Any]] | None = None,
    s3_buckets: list[dict[str, Any]] | None = None,
    ecs_clusters: list[str] | None = None,
    client_error: Exception | None = None,
) -> MagicMock:
    """Build a mock boto3 Session with configurable responses."""
    session = MagicMock()

    def _make_client(service: str, **_kwargs: Any) -> MagicMock:
        client = MagicMock()

        if service == "sts":
            if sts_error is not None:
                client.get_caller_identity.side_effect = sts_error
            else:
                client.get_caller_identity.return_value = {
                    "Account": account_id,
                    "UserId": "AIDAIOSFODNN7EXAMPLE",
                    "Arn": f"arn:aws:iam::{account_id}:user/test",
                }

        elif service == "ec2":
            if client_error is not None:
                client.describe_vpcs.side_effect = client_error
                client.describe_subnets.side_effect = client_error
                client.describe_instances.side_effect = client_error
                client.describe_security_groups.side_effect = client_error
            else:
                vpcs = ec2_vpcs or []
                client.describe_vpcs.return_value = {"Vpcs": vpcs}
                client.describe_subnets.return_value = {"Subnets": []}
                instances = ec2_instances or []
                reservations = (
                    [{"ReservationId": "r-00000001", "Instances": instances}]
                    if instances
                    else []
                )
                client.describe_instances.return_value = {"Reservations": reservations}
                client.describe_security_groups.return_value = {"SecurityGroups": []}

        elif service == "s3":
            if client_error is not None:
                client.list_buckets.side_effect = client_error
            else:
                buckets = s3_buckets or []
                client.list_buckets.return_value = {"Buckets": buckets}

        elif service == "ecs":
            if client_error is not None:
                client.list_clusters.side_effect = client_error
            else:
                arns = ecs_clusters or []
                client.list_clusters.return_value = {"clusterArns": arns}

        else:
            if client_error is not None:
                for attr in dir(client):
                    if not attr.startswith("_"):
                        with contextlib.suppress(AttributeError):
                            getattr(client, attr).side_effect = client_error

        return client

    session.client.side_effect = _make_client
    return session