Rolling update fixture (#18430)

4 files changed, 150 insertions, 0 deletions

diff --git a/ydb/tests/compatibility/test_rolling.py b/ydb/tests/compatibility/test_rolling.py
new file mode 100644
index 00000000000..ce7fbb260e9
--- /dev/null
+++ b/ydb/tests/compatibility/test_rolling.py
@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+import os
+import time
+
+import pytest
+
+import yatest
+
+from ydb.tests.library.compatibility.fixtures import RollingUpdateFixture
+
+
+class TestRolling(RollingUpdateFixture):
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        output_path = yatest.common.test_output_path()
+        self.output_f = open(os.path.join(output_path, "out.log"), "w")
+        yield from self.setup_cluster(
+            extra_feature_flags={
+                "enable_column_store": True,
+            },
+
+            column_shard_config={
+                "disabled_on_scheme_shard": False,
+            },
+        )
+
+    @pytest.mark.parametrize("store_type", ["row", "column"])
+    def test_kv(self, store_type):
+        init_command_prefix = [
+            yatest.common.binary_path(os.getenv("YDB_CLI_BINARY")),
+            "--verbose",
+            "--endpoint",
+            "grpc://localhost:%d" % self.cluster.nodes[1].grpc_port,
+            "--database=/Root",
+            "workload",
+            "kv",
+            "init",
+            "--min-partitions",
+            "10",
+            "--partition-size",
+            "10",
+            "--auto-partition",
+            "0",
+            "--init-upserts",
+            "0",
+            "--cols",
+            "5",
+            "--int-cols",
+            "2",
+            "--key-cols",
+            "3",
+        ]
+
+        run_command_prefix = [
+            yatest.common.binary_path(os.getenv("YDB_CLI_BINARY")),
+            "--verbose",
+            "--endpoint",
+            "grpc://localhost:%d" % self.cluster.nodes[1].grpc_port,
+            "--database=/Root",
+            "workload",
+            "kv",
+            "run",
+            "mixed",
+            "--seconds",
+            "10000000",  # infinity
+            "--threads",
+            "10",
+            "--cols",
+            "5",
+            "--len",
+            "200",
+            "--int-cols",
+            "2",
+            "--key-cols",
+            "3",
+        ]
+
+        init_command = init_command_prefix
+        init_command.extend(
+            [
+                "--path",
+                store_type,
+                "--store",
+                store_type,
+            ]
+        )
+        run_command = run_command_prefix
+        run_command.extend(
+            [
+                "--path",
+                store_type,
+            ]
+        )
+        yatest.common.execute(init_command, wait=True, stdout=self.output_f, stderr=self.output_f)
+        run = yatest.common.execute(run_command, wait=False, stdout=self.output_f, stderr=self.output_f)
+
+        for _ in self.roll():
+            time.sleep(5)
+
+        run.kill()
diff --git a/ydb/tests/compatibility/ya.make b/ydb/tests/compatibility/ya.make
index 5e2dce1857a..c3cd6fba37a 100644
--- a/ydb/tests/compatibility/ya.make
+++ b/ydb/tests/compatibility/ya.make
@@ -12,6 +12,7 @@ TEST_SRCS(
     test_followers.py
     test_compatibility.py
     test_stress.py
+    test_rolling.py
 )
 
 SIZE(LARGE)
diff --git a/ydb/tests/library/compatibility/fixtures.py b/ydb/tests/library/compatibility/fixtures.py
index c2ce8672ae7..73c6f2edbf1 100644
--- a/ydb/tests/library/compatibility/fixtures.py
+++ b/ydb/tests/library/compatibility/fixtures.py
@@ -109,3 +109,47 @@ class MixedClusterFixture:
         self.driver.wait()
         yield
         self.cluster.stop()
+
+
+class RollingUpdateFixture:
+    @pytest.fixture(autouse=True)
+    def base_setup(self):
+        self.all_binary_paths = [last_stable_binary_path]
+
+    def setup_cluster(self, **kwargs):
+        self.config = KikimrConfigGenerator(
+            erasure=Erasure.MIRROR_3_DC,
+            binary_paths=self.all_binary_paths,
+            **kwargs,
+        )
+
+        self.cluster = KiKiMR(self.config)
+        self.cluster.start()
+        self.endpoint = "grpc://%s:%s" % ('localhost', self.cluster.nodes[1].port)
+
+        self.driver = ydb.Driver(
+            ydb.DriverConfig(
+                database='/Root',
+                endpoint=self.endpoint
+            )
+        )
+        self.driver.wait()
+        yield
+        self.cluster.stop()
+
+    def roll(self):
+        # from old to new
+        for node_id, node in self.cluster.nodes.items():
+            node.stop()
+            node.binary_path = current_binary_path
+            node.start()
+            yield
+
+        # from new to old
+        for node_id, node in self.cluster.nodes.items():
+            node.stop()
+            node.binary_path = last_stable_binary_path
+            node.start()
+            yield
+
+        yield
diff --git a/ydb/tests/library/harness/kikimr_runner.py b/ydb/tests/library/harness/kikimr_runner.py
index 82eb8dbd428..e1448798f6d 100644
--- a/ydb/tests/library/harness/kikimr_runner.py
+++ b/ydb/tests/library/harness/kikimr_runner.py
@@ -154,6 +154,10 @@ class KiKiMRNode(daemon.Daemon, kikimr_node_interface.NodeInterface):
     def binary_path(self):
         return self.__binary_path
 
+    @binary_path.setter
+    def binary_path(self, value):
+        self.__binary_path = value
+
     @property
     def command(self):
         return self.__make_run_command()
@@ -277,6 +281,7 @@ class KiKiMRNode(daemon.Daemon, kikimr_node_interface.NodeInterface):
 
     def start(self):
         try:
+            self.update_command(self.__make_run_command())
             super(KiKiMRNode, self).start()
         finally:
             logger.info("Started node %s", self)