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.")
 
@@ -1949,6 +1953,26 @@ class TestGenStartStep(Step):
     label = "Starting docker compose application"
 
     def execute(self, action, args):
+        if action.args_cmd == "upgrade":
+            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 Exception:
+                pass
+
         action.run_cmd(
             "docker",
             "compose",
@@ -2034,7 +2058,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)}")
 
 

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

@@ -18,14 +18,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
     ]
 
@@ -59,7 +59,7 @@ def get_compose_content(*extra_vars):
 
 @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 +68,9 @@ 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):
+def test_tg_upgrade(tg_upgrade_action, compose_path, start_cmd_mock, tg_upgrade_stdout_side_effect, args_mock, version_check_mock):
+    version_check_mock.return_value.code = 200
+    version_check_mock.return_value.read.return_value = b"{\"docker\": {\"datakitchen/dataops-testgen\": \"1.1.0\"}}"
     compose_path.write_text(get_compose_content())
 
     tg_upgrade_action.execute(args_mock)
@@ -93,21 +95,17 @@ 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
+    version_check_mock.return_value.code = 200
+    version_check_mock.return_value.read.return_value = b"{\"docker\": {\"datakitchen/dataops-testgen\": \"1.0.0\"}}"
     initial_compose_content = get_compose_content("TG_INSTANCE_ID: test-instance-id")
     compose_path.write_text(initial_compose_content)
 
@@ -116,7 +114,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 +128,10 @@ 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"
+    version_check_mock.return_value.code = 200
+    version_check_mock.return_value.read.return_value = b"{\"docker\": {\"datakitchen/dataops-testgen\": \"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 +153,11 @@ 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"
+    version_check_mock.return_value.code = 200
+    version_check_mock.return_value.read.return_value = b"{\"docker\": {\"datakitchen/dataops-testgen\": \"1.0.0\"}}"
     compose_path.write_text(
         get_compose_content("TG_INSTANCE_ID: test-instance-id", "TG_ANALYTICS: yes" if explicitly_enabled else "")
     )