diff --git a/e2e-tests/fixtures/import/cleanup_resources.py b/e2e-tests/fixtures/import/cleanup_resources.py index 120ced18f..eec1d41ec 100644 --- a/e2e-tests/fixtures/import/cleanup_resources.py +++ b/e2e-tests/fixtures/import/cleanup_resources.py @@ -10,12 +10,16 @@ import json import os import sys +import time sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from common import REGION, RESOURCES_FILE, get_control_client, get_account_id import boto3 +DELETE_TIMEOUT_SECONDS = 120 +DELETE_POLL_SECONDS = 5 + def cleanup_s3_code_objects(): """Delete uploaded code.zip objects from the bugbash S3 bucket.""" @@ -36,6 +40,91 @@ def cleanup_s3_code_objects(): print(f"Could not clean up S3 objects: {e}") +def is_not_found(error): + """Return whether an AWS SDK error means the resource is already gone.""" + response = getattr(error, "response", {}) + code = response.get("Error", {}).get("Code") + return code in ("ResourceNotFoundException", "NotFoundException") + + +def wait_for_gateway_target_deleted( + client, + gateway_id, + target_id, + timeout=DELETE_TIMEOUT_SECONDS, + poll_interval=DELETE_POLL_SECONDS, +): + """Wait until a target is gone so its parent gateway can be deleted.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + client.get_gateway_target( + gatewayIdentifier=gateway_id, + targetId=target_id, + ) + except Exception as e: + if is_not_found(e): + return + raise + time.sleep(poll_interval) + raise TimeoutError(f"Gateway target {target_id} was not deleted within {timeout}s") + + +def delete_tracked_resources(client, resources): + """Delete tracked resources and return keys that still require cleanup.""" + gateways_by_arn = { + value.get("arn"): value.get("id") + for key, value in resources.items() + if key == "gateway" + } + + def delete_order(item): + key = item[0] + if key.startswith("gateway-target"): + return 0 + if key == "gateway": + return 2 + return 1 + + failed = [] + for key, value in sorted(resources.items(), key=delete_order): + resource_id = value.get("id") + if not resource_id: + print(f"Could not delete {key}: tracked resource has no id") + failed.append(key) + continue + + try: + if key.startswith("gateway-target"): + gateway_id = gateways_by_arn.get(value.get("arn")) + if not gateway_id: + raise ValueError("tracked gateway target has no matching parent gateway") + client.delete_gateway_target( + gatewayIdentifier=gateway_id, + targetId=resource_id, + ) + wait_for_gateway_target_deleted(client, gateway_id, resource_id) + elif key == "gateway": + client.delete_gateway(gatewayIdentifier=resource_id) + elif key.startswith("runtime-"): + client.delete_agent_runtime(agentRuntimeId=resource_id) + elif key.startswith("memory-"): + client.delete_memory(memoryId=resource_id) + elif key.startswith("evaluator-"): + client.delete_evaluator(evaluatorId=resource_id) + else: + raise ValueError(f"unsupported tracked resource key: {key}") + print(f"Deleted {key}: {resource_id}") + except Exception as e: + if is_not_found(e): + print(f"Already deleted {key}: {resource_id}") + continue + print(f"Could not delete {key} ({resource_id}): {e}") + failed.append(key) + + return failed + + def main(): if not os.path.exists(RESOURCES_FILE): print("No bugbash-resources.json found, nothing to clean up") @@ -45,23 +134,7 @@ def main(): resources = json.load(f) client = get_control_client() - - failed = [] - for key, val in resources.items(): - rid = val.get("id") - if not rid: - continue - try: - if "runtime" in key: - client.delete_agent_runtime(agentRuntimeId=rid) - elif "memory" in key: - client.delete_memory(memoryId=rid) - elif "evaluator" in key: - client.delete_evaluator(evaluatorId=rid) - print(f"Deleted {key}: {rid}") - except Exception as e: - print(f"Could not delete {key} ({rid}): {e}") - failed.append(key) + failed = delete_tracked_resources(client, resources) if failed: remaining = {k: v for k, v in resources.items() if k in failed}