From e579f9bf43bfd75569a7918d2c967c0ae876b891 Mon Sep 17 00:00:00 2001
From: Sait Talha Nisanci <s.talhanisanci@gmail.com>
Date: Tue, 19 Oct 2021 16:27:01 +0300
Subject: [PATCH] run on coordinator and a random worker

---
 .../arbitrary_configs/citus_arbitrary_configs.py   | 14 +++++++++-----
 src/test/regress/citus_tests/config.py             | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/test/regress/citus_tests/arbitrary_configs/citus_arbitrary_configs.py b/src/test/regress/citus_tests/arbitrary_configs/citus_arbitrary_configs.py
index b5f71f945..95f526c93 100755
--- a/src/test/regress/citus_tests/arbitrary_configs/citus_arbitrary_configs.py
+++ b/src/test/regress/citus_tests/arbitrary_configs/citus_arbitrary_configs.py
@@ -84,7 +84,7 @@ def run_for_config(config, lock, sql_schedule_name):
     extra_tests = os.getenv("EXTRA_TESTS", "")
     if config.is_mx and config.worker_amount > 0:
         exitCode |= _run_pg_regress_on_port(
-            config, config.random_port(), sql_schedule_name, extra_tests=extra_tests
+            config, config.sql_port, sql_schedule_name, extra_tests=extra_tests
         )
     else:
         exitCode |= _run_pg_regress_on_port(
@@ -171,9 +171,13 @@ def read_configs(docoptRes):
     configs = []
     # We fill the configs from all of the possible classes in config.py so that if we add a new config,
     # we don't need to add it here. And this avoids the problem where we forget to add it here
-    for x in cfg.__dict__.values():
-        if cfg.should_include_config(x):
-            configs.append(x(docoptRes))
+    for class_name in cfg.__dict__.values():
+        if cfg.should_include_config(class_name):
+            if issubclass(class_name, cfg.CitusMXBaseClusterConfig):
+                arguments = dict(docoptRes)
+                arguments["run_test_on_coordinator"] = True
+                configs.append(class_name(arguments))
+            configs.append(class_name(docoptRes))
     return configs
 
 
@@ -195,7 +199,7 @@ def read_arguments(docoptRes):
         given_configs = docoptRes["--configs"].split(",")
         new_configs = []
         for config in configs:
-            if config.name in given_configs:
+            if config.class_name in given_configs:
                 new_configs.append(config)
         if len(new_configs) > 0:
             configs = new_configs
diff --git a/src/test/regress/citus_tests/config.py b/src/test/regress/citus_tests/config.py
index b4bdb9991..365c484fe 100644
--- a/src/test/regress/citus_tests/config.py
+++ b/src/test/regress/citus_tests/config.py
@@ -99,7 +99,14 @@ class CitusBaseClusterConfig(object, metaclass=NewInitCaller):
         self.user = REGULAR_USER_NAME
         self.is_mx = False
         self.is_citus = True
+        self.run_test_on_coordinator = (
+            True if "run_test_on_coordinator" in arguments else False
+        )
+        self.class_name = type(self).__name__
         self.name = type(self).__name__
+        if self.run_test_on_coordinator:
+            self.name += "_coordinator"
+
         self.settings = {
             "shared_preload_libraries": "citus",
             "log_error_verbosity": "terse",
@@ -122,6 +129,11 @@ class CitusBaseClusterConfig(object, metaclass=NewInitCaller):
         if self.worker_amount > 0:
             self.chosen_random_worker_port = self.random_worker_port()
         self.settings.update(self.new_settings)
+        self.sql_port = (
+            self.coordinator_port()
+            if self.run_test_on_coordinator
+            else self.random_worker_port()
+        )
 
     def coordinator_port(self):
         return self.node_name_to_ports[COORDINATOR_NAME]
@@ -130,6 +142,8 @@ class CitusBaseClusterConfig(object, metaclass=NewInitCaller):
         pass
 
     def random_worker_port(self):
+        if len(self.worker_ports) == 0:
+            return self.coordinator_port()
         return random.choice(self.worker_ports)
 
     def random_port(self):