From 61bfff62bed7e2e4c914d0c721ef4277b2c8f84a Mon Sep 17 00:00:00 2001
From: Daniil Komarevtsev <komarevtsev@yandex-team.ru>
Date: Fri, 12 Jul 2024 14:40:51 +0000
Subject: [PATCH] review fixes

---
 .../disk_registry/disk_registry_state.cpp     | 33 ++++++++-----------
 cloud/blockstore/tests/infra-cms/test.py      |  1 -
 2 files changed, 14 insertions(+), 20 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 67c44d02dac..44e0af44365 100644
--- a/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.cpp
+++ b/cloud/blockstore/libs/storage/disk_registry/disk_registry_state.cpp
@@ -4694,17 +4694,10 @@ NProto::TError TDiskRegistryState::UpdateAgentState(
     // an idle agent. DR waits for it to become online and blocks host
     // deployment.
     if (StorageConfig->GetCleanupDRConfigOnCMSActions() &&
-        newState == NProto::EAgentState::AGENT_STATE_UNAVAILABLE &&
+        newState == NProto::AGENT_STATE_UNAVAILABLE &&
         agent->GetDevices().empty())
     {
-        bool success = RemoveAgent(db, agent->GetNodeId());
-        if (!success) {
-            return MakeError(
-                E_FAIL,
-                TStringBuilder() << "Couldn't remove the agent "
-                                 << agent->GetAgentId().Quote());
-        }
-
+        RemoveAgent(db, *agent);
         return {};
     }
 
@@ -5060,17 +5053,19 @@ NProto::TError TDiskRegistryState::UpdateCmsHostState(
     ApplyAgentStateChange(db, *agent, now, affectedDisks);
 
     if (newState != NProto::AGENT_STATE_ONLINE && !HasError(result)) {
-        auto error = TryToRemoveAgentDevices(db, agent->GetAgentId());
-        if (!HasError(error)) {
-            return result;
-        }
+        if (StorageConfig->GetCleanupDRConfigOnCMSActions()) {
+            auto error = TryToRemoveAgentDevices(db, agent->GetAgentId());
+            if (!HasError(error)) {
+                return result;
+            }
 
-        // Do not return the error from "TryToRemoveAgentDevices()" since it's
-        // internal and shouldn't block node removal.
-        STORAGE_WARN(
-            "Could not remove device from agent %s: %s",
-            agent->GetAgentId().Quote().c_str(),
-            FormatError(error).c_str());
+            // Do not return the error from "TryToRemoveAgentDevices()" since
+            // it's internal and shouldn't block node removal.
+            STORAGE_WARN(
+                "Could not remove device from agent %s: %s",
+                agent->GetAgentId().Quote().c_str(),
+                FormatError(error).c_str());
+        }
 
         SuspendLocalDevices(db, *agent);
     }
diff --git a/cloud/blockstore/tests/infra-cms/test.py b/cloud/blockstore/tests/infra-cms/test.py
index 234515fbe7a..20253a1f343 100644
--- a/cloud/blockstore/tests/infra-cms/test.py
+++ b/cloud/blockstore/tests/infra-cms/test.py
@@ -338,7 +338,6 @@ def wait_for_stats(self, **kwargs):
 
             for name, expected in kwargs.items():
                 val = get_sensor_by_name(sensors, 'disk_registry', name)
-                logging.info('looking for {}, current val = {}, waitong for = {}'.format(name, val, expected))
                 if val != expected:
                     break
                 satisfied.add(name)