From 889e108be9d3fb923eb7d43bbd4c0cc0283b03a8 Mon Sep 17 00:00:00 2001
From: Andrey Berezin <a.berezin@yadro.com>
Date: Tue, 27 Jun 2023 11:47:42 +0300
Subject: [PATCH] Update epoch align check

Signed-off-by: Andrey Berezin <a.berezin@yadro.com>
---
 src/frostfs_testlib/steps/epoch.py             | 18 ++++++++----------
 .../controllers/cluster_state_controller.py    | 11 -----------
 2 files changed, 8 insertions(+), 21 deletions(-)

diff --git a/src/frostfs_testlib/steps/epoch.py b/src/frostfs_testlib/steps/epoch.py
index 0d40f8d..a589569 100644
--- a/src/frostfs_testlib/steps/epoch.py
+++ b/src/frostfs_testlib/steps/epoch.py
@@ -54,16 +54,14 @@ def ensure_fresh_epoch(
     return epoch
 
 
-@reporter.step_deco("Wait for epochs align in whole cluster")
-@wait_for_success(60, 5)
-def wait_for_epochs_align(shell: Shell, cluster: Cluster) -> None:
-    epochs = []
-    for node in cluster.services(StorageNode):
-        epochs.append(get_epoch(shell, cluster, node))
-    unique_epochs = list(set(epochs))
-    assert (
-        len(unique_epochs) == 1
-    ), f"unaligned epochs found,  {epochs}, count of unique epochs {len(unique_epochs)}"
+@reporter.step_deco("Wait up to {timeout} seconds for nodes on cluster to align epochs")
+def wait_for_epochs_align(shell: Shell, cluster: Cluster, timeout=60):
+    @wait_for_success(timeout, 5, None, True)
+    def check_epochs():
+        epochs_by_node = get_epochs_from_nodes(shell, cluster)
+        assert len(set(epochs_by_node.values())) == 1, f"unaligned epochs found: {epochs_by_node}"
+
+    check_epochs()
 
 
 @reporter.step_deco("Get Epoch")
diff --git a/src/frostfs_testlib/storage/controllers/cluster_state_controller.py b/src/frostfs_testlib/storage/controllers/cluster_state_controller.py
index 6126f9d..c73a8f4 100644
--- a/src/frostfs_testlib/storage/controllers/cluster_state_controller.py
+++ b/src/frostfs_testlib/storage/controllers/cluster_state_controller.py
@@ -220,17 +220,6 @@ class ClusterStateController:
             wait_for_host_online(self.shell, node.storage_node)
             wait_for_node_online(node.storage_node)
 
-    @reporter.step_deco("Wait up to {timeout} seconds for nodes on cluster to align epochs")
-    def wait_for_epochs_align(self, timeout=60):
-        @wait_for_success(timeout, 5, None, True)
-        def check_epochs():
-            epochs_by_node = epoch.get_epochs_from_nodes(self.shell, self.cluster)
-            assert (
-                len(set(epochs_by_node.values())) == 1
-            ), f"unaligned epochs found: {epochs_by_node}"
-
-        check_epochs()
-
     def _get_disk_controller(
         self, node: StorageNode, device: str, mountpoint: str
     ) -> DiskController: