From 47b6fe9dfd48fcb55be51cde9221365b6545a920 Mon Sep 17 00:00:00 2001
From: Sang Jun Bak <sangjun.bak99@gmail.com>
Date: Tue, 14 Oct 2025 23:17:09 -0400
Subject: [PATCH] Add preflight check scenarios for upgrading to newer and
 older versions

This commit introduces two new scenarios: PreflightCheckThenUpgradeToNewerVersion and PreflightCheckThenUpgradeToOlderVersion. Each scenario performs a preflight check before attempting to upgrade to a specified version, handling both newer and older version upgrades.
---
 .../materialize/checks/scenarios_upgrade.py   | 78 +++++++++++++++++++
 1 file changed, 78 insertions(+)

diff --git a/misc/python/materialize/checks/scenarios_upgrade.py b/misc/python/materialize/checks/scenarios_upgrade.py
index 0c408e5045852..d89b157ef21a9 100644
--- a/misc/python/materialize/checks/scenarios_upgrade.py
+++ b/misc/python/materialize/checks/scenarios_upgrade.py
@@ -509,3 +509,81 @@ def actions(self) -> list[Action]:
             ),
             Validate(self),
         ]
+
+
+class PreflightCheckThenUpgradeToNewerVersion(Scenario):
+    """Preflight check, then upgrade to a newer version than the last upgrade version"""
+
+    def base_version(self) -> MzVersion:
+        return get_previous_version()
+
+    def actions(self) -> list[Action]:
+        print(f"Upgrading from tag {self.base_version()}")
+        return [
+            StartMz(
+                self,
+                tag=self.base_version(),
+                mz_service="mz_1",
+            ),
+            Initialize(self, mz_service="mz_1"),
+            Manipulate(self, phase=1, mz_service="mz_1"),
+            Manipulate(self, phase=2, mz_service="mz_1"),
+            # Try to upgrade to a new version
+            start_mz_read_only(
+                self,
+                tag=get_last_version(),
+                deploy_generation=1,
+                mz_service="mz_2",
+            ),
+            WaitReadyMz(mz_service="mz_2"),
+            KillMz(capture_logs=True, mz_service="mz_2"),
+            # Try to upgrade to a newer version than the last upgrade version
+            start_mz_read_only(
+                self,
+                tag=None,
+                deploy_generation=1,
+                mz_service="mz_3",
+            ),
+            WaitReadyMz(mz_service="mz_3"),
+            PromoteMz(mz_service="mz_3"),
+            Validate(self),
+        ]
+
+
+class PreflightCheckThenUpgradeToOlderVersion(Scenario):
+    """Preflight check, then upgrade"""
+
+    def base_version(self) -> MzVersion:
+        return get_previous_version()
+
+    def actions(self) -> list[Action]:
+        print(f"Upgrading from tag {self.base_version()}")
+        return [
+            StartMz(
+                self,
+                tag=self.base_version(),
+                mz_service="mz_1",
+            ),
+            Initialize(self, mz_service="mz_1"),
+            Manipulate(self, phase=1, mz_service="mz_1"),
+            Manipulate(self, phase=2, mz_service="mz_1"),
+            # Try to upgrade to a new version
+            start_mz_read_only(
+                self,
+                tag=None,
+                deploy_generation=1,
+                mz_service="mz_2",
+            ),
+            WaitReadyMz(mz_service="mz_2"),
+            KillMz(capture_logs=True, mz_service="mz_2"),
+            # Try to upgrade to a newer version than the last upgrade version
+            start_mz_read_only(
+                self,
+                tag=get_last_version(),
+                deploy_generation=1,
+                mz_service="mz_3",
+            ),
+            WaitReadyMz(mz_service="mz_3"),
+            PromoteMz(mz_service="mz_3"),
+            Validate(self),
+        ]