Dynamically adjust test module's metadata and yaml file during CLI st…

…arts and prevent unnecessary restarts and waits for tests that do not require it

src/lib/modules/catalog/test/resources/bootstrap/bootstrap-test.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 
 set -euxo pipefail
-echo "hello world" > ~/test_bootstrap.txt
+echo "hello world" > /tmp/bootstrap.txt

src/lib/modules/catalog/test/resources/bootstrap/bootstrap-trino.sh

@@ -1,5 +1,4 @@
 #!/usr/bin/env bash
 
 set -euxo pipefail
-
-touch /etc/starburst/test_bootstrap.txt
+echo "hello world" > /tmp/bootstrap.txt

src/lib/modules/catalog/test/test.yaml

@@ -2,23 +2,16 @@
 services:
 
   trino:
-    environment:
-      MINITRINO_BOOTSTRAP: bootstrap-trino.sh
-      CONFIG_PROPERTIES: |-
-        query.max-stage-count=85
-        query.max-execution-time=1h
-      JVM_CONFIG: |-
-        -Xmx2G
-        -Xms1G
+    labels:
+      - com.starburst.tests=minitrino
+      - com.starburst.tests.module.test=catalog-test
 
   test:
     build:
       context: ./modules/catalog/test/
       labels:
         - com.starburst.tests=minitrino
         - com.starburst.tests.module.test=catalog-test
-    environment:
-      MINITRINO_BOOTSTRAP: bootstrap-test.sh
     command: [tail, -f, /dev/null]  # Keep alive
     image: minitrino/test:latest
     container_name: test

src/test/setup.py

@@ -7,5 +7,5 @@
     name="minitrino-tests",
     version="0.0",
     packages=["src", "src/cli", "src/lib"],
-    install_requires=["minitrino", "jsonschema", "requests>=2.32.2"],
+    install_requires=["minitrino", "jsonschema", "PyYAML", "requests>=2.32.2"],
 )

src/test/src/cli/test_cmd_provision.py

@@ -1,10 +1,8 @@
 #!usr/bin/env/python3
 # -*- coding: utf-8 -*-
 
-# TODO: Test no rollback
-# TODO: Test invalid user config (Trino/JVM)
-
 import docker
+import yaml
 
 import src.common as common
 import src.cli.utils as utils
@@ -24,6 +22,7 @@ def main():
     test_bad_sep_version()
     test_invalid_module()
     test_docker_native()
+    test_bootstrap()
     test_enterprise()
     test_version_requirements()
     test_valid_user_config()
@@ -95,9 +94,7 @@ def test_invalid_module():
 
 def test_docker_native():
     """Ensures that native Docker Compose command options can be appended to the
-    provisioning command.
-
-    This function also calls the bootstrap script test functions."""
+    provisioning command."""
 
     common.log_status(cast(FrameType, currentframe()).f_code.co_name)
 
@@ -110,8 +107,6 @@ def test_docker_native():
     assert len(containers) == 2
 
     common.log_success(cast(FrameType, currentframe()).f_code.co_name)
-
-    test_bootstrap_script(result)
     cleanup()
 
 
@@ -120,7 +115,7 @@ def test_enterprise():
 
     common.log_status(cast(FrameType, currentframe()).f_code.co_name)
 
-    utils.update_metadata_json([{"enterprise": True}])
+    utils.update_metadata_json("test", [{"enterprise": True}])
     result = utils.execute_cli_cmd(["-v", "provision", "--module", "test"])
 
     assert result.exit_code == 2
@@ -143,7 +138,7 @@ def test_enterprise():
     )
 
     assert "LIC_PATH" and "/tmp/dummy.license" in result.output
-    utils.reset_metadata_json()
+    utils.reset_test_metadata_json()
     cleanup()
 
     # Ensure default dummy license satisfies Compose env vars
@@ -162,7 +157,7 @@ def test_version_requirements():
 
     common.log_status(cast(FrameType, currentframe()).f_code.co_name)
 
-    utils.update_metadata_json([{"versions": [MIN_SEP_VER + 1, 998]}])
+    utils.update_metadata_json("test", [{"versions": [MIN_SEP_VER + 1, 998]}])
 
     # Should fail - lower bound
     result = utils.execute_cli_cmd(
@@ -189,33 +184,79 @@ def test_version_requirements():
     assert "maximum required" in result.output
 
     common.log_success(cast(FrameType, currentframe()).f_code.co_name)
-    utils.reset_metadata_json()
+    utils.reset_test_metadata_json()
     cleanup()
 
 
-def test_bootstrap_script(result):
+def test_bootstrap():
     """Ensures that bootstrap scripts properly execute in containers."""
 
     common.log_status(cast(FrameType, currentframe()).f_code.co_name)
 
+    def add_yaml_bootstrap(yaml_path=""):
+        # In case left over from previous run
+        del_yaml_bootstrap(yaml_path)
+
+        with open(yaml_path, "r") as file:
+            data = yaml.safe_load(file)
+
+        for svc_name, svc_content in data.get("services", {}).items():
+            if "environment" not in svc_content:
+                svc_content["environment"] = {
+                    "MINITRINO_BOOTSTRAP": f"bootstrap-{svc_name}.sh"
+                }
+
+        with open(yaml_path, "w") as file:
+            yaml.dump(data, file, default_flow_style=False)
+
+    def del_yaml_bootstrap(yaml_path=""):
+        with open(yaml_path, "r") as file:
+            data = yaml.safe_load(file)
+
+        try:
+            for _, svc_content in data.get("services", {}).items():
+                if "environment" in svc_content:
+                    del svc_content["environment"]
+            with open(yaml_path, "w") as file:
+                yaml.dump(data, file, default_flow_style=False)
+        except:
+            pass
+
+    yaml_path = utils.get_module_yaml_path("test")
+    add_yaml_bootstrap(yaml_path)
+
+    result = utils.execute_cli_cmd(
+        [
+            "-v",
+            "provision",
+            "--module",
+            "test",
+        ]
+    )
+
     assert all(
         (
-            "Successfully executed bootstrap script in container: 'trino'",
-            "Successfully executed bootstrap script in container: 'test'",
+            "Successfully executed bootstrap script in container 'trino'"
+            in result.output,
+            "Successfully executed bootstrap script in container 'test'"
+            in result.output,
         )
     )
 
-    trino_bootstrap_check = common.execute_command(
-        f"docker exec -i trino ls /etc/starburst/"
+    trino_bootstrap = common.execute_command(
+        f"docker exec -i trino cat /tmp/bootstrap.txt"
     )
-    test_bootstrap_check = common.execute_command(
-        f"docker exec -i test cat /root/test_bootstrap.txt"
+    test_bootstrap = common.execute_command(
+        f"docker exec -i test cat /tmp/bootstrap.txt"
     )
 
-    assert "test_bootstrap.txt" in trino_bootstrap_check.get("output", "")
-    assert "hello world" in test_bootstrap_check.get("output", "")
+    assert "hello world" in trino_bootstrap.get("output", "") and test_bootstrap.get(
+        "output", ""
+    )
 
     common.log_success(cast(FrameType, currentframe()).f_code.co_name)
+    del_yaml_bootstrap(yaml_path)
+    cleanup()
 
 
 def test_valid_user_config():
@@ -227,6 +268,10 @@ def test_valid_user_config():
     result = utils.execute_cli_cmd(
         [
             "-v",
+            "--env",
+            "CONFIG_PROPERTIES=query.max-stage-count=85\nquery.max-execution-time=1h",
+            "--env",
+            "JVM_CONFIG=-Xms1G\n-Xmx2G",
             "provision",
             "--module",
             "test",

src/test/src/cli/utils.py

@@ -61,11 +61,11 @@ def make_sample_config():
     common.execute_command(cmd)
 
 
-def update_metadata_json(updates=[]):
-    """Updates the test module's metadata.json file with the provided
+def update_metadata_json(module="", updates=[]):
+    """Updates a given module's metadata.json file with the provided
     list of dicts."""
 
-    path = get_metadata_json("test")
+    path = get_metadata_json_path(module)
 
     with open(path, "r") as f:
         data = json.load(f)
@@ -78,7 +78,7 @@ def update_metadata_json(updates=[]):
         json.dump(data, f, indent=4)
 
 
-def reset_metadata_json():
+def reset_test_metadata_json():
     """Resets the test module's metadata.json file to default values."""
 
     default = {
@@ -88,16 +88,31 @@ def reset_metadata_json():
         "versions": [],
         "enterprise": False,
     }
-    path = get_metadata_json("test")
+    path = get_metadata_json_path("test")
 
     with open(path, "w") as f:
-        json.dump(default, f, indent=4)
+        json.dump(default, f, indent=2)
+        f.write("\n")
 
 
-def get_metadata_json(module=""):
+def get_metadata_json_path(module=""):
     """Fetches the metadata.json file path for a given module."""
 
-    result = execute_cli_cmd(["modules", "-m", module, "--json"])
-    output = json.loads(result.output)
+    metadata = get_module_metadata(module)
+    return os.path.abspath(
+        os.path.join(metadata[module]["module_dir"], "metadata.json")
+    )
+
+
+def get_module_metadata(module=""):
+    """Fetches (all) module metadata for a given module."""
+
+    metadata = execute_cli_cmd(["modules", "-m", module, "--json"])
+    return json.loads(metadata.output)
+
+
+def get_module_yaml_path(module=""):
+    """Fetches the module.yaml file path for a given module."""
 
-    return os.path.join(output[module]["module_dir"], "metadata.json")
+    metadata = get_module_metadata(module)
+    return os.path.abspath(metadata[module]["yaml_file"])