"""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