Merge pull request #74 from DataKitchen/volume-owner

fix: attempt to set volume owner on tg upgrade

Author: aarthy-dk

dk-installer.py

@@ -77,7 +77,9 @@
 TESTGEN_PULL_TIMEOUT = 5
 TESTGEN_PULL_RETRIES = 3
 TESTGEN_DEFAULT_PORT = 8501
-TESTGEN_LATEST_VERSIONS_URL = "https://dk-support-external.s3.us-east-1.amazonaws.com/testgen-observability/testgen-latest-versions.json"
+TESTGEN_LATEST_VERSIONS_URL = (
+    "https://dk-support-external.s3.us-east-1.amazonaws.com/testgen-observability/testgen-latest-versions.json"
+)
 
 MIXPANEL_TOKEN = "4eff51580bc1685b8ffe79ffb22d2704"
 MIXPANEL_URL = "https://api.mixpanel.com"
@@ -1636,7 +1638,7 @@ class UpdateComposeFileStep(Step):
     label = "Updating the Docker compose file"
 
     def __init__(self):
-        self.update_version = None
+        self.update_version = False
         self.update_analytics = False
         self.update_token = False
         super().__init__()
@@ -1647,7 +1649,9 @@ def pre_execute(self, action, args):
         CONSOLE.space()
 
         contents = action.docker_compose_file_path.read_text()
-        if not args.skip_verify:
+        if args.skip_verify:
+            self.update_version = True
+        else:
             try:
                 output = action.run_cmd(
                     "docker",
@@ -1660,30 +1664,30 @@ def pre_execute(self, action, args):
                     "--help",
                     capture_text=True,
                 )
-                version_match = re.search(r"TestGen\s([0-9.]*)", output)
+                version_match = re.search(r"TestGen\s(?:[a-zA-Z]+\s)*([0-9.]*)", output)
                 current_version = version_match.group(1)
 
                 image_match = re.search(r"image:\s*(datakitchen.+):.+\n", contents)
                 docker_image = image_match.group(1)
                 latest_version = "unknown"
-                
+
                 ssl_context = ssl.create_default_context()
                 ssl_context.check_hostname = False
                 ssl_context.verify_mode = ssl.CERT_NONE
                 resp = urllib.request.urlopen(TESTGEN_LATEST_VERSIONS_URL, timeout=3, context=ssl_context)
                 if resp.code == 200:
-                    json_data = json.loads(resp.read().decode('utf-8'))
+                    json_data = json.loads(resp.read().decode("utf-8"))
                     latest_version = json_data.get("docker", {}).get(docker_image)
             except Exception:
                 CONSOLE.msg("Current version: unknown")
                 CONSOLE.msg("Latest version: unknown")
-                pass
+                self.update_version = True
             else:
                 CONSOLE.msg(f"Current version: {current_version}")
                 CONSOLE.msg(f"Latest version: {latest_version}")
 
                 if current_version != latest_version:
-                    self.update_version = latest_version
+                    self.update_version = True
                 else:
                     CONSOLE.msg("Application is already up-to-date.")
 
@@ -1983,6 +1987,30 @@ def execute(self, action, args):
         )
 
 
+class TestGenUpdateVolumeStep(Step):
+    label = "Updating docker volume"
+
+    def execute(self, action, args):
+        try:
+            # Set testgen user as volume owner in case UID changes
+            action.run_cmd(
+                "docker",
+                "compose",
+                "-f",
+                action.docker_compose_file_path,
+                "run",
+                "--entrypoint",
+                "/bin/sh -c",
+                "--user",
+                "root",
+                "--rm",
+                "engine",
+                "chown -R testgen:testgen /var/lib/testgen",
+            )
+        except CommandFailed:
+            raise SkipStep
+
+
 class TestGenSetupDatabaseStep(Step):
     label = "Initializing the platform database"
 
@@ -2034,7 +2062,7 @@ def on_action_success(self, action, args):
             capture_text=True,
         )
 
-        match = re.search("This version:(.*)", output)
+        match = re.search(r"TestGen\s(?:[a-zA-Z]+\s)*([0-9.]*)", output)
         CONSOLE.msg(f"Application version: {match.group(1)}")
 
 
@@ -2119,6 +2147,7 @@ class TestgenUpgradeAction(TestgenActionMixin, AnalyticsMultiStepAction):
         UpdateComposeFileStep,
         TestGenStopStep,
         TestGenPullImagesStep,
+        TestGenUpdateVolumeStep,
         TestGenStartStep,
         TestGenUpgradeDatabaseStep,
     ]

tests/conftest.py

@@ -72,6 +72,12 @@ def stream_iter_mock():
         yield si_mock
 
 
+@pytest.fixture
+def version_check_mock():
+    with patch("urllib.request.urlopen") as mock:
+        yield mock
+
+
 @pytest.fixture
 def analytics_mock():
     with patch("tests.installer.AnalyticsWrapper") as mock:

tests/test_tg_upgrade.py

@@ -1,3 +1,4 @@
+import json
 import textwrap
 from unittest.mock import call
 
@@ -18,14 +19,14 @@ def tg_upgrade_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock, re
 def tg_upgrade_stdout_side_effect(stdout_mock):
     side_effect = [
         # Pre-execute calls
-        [b"This version: 1.0.0\n", b"Latest version: 1.1.0\n"],  # Version check
+        [b"TestGen 1.0.0\n"],  # Version check
         # Execute calls
         [],  # Down
         [],  # Pull
         [],  # Up
         [],  # Upgrade DB
         # Post-execute calls
-        [b"This version: 1.1.0\n", b"Latest version: 1.1.0\n"],  # Confirmation version check
+        [b"TestGen 1.1.0\n"],  # Confirmation version check
         [b"[]"],  # Image data collection
     ]
 
@@ -57,9 +58,15 @@ def get_compose_content(*extra_vars):
     return template.format(textwrap.indent("\n".join(extra_vars), "  "))
 
 
+def set_version_check_mock(version_check_mock, latest_version):
+    version_check_mock.return_value.code = 200
+    version_values = { "docker": {"datakitchen/dataops-testgen": latest_version } }
+    version_check_mock.return_value.read.return_value = json.dumps(version_values).encode("utf-8")
+
+
 @pytest.mark.integration
 def test_tg_upgrade_compose_missing(tg_upgrade_action, args_mock, start_cmd_mock, console_msg_mock):
-    start_cmd_mock.__exit__.side_effect = [None, None, None, CommandFailed]
+    start_cmd_mock.__exit__.side_effect = [None, None, CommandFailed]
 
     with pytest.raises(AbortAction, match=""):
         tg_upgrade_action._check_requirements(args_mock)
@@ -68,7 +75,22 @@ def test_tg_upgrade_compose_missing(tg_upgrade_action, args_mock, start_cmd_mock
 
 
 @pytest.mark.integration
-def test_tg_upgrade(tg_upgrade_action, compose_path, start_cmd_mock, tg_upgrade_stdout_side_effect, args_mock):
+@pytest.mark.parametrize(
+    "skip_verify, latest_version",
+    ((True, "1.0.0"), (False, "1.1.0")),
+)
+def test_tg_upgrade(
+    skip_verify,
+    latest_version,
+    tg_upgrade_action,
+    compose_path,
+    start_cmd_mock,
+    tg_upgrade_stdout_side_effect,
+    args_mock,
+    version_check_mock,
+):
+    args_mock.skip_verify = skip_verify
+    set_version_check_mock(version_check_mock, latest_version)
     compose_path.write_text(get_compose_content())
 
     tg_upgrade_action.execute(args_mock)
@@ -93,21 +115,16 @@ def test_tg_upgrade(tg_upgrade_action, compose_path, start_cmd_mock, tg_upgrade_
 
 
 @pytest.mark.integration
-@pytest.mark.parametrize(
-    "skip_verify, latest_version",
-    ((True, b"1.0.0"), (True, b"1.1.0"), (False, b"1.0.0")),
-)
 def test_tg_upgrade_abort(
-    skip_verify,
-    latest_version,
     tg_upgrade_action,
     compose_path,
     start_cmd_mock,
     tg_upgrade_stdout_side_effect,
     args_mock,
+    version_check_mock,
 ):
-    args_mock.skip_verify = skip_verify
-    tg_upgrade_stdout_side_effect[0][1] = b"Latest version: %b\n" % latest_version
+    args_mock.skip_verify = False
+    set_version_check_mock(version_check_mock, "1.0.0")
     initial_compose_content = get_compose_content("TG_INSTANCE_ID: test-instance-id")
     compose_path.write_text(initial_compose_content)
 
@@ -116,7 +133,7 @@ def test_tg_upgrade_abort(
 
     compose_content = compose_path.read_text()
     assert compose_content == initial_compose_content
-    assert start_cmd_mock.call_count == 0 if skip_verify else 1
+    assert start_cmd_mock.call_count == 1
 
 
 @pytest.mark.integration
@@ -130,8 +147,9 @@ def test_tg_upgrade_enable_analytics(
     args_mock,
     console_msg_mock,
     analytics_mock,
+    version_check_mock,
 ):
-    tg_upgrade_stdout_side_effect[0][1] = b"Latest version: 1.0.0\n"
+    set_version_check_mock(version_check_mock, "1.0.0")
     compose_path.write_text(get_compose_content("TG_ANALYTICS: no" if re_enable else ""))
     analytics_mock.get_instance_id.return_value = "test-instance-id"
 
@@ -153,9 +171,10 @@ def test_tg_upgrade_disable_analytics(
     tg_upgrade_stdout_side_effect,
     args_mock,
     console_msg_mock,
+    version_check_mock,
 ):
     args_mock.send_analytics_data = False
-    tg_upgrade_stdout_side_effect[0][1] = b"Latest version: 1.0.0\n"
+    set_version_check_mock(version_check_mock, "1.0.0")
     compose_path.write_text(
         get_compose_content("TG_INSTANCE_ID: test-instance-id", "TG_ANALYTICS: yes" if explicitly_enabled else "")
     )