From 4286dea5c7627b919e1a0e1874fad45b80bc1a94 Mon Sep 17 00:00:00 2001
From: Siyao Meng <50227127+smengcl@users.noreply.github.com>
Date: Sat, 8 Aug 2026 20:19:57 -0700
Subject: [PATCH] HDDS-16113. Add regression tests for snapshot defragmentation
 cleanup

---
 .../defrag/TestSnapshotDefragService.java     | 38 ++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/defrag/TestSnapshotDefragService.java b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/defrag/TestSnapshotDefragService.java
index 57e2d36b92d..036226d81cf 100644
--- a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/defrag/TestSnapshotDefragService.java
+++ b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/defrag/TestSnapshotDefragService.java
@@ -634,6 +634,41 @@ public void close() {
     }
   }
 
+  @Test
+  public void testCreateCheckpointCleansUpAfterOpenFailure() throws Exception {
+    SnapshotInfo snapshotInfo = createMockSnapshotInfo(UUID.randomUUID(), "vol1", "bucket1", "snap1");
+    Path checkpointPath = tempDir.resolve("tmp_defrag").resolve("checkpoint-" + UUID.randomUUID());
+    Files.createDirectories(checkpointPath);
+    DBCheckpoint checkpoint = new RocksDBCheckpoint(checkpointPath);
+    DBStore snapshotStore = mock(DBStore.class);
+    when(snapshotStore.getCheckpoint(anyString(), eq(true))).thenReturn(checkpoint);
+    createMockSnapshot(snapshotInfo, snapshotStore);
+
+    SnapshotDefragService spyDefragService = Mockito.spy(defragService);
+    Mockito.doThrow(new IOException("Failed to open checkpoint"))
+        .when(spyDefragService).createDefragCheckpointMetadataManager(eq(checkpoint), eq(false));
+
+    org.junit.jupiter.api.Assertions.assertThrows(IOException.class,
+        () -> spyDefragService.createCheckpoint(snapshotInfo, COLUMN_FAMILIES_TO_TRACK_IN_SNAPSHOT));
+    assertFalse(Files.exists(checkpointPath));
+  }
+
+  @Test
+  public void testCreateCheckpointCleansUpWhenStoreReturnsNull() throws Exception {
+    SnapshotInfo snapshotInfo = createMockSnapshotInfo(UUID.randomUUID(), "vol1", "bucket1", "snap1");
+    String dbName = "snapshot-" + UUID.randomUUID();
+    Path partialCheckpoint = tempDir.resolve("tmp_defrag").resolve(dbName + "_checkpoint_1");
+    Files.createDirectories(partialCheckpoint);
+    DBStore snapshotStore = mock(DBStore.class);
+    when(snapshotStore.getCheckpoint(anyString(), eq(true))).thenReturn(null);
+    when(snapshotStore.getDbLocation()).thenReturn(tempDir.resolve(dbName).toFile());
+    createMockSnapshot(snapshotInfo, snapshotStore);
+
+    org.junit.jupiter.api.Assertions.assertThrows(IOException.class,
+        () -> defragService.createCheckpoint(snapshotInfo, COLUMN_FAMILIES_TO_TRACK_IN_SNAPSHOT));
+    assertFalse(Files.exists(partialCheckpoint));
+  }
+
   @Test
   public void testDefragCheckpointMetadataManagerSkipsRocksDBMetrics() throws Exception {
     Path checkpointPath = tempDir.resolve("defrag-metrics-" + UUID.randomUUID());
@@ -1036,9 +1071,10 @@ public void testCheckAndDefragAlreadyDefraggedSnapshot() throws IOException {
     try (MockedStatic<SnapshotUtils> mockedStatic = Mockito.mockStatic(SnapshotUtils.class)) {
       mockedStatic.when(() -> SnapshotUtils.getSnapshotInfo(eq(ozoneManager), eq(chainManager),
           eq(snapshotInfo.getSnapshotId()))).thenReturn(snapshotInfo);
-      doReturn(Pair.of(false, 0)).when(spyDefragService).needsDefragmentation(eq(snapshotInfo));
+      doReturn(Pair.of(false, 2)).when(spyDefragService).needsDefragmentation(eq(snapshotInfo));
       assertFalse(spyDefragService.checkAndDefragSnapshot(chainManager, snapshotInfo.getSnapshotId()));
       verify(snapshotMetrics).incNumSnapshotDefragSnapshotSkipped();
+      verify(omSnapshotManager).deleteSnapshotCheckpointDirectories(snapshotInfo.getSnapshotId(), 1);
     }
   }