add test for upgrading orchestratord

Author: doy-materialize

ci/nightly/pipeline.template.yml

@@ -2504,3 +2504,16 @@ steps:
               ci-builder: stable
         agents:
           queue: hetzner-aarch64-16cpu-32gb
+
+      - id: orchestratord-upgrade-operator
+        label: "Orchestratord operator upgrade tests"
+        artifact_paths: ["mz_debug_*.zip"]
+        depends_on: devel-docker-tags
+        timeout_in_minutes: 120
+        plugins:
+          - ./ci/plugins/mzcompose:
+              composition: orchestratord
+              run: orchestratord-upgrade
+              ci-builder: stable
+        agents:
+          queue: hetzner-aarch64-16cpu-32gb

test/orchestratord/mzcompose.py

@@ -183,6 +183,12 @@ def get_environmentd_data() -> dict[str, Any]:
     )
 
 
+def get_clusterd_data() -> dict[str, Any]:
+    return get_pod_data(
+        labels={"environmentd.materialize.cloud/namespace": "cluster"},
+    )
+
+
 @contextmanager
 def port_forward_environmentd(
     port: int = 6875,
@@ -1990,6 +1996,208 @@ def workflow_balancer(c: Composition, parser: WorkflowArgumentParser) -> None:
     run_balancer(definition, False)
 
 
+def workflow_orchestratord_upgrade(
+    c: Composition,
+    parser: WorkflowArgumentParser,
+) -> None:
+    # TODO: ideally we'd just be able to compare the images directly, but
+    # this test isn't consistently handling ghcr overrides yet - remove this
+    # and go back to full image comparisons once we fix that
+    def image_tag(image: str):
+        return image.rsplit(":", 1)[1]
+
+    def check_orchestratord_version(version: MzVersion):
+        def check():
+            data = get_orchestratord_data()
+            assert len(data["items"]) == 1, f"got {len(data['items'])} items"
+
+            got_image = data["items"][0]["spec"]["containers"][0]["image"]
+            expected_image = get_image(
+                c.compose["services"]["orchestratord"]["image"],
+                str(version),
+            )
+            assert image_tag(got_image) == image_tag(
+                expected_image
+            ), f"{got_image} != {expected_image}"
+
+        retry(check, 60)
+
+    def check_environmentd_version(version: MzVersion):
+        def check():
+            data = get_environmentd_data()
+            assert len(data["items"]) == 1, f"got {len(data['items'])} items"
+
+            got_image = data["items"][0]["spec"]["containers"][0]["image"]
+            expected_image = get_image(
+                c.compose["services"]["environmentd"]["image"],
+                str(version),
+            )
+            assert image_tag(got_image) == image_tag(
+                expected_image
+            ), f"{got_image} != {expected_image}"
+
+        retry(check, 60)
+
+    def check_clusterd_version(version: MzVersion):
+        def check():
+            data = get_clusterd_data()
+            assert len(data["items"]) == 2, f"got {len(data['items'])} items"
+
+            got_image = data["items"][0]["spec"]["containers"][0]["image"]
+            expected_image = get_image(
+                c.compose["services"]["clusterd"]["image"],
+                str(version),
+            )
+            assert image_tag(got_image) == image_tag(
+                expected_image
+            ), f"{got_image} != {expected_image}"
+
+        retry(check, 60)
+
+    def check_balancerd_version(version: MzVersion):
+        def check():
+            data = get_balancerd_data()
+            assert len(data["items"]) == 2, f"got {len(data['items'])} items"
+
+            got_image = data["items"][0]["spec"]["containers"][0]["image"]
+            expected_image = get_image(
+                c.compose["services"]["balancerd"]["image"],
+                str(version),
+            )
+            assert image_tag(got_image) == image_tag(
+                expected_image
+            ), f"{got_image} != {expected_image}"
+
+        # balancerd pods have a 60s draining period after termination, so we
+        # have to wait longer to ensure that the old ones are gone
+        retry(check, 180)
+
+    parser.add_argument(
+        "--recreate-cluster",
+        action=argparse.BooleanOptionalAction,
+        help="Recreate cluster if it exists already",
+    )
+    parser.add_argument(
+        "--tag",
+        type=str,
+        help="Custom version tag to use",
+    )
+    parser.add_argument(
+        "--orchestratord-override",
+        default=True,
+        action=argparse.BooleanOptionalAction,
+        help="Override orchestratord tag",
+    )
+    args = parser.parse_args()
+
+    definition = setup(c, args)
+    versions = get_all_self_managed_versions()
+    versions.append(get_version(args.tag))
+
+    print(f"running orchestratord {versions[-3]}")
+    definition["operator"]["operator"]["image"]["tag"] = str(versions[-3])
+    init(definition)
+    check_orchestratord_version(versions[-3])
+
+    print(f"running environmentd {versions[-3]}")
+    definition["materialize"]["spec"]["environmentdImageRef"] = get_image(
+        c.compose["services"]["environmentd"]["image"],
+        str(versions[-3]),
+    )
+    run(definition, False)
+    check_environmentd_version(versions[-3])
+    check_clusterd_version(versions[-3])
+    check_balancerd_version(versions[-3])
+
+    for version in versions[-2:]:
+        print(f"running orchestratord {version}")
+        definition["operator"]["operator"]["image"]["tag"] = str(version)
+        spawn.runv(
+            [
+                "helm",
+                "upgrade",
+                "operator",
+                MZ_ROOT / "misc" / "helm-charts" / "operator",
+                "--namespace=materialize",
+                "--create-namespace",
+                "--version",
+                "v26.0.0",
+                "--wait",
+                "-f",
+                "-",
+            ],
+            stdin=yaml.dump(definition["operator"]).encode(),
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        check_orchestratord_version(version)
+
+        print(f"running environmentd {version}")
+        definition["materialize"]["spec"]["environmentdImageRef"] = get_image(
+            c.compose["services"]["environmentd"]["image"],
+            str(version),
+        )
+        definition["materialize"]["spec"]["requestRollout"] = str(uuid.uuid4())
+        run(definition, False)
+        check_environmentd_version(version)
+        check_clusterd_version(version)
+        # balancerd is broken in the upgrade to 26.4.0
+        if str(version) != "v26.4.0":
+            check_balancerd_version(version)
+
+    definition = setup(c, args)
+
+    print(f"running orchestratord {versions[-3]}")
+    definition["operator"]["operator"]["image"]["tag"] = str(versions[-3])
+    init(definition)
+    check_orchestratord_version(versions[-3])
+
+    print(f"running environmentd {versions[-3]}")
+    definition["materialize"]["spec"]["environmentdImageRef"] = get_image(
+        c.compose["services"]["environmentd"]["image"],
+        str(versions[-3]),
+    )
+    run(definition, False)
+    check_environmentd_version(versions[-3])
+    check_clusterd_version(versions[-3])
+    check_balancerd_version(versions[-3])
+
+    print(f"running orchestratord {versions[-1]}")
+    definition["operator"]["operator"]["image"]["tag"] = str(versions[-1])
+    spawn.runv(
+        [
+            "helm",
+            "upgrade",
+            "operator",
+            MZ_ROOT / "misc" / "helm-charts" / "operator",
+            "--namespace=materialize",
+            "--create-namespace",
+            "--version",
+            "v26.0.0",
+            "--wait",
+            "-f",
+            "-",
+        ],
+        stdin=yaml.dump(definition["operator"]).encode(),
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    check_orchestratord_version(versions[-1])
+
+    print(f"running environmentd {versions[-1]}")
+    definition["materialize"]["spec"]["environmentdImageRef"] = get_image(
+        c.compose["services"]["environmentd"]["image"],
+        str(versions[-1]),
+    )
+    definition["materialize"]["spec"]["requestRollout"] = str(uuid.uuid4())
+    run(definition, False)
+    check_environmentd_version(versions[-1])
+    check_clusterd_version(versions[-1])
+    # balancerd is broken in the upgrade to 26.4.0
+    if str(versions[-1]) != "v26.4.0":
+        check_balancerd_version(versions[-1])
+
+
 def workflow_default(c: Composition, parser: WorkflowArgumentParser) -> None:
     parser.add_argument(
         "--recreate-cluster",
@@ -2420,6 +2628,7 @@ def init(definition: dict[str, Any]) -> None:
             "--create-namespace",
             "--version",
             "v26.0.0",
+            "--wait",
             "-f",
             "-",
         ],