fix: correct unset_maintenance_mode behavior

GreptimeTeam · Nov 18, 2024 · bcc34c0 · bcc34c0
1 parent 485782a
commit bcc34c0
Showing 1 changed file with 29 additions and 6 deletions.
diff --git 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());
+    }
+}