From e4ce28043b811a74f4f5839dec8fcc8eab323dc7 Mon Sep 17 00:00:00 2001
From: sharpeye <misko.pawel@gmail.com>
Date: Thu, 28 Mar 2024 09:17:11 +0000
Subject: [PATCH] add ut

---
 .../disk_registry/disk_registry_state.cpp     | 18 ++--
 .../disk_registry/disk_registry_ut_cms.cpp    | 91 +++++++++++++++++++
 2 files changed, 99 insertions(+), 10 deletions(-)

diff --git a/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.cpp b/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.cpp
index 3fd1c38eade..24903c04da5 100644
--- a/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.cpp
+++ b/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.cpp
@@ -4842,28 +4842,26 @@ NProto::TError TDiskRegistryState::UpdateCmsHostState(
     }
 
     NProto::TError result;
-    if (timeout == TDuration::Zero()) {
-        result = MakeError(S_OK);
-        cmsTs = TInstant::Zero();
-    } else {
-        result = MakeError(
-            E_TRY_AGAIN,
-            TStringBuilder() << "time remaining: " << timeout);
-    }
 
     // Agent can return from 'unavailable' state only when it is reconnected to
     // the cluster.
     if (agent->GetState() == NProto::AGENT_STATE_UNAVAILABLE &&
         newState == NProto::AGENT_STATE_ONLINE)
     {
-        result = MakeError(E_TRY_AGAIN, "agent currently unavailable");
         timeout = cmsTs + CMS_UPDATE_STATE_TO_ONLINE_TIMEOUT - now;
-
         if (!timeout) {
             result.SetCode(E_INVALID_STATE);
         }
     }
 
+    if (timeout) {
+        result = MakeError(
+            E_TRY_AGAIN,
+            TStringBuilder() << "time remaining: " << timeout);
+    } else {
+        cmsTs = TInstant::Zero();
+    }
+
     if (dryRun) {
         return result;
     }
diff --git a/cloud/blockstore/libs/storage/disk_registry/disk_registry_ut_cms.cpp b/cloud/blockstore/libs/storage/disk_registry/disk_registry_ut_cms.cpp
index 18e40ec1b9a..d6d574038d7 100644
--- a/cloud/blockstore/libs/storage/disk_registry/disk_registry_ut_cms.cpp
+++ b/cloud/blockstore/libs/storage/disk_registry/disk_registry_ut_cms.cpp
@@ -20,11 +20,14 @@
 #include <util/datetime/base.h>
 #include <util/generic/size_literals.h>
 
+#include <chrono>
+
 namespace NCloud::NBlockStore::NStorage {
 
 using namespace NActors;
 using namespace NKikimr;
 using namespace NDiskRegistryTest;
+using namespace std::chrono_literals;
 
 namespace {
 
@@ -1112,6 +1115,94 @@ Y_UNIT_TEST_SUITE(TDiskRegistryTest)
             UNIT_ASSERT_VALUES_EQUAL(0, timeout);
         }
     }
+
+    Y_UNIT_TEST(ShouldRejectAddHostForUnavailableAgent)
+    {
+        const TVector agents{CreateAgentConfig(
+            "agent-1",
+            {Device("dev-1", "uuid-1", "rack-1", 10_GB),
+             Device("dev-2", "uuid-2", "rack-1", 10_GB)})};
+
+        auto runtime = TTestRuntimeBuilder().WithAgents(agents).Build();
+
+        TDiskRegistryClient diskRegistry(*runtime);
+        diskRegistry.SetWritableState(true);
+        diskRegistry.WaitReady();
+        diskRegistry.UpdateConfig(CreateRegistryConfig(0, agents));
+
+        RegisterAgents(*runtime, 1);
+        WaitForAgents(*runtime, 1);
+        WaitForSecureErase(*runtime, agents);
+
+        {
+            auto [error, timeout] = RemoveHost(diskRegistry, "agent-1");
+
+            UNIT_ASSERT_VALUES_EQUAL(S_OK, error.GetCode());
+            UNIT_ASSERT_VALUES_EQUAL(0, timeout);
+        }
+
+        runtime->AdvanceCurrentTime(15min);
+
+        diskRegistry.ChangeAgentState(
+            "agent-1",
+            NProto::EAgentState::AGENT_STATE_UNAVAILABLE);
+
+        runtime->AdvanceCurrentTime(15min);
+
+        {
+            auto [error, timeout] = AddHost(diskRegistry, "agent-1");
+
+            UNIT_ASSERT_VALUES_EQUAL(E_TRY_AGAIN, error.GetCode());
+            UNIT_ASSERT_VALUES_EQUAL(
+                TDuration{5min},
+                TDuration::Seconds(timeout));
+        }
+
+        runtime->AdvanceCurrentTime(2min);
+
+        {
+            auto [error, timeout] = AddHost(diskRegistry, "agent-1");
+
+            UNIT_ASSERT_VALUES_EQUAL(E_TRY_AGAIN, error.GetCode());
+            UNIT_ASSERT_LE(TDuration::Seconds(timeout), TDuration{3min});
+            UNIT_ASSERT_GT(TDuration::Seconds(timeout), TDuration{2min});
+        }
+
+        runtime->AdvanceCurrentTime(10min);
+
+        {
+            auto [error, timeout] = AddHost(diskRegistry, "agent-1");
+
+            UNIT_ASSERT_VALUES_EQUAL(E_INVALID_STATE, error.GetCode());
+            UNIT_ASSERT_VALUES_EQUAL(0, timeout);
+        }
+
+        runtime->AdvanceCurrentTime(10min);
+
+        {
+            auto [error, timeout] = AddHost(diskRegistry, "agent-1");
+
+            UNIT_ASSERT_VALUES_EQUAL(E_TRY_AGAIN, error.GetCode());
+            UNIT_ASSERT_VALUES_EQUAL(
+                TDuration{5min},
+                TDuration::Seconds(timeout));
+        }
+
+        runtime->AdvanceCurrentTime(1min);
+
+        diskRegistry.ChangeAgentState(
+            "agent-1",
+            NProto::EAgentState::AGENT_STATE_WARNING);
+
+        runtime->AdvanceCurrentTime(1min);
+
+        {
+            auto [error, timeout] = AddHost(diskRegistry, "agent-1");
+
+            UNIT_ASSERT_VALUES_EQUAL(S_OK, error.GetCode());
+            UNIT_ASSERT_VALUES_EQUAL(0, timeout);
+        }
+    }
 }
 
 }   // namespace NCloud::NBlockStore::NStorage