diff --git a/src/python_testing/TC_TSTAT_2_2.py b/src/python_testing/TC_TSTAT_2_2.py
index f4e1dfbc2b70d3..02fc9c2ff3c4df 100644
--- a/src/python_testing/TC_TSTAT_2_2.py
+++ b/src/python_testing/TC_TSTAT_2_2.py
@@ -18,27 +18,19 @@
 # for details about the block below.
 #
 # === BEGIN CI TEST ARGUMENTS ===
-# test-runner-runs:
-#   run1:
-#     app: ${ALL_CLUSTERS_APP}
-#     app-args: --discriminator 1234 --KVS kvs1 --trace-to json:${TRACE_APP}.json
-#     script-args: >
-#       --storage-path admin_storage.json
-#       --commissioning-method on-network
-#       --discriminator 1234
-#       --passcode 20202021
-#       --endpoint 1
-#       --trace-to json:${TRACE_TEST_JSON}.json
-#       --trace-to perfetto:${TRACE_TEST_PERFETTO}.perfetto
-#     factory-reset: true
-#     quiet: true
+# test-runner-runs: run1
+# test-runner-run/run1/app: ${ALL_CLUSTERS_APP}
+# test-runner-run/run1/app-args: --discriminator 1234 --KVS kvs1 --trace-to json:${TRACE_APP}.json
+# test-runner-run/run1/script-args: --storage-path admin_storage.json --commissioning-method on-network --discriminator 1234 --passcode 20202021 --endpoint 1 --trace-to json:${TRACE_TEST_JSON}.json --trace-to perfetto:${TRACE_TEST_PERFETTO}.perfetto
+# test-runner-run/run1/factoryreset: True
+# test-runner-run/run1/quiet: True
 # === END CI TEST ARGUMENTS ===
 
 import logging
 
 import chip.clusters as Clusters
 from chip.interaction_model import Status
-from chip.testing.matter_testing import MatterBaseTest, TestStep, async_test_body, default_matter_test_main
+from matter_testing_support import MatterBaseTest, TestStep, async_test_body, default_matter_test_main
 from mobly import asserts
 
 logger = logging.getLogger(__name__)
diff --git a/src/python_testing/matter_testing_support.py b/src/python_testing/matter_testing_support.py
index 96f76479fc8aa6..1471da5e7f8246 100644
--- a/src/python_testing/matter_testing_support.py
+++ b/src/python_testing/matter_testing_support.py
@@ -1201,6 +1201,9 @@ def certificate_authority_manager(self) -> chip.CertificateAuthority.Certificate
     def dut_node_id(self) -> int:
         return self.matter_test_config.dut_node_ids[0]
 
+    def get_endpoint(self, default: Optional[int] = 0) -> int:
+        return self.matter_test_config.endpoint if self.matter_test_config.endpoint is not None else default
+
     def setup_class(self):
         super().setup_class()