Enable Session drain test for Python client (grpc#64)

The docker image used for testing was pushed to Container Registry and
can be tested locally using those flags:
*
`--client_image=gcr.io/grpc-testing/xds-interop/python-client:test_ssa_interop_python`
*
`--server_image=gcr.io/grpc-testing/xds-interop/python-server:test_ssa_interop_python`

Tested locally use:
- [x] `python -m tests.gamma.affinity_session_drain_test
--flagfile="config/local-dev-gamma.cfg"`
- [x] `python -m tests.gamma.gamma_baseline_test
--flagfile="config/local-dev-gamma.cfg"`
- [x] `python -m tests.gamma.affinity_test
--flagfile="config/local-dev-gamma.cfg"`

config/common-csm.cfg

@@ -1,3 +1,4 @@
 # Common config file for PSM CSM tests.
 --resource_prefix=psm-csm
 --noenable_workload_identity
+--csm_server_image_canonical=gcr.io/grpc-testing/xds-interop/cpp-server:v1.62.x

framework/xds_gamma_testcase.py

@@ -15,8 +15,10 @@
 import logging
 from typing import List
 
+from absl import flags
 from typing_extensions import override
 
+from framework import xds_k8s_flags
 from framework.infrastructure import k8s
 import framework.infrastructure.traffic_director_gamma as td_gamma
 from framework.test_app import client_app
@@ -31,13 +33,21 @@
 XdsTestServer = server_app.XdsTestServer
 
 logger = logging.getLogger(__name__)
+flags.adopt_module_key_flags(xds_k8s_flags)
 
 
 # TODO(sergiitk): [GAMMA] Move into framework/test_cases
 class GammaXdsKubernetesTestCase(xds_k8s_testcase.RegularXdsKubernetesTestCase):
     server_runner: GammaServerRunner
     frontend_service_name: str
 
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.csm_server_image_canonical = (
+            xds_k8s_flags.CSM_SERVER_IMAGE_CANONICAL.value
+        )
+
     def setUp(self):
         """Hook method for setting up the test fixture before exercising it."""
         # TODO(sergiitk): [GAMMA] Remove when refactored to be TD-manager-less.

framework/xds_k8s_flags.py

@@ -46,6 +46,15 @@
         "does not exist, or missing a feature required for the test."
     ),
 )
+CSM_SERVER_IMAGE_CANONICAL = flags.DEFINE_string(
+    "csm_server_image_canonical",
+    default=None,
+    help=(
+        "The canonical implementation of the CSM xDS test server.\n"
+        "Can be used in tests where language-specific xDS test server"
+        "does not exist, or missing a feature required for the test."
+    ),
+)
 CLIENT_IMAGE = flags.DEFINE_string(
     "client_image", default=None, help="Client Docker image name"
 )

tests/app_net_ssa_test.py

@@ -40,7 +40,7 @@
 class AppNetSsaTest(xds_k8s_testcase.AppNetXdsKubernetesTestCase):
     @staticmethod
     def is_supported(config: skips.TestConfig) -> bool:
-        if config.client_lang is _Lang.CPP:
+        if config.client_lang in _Lang.CPP | _Lang.PYTHON:
             return config.version_gte("v1.62.x")
         return False
 

tests/gamma/affinity_session_drain_test.py

@@ -50,10 +50,18 @@ class AffinitySessionDrainTest(  # pylint: disable=too-many-ancestors
     xds_gamma_testcase.GammaXdsKubernetesTestCase,
     session_affinity_mixin.SessionAffinityMixin,
 ):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        # Force the python client to use the reference server image (C++)
+        # because the python server doesn't yet support session drain test.
+        if cls.lang_spec.client_lang == _Lang.PYTHON:
+            cls.server_image = cls.csm_server_image_canonical
+
     @staticmethod
     @override
     def is_supported(config: skips.TestConfig) -> bool:
-        if config.client_lang == _Lang.CPP and config.server_lang == _Lang.CPP:
+        if config.client_lang in _Lang.CPP | _Lang.PYTHON:
             return config.version_gte("v1.62.x")
         return False