From 81c93fcacec22c559e09fd6fc44481802c4ec5fc Mon Sep 17 00:00:00 2001
From: AgraVator <reachabhishekag@gmail.com>
Date: Mon, 17 Feb 2025 15:09:00 +0530
Subject: [PATCH] Updates resource cleanup script to support gamma resources
 for Prod

---
 bin/cleanup/cleanup.py                        | 61 +++++++++++++++++--
 .../runners/k8s/gamma_server_runner.py        |  4 +-
 2 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/bin/cleanup/cleanup.py b/bin/cleanup/cleanup.py
index a778610e..cc58b035 100644
--- a/bin/cleanup/cleanup.py
+++ b/bin/cleanup/cleanup.py
@@ -44,12 +44,13 @@
 from framework.infrastructure import k8s
 from framework.infrastructure import traffic_director
 from framework.test_app.runners.k8s import k8s_xds_client_runner
-from framework.test_app.runners.k8s import k8s_xds_server_runner
+from framework.test_app.runners.k8s import k8s_xds_server_runner, gamma_server_runner
 
 logger = logging.getLogger(__name__)
 Json = Any
 _KubernetesClientRunner = k8s_xds_client_runner.KubernetesClientRunner
 _KubernetesServerRunner = k8s_xds_server_runner.KubernetesServerRunner
+_GammaServerRunner = gamma_server_runner.GammaServerRunner
 
 GCLOUD = os.environ.get("GCLOUD", "gcloud")
 GCLOUD_CMD_TIMEOUT_S = datetime.timedelta(seconds=5).total_seconds()
@@ -68,6 +69,7 @@
 LEGACY_DRIVER_SECONDARY_ZONE = "us-west1-b"
 
 PSM_INTEROP_PREFIX = "psm-interop"  # Prefix for gke resources to delete.
+GAMMA_PREFIX = "csm-psm" # Prefix for csm gke resources to delete.
 URL_MAP_TEST_PREFIX = (
     "interop-psm-url-map"  # Prefix for url-map test resources to delete.
 )
@@ -95,7 +97,7 @@
 )
 SERVER_PREFIXES = flags.DEFINE_list(
     "server_prefixes",
-    default=[PSM_INTEROP_PREFIX],
+    default=[PSM_INTEROP_PREFIX, GAMMA_PREFIX],
     help=(
         "a comma-separated list of prefixes for which the leaked servers will"
         " be deleted"
@@ -103,7 +105,7 @@
 )
 CLIENT_PREFIXES = flags.DEFINE_list(
     "client_prefixes",
-    default=[PSM_INTEROP_PREFIX, URL_MAP_TEST_PREFIX],
+    default=[PSM_INTEROP_PREFIX, URL_MAP_TEST_PREFIX, GAMMA_PREFIX],
     help=(
         "a comma-separated list of prefixes for which the leaked clients will"
         " be deleted"
@@ -412,6 +414,48 @@ def cleanup_client(
         )
         raise
 
+# cleanup_gamma_server creates a server runner, and calls its cleanup() method.
+def cleanup_gamma_server(
+        project,
+        network,
+        k8s_api_manager,
+        server_namespace,
+        gcp_api_manager,
+        gcp_service_account,
+        *,
+        suffix: Optional[str] = "",
+):
+    deployment_name = xds_flags.SERVER_NAME.value
+    if suffix:
+        deployment_name = f"{deployment_name}-{suffix}"
+
+    ns = k8s.KubernetesNamespace(k8s_api_manager, server_namespace)
+    # Shorten the timeout to avoid waiting for the stuck namespaces.
+    # Normal ns deletion during the cleanup takes less two minutes.
+    ns.wait_for_namespace_deleted_timeout_sec = 5 * 60
+    server_runner = _GammaServerRunner(
+        k8s_namespace=ns,
+        deployment_name=deployment_name,
+        gcp_project=project,
+        network=network,
+        gcp_service_account=gcp_service_account,
+        gcp_api_manager=gcp_api_manager,
+        image_name="",
+        td_bootstrap_image="",
+    )
+
+    logger.info("Cleanup gamma server")
+    try:
+        server_runner.cleanup(force=True, force_namespace=True)
+    except retryers.RetryError as err:
+        logger.error(
+            "Timeout waiting for namespace %s deletion. "
+            "Failed resource status:\n\n%s",
+            ns.name,
+            ns.pretty_format_status(err.result()),
+        )
+        raise
+
 
 # cleanup_server creates a server runner, and calls its cleanup() method.
 def cleanup_server(
@@ -580,9 +624,14 @@ def find_and_remove_leaked_k8s_resources(
             K8sResourceRule(f"{prefix}-client-(.*)", cleanup_client)
         )
     for prefix in SERVER_PREFIXES.value:
-        k8s_resource_rules.append(
-            K8sResourceRule(f"{prefix}-server-(.*)", cleanup_server)
-        )
+        if prefix == GAMMA_PREFIX:
+            k8s_resource_rules.append(
+                K8sResourceRule(f"{prefix}-server-(.*)", cleanup_gamma_server)
+            )
+        else:
+            k8s_resource_rules.append(
+                K8sResourceRule(f"{prefix}-server-(.*)", cleanup_server)
+            )
 
     # Delete leaked k8s namespaces, those usually mean there are leaked testing
     # client/servers from the gke framework.
diff --git a/framework/test_app/runners/k8s/gamma_server_runner.py b/framework/test_app/runners/k8s/gamma_server_runner.py
index 4971a476..eae1e706 100644
--- a/framework/test_app/runners/k8s/gamma_server_runner.py
+++ b/framework/test_app/runners/k8s/gamma_server_runner.py
@@ -58,7 +58,7 @@ class GammaServerRunner(KubernetesServerRunner):
     def __init__(
         self,
         k8s_namespace: k8s.KubernetesNamespace,
-        frontend_service_name: str,
+        frontend_service_name: str, #
         *,
         deployment_name: str,
         image_name: str,
@@ -71,7 +71,7 @@ def __init__(
         gcp_service_account: str,
         service_account_name: Optional[str] = None,
         service_name: Optional[str] = None,
-        route_kind: k8s.RouteKind = k8s.RouteKind.HTTP,
+        route_kind: k8s.RouteKind = k8s.RouteKind.HTTP, #
         neg_name: Optional[str] = None,
         deployment_template: str = "server.deployment.yaml",
         service_account_template: str = "service-account.yaml",