[PSM Interop] URL Map: register the cleanup hook early (#33137)

I've noticed we add the cleanup hook after setting up the
infrastructure. Thus, if infra setup failed, the cleanup won't work.
This fixes it, and adds extra checks to not call
`cls.test_client_runner` if it's not set.

tools/run_tests/xds_k8s_test_driver/framework/xds_url_map_testcase.py

@@ -36,6 +36,7 @@ from framework.helpers import retryers
 from framework.helpers import skips
 from framework.infrastructure import k8s
 from framework.test_app import client_app
+from framework.test_app.runners.k8s import k8s_xds_client_runner
 
 # Load existing flags
 flags.adopt_module_key_flags(xds_k8s_testcase)
@@ -57,6 +58,7 @@ XdsTestClient = client_app.XdsTestClient
 GcpResourceManager = xds_url_map_test_resources.GcpResourceManager
 HostRule = xds_url_map_test_resources.HostRule
 PathMatcher = xds_url_map_test_resources.PathMatcher
+_KubernetesClientRunner = k8s_xds_client_runner.KubernetesClientRunner
 JsonType = Any
 _timedelta = datetime.timedelta
 
@@ -248,6 +250,8 @@ class XdsUrlMapTestCase(absltest.TestCase, metaclass=_MetaXdsUrlMapTestCase):
     - rpc_distribution_validate: Validates if the routing behavior is correct
     """
 
+    test_client_runner: Optional[_KubernetesClientRunner] = None
+
     @staticmethod
     def is_supported(config: skips.TestConfig) -> bool:
         """Allow the test case to decide whether it supports the given config.
@@ -339,15 +343,15 @@ class XdsUrlMapTestCase(absltest.TestCase, metaclass=_MetaXdsUrlMapTestCase):
         # support current test case.
         skips.evaluate_test_config(cls.is_supported)
 
+        # Configure cleanup to run after all tests regardless of
+        # whether setUpClass failed.
+        cls.addClassCleanup(cls.cleanupAfterTests)
+
         if not cls.started_test_cases:
             # Create the GCP resource once before the first test start
             GcpResourceManager().setup(cls.test_case_classes)
         cls.started_test_cases.add(cls.__name__)
 
-        # Configure cleanup to run after all tests regardless of whether or not
-        # this setUpClass failed
-        cls.addClassCleanup(cls.cleanupAfterTests)
-
         # Create the test case's own client runner with it's own namespace,
         # enables concurrent running with other test cases.
         cls.test_client_runner = GcpResourceManager().create_test_client_runner(
@@ -365,17 +369,15 @@ class XdsUrlMapTestCase(absltest.TestCase, metaclass=_MetaXdsUrlMapTestCase):
     @classmethod
     def cleanupAfterTests(cls):
         logging.info('----- TestCase %s teardown -----', cls.__name__)
-        logging.debug('Getting pods restart times')
         client_restarts: int = 0
-        try:
-            client_restarts = cls.test_client_runner.get_pod_restarts(
-                cls.test_client_runner.deployment)
-        except (retryers.RetryError, k8s.NotFound) as e:
-            logging.exception(e)
+        if cls.test_client_runner:
+            try:
+                logging.debug('Getting pods restart times')
+                client_restarts = cls.test_client_runner.get_pod_restarts(
+                    cls.test_client_runner.deployment)
+            except (retryers.RetryError, k8s.NotFound) as e:
+                logging.exception(e)
 
-        retryer = retryers.constant_retryer(wait_fixed=_timedelta(seconds=10),
-                                            attempts=3,
-                                            log_level=logging.INFO)
         cls.finished_test_cases.add(cls.__name__)
         # Whether to clean up shared pre-provisioned infrastructure too.
         # We only do it after all tests are finished.
@@ -383,6 +385,9 @@ class XdsUrlMapTestCase(absltest.TestCase, metaclass=_MetaXdsUrlMapTestCase):
 
         # Graceful cleanup: try three times, and don't fail the test on
         # a cleanup failure.
+        retryer = retryers.constant_retryer(wait_fixed=_timedelta(seconds=10),
+                                            attempts=3,
+                                            log_level=logging.INFO)
         try:
             retryer(cls._cleanup, cleanup_all)
         except retryers.RetryError:
@@ -397,7 +402,8 @@ class XdsUrlMapTestCase(absltest.TestCase, metaclass=_MetaXdsUrlMapTestCase):
 
     @classmethod
     def _cleanup(cls, cleanup_all: bool = False):
-        cls.test_client_runner.cleanup(force=True, force_namespace=True)
+        if cls.test_client_runner:
+            cls.test_client_runner.cleanup(force=True, force_namespace=True)
         if cleanup_all:
             GcpResourceManager().cleanup()