Make run_test.py and create_test.py importable without errors (#6736)

Allowing scripts to be importable is good practice in general and it's
required for the pytest testing framework that I'm adding in a follow up
PR.

src/test/regress/bin/create_test.py

@@ -4,25 +4,26 @@ import os
 import random
 import sys
 
-if len(sys.argv) != 2:
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
         print(
             "ERROR: Expected the name of the new test as an argument, such as:\n"
             "src/test/regress/bin/create_test.py my_awesome_test"
         )
         sys.exit(1)
 
-test_name = sys.argv[1]
+    test_name = sys.argv[1]
 
-regress_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+    regress_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-filename = os.path.join(regress_dir, "sql", f"{test_name}.sql")
+    filename = os.path.join(regress_dir, "sql", f"{test_name}.sql")
 
-if os.path.isfile(filename):
+    if os.path.isfile(filename):
         print(f"ERROR: test file '{filename}' already exists")
         sys.exit(1)
 
-shard_id = random.randint(1, 999999) * 100
+    shard_id = random.randint(1, 999999) * 100
 
-contents = f"""CREATE SCHEMA {test_name};
+    contents = f"""CREATE SCHEMA {test_name};
 SET search_path TO {test_name};
 SET citus.shard_count TO 4;
 SET citus.shard_replication_factor TO 1;
@@ -34,9 +35,8 @@ SET client_min_messages TO WARNING;
 DROP SCHEMA {test_name} CASCADE;
 """
 
-
+    with open(filename, "w") as f:
-with open(filename, "w") as f:
         f.write(contents)
 
-print(f"Created {filename}")
+    print(f"Created {filename}")
-print(f"Don't forget to add '{test_name}' in multi_schedule somewhere")
+    print(f"Don't forget to add '{test_name}' in multi_schedule somewhere")

src/test/regress/citus_tests/run_test.py

@@ -15,49 +15,51 @@ import common
 
 import config
 
-args = argparse.ArgumentParser()
+if __name__ == "__main__":
-args.add_argument(
+    args = argparse.ArgumentParser()
+    args.add_argument(
         "test_name", help="Test name (must be included in a schedule.)", nargs="?"
-)
+    )
-args.add_argument(
+    args.add_argument(
         "-p",
         "--path",
         required=False,
         help="Relative path for test file (must have a .sql or .spec extension)",
         type=pathlib.Path,
-)
+    )
-args.add_argument("-r", "--repeat", help="Number of test to run", type=int, default=1)
+    args.add_argument(
-args.add_argument(
+        "-r", "--repeat", help="Number of test to run", type=int, default=1
+    )
+    args.add_argument(
         "-b",
         "--use-base-schedule",
         required=False,
         help="Choose base-schedules rather than minimal-schedules",
         action="store_true",
-)
+    )
-args.add_argument(
+    args.add_argument(
         "-w",
         "--use-whole-schedule-line",
         required=False,
         help="Use the whole line found in related schedule",
         action="store_true",
-)
+    )
-args.add_argument(
+    args.add_argument(
         "--valgrind",
         required=False,
         help="Run the test with valgrind enabled",
         action="store_true",
-)
+    )
 
-args = vars(args.parse_args())
+    args = vars(args.parse_args())
 
-regress_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+    regress_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-test_file_path = args["path"]
+    test_file_path = args["path"]
-test_file_name = args["test_name"]
+    test_file_name = args["test_name"]
-use_base_schedule = args["use_base_schedule"]
+    use_base_schedule = args["use_base_schedule"]
-use_whole_schedule_line = args["use_whole_schedule_line"]
+    use_whole_schedule_line = args["use_whole_schedule_line"]
 
-
+    class TestDeps:
-class TestDeps:
         schedule: Optional[str]
         direct_extra_tests: list[str]
 
@@ -76,8 +78,7 @@ class TestDeps:
 
             return list(all_deps.keys())
 
-
+    deps = {
-deps = {
         "multi_cluster_management": TestDeps(
             None, ["multi_test_helpers_superuser"], repeatable=False
         ),
@@ -86,14 +87,13 @@ deps = {
         "multi_extension": TestDeps(None, repeatable=False),
         "multi_test_helpers": TestDeps(None),
         "multi_insert_select": TestDeps("base_schedule"),
-}
+    }
 
-if not (test_file_name or test_file_path):
+    if not (test_file_name or test_file_path):
         print("FATAL: No test given.")
         sys.exit(2)
 
-
+    if test_file_path:
-if test_file_path:
         test_file_path = os.path.join(os.getcwd(), args["path"])
 
         if not os.path.isfile(test_file_path):
@@ -109,11 +109,11 @@ if test_file_path:
             )
             sys.exit(1)
 
-test_schedule = ""
+    test_schedule = ""
-dependencies = []
+    dependencies = []
 
-# find related schedule
+    # find related schedule
-for schedule_file_path in sorted(glob(os.path.join(regress_dir, "*_schedule"))):
+    for schedule_file_path in sorted(glob(os.path.join(regress_dir, "*_schedule"))):
         for schedule_line in open(schedule_file_path, "r"):
             if re.search(r"\b" + test_file_name + r"\b", schedule_line):
                 test_schedule = pathlib.Path(schedule_file_path).stem
@@ -125,11 +125,10 @@ for schedule_file_path in sorted(glob(os.path.join(regress_dir, "*_schedule"))):
         else:
             continue
         break
-else:
+    else:
         raise Exception("Test could not be found in any schedule")
 
-
+    def default_base_schedule(test_schedule):
-def default_base_schedule(test_schedule):
         if "isolation" in test_schedule:
             return "base_isolation_schedule"
 
@@ -151,29 +150,32 @@ def default_base_schedule(test_schedule):
             return "minimal_schedule"
 
         if test_schedule in config.ARBITRARY_SCHEDULE_NAMES:
-        print(f"WARNING: Arbitrary config schedule ({test_schedule}) is not supported.")
+            print(
+                f"WARNING: Arbitrary config schedule ({test_schedule}) is not supported."
+            )
             sys.exit(0)
 
         if use_base_schedule:
             return "base_schedule"
         return "minimal_schedule"
 
-
+    if test_file_name in deps:
-if test_file_name in deps:
         dependencies = deps[test_file_name]
-else:
+    else:
         dependencies = TestDeps(default_base_schedule(test_schedule))
 
-# copy base schedule to a temp file and append test_schedule_line
+    # copy base schedule to a temp file and append test_schedule_line
-# to be able to run tests in parallel (if test_schedule_line is a parallel group.)
+    # to be able to run tests in parallel (if test_schedule_line is a parallel group.)
-tmp_schedule_path = os.path.join(
+    tmp_schedule_path = os.path.join(
         regress_dir, f"tmp_schedule_{ random.randint(1, 10000)}"
-)
+    )
-# some tests don't need a schedule to run
+    # some tests don't need a schedule to run
-# e.g tests that are in the first place in their own schedule
+    # e.g tests that are in the first place in their own schedule
-if dependencies.schedule:
+    if dependencies.schedule:
-    shutil.copy2(os.path.join(regress_dir, dependencies.schedule), tmp_schedule_path)
+        shutil.copy2(
-with open(tmp_schedule_path, "a") as myfile:
+            os.path.join(regress_dir, dependencies.schedule), tmp_schedule_path
+        )
+    with open(tmp_schedule_path, "a") as myfile:
         for dependency in dependencies.extra_tests():
             myfile.write(f"test: {dependency}\n")
 
@@ -184,25 +186,24 @@ with open(tmp_schedule_path, "a") as myfile:
         for _ in range(repetition_cnt):
             myfile.write(test_schedule_line)
 
-
+    # find suitable make recipe
-# find suitable make recipe
+    if dependencies.schedule == "base_isolation_schedule":
-if dependencies.schedule == "base_isolation_schedule":
         make_recipe = "check-isolation-custom-schedule"
-elif dependencies.schedule == "failure_base_schedule":
+    elif dependencies.schedule == "failure_base_schedule":
         make_recipe = "check-failure-custom-schedule"
-else:
+    else:
         make_recipe = "check-custom-schedule"
 
-if args["valgrind"]:
+    if args["valgrind"]:
         make_recipe += "-vg"
 
-# prepare command to run tests
+    # prepare command to run tests
-test_command = f"make -C {regress_dir} {make_recipe} SCHEDULE='{pathlib.Path(tmp_schedule_path).stem}'"
+    test_command = f"make -C {regress_dir} {make_recipe} SCHEDULE='{pathlib.Path(tmp_schedule_path).stem}'"
 
-# run test command n times
+    # run test command n times
-try:
+    try:
         print(f"Executing.. {test_command}")
         result = common.run(test_command)
-finally:
+    finally:
         # remove temp schedule file
         os.remove(tmp_schedule_path)