[Disk Manager] Fix flappy test for snapshot service

cloud/disk_manager/internal/pkg/facade/snapshot_service_test/snapshot_service_test.go

@@ -598,7 +598,7 @@ func TestSnapshotServiceDeleteIncrementalSnapshotWhileCreating(t *testing.T) {
 	require.NotEmpty(t, createOperation)
 
 	// Need to add some variance for better testing.
-	common.WaitForRandomDuration(1*time.Second, 3*time.Second)
+	common.WaitForRandomDuration(1*time.Millisecond, 3*time.Second)
 
 	reqCtx = testcommon.GetRequestContext(t, ctx)
 	deleteOperation, err := client.DeleteSnapshot(reqCtx, &disk_manager.DeleteSnapshotRequest{
@@ -612,22 +612,25 @@ 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 {
+	// If snapshot creation ends up successfuly, there may be two cases:
+	// Snapshot was created and then deleted, so should be no checkpoints left
+	// or snapshot deletion ended up before creation, snapshot was not deleted,
+	// so there should be a checkpoint.
+	if creationErr != nil {
 		snapshotID, _, err := testcommon.GetIncremental(ctx, &types.Disk{
 			ZoneId: "zone-a",
 			DiskId: diskID,
 		})
 		require.NoError(t, err)
 
+		// Should wait here because checkpoint is deleted on |createOperation|
+		// operation cancel (and exact time of this event is unknown).
+		testcommon.WaitForCheckpointsDoNotExist(t, ctx, diskID, snapshotID1)
 		// 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.
+		// on the disk. Otherwise base snapshot checkpoint should exist.
 		if len(snapshotID) > 0 {
+			require.Equal(t, snapshotID, baseSnapshotID)
 			testcommon.RequireCheckpoint(t, ctx, diskID, baseSnapshotID)
 		} else {
 			testcommon.RequireCheckpointsDoNotExist(t, ctx, diskID)

cloud/disk_manager/internal/pkg/facade/testcommon/common.go

@@ -369,6 +369,7 @@ func WaitForCheckpointsDoNotExist(
 	t *testing.T,
 	ctx context.Context,
 	diskID string,
+	awaitedCheckpoints ...string,
 ) {
 
 	nbsClient := NewNbsTestingClient(t, ctx, "zone-a")
@@ -377,7 +378,21 @@ func WaitForCheckpointsDoNotExist(
 		checkpoints, err := nbsClient.GetCheckpoints(ctx, diskID)
 		require.NoError(t, err)
 
-		if len(checkpoints) == 0 {
+		found := false
+
+		for _, awaitedCheckpoint := range awaitedCheckpoints {
+			for _, checkpoint := range checkpoints {
+				if awaitedCheckpoint == checkpoint {
+					found = true
+					break
+				}
+			}
+			if found {
+				break
+			}
+		}
+
+		if !found {
 			return
 		}