"""ECS cluster, task definition, service creation.""" import time from contextlib import suppress from typing import Any from botocore.exceptions import ClientError from tests.shared.infrastructure_sdk.deployer import ( DEFAULT_REGION, get_boto3_client, get_standard_tags_ecs, ) def create_cluster(name: str, stack_name: str, region: str = DEFAULT_REGION) -> dict[str, Any]: """Create ECS cluster. Args: name: Cluster name. stack_name: Stack name for tagging. region: AWS region. Returns: Dictionary with cluster info: arn, name. """ ecs_client = get_boto3_client("ecs", region) # First, check if cluster already exists and is ACTIVE with suppress(ClientError): response = ecs_client.describe_clusters(clusters=[name]) if response.get("clusters"): cluster = response["clusters"][0] if cluster.get("status") == "ACTIVE": return { "arn": cluster["clusterArn"], "name": name, } # If INACTIVE, wait for it to be fully deleted if cluster.get("status") == "INACTIVE": time.sleep(5) # Give time for cleanup # Create new cluster try: response = ecs_client.create_cluster( clusterName=name, tags=get_standard_tags_ecs(stack_name), capacityProviders=["FARGATE", "FARGATE_SPOT"], defaultCapacityProviderStrategy=[ {"capacityProvider": "FARGATE", "weight": 1}, ], ) cluster_arn = response["cluster"]["clusterArn"] except ClientError as e: error_code = e.response.get("Error", {}).get("Code", "") # Handle cluster already exists or idempotent request conflict if "ClusterAlreadyExists" in str(e) or error_code == "InvalidParameterException": # Get existing cluster ARN - it might have become active response = ecs_client.describe_clusters(clusters=[name]) if response.get("clusters") and response["clusters"][0].get("status") == "ACTIVE": cluster_arn = response["clusters"][0]["clusterArn"] else: raise else: raise return { "arn": cluster_arn, "name": name, } def create_task_definition( family: str, container_definitions: list[dict[str, Any]], task_role_arn: str, execution_role_arn: str, cpu: str = "256", memory: str = "512", runtime_platform: dict[str, str] | None = None, stack_name: str | None = None, region: str = DEFAULT_REGION, ) -> dict[str, Any]: """Register task definition. Args: family: Task definition family name. container_definitions: List of container definitions. task_role_arn: ARN of the task role. execution_role_arn: ARN of the execution role. cpu: CPU units (256, 512, 1024, 2048, 4096). memory: Memory in MB (512, 1024, 2048, etc.). runtime_platform: Platform config (e.g., {"cpuArchitecture": "ARM64"}). stack_name: Stack name for tagging. region: AWS region. Returns: Dictionary with task definition info: arn, family, revision. """ ecs_client = get_boto3_client("ecs", region) config: dict[str, Any] = { "family": family, "containerDefinitions": container_definitions, "taskRoleArn": task_role_arn, "executionRoleArn": execution_role_arn, "cpu": cpu, "memory": memory, "networkMode": "awsvpc", "requiresCompatibilities": ["FARGATE"], } if runtime_platform: config["runtimePlatform"] = { "operatingSystemFamily": runtime_platform.get("operatingSystemFamily", "LINUX"), "cpuArchitecture": runtime_platform.get("cpuArchitecture", "X86_64"), } if stack_name: config["tags"] = get_standard_tags_ecs(stack_name) response = ecs_client.register_task_definition(**config) task_def = response["taskDefinition"] return { "arn": task_def["taskDefinitionArn"], "family": task_def["family"], "revision": task_def["revision"], } def create_service( cluster: str, name: str, task_definition: str, desired_count: int = 1, subnets: list[str] | None = None, security_groups: list[str] | None = None, assign_public_ip: bool = True, stack_name: str | None = None, region: str = DEFAULT_REGION, ) -> dict[str, Any]: """Create Fargate service. Args: cluster: Cluster name or ARN. name: Service name. task_definition: Task definition family:revision or ARN. desired_count: Number of tasks to run. subnets: List of subnet IDs. security_groups: List of security group IDs. assign_public_ip: Whether to assign public IP. stack_name: Stack name for tagging. region: AWS region. Returns: Dictionary with service info: arn, name. """ from tests.shared.infrastructure_sdk.resources.vpc import get_default_vpc, get_public_subnets ecs_client = get_boto3_client("ecs", region) # Get default VPC resources if not provided if not subnets: vpc = get_default_vpc(region) subnets = get_public_subnets(vpc["vpc_id"], region) network_config: dict[str, Any] = { "awsvpcConfiguration": { "subnets": subnets, "assignPublicIp": "ENABLED" if assign_public_ip else "DISABLED", } } if security_groups: network_config["awsvpcConfiguration"]["securityGroups"] = security_groups config: dict[str, Any] = { "cluster": cluster, "serviceName": name, "taskDefinition": task_definition, "desiredCount": desired_count, "launchType": "FARGATE", "networkConfiguration": network_config, } if stack_name: config["tags"] = get_standard_tags_ecs(stack_name) config["propagateTags"] = "SERVICE" try: response = ecs_client.create_service(**config) service_arn = response["service"]["serviceArn"] except ClientError as e: if ( "ServiceAlreadyExists" in str(e) or e.response["Error"]["Code"] == "InvalidParameterException" ): # Update existing service ecs_client.update_service( cluster=cluster, service=name, taskDefinition=task_definition, desiredCount=desired_count, ) response = ecs_client.describe_services(cluster=cluster, services=[name]) service_arn = response["services"][0]["serviceArn"] else: raise return { "arn": service_arn, "name": name, } def run_task( cluster: str, task_definition: str, subnets: list[str], security_groups: list[str] | None = None, assign_public_ip: bool = True, overrides: dict[str, Any] | None = None, count: int = 1, region: str = DEFAULT_REGION, ) -> dict[str, Any]: """Run one-shot ECS task. Args: cluster: Cluster name or ARN. task_definition: Task definition family:revision or ARN. subnets: List of subnet IDs. security_groups: List of security group IDs. assign_public_ip: Whether to assign public IP. overrides: Task overrides (environment, command, etc.). count: Number of tasks to run. region: AWS region. Returns: Dictionary with task info: task_arns, cluster. """ ecs_client = get_boto3_client("ecs", region) network_config: dict[str, Any] = { "awsvpcConfiguration": { "subnets": subnets, "assignPublicIp": "ENABLED" if assign_public_ip else "DISABLED", } } if security_groups: network_config["awsvpcConfiguration"]["securityGroups"] = security_groups config: dict[str, Any] = { "cluster": cluster, "taskDefinition": task_definition, "launchType": "FARGATE", "networkConfiguration": network_config, "count": count, } if overrides: config["overrides"] = overrides response = ecs_client.run_task(**config) return { "task_arns": [t["taskArn"] for t in response.get("tasks", [])], "cluster": cluster, "failures": response.get("failures", []), } def wait_for_tasks_stopped( cluster: str, task_arns: list[str], region: str = DEFAULT_REGION, timeout_seconds: int = 600, ) -> list[dict[str, Any]]: """Wait for tasks to stop and return their exit codes. Args: cluster: Cluster name or ARN. task_arns: List of task ARNs. region: AWS region. timeout_seconds: Maximum time to wait. Returns: List of task details with exit codes. """ ecs_client = get_boto3_client("ecs", region) waiter = ecs_client.get_waiter("tasks_stopped") waiter.wait( cluster=cluster, tasks=task_arns, WaiterConfig={ "Delay": 10, "MaxAttempts": timeout_seconds // 10, }, ) response = ecs_client.describe_tasks(cluster=cluster, tasks=task_arns) results = [] for task in response.get("tasks", []): container_results = [] for container in task.get("containers", []): container_results.append( { "name": container["name"], "exit_code": container.get("exitCode"), "reason": container.get("reason"), } ) results.append( { "task_arn": task["taskArn"], "last_status": task["lastStatus"], "stopped_reason": task.get("stoppedReason"), "containers": container_results, } ) return results def delete_cluster(name: str, region: str = DEFAULT_REGION) -> None: """Delete cluster (must have no services). Args: name: Cluster name. region: AWS region. """ ecs_client = get_boto3_client("ecs", region) try: ecs_client.delete_cluster(cluster=name) except ClientError as e: if e.response["Error"]["Code"] != "ClusterNotFoundException": raise def delete_service(cluster: str, service: str, region: str = DEFAULT_REGION) -> None: """Delete service. Args: cluster: Cluster name or ARN. service: Service name or ARN. region: AWS region. """ ecs_client = get_boto3_client("ecs", region) with suppress(ClientError): # Scale to 0 first ecs_client.update_service(cluster=cluster, service=service, desiredCount=0) time.sleep(5) try: ecs_client.delete_service(cluster=cluster, service=service, force=True) except ClientError as e: if e.response["Error"]["Code"] != "ServiceNotFoundException": raise def get_cluster(name: str, region: str = DEFAULT_REGION) -> dict[str, Any] | None: """Get cluster details. Args: name: Cluster name. region: AWS region. Returns: Cluster details or None if not found. """ ecs_client = get_boto3_client("ecs", region) try: response = ecs_client.describe_clusters(clusters=[name]) if response["clusters"]: cluster = response["clusters"][0] return { "arn": cluster["clusterArn"], "name": cluster["clusterName"], "status": cluster["status"], } return None except ClientError: return None def deregister_task_definition(task_definition: str, region: str = DEFAULT_REGION) -> None: """Deregister a task definition. Args: task_definition: Task definition ARN or family:revision. region: AWS region. """ ecs_client = get_boto3_client("ecs", region) try: ecs_client.deregister_task_definition(taskDefinition=task_definition) except ClientError as e: if e.response["Error"]["Code"] != "InvalidParameterException": raise def build_container_definition( name: str, image: str, cpu: int = 256, memory: int = 512, port_mappings: list[dict[str, Any]] | None = None, environment: dict[str, str] | None = None, secrets: list[dict[str, str]] | None = None, command: list[str] | None = None, entry_point: list[str] | None = None, log_group: str | None = None, region: str = DEFAULT_REGION, essential: bool = True, ) -> dict[str, Any]: """Build a container definition dict. Args: name: Container name. image: Docker image URI. cpu: CPU units for this container. memory: Memory limit in MB. port_mappings: List of port mappings. environment: Environment variables dict. secrets: List of secrets (name, valueFrom ARN). command: Container command override. entry_point: Container entry point override. log_group: CloudWatch log group name. region: AWS region. essential: Whether container is essential. Returns: Container definition dict ready for task definition. """ container: dict[str, Any] = { "name": name, "image": image, "cpu": cpu, "memory": memory, "essential": essential, } if port_mappings: container["portMappings"] = port_mappings if environment: container["environment"] = [{"name": k, "value": v} for k, v in environment.items()] if secrets: container["secrets"] = secrets if command: container["command"] = command if entry_point: container["entryPoint"] = entry_point if log_group: container["logConfiguration"] = { "logDriver": "awslogs", "options": { "awslogs-group": log_group, "awslogs-region": region, "awslogs-stream-prefix": name, }, } return container