[Disk Manager] Fix disk manager tests after updating StatVolume behaviour

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 {

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)
 }

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")
+}

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
+}

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,

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)

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)
 }

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)
 }
 
@@ -758,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)
 }

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: