From 324b8272ab5a1f74ea82db9a9838cb93337dd353 Mon Sep 17 00:00:00 2001
From: Kirill Pleshivtsev <drbasic@yandex-team.ru>
Date: Thu, 19 Sep 2024 18:12:01 +0700
Subject: [PATCH] Fix vhost test ShouldStatAllZeroesBlocks flakiness (#2070)

---
 cloud/blockstore/vhost-server/server_ut.cpp | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/cloud/blockstore/vhost-server/server_ut.cpp b/cloud/blockstore/vhost-server/server_ut.cpp
index e9af73d05b..2ccafe4d0c 100644
--- a/cloud/blockstore/vhost-server/server_ut.cpp
+++ b/cloud/blockstore/vhost-server/server_ut.cpp
@@ -241,6 +241,14 @@ class TServerTest
     }
 
     TCompleteStats GetStats(ui64 expectedCompleted) const
+    {
+        return GetStats(
+            [expectedCompleted](const TCompleteStats& stats)
+            { return stats.SimpleStats.Completed == expectedCompleted; });
+    }
+
+    TCompleteStats GetStats(
+        std::function<bool(const TCompleteStats&)> func) const
     {
         // Without I/O, stats are synced every second and only if there is a
         // pending GetStats call. The first call to GetStats might not bring the
@@ -251,7 +259,7 @@ class TServerTest
         TCompleteStats stats;
         for (int i = 0; i != 5; ++i) {
             stats = Server->GetStats(prevStats);
-            if (stats.SimpleStats.Completed == expectedCompleted) {
+            if (func(stats)) {
                 break;
             }
             Sleep(TDuration::Seconds(1));
@@ -959,7 +967,9 @@ TEST_P(TServerTest, ShouldStatAllZeroesBlocks)
     }
 
     // validate stats
-    const auto completeStats = GetStats(0);
+    const auto completeStats =
+        GetStats([](const TCompleteStats& stats)
+                 { return stats.CriticalEvents.size() != 0; });
     const auto& stats = completeStats.SimpleStats;
 
     EXPECT_EQ(0u, stats.CompFailed);