From bcc34c0727fdd8ef3aa33ceca3725cc1a0215bd5 Mon Sep 17 00:00:00 2001
From: WenyXu <wenymedia@gmail.com>
Date: Mon, 18 Nov 2024 03:20:45 +0000
Subject: [PATCH] fix: correct `unset_maintenance_mode` behavior

---
 src/common/meta/src/key/maintenance.rs | 35 +++++++++++++++++++++-----
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/common/meta/src/key/maintenance.rs b/src/common/meta/src/key/maintenance.rs
index b8ee760eb8e3..c1cb93d76ef7 100644
--- a/src/common/meta/src/key/maintenance.rs
+++ b/src/common/meta/src/key/maintenance.rs
@@ -47,12 +47,9 @@ impl MaintenanceModeManager {
 
     /// Unsets maintenance mode.
     pub async fn unset_maintenance_mode(&self) -> Result<()> {
-        let req = PutRequest {
-            key: Vec::from(MAINTENANCE_KEY),
-            value: vec![],
-            prev_kv: false,
-        };
-        self.kv_backend.put(req).await?;
+        self.kv_backend
+            .delete(MAINTENANCE_KEY.as_bytes(), false)
+            .await?;
         Ok(())
     }
 
@@ -61,3 +58,29 @@ impl MaintenanceModeManager {
         self.kv_backend.exists(MAINTENANCE_KEY.as_bytes()).await
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use crate::key::maintenance::MaintenanceModeManager;
+    use crate::kv_backend::memory::MemoryKvBackend;
+
+    #[tokio::test]
+    async fn test_maintenance_mode_manager() {
+        let maintenance_mode_manager = Arc::new(MaintenanceModeManager::new(Arc::new(
+            MemoryKvBackend::new(),
+        )));
+        assert!(!maintenance_mode_manager.maintenance_mode().await.unwrap());
+        maintenance_mode_manager
+            .set_maintenance_mode()
+            .await
+            .unwrap();
+        assert!(maintenance_mode_manager.maintenance_mode().await.unwrap());
+        maintenance_mode_manager
+            .unset_maintenance_mode()
+            .await
+            .unwrap();
+        assert!(!maintenance_mode_manager.maintenance_mode().await.unwrap());
+    }
+}