From 74ce7f2f180a81a209af3424bddbe7ab2c342495 Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Fri, 24 Jan 2025 17:30:54 +0000
Subject: [PATCH 01/11] edit volume actor statvolume and fix it's unit tests

---
 .../libs/storage/volume/volume_ut_checkpoint.cpp          | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp b/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp
index 8d464a9944..1aad1527f8 100644
--- a/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp
+++ b/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp
@@ -1933,6 +1933,14 @@ Y_UNIT_TEST_SUITE(TVolumeCheckpointTest)
 
         volume.DeleteCheckpointData("c1");
 
+        {
+            auto stat = volume.StatVolume();
+            const auto& cp = stat->Record.GetCheckpoints();
+            UNIT_ASSERT_VALUES_EQUAL(1, cp.size());
+        }
+
+        volume.DeleteCheckpoint("c1");
+
         {
             auto stat = volume.StatVolume();
             const auto& cp = stat->Record.GetCheckpoints();

From 7730c9f080f97adfd643af364a67a3b07fa5709f Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Mon, 3 Feb 2025 12:53:26 +0000
Subject: [PATCH 02/11] add additional checks in UT

---
 cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp b/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp
index 1aad1527f8..8fa0b3a4db 100644
--- a/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp
+++ b/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp
@@ -1937,6 +1937,7 @@ Y_UNIT_TEST_SUITE(TVolumeCheckpointTest)
             auto stat = volume.StatVolume();
             const auto& cp = stat->Record.GetCheckpoints();
             UNIT_ASSERT_VALUES_EQUAL(1, cp.size());
+            UNIT_ASSERT_VALUES_EQUAL("c1", stat->Record.GetCheckpoints(0));
         }
 
         volume.DeleteCheckpoint("c1");

From 345b77b4b259589b23cd9fb94d118d14d8d237bc Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Tue, 4 Feb 2025 16:06:38 +0000
Subject: [PATCH 03/11] fix dm tests

---
 .../snapshot/storage/mocks/storage_mock.go    |  9 ++++++
 .../pkg/dataplane/snapshot/storage/storage.go |  5 ++++
 .../snapshot/storage/storage_legacy.go        |  8 +++++
 .../dataplane/snapshot/storage/storage_ydb.go | 19 ++++++++++++
 .../snapshot/storage/storage_ydb_impl.go      | 27 ++++++++++++++++-
 .../snapshot/storage/storage_ydb_test.go      | 30 +++++++++++++++++++
 .../image_service_test/image_service_test.go  |  8 ++---
 .../snapshot_service_test.go                  | 27 ++++++++++-------
 .../internal/pkg/facade/testcommon/common.go  | 15 +++++++++-
 9 files changed, 131 insertions(+), 17 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/mocks/storage_mock.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/mocks/storage_mock.go
index e7f928d7af..75a6fab928 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/mocks/storage_mock.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/mocks/storage_mock.go
@@ -234,6 +234,15 @@ func (s *StorageMock) GetSnapshotMeta(
 	return args.Get(0).(*storage.SnapshotMeta), args.Error(1)
 }
 
+func (s *StorageMock) GetIncremental(
+	ctx context.Context,
+	snapshotID string,
+) (string, string, error) {
+
+	args := s.Called(ctx, snapshotID)
+	return args.String(0), args.String(1), args.Error(2)
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 func NewStorageMock() *StorageMock {
diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage.go
index a08d047afb..94441f89a2 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage.go
@@ -138,4 +138,9 @@ type Storage interface {
 		ctx context.Context,
 		snapshotID string,
 	) (*SnapshotMeta, error)
+
+	GetIncremental(
+		ctx context.Context,
+		disk *types.Disk,
+	) (string, string, error)
 }
diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_legacy.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_legacy.go
index c12b9e2441..da398631f1 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_legacy.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_legacy.go
@@ -277,3 +277,11 @@ func (s *legacyStorage) GetSnapshotMeta(
 
 	return nil, task_errors.NewNonRetriableErrorf("not implemented")
 }
+
+func (s *legacyStorage) GetIncremental(
+	ctx context.Context,
+	disk *types.Disk,
+) (string, string, error) {
+
+	return "", "", task_errors.NewNonRetriableErrorf("not implemented")
+}
diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb.go
index 572b2872e6..a321a9d183 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb.go
@@ -238,3 +238,22 @@ func (s *storageYDB) GetSnapshotMeta(
 	)
 	return snapshotMeta, err
 }
+
+func (s *storageYDB) GetIncremental(
+	ctx context.Context,
+	disk *types.Disk,
+) (snapshotID string, checkpointID string, err error) {
+
+	err = s.db.Execute(
+		ctx,
+		func(ctx context.Context, session *persistence.Session) (err error) {
+			snapshotID, checkpointID, err = s.getIncremental(
+				ctx,
+				session,
+				disk,
+			)
+			return err
+		},
+	)
+	return snapshotID, checkpointID, err
+}
diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_impl.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_impl.go
index bf382c08a6..4892b222f9 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_impl.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_impl.go
@@ -37,6 +37,31 @@ func makeShardID(s string) uint64 {
 ////////////////////////////////////////////////////////////////////////////////
 
 func (s *storageYDB) getIncremental(
+	ctx context.Context,
+	session *persistence.Session,
+	disk *types.Disk,
+) (snapshotID string, checkpointID string, err error) {
+
+	tx, err := session.BeginRWTransaction(ctx)
+	if err != nil {
+		return "", "", err
+	}
+	defer tx.Rollback(ctx)
+
+	snapshotID, checkpointID, err = s.getIncrementalTx(ctx, tx, disk)
+	if err != nil {
+		return "", "", err
+	}
+
+	err = tx.Commit(ctx)
+	if err != nil {
+		return "", "", err
+	}
+
+	return snapshotID, checkpointID, err
+}
+
+func (s *storageYDB) getIncrementalTx(
 	ctx context.Context,
 	tx *persistence.Transaction,
 	disk *types.Disk,
@@ -142,7 +167,7 @@ func (s *storageYDB) createSnapshot(
 		state.diskID = snapshotMeta.Disk.DiskId
 		state.checkpointID = snapshotMeta.CheckpointID
 
-		baseSnapshotID, baseCheckpointID, err := s.getIncremental(
+		baseSnapshotID, baseCheckpointID, err := s.getIncrementalTx(
 			ctx,
 			tx,
 			snapshotMeta.Disk,
diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_test.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_test.go
index 361c9e2fdf..d72bc60c1c 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_test.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_test.go
@@ -415,6 +415,20 @@ func testCases() []differentChunkStorageTestCase {
 
 ////////////////////////////////////////////////////////////////////////////////
 
+func checkIncremental(
+	t *testing.T,
+	f *fixture,
+	expectedSnapshotID string,
+	expectedCheckpointID string,
+	disk *types.Disk,
+) {
+
+	snapshotID, checkpointID, err := f.storage.GetIncremental(f.ctx, disk)
+	require.NoError(t, err)
+	require.Equal(t, expectedSnapshotID, snapshotID)
+	require.Equal(t, expectedCheckpointID, checkpointID)
+}
+
 func checkBaseSnapshot(
 	t *testing.T,
 	ctx context.Context,
@@ -508,6 +522,10 @@ func TestSnapshotsCreateIncrementalSnapshot(t *testing.T) {
 			require.NotNil(t, created)
 			require.Empty(t, created.BaseSnapshotID)
 			require.Empty(t, created.BaseCheckpointID)
+			checkIncremental(t, f, "", "", &types.Disk{
+				ZoneId: "zone",
+				DiskId: "disk",
+			})
 
 			created, err = f.storage.CreateSnapshot(f.ctx, SnapshotMeta{
 				ID: "snapshot2",
@@ -525,6 +543,10 @@ func TestSnapshotsCreateIncrementalSnapshot(t *testing.T) {
 
 			err = f.storage.SnapshotCreated(f.ctx, snapshot1.ID, 0, 0, 0, nil)
 			require.NoError(t, err)
+			checkIncremental(t, f, "snapshot1", "checkpoint1", &types.Disk{
+				ZoneId: "zone",
+				DiskId: "disk",
+			})
 
 			snapshot3 := SnapshotMeta{
 				ID: "snapshot3",
@@ -545,6 +567,10 @@ func TestSnapshotsCreateIncrementalSnapshot(t *testing.T) {
 
 			err = f.storage.SnapshotCreated(f.ctx, snapshot3.ID, 0, 0, 0, nil)
 			require.NoError(t, err)
+			checkIncremental(t, f, "snapshot3", "checkpoint3", &types.Disk{
+				ZoneId: "zone",
+				DiskId: "disk",
+			})
 
 			snapshot4 := SnapshotMeta{
 				ID: "snapshot4",
@@ -565,6 +591,10 @@ func TestSnapshotsCreateIncrementalSnapshot(t *testing.T) {
 
 			err = f.storage.SnapshotCreated(f.ctx, snapshot4.ID, 0, 0, 0, nil)
 			require.NoError(t, err)
+			checkIncremental(t, f, "snapshot4", "checkpoint4", &types.Disk{
+				ZoneId: "zone",
+				DiskId: "disk",
+			})
 
 			_, err = f.storage.DeletingSnapshot(f.ctx, snapshot1.ID, "delete1")
 			require.NoError(t, err)
diff --git a/cloud/disk_manager/internal/pkg/facade/image_service_test/image_service_test.go b/cloud/disk_manager/internal/pkg/facade/image_service_test/image_service_test.go
index 60f8e27a61..eb6dbe432a 100644
--- a/cloud/disk_manager/internal/pkg/facade/image_service_test/image_service_test.go
+++ b/cloud/disk_manager/internal/pkg/facade/image_service_test/image_service_test.go
@@ -238,7 +238,7 @@ func testImageServiceCreateImageFromDiskWithKind(
 	require.NoError(t, err)
 	require.Equal(t, float64(1), meta.Progress)
 
-	testcommon.RequireCheckpointsAreEmpty(t, ctx, diskID)
+	testcommon.RequireCheckpoint(t, ctx, diskID, imageID)
 
 	checkUnencryptedImage(
 		t,
@@ -1144,7 +1144,7 @@ func TestImageServiceCreateIncrementalImageFromDisk(t *testing.T) {
 	err = internal_client.GetOperationMetadata(ctx, client, operation.Id, &meta)
 	require.NoError(t, err)
 	require.Equal(t, float64(1), meta.Progress)
-	testcommon.RequireCheckpointsAreEmpty(t, ctx, diskID1)
+	testcommon.RequireCheckpoint(t, ctx, diskID1, imageID1)
 
 	nbsClient := testcommon.NewNbsTestingClient(t, ctx, "zone-a")
 	waitForWrite, err := nbsClient.GoWriteRandomBlocksToNbsDisk(ctx, diskID1)
@@ -1177,7 +1177,7 @@ func TestImageServiceCreateIncrementalImageFromDisk(t *testing.T) {
 	err = internal_client.GetOperationMetadata(ctx, client, operation.Id, &meta)
 	require.NoError(t, err)
 	require.Equal(t, float64(1), meta.Progress)
-	testcommon.RequireCheckpointsAreEmpty(t, ctx, diskID1)
+	testcommon.RequireCheckpoint(t, ctx, diskID1, imageID2)
 
 	testcommon.CheckBaseSnapshot(t, ctx, imageID2, imageID1)
 
@@ -1223,6 +1223,6 @@ func TestImageServiceCreateIncrementalImageFromDisk(t *testing.T) {
 	err = internal_client.WaitOperation(ctx, client, operation.Id)
 	require.NoError(t, err)
 
-	testcommon.RequireCheckpointsAreEmpty(t, ctx, diskID1)
+	testcommon.RequireNoCheckpoints(t, ctx, diskID1)
 	testcommon.CheckConsistency(t, ctx)
 }
diff --git a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
index 3b28520a8e..73fbb7258d 100644
--- a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
+++ b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/clients/nbs"
 	"github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/common"
 	"github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/facade/testcommon"
+	"github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/types"
 )
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -77,7 +78,7 @@ func testCreateSnapshotFromDisk(
 	require.NoError(t, err)
 	require.Equal(t, float64(1), meta.Progress)
 
-	testcommon.RequireCheckpointsAreEmpty(t, ctx, diskID)
+	testcommon.RequireCheckpoint(t, ctx, diskID, snapshotID)
 
 	reqCtx = testcommon.GetRequestContext(t, ctx)
 	operation, err = client.DeleteSnapshot(reqCtx, &disk_manager.DeleteSnapshotRequest{
@@ -604,17 +605,21 @@ func TestSnapshotServiceDeleteIncrementalSnapshotWhileCreating(t *testing.T) {
 	err = internal_client.WaitOperation(ctx, client, deleteOperation.Id)
 	require.NoError(t, err)
 
-	//nolint:sa9003
-	// TODO: remove line above after
-	// https://github.com/ydb-platform/nbs/issues/2008
 	if creationErr == nil {
-		testcommon.RequireCheckpointsAreEmpty(t, ctx, diskID)
+		testcommon.RequireNoCheckpoints(t, ctx, diskID)
 	} else {
-		// Checkpoint that corresponds to base snapshot should not be deleted.
-		// NOTE: we use snapshot id as checkpoint id.
-		// TODO: enable this check after resolving issue
-		// https://github.com/ydb-platform/nbs/issues/2008.
-		// testcommon.RequireCheckpoint(t, ctx, diskID, baseSnapshotID)
+		snapshotID, _, err := testcommon.GetIncremental(
+			ctx,
+			&types.Disk{
+				ZoneId: "zone-a",
+				DiskId: diskID,
+			},
+		)
+
+		require.NoError(t, err)
+		if snapshotID != "" {
+			testcommon.RequireCheckpoint(t, ctx, diskID, baseSnapshotID)
+		}
 	}
 
 	snapshotID2 := t.Name() + "2"
@@ -702,7 +707,7 @@ func TestSnapshotServiceDeleteSnapshot(t *testing.T) {
 	err = internal_client.WaitOperation(ctx, client, operation2.Id)
 	require.NoError(t, err)
 
-	testcommon.RequireCheckpointsAreEmpty(t, ctx, diskID)
+	testcommon.RequireNoCheckpoints(t, ctx, diskID)
 	testcommon.CheckConsistency(t, ctx)
 }
 
diff --git a/cloud/disk_manager/internal/pkg/facade/testcommon/common.go b/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
index 1afe621460..92b07b424b 100644
--- a/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
+++ b/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
@@ -353,7 +353,7 @@ func RequireCheckpoint(
 	require.EqualValues(t, checkpointID, checkpoints[0])
 }
 
-func RequireCheckpointsAreEmpty(
+func RequireNoCheckpoints(
 	t *testing.T,
 	ctx context.Context,
 	diskID string,
@@ -626,6 +626,19 @@ func CheckConsistency(t *testing.T, ctx context.Context) {
 
 ////////////////////////////////////////////////////////////////////////////////
 
+func GetIncremental(
+	ctx context.Context,
+	disk *types.Disk,
+) (string, string, error) {
+
+	storage, err := newSnapshotStorage(ctx)
+	if err != nil {
+		return "", "", err
+	}
+
+	return storage.GetIncremental(ctx, disk)
+}
+
 func GetEncryptionKeyHash(encryptionDesc *types.EncryptionDesc) ([]byte, error) {
 	switch key := encryptionDesc.Key.(type) {
 	case *types.EncryptionDesc_KeyHash:

From 15301c708787a73cce1bfd907282c6c43463eb53 Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Wed, 5 Feb 2025 11:48:27 +0000
Subject: [PATCH 04/11] fix
 TestSnapshotServiceDeleteSnapshotWhenCreationIsInFlight

---
 .../facade/snapshot_service_test/snapshot_service_test.go | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
index 73fbb7258d..844bca351d 100644
--- a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
+++ b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
@@ -763,13 +763,13 @@ func TestSnapshotServiceDeleteSnapshotWhenCreationIsInFlight(t *testing.T) {
 	err = internal_client.WaitOperation(ctx, client, operation.Id)
 	require.NoError(t, err)
 
-	_ = internal_client.WaitOperation(ctx, client, createOp.Id)
+	createErr := internal_client.WaitOperation(ctx, client, createOp.Id)
 
 	// Should wait here because checkpoint is deleted on |createOp| operation
 	// cancel (and exact time of this event is unknown).
-	// TODO: enable this check after resolving issue
-	// https://github.com/ydb-platform/nbs/issues/2008.
-	// testcommon.WaitForCheckpointsAreEmpty(t, ctx, diskID)
+	if createErr != nil {
+		testcommon.WaitForCheckpointsAreEmpty(t, ctx, diskID)
+	}
 
 	testcommon.CheckConsistency(t, ctx)
 }

From 0d5ca4d687ed5d1fe5fcd6d1cfc523ad6ee1c5d2 Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Wed, 5 Feb 2025 15:30:52 +0000
Subject: [PATCH 05/11] remove blockstore changes

---
 .../libs/storage/volume/volume_ut_checkpoint.cpp         | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp b/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp
index 8fa0b3a4db..8d464a9944 100644
--- a/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp
+++ b/cloud/blockstore/libs/storage/volume/volume_ut_checkpoint.cpp
@@ -1940,15 +1940,6 @@ Y_UNIT_TEST_SUITE(TVolumeCheckpointTest)
             UNIT_ASSERT_VALUES_EQUAL("c1", stat->Record.GetCheckpoints(0));
         }
 
-        volume.DeleteCheckpoint("c1");
-
-        {
-            auto stat = volume.StatVolume();
-            const auto& cp = stat->Record.GetCheckpoints();
-            UNIT_ASSERT_VALUES_EQUAL(1, cp.size());
-            UNIT_ASSERT_VALUES_EQUAL("c1", stat->Record.GetCheckpoints(0));
-        }
-
         volume.RebootTablet();
         volume.AddClient(clientInfo);
         volume.WaitReady();

From d3447adc753927d0a1d2a8f1ca8086d102f80c57 Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Wed, 5 Feb 2025 15:53:17 +0000
Subject: [PATCH 06/11] refactor

---
 .../snapshot/storage/storage_ydb_test.go      | 34 +++++++------------
 .../image_service_test/image_service_test.go  |  2 +-
 .../snapshot_service_test.go                  | 19 +++++------
 .../internal/pkg/facade/testcommon/common.go  |  2 +-
 4 files changed, 22 insertions(+), 35 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_test.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_test.go
index d72bc60c1c..0e4750dfae 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_test.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_ydb_test.go
@@ -418,9 +418,9 @@ func testCases() []differentChunkStorageTestCase {
 func checkIncremental(
 	t *testing.T,
 	f *fixture,
+	disk *types.Disk,
 	expectedSnapshotID string,
 	expectedCheckpointID string,
-	disk *types.Disk,
 ) {
 
 	snapshotID, checkpointID, err := f.storage.GetIncremental(f.ctx, disk)
@@ -507,12 +507,14 @@ func TestSnapshotsCreateIncrementalSnapshot(t *testing.T) {
 			f := createFixture(t)
 			defer f.teardown()
 
+			disk := types.Disk{
+				ZoneId: "zone",
+				DiskId: "disk",
+			}
+
 			snapshot1 := SnapshotMeta{
-				ID: "snapshot1",
-				Disk: &types.Disk{
-					ZoneId: "zone",
-					DiskId: "disk",
-				},
+				ID:           "snapshot1",
+				Disk:         &disk,
 				CheckpointID: "checkpoint1",
 				CreateTaskID: "create1",
 			}
@@ -522,10 +524,7 @@ func TestSnapshotsCreateIncrementalSnapshot(t *testing.T) {
 			require.NotNil(t, created)
 			require.Empty(t, created.BaseSnapshotID)
 			require.Empty(t, created.BaseCheckpointID)
-			checkIncremental(t, f, "", "", &types.Disk{
-				ZoneId: "zone",
-				DiskId: "disk",
-			})
+			checkIncremental(t, f, &disk, "", "")
 
 			created, err = f.storage.CreateSnapshot(f.ctx, SnapshotMeta{
 				ID: "snapshot2",
@@ -543,10 +542,7 @@ func TestSnapshotsCreateIncrementalSnapshot(t *testing.T) {
 
 			err = f.storage.SnapshotCreated(f.ctx, snapshot1.ID, 0, 0, 0, nil)
 			require.NoError(t, err)
-			checkIncremental(t, f, "snapshot1", "checkpoint1", &types.Disk{
-				ZoneId: "zone",
-				DiskId: "disk",
-			})
+			checkIncremental(t, f, &disk, "snapshot1", "checkpoint1")
 
 			snapshot3 := SnapshotMeta{
 				ID: "snapshot3",
@@ -567,10 +563,7 @@ func TestSnapshotsCreateIncrementalSnapshot(t *testing.T) {
 
 			err = f.storage.SnapshotCreated(f.ctx, snapshot3.ID, 0, 0, 0, nil)
 			require.NoError(t, err)
-			checkIncremental(t, f, "snapshot3", "checkpoint3", &types.Disk{
-				ZoneId: "zone",
-				DiskId: "disk",
-			})
+			checkIncremental(t, f, &disk, "snapshot3", "checkpoint3")
 
 			snapshot4 := SnapshotMeta{
 				ID: "snapshot4",
@@ -591,10 +584,7 @@ func TestSnapshotsCreateIncrementalSnapshot(t *testing.T) {
 
 			err = f.storage.SnapshotCreated(f.ctx, snapshot4.ID, 0, 0, 0, nil)
 			require.NoError(t, err)
-			checkIncremental(t, f, "snapshot4", "checkpoint4", &types.Disk{
-				ZoneId: "zone",
-				DiskId: "disk",
-			})
+			checkIncremental(t, f, &disk, "snapshot4", "checkpoint4")
 
 			_, err = f.storage.DeletingSnapshot(f.ctx, snapshot1.ID, "delete1")
 			require.NoError(t, err)
diff --git a/cloud/disk_manager/internal/pkg/facade/image_service_test/image_service_test.go b/cloud/disk_manager/internal/pkg/facade/image_service_test/image_service_test.go
index eb6dbe432a..a4bf6f23cd 100644
--- a/cloud/disk_manager/internal/pkg/facade/image_service_test/image_service_test.go
+++ b/cloud/disk_manager/internal/pkg/facade/image_service_test/image_service_test.go
@@ -1223,6 +1223,6 @@ func TestImageServiceCreateIncrementalImageFromDisk(t *testing.T) {
 	err = internal_client.WaitOperation(ctx, client, operation.Id)
 	require.NoError(t, err)
 
-	testcommon.RequireNoCheckpoints(t, ctx, diskID1)
+	testcommon.RequireCheckpointsDoNotExist(t, ctx, diskID1)
 	testcommon.CheckConsistency(t, ctx)
 }
diff --git a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
index 844bca351d..12e24a429d 100644
--- a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
+++ b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
@@ -606,18 +606,15 @@ func TestSnapshotServiceDeleteIncrementalSnapshotWhileCreating(t *testing.T) {
 	require.NoError(t, err)
 
 	if creationErr == nil {
-		testcommon.RequireNoCheckpoints(t, ctx, diskID)
+		testcommon.RequireCheckpointsDoNotExist(t, ctx, diskID)
 	} else {
-		snapshotID, _, err := testcommon.GetIncremental(
-			ctx,
-			&types.Disk{
-				ZoneId: "zone-a",
-				DiskId: diskID,
-			},
-		)
-
+		snapshotID, _, err := testcommon.GetIncremental(ctx, &types.Disk{
+			ZoneId: "zone-a",
+			DiskId: diskID,
+		})
 		require.NoError(t, err)
-		if snapshotID != "" {
+
+		if len(snapshotID) > 0 {
 			testcommon.RequireCheckpoint(t, ctx, diskID, baseSnapshotID)
 		}
 	}
@@ -707,7 +704,7 @@ func TestSnapshotServiceDeleteSnapshot(t *testing.T) {
 	err = internal_client.WaitOperation(ctx, client, operation2.Id)
 	require.NoError(t, err)
 
-	testcommon.RequireNoCheckpoints(t, ctx, diskID)
+	testcommon.RequireCheckpointsDoNotExist(t, ctx, diskID)
 	testcommon.CheckConsistency(t, ctx)
 }
 
diff --git a/cloud/disk_manager/internal/pkg/facade/testcommon/common.go b/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
index 92b07b424b..46bb33871f 100644
--- a/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
+++ b/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
@@ -353,7 +353,7 @@ func RequireCheckpoint(
 	require.EqualValues(t, checkpointID, checkpoints[0])
 }
 
-func RequireNoCheckpoints(
+func RequireCheckpointsDoNotExist(
 	t *testing.T,
 	ctx context.Context,
 	diskID string,

From 4b577ca5950f9042a7fb2810d000f9b03caedfac Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Wed, 5 Feb 2025 19:19:16 +0000
Subject: [PATCH 07/11] edit comments

---
 .../facade/snapshot_service_test/snapshot_service_test.go | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
index 12e24a429d..8d40b4b2a6 100644
--- a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
+++ b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
@@ -614,8 +614,12 @@ func TestSnapshotServiceDeleteIncrementalSnapshotWhileCreating(t *testing.T) {
 		})
 		require.NoError(t, err)
 
+		// If there is a record about this disk left in incrementality table,
+		// checkpoint that corresponds to base snapshot should not be deleted.
 		if len(snapshotID) > 0 {
 			testcommon.RequireCheckpoint(t, ctx, diskID, baseSnapshotID)
+		} else {
+			testcommon.RequireCheckpointsDoNotExist(t, ctx, diskID)
 		}
 	}
 
@@ -762,9 +766,9 @@ func TestSnapshotServiceDeleteSnapshotWhenCreationIsInFlight(t *testing.T) {
 
 	createErr := internal_client.WaitOperation(ctx, client, createOp.Id)
 
-	// Should wait here because checkpoint is deleted on |createOp| operation
-	// cancel (and exact time of this event is unknown).
 	if createErr != nil {
+		// Should wait here because checkpoint is deleted on |createOp| operation
+		// cancel (and exact time of this event is unknown).
 		testcommon.WaitForCheckpointsAreEmpty(t, ctx, diskID)
 	}
 

From a49be04b1d90a500fb1b0cb4d4c3dded0572e4d6 Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 13 Feb 2025 11:49:46 +0000
Subject: [PATCH 08/11] rename method, add comments

---
 .../snapshot_service_test.go                  | 24 ++++++++++++-------
 .../internal/pkg/facade/testcommon/common.go  |  4 ++--
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
index 8d40b4b2a6..fd713682ef 100644
--- a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
+++ b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
@@ -171,7 +171,7 @@ func TestSnapshotServiceCancelCreateSnapshotFromDisk(t *testing.T) {
 
 	// Should wait here because checkpoint is deleted on operation cancel (and
 	// exact time of this event is unknown).
-	testcommon.WaitForCheckpointsAreEmpty(t, ctx, diskID)
+	testcommon.WaitForCheckpointsDoNotExist(t, ctx, diskID)
 
 	testcommon.CheckConsistency(t, ctx)
 }
@@ -605,6 +605,9 @@ func TestSnapshotServiceDeleteIncrementalSnapshotWhileCreating(t *testing.T) {
 	err = internal_client.WaitOperation(ctx, client, deleteOperation.Id)
 	require.NoError(t, err)
 
+	// In case of successful snapshot1 creation, snapshot1 should be deleted
+	// and checkpoint should be deleted too.
+	// Otherwise we should check incremental table.
 	if creationErr == nil {
 		testcommon.RequireCheckpointsDoNotExist(t, ctx, diskID)
 	} else {
@@ -614,8 +617,9 @@ func TestSnapshotServiceDeleteIncrementalSnapshotWhileCreating(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		// If there is a record about this disk left in incrementality table,
-		// checkpoint that corresponds to base snapshot should not be deleted.
+		// In case of snapshot creation failure baseSnapshot may be already
+		// deleted from incremental table and then checkpoint should not exist
+		// on the disk. Otherwise checkpoint should exist.
 		if len(snapshotID) > 0 {
 			testcommon.RequireCheckpoint(t, ctx, diskID, baseSnapshotID)
 		} else {
@@ -764,12 +768,16 @@ func TestSnapshotServiceDeleteSnapshotWhenCreationIsInFlight(t *testing.T) {
 	err = internal_client.WaitOperation(ctx, client, operation.Id)
 	require.NoError(t, err)
 
-	createErr := internal_client.WaitOperation(ctx, client, createOp.Id)
+	err = internal_client.WaitOperation(ctx, client, createOp.Id)
 
-	if createErr != nil {
-		// Should wait here because checkpoint is deleted on |createOp| operation
-		// cancel (and exact time of this event is unknown).
-		testcommon.WaitForCheckpointsAreEmpty(t, ctx, diskID)
+	// In case snapshot creation ends up successfuly,
+	// there should be a checkpoint without data.
+	if err == nil {
+		testcommon.RequireCheckpoint(t, ctx, diskID, snapshotID)
+	} else {
+		// Should wait here because checkpoint is deleted on |createOp|
+		// operation cancel (and exact time of this event is unknown).
+		testcommon.WaitForCheckpointsDoNotExist(t, ctx, diskID)
 	}
 
 	testcommon.CheckConsistency(t, ctx)
diff --git a/cloud/disk_manager/internal/pkg/facade/testcommon/common.go b/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
index 46bb33871f..74d596417d 100644
--- a/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
+++ b/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
@@ -367,7 +367,7 @@ func RequireCheckpointsDoNotExist(
 	require.Empty(t, checkpoints)
 }
 
-func WaitForCheckpointsAreEmpty(
+func WaitForCheckpointsDoNotExist(
 	t *testing.T,
 	ctx context.Context,
 	diskID string,
@@ -385,7 +385,7 @@ func WaitForCheckpointsAreEmpty(
 
 		logging.Warn(
 			ctx,
-			"waitForCheckpointsAreEmpty proceeding to next iteration",
+			"WaitForCheckpointsDoNotExist proceeding to next iteration",
 		)
 
 		<-time.After(100 * time.Millisecond)

From ec52f94663bede471889eb79ed6e914719b9dac3 Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 13 Feb 2025 11:53:21 +0000
Subject: [PATCH 09/11] enable RequireCheckpointsDoNotExist

---
 cloud/disk_manager/internal/pkg/facade/testcommon/common.go | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/facade/testcommon/common.go b/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
index 74d596417d..258a03a309 100644
--- a/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
+++ b/cloud/disk_manager/internal/pkg/facade/testcommon/common.go
@@ -358,9 +358,7 @@ func RequireCheckpointsDoNotExist(
 	ctx context.Context,
 	diskID string,
 ) {
-	// TODO: enable this method after resolving this issue
-	// https://github.com/ydb-platform/nbs/issues/2008.
-	return
+
 	nbsClient := NewNbsTestingClient(t, ctx, "zone-a")
 	checkpoints, err := nbsClient.GetCheckpoints(ctx, diskID)
 	require.NoError(t, err)

From eec2dc3a1a44d2b9f26aaf03cd7942721dc0389c Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 13 Feb 2025 16:48:18 +0000
Subject: [PATCH 10/11] make arguments of storage*.GetIncremental named

---
 .../pkg/dataplane/snapshot/storage/mocks/storage_mock.go        | 2 +-
 .../internal/pkg/dataplane/snapshot/storage/storage.go          | 2 +-
 .../internal/pkg/dataplane/snapshot/storage/storage_legacy.go   | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/mocks/storage_mock.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/mocks/storage_mock.go
index 75a6fab928..913f543ef5 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/mocks/storage_mock.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/mocks/storage_mock.go
@@ -237,7 +237,7 @@ func (s *StorageMock) GetSnapshotMeta(
 func (s *StorageMock) GetIncremental(
 	ctx context.Context,
 	snapshotID string,
-) (string, string, error) {
+) (snapshotID string, checkpointID string, err error) {
 
 	args := s.Called(ctx, snapshotID)
 	return args.String(0), args.String(1), args.Error(2)
diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage.go
index 94441f89a2..7532f77ca0 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage.go
@@ -142,5 +142,5 @@ type Storage interface {
 	GetIncremental(
 		ctx context.Context,
 		disk *types.Disk,
-	) (string, string, error)
+	) (snapshotID string, checkpointID string, err error)
 }
diff --git a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_legacy.go b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_legacy.go
index da398631f1..c3cac3d1ff 100644
--- a/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_legacy.go
+++ b/cloud/disk_manager/internal/pkg/dataplane/snapshot/storage/storage_legacy.go
@@ -281,7 +281,7 @@ func (s *legacyStorage) GetSnapshotMeta(
 func (s *legacyStorage) GetIncremental(
 	ctx context.Context,
 	disk *types.Disk,
-) (string, string, error) {
+) (snapshotID string, checkpointID string, err error {
 
 	return "", "", task_errors.NewNonRetriableErrorf("not implemented")
 }

From 3f4bdf47127c349915f9baba52e51894a9218bda Mon Sep 17 00:00:00 2001
From: Rattysed <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 13 Feb 2025 17:46:11 +0000
Subject: [PATCH 11/11] fix comments

---
 .../pkg/facade/snapshot_service_test/snapshot_service_test.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
index fd713682ef..adb8bc9733 100644
--- a/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
+++ b/cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go
@@ -617,7 +617,7 @@ func TestSnapshotServiceDeleteIncrementalSnapshotWhileCreating(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		// In case of snapshot creation failure baseSnapshot may be already
+		// In case of snapshot1 creation failure base snapshot may be already
 		// deleted from incremental table and then checkpoint should not exist
 		// on the disk. Otherwise checkpoint should exist.
 		if len(snapshotID) > 0 {
@@ -770,7 +770,7 @@ func TestSnapshotServiceDeleteSnapshotWhenCreationIsInFlight(t *testing.T) {
 
 	err = internal_client.WaitOperation(ctx, client, createOp.Id)
 
-	// In case snapshot creation ends up successfuly,
+	// If snapshot creation ends up successfuly,
 	// there should be a checkpoint without data.
 	if err == nil {
 		testcommon.RequireCheckpoint(t, ctx, diskID, snapshotID)