From ce6c4c3df10a643c3308021994194a4739a3b7f4 Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Wed, 26 Nov 2025 16:15:46 +0000
Subject: [PATCH 01/11] remove clusterCapacityState

---
 .../cells/collect_cluster_capacity_task.go    |   3 +
 .../collect_cluster_capacity_task_test.go     | 153 +++++++++++-------
 .../internal/pkg/cells/storage/storage.go     |  58 +++----
 .../internal/pkg/cells/storage/storage_ydb.go |  14 +-
 4 files changed, 121 insertions(+), 107 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task.go b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task.go
index c176a238a90..76647141f2d 100644
--- a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task.go
+++ b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task.go
@@ -143,6 +143,8 @@ func (t *collectClusterCapacityTask) updateCellCapacity(
 		return err
 	}
 
+	createdAt := time.Now()
+
 	var capacities []storage.ClusterCapacity
 	for _, info := range capacityInfos {
 		capacities = append(capacities, storage.ClusterCapacity{
@@ -151,6 +153,7 @@ func (t *collectClusterCapacityTask) updateCellCapacity(
 			Kind:       info.DiskKind,
 			FreeBytes:  info.FreeBytes,
 			TotalBytes: info.TotalBytes,
+			CreatedAt:  createdAt,
 		})
 	}
 
diff --git a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
index d406f43c681..78ebeaa5f3b 100644
--- a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
@@ -19,6 +19,26 @@ import (
 
 ////////////////////////////////////////////////////////////////////////////////
 
+func matchClusterCapacities(t *testing.T, want []cells_storage.ClusterCapacity) interface{} {
+	return mock.MatchedBy(func(actual []cells_storage.ClusterCapacity) bool {
+		require.Equal(t, len(want), len(actual))
+		for i := range want {
+			expectd := want[i]
+			got := actual[i]
+
+			require.Equal(t, expectd.ZoneID, got.ZoneID)
+			require.Equal(t, expectd.CellID, got.CellID)
+			require.Equal(t, expectd.Kind, got.Kind)
+			require.Equal(t, expectd.FreeBytes, got.FreeBytes)
+			require.Equal(t, expectd.TotalBytes, got.TotalBytes)
+			require.WithinDuration(t, time.Now(), got.CreatedAt, 15*time.Minute)
+		}
+		return true
+	})
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 func TestCollectClusterCapacityTask(t *testing.T) {
 	ctx := newContext()
 	execCtx := tasks_mocks.NewExecutionContextMock()
@@ -56,41 +76,50 @@ func TestCollectClusterCapacityTask(t *testing.T) {
 
 	storage.On("UpdateClusterCapacities",
 		mock.Anything, // ctx.
-		[]cells_storage.ClusterCapacity{
-			{
-				ZoneID:     "zone-a",
-				CellID:     "zone-a-cell1",
-				Kind:       types.DiskKind_DISK_KIND_SSD,
-				FreeBytes:  1024,
-				TotalBytes: 2048,
+		matchClusterCapacities(
+			t,
+			[]cells_storage.ClusterCapacity{
+				{
+					ZoneID:     "zone-a",
+					CellID:     "zone-a-cell1",
+					Kind:       types.DiskKind_DISK_KIND_SSD,
+					FreeBytes:  1024,
+					TotalBytes: 2048,
+				},
 			},
-		},
+		),
 		deleteOlderThanExpectation,
 	).Return(nil).Once()
 	storage.On("UpdateClusterCapacities",
 		mock.Anything, // ctx.
-		[]cells_storage.ClusterCapacity{
-			{
-				ZoneID:     "zone-a",
-				CellID:     "zone-a",
-				Kind:       types.DiskKind_DISK_KIND_SSD,
-				FreeBytes:  1024,
-				TotalBytes: 2048,
+		matchClusterCapacities(
+			t,
+			[]cells_storage.ClusterCapacity{
+				{
+					ZoneID:     "zone-a",
+					CellID:     "zone-a",
+					Kind:       types.DiskKind_DISK_KIND_SSD,
+					FreeBytes:  1024,
+					TotalBytes: 2048,
+				},
 			},
-		},
+		),
 		deleteOlderThanExpectation,
 	).Return(nil).Once()
 	storage.On("UpdateClusterCapacities",
 		mock.Anything, // ctx.
-		[]cells_storage.ClusterCapacity{
-			{
-				ZoneID:     "zone-b",
-				CellID:     "zone-b",
-				Kind:       types.DiskKind_DISK_KIND_SSD,
-				FreeBytes:  1024,
-				TotalBytes: 2048,
+		matchClusterCapacities(
+			t,
+			[]cells_storage.ClusterCapacity{
+				{
+					ZoneID:     "zone-b",
+					CellID:     "zone-b",
+					Kind:       types.DiskKind_DISK_KIND_SSD,
+					FreeBytes:  1024,
+					TotalBytes: 2048,
+				},
 			},
-		},
+		),
 		deleteOlderThanExpectation,
 	).Return(nil).Once()
 
@@ -163,15 +192,18 @@ func TestCollectClusterCapacityFailureNbsReturnsError(t *testing.T) {
 	// Only the successful cell should be updated
 	storage.On("UpdateClusterCapacities",
 		mock.Anything, // ctx.
-		[]cells_storage.ClusterCapacity{
-			{
-				ZoneID:     "zone-a",
-				CellID:     "zone-a",
-				Kind:       types.DiskKind_DISK_KIND_SSD,
-				FreeBytes:  1024,
-				TotalBytes: 2048,
+		matchClusterCapacities(
+			t,
+			[]cells_storage.ClusterCapacity{
+				{
+					ZoneID:     "zone-a",
+					CellID:     "zone-a",
+					Kind:       types.DiskKind_DISK_KIND_SSD,
+					FreeBytes:  1024,
+					TotalBytes: 2048,
+				},
 			},
-		},
+		),
 		mock.Anything, // deleteOlderThan.
 	).Return(nil).Once()
 
@@ -232,28 +264,34 @@ func TestCollectClusterCapacityFailureStorageReturnsError(t *testing.T) {
 
 	storage.On("UpdateClusterCapacities",
 		mock.Anything, // ctx.
-		[]cells_storage.ClusterCapacity{
-			{
-				ZoneID:     "zone-a",
-				CellID:     "zone-a",
-				Kind:       types.DiskKind_DISK_KIND_SSD,
-				FreeBytes:  1024,
-				TotalBytes: 2048,
+		matchClusterCapacities(
+			t,
+			[]cells_storage.ClusterCapacity{
+				{
+					ZoneID:     "zone-a",
+					CellID:     "zone-a",
+					Kind:       types.DiskKind_DISK_KIND_SSD,
+					FreeBytes:  1024,
+					TotalBytes: 2048,
+				},
 			},
-		},
+		),
 		mock.Anything, // deleteOlderThan.
 	).Return(nil).Once()
 	storage.On("UpdateClusterCapacities",
 		mock.Anything,
-		[]cells_storage.ClusterCapacity{
-			{
-				ZoneID:     "zone-a",
-				CellID:     "zone-a-cell1",
-				Kind:       types.DiskKind_DISK_KIND_SSD,
-				FreeBytes:  1024,
-				TotalBytes: 2048,
+		matchClusterCapacities(
+			t,
+			[]cells_storage.ClusterCapacity{
+				{
+					ZoneID:     "zone-a",
+					CellID:     "zone-a-cell1",
+					Kind:       types.DiskKind_DISK_KIND_SSD,
+					FreeBytes:  1024,
+					TotalBytes: 2048,
+				},
 			},
-		},
+		),
 		mock.Anything, // deleteOlderThan.
 	).Return(assert.AnError).Once()
 
@@ -313,15 +351,18 @@ func TestCollectClusterCapacityOneCellHasAlreadyBeenProcessed(t *testing.T) {
 
 	storage.On("UpdateClusterCapacities",
 		mock.Anything, // ctx.
-		[]cells_storage.ClusterCapacity{
-			{
-				ZoneID:     "zone-a",
-				CellID:     "zone-a",
-				Kind:       types.DiskKind_DISK_KIND_SSD,
-				FreeBytes:  1024,
-				TotalBytes: 2048,
+		matchClusterCapacities(
+			t,
+			[]cells_storage.ClusterCapacity{
+				{
+					ZoneID:     "zone-a",
+					CellID:     "zone-a",
+					Kind:       types.DiskKind_DISK_KIND_SSD,
+					FreeBytes:  1024,
+					TotalBytes: 2048,
+				},
 			},
-		},
+		),
 		mock.Anything, // deleteOlderThan.
 	).Return(nil).Once()
 
diff --git a/cloud/disk_manager/internal/pkg/cells/storage/storage.go b/cloud/disk_manager/internal/pkg/cells/storage/storage.go
index e4c0f54b8c4..bf64291d6a0 100644
--- a/cloud/disk_manager/internal/pkg/cells/storage/storage.go
+++ b/cloud/disk_manager/internal/pkg/cells/storage/storage.go
@@ -19,42 +19,17 @@ type ClusterCapacity struct {
 	Kind       types.DiskKind
 	TotalBytes uint64
 	FreeBytes  uint64
-}
-
-////////////////////////////////////////////////////////////////////////////////
-
-type clusterCapacityState struct {
-	ZoneID     string
-	CellID     string
-	Kind       string
-	TotalBytes uint64
-	FreeBytes  uint64
 	CreatedAt  time.Time
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 
-func (c *clusterCapacityState) toClusterCapacity() (ClusterCapacity, error) {
-	kind, err := common.DiskKindFromString(c.Kind)
-	if err != nil {
-		return ClusterCapacity{}, err
-	}
-
-	return ClusterCapacity{
-		ZoneID:     c.ZoneID,
-		CellID:     c.CellID,
-		Kind:       kind,
-		TotalBytes: c.TotalBytes,
-		FreeBytes:  c.FreeBytes,
-	}, nil
-}
-
 // Returns ydb entity of the node object.
-func (c *clusterCapacityState) structValue() persistence.Value {
+func (c *ClusterCapacity) structValue() persistence.Value {
 	return persistence.StructValue(
 		persistence.StructFieldValue("zone_id", persistence.UTF8Value(c.ZoneID)),
 		persistence.StructFieldValue("cell_id", persistence.UTF8Value(c.CellID)),
-		persistence.StructFieldValue("kind", persistence.UTF8Value(c.Kind)),
+		persistence.StructFieldValue("kind", persistence.UTF8Value(common.DiskKindToString(c.Kind))),
 		persistence.StructFieldValue("total_bytes", persistence.Uint64Value(c.TotalBytes)),
 		persistence.StructFieldValue("free_bytes", persistence.Uint64Value(c.FreeBytes)),
 		persistence.StructFieldValue("created_at", persistence.TimestampValue(c.CreatedAt)),
@@ -63,7 +38,7 @@ func (c *clusterCapacityState) structValue() persistence.Value {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-func clusterCapacityStateStructTypeString() string {
+func clusterCapacityStructTypeString() string {
 	return `Struct<
 		zone_id: Utf8,
 		cell_id: Utf8,
@@ -73,7 +48,7 @@ func clusterCapacityStateStructTypeString() string {
 		created_at: Timestamp>`
 }
 
-func clusterCapacityStateTableDescription() persistence.CreateTableDescription {
+func clusterCapacityTableDescription() persistence.CreateTableDescription {
 	return persistence.NewCreateTableDescription(
 		persistence.WithColumn("zone_id", persistence.Optional(persistence.TypeUTF8)),
 		persistence.WithColumn("cell_id", persistence.Optional(persistence.TypeUTF8)),
@@ -89,17 +64,27 @@ func clusterCapacityStateTableDescription() persistence.CreateTableDescription {
 
 func scanClusterCapacity(
 	res persistence.Result,
-) (capacity clusterCapacityState, err error) {
+) (capacity ClusterCapacity, err error) {
 
+	var kind string
 	err = res.ScanNamed(
 		persistence.OptionalWithDefault("zone_id", &capacity.ZoneID),
 		persistence.OptionalWithDefault("cell_id", &capacity.CellID),
-		persistence.OptionalWithDefault("kind", &capacity.Kind),
+		persistence.OptionalWithDefault("kind", kind),
 		persistence.OptionalWithDefault("total_bytes", &capacity.TotalBytes),
 		persistence.OptionalWithDefault("free_bytes", &capacity.FreeBytes),
+		persistence.OptionalWithDefault("created_at", &capacity.CreatedAt),
 	)
+	if err != nil {
+		return ClusterCapacity{}, err
+	}
+
+	capacity.Kind, err = common.DiskKindFromString(kind)
+	if err != nil {
+		return ClusterCapacity{}, err
+	}
 
-	return capacity, err
+	return capacity, nil
 }
 
 func scanClusterCapacities(
@@ -111,12 +96,7 @@ func scanClusterCapacities(
 
 	for res.NextResultSet(ctx) {
 		for res.NextRow() {
-			capacityState, err := scanClusterCapacity(res)
-			if err != nil {
-				return nil, err
-			}
-
-			capacity, err := capacityState.toClusterCapacity()
+			capacity, err := scanClusterCapacity(res)
 			if err != nil {
 				return nil, err
 			}
@@ -163,7 +143,7 @@ func CreateYDBTables(
 		ctx,
 		config.GetStorageFolder(),
 		"cluster_capacity",
-		clusterCapacityStateTableDescription(),
+		clusterCapacityTableDescription(),
 		dropUnusedColumns,
 	)
 	if err != nil {
diff --git a/cloud/disk_manager/internal/pkg/cells/storage/storage_ydb.go b/cloud/disk_manager/internal/pkg/cells/storage/storage_ydb.go
index 73cf7114b10..e85f083b21a 100644
--- a/cloud/disk_manager/internal/pkg/cells/storage/storage_ydb.go
+++ b/cloud/disk_manager/internal/pkg/cells/storage/storage_ydb.go
@@ -105,19 +105,9 @@ func (s *storageYDB) updateClusterCapacities(ctx context.Context,
 	}
 
 	var values []persistence.Value
-	createdAt := time.Now()
 
 	for _, capacity := range capacities {
-		capacityState := &clusterCapacityState{
-			ZoneID:     capacity.ZoneID,
-			CellID:     capacity.CellID,
-			Kind:       common.DiskKindToString(capacity.Kind),
-			TotalBytes: capacity.TotalBytes,
-			FreeBytes:  capacity.FreeBytes,
-			CreatedAt:  createdAt,
-		}
-
-		values = append(values, capacityState.structValue())
+		values = append(values, capacity.structValue())
 	}
 
 	_, err = tx.Execute(ctx, fmt.Sprintf(`
@@ -128,7 +118,7 @@ func (s *storageYDB) updateClusterCapacities(ctx context.Context,
 		upsert into cluster_capacity
 		select *
 		from AS_TABLE($capacities)
-	`, s.tablesPath, clusterCapacityStateStructTypeString()),
+	`, s.tablesPath, clusterCapacityStructTypeString()),
 		persistence.ValueParam("$capacities", persistence.ListValue(values...)),
 	)
 	if err != nil {

From c7a9482e595cdd601457e2ebb39cbfccb61cc3fc Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Wed, 26 Nov 2025 21:39:08 +0000
Subject: [PATCH 02/11] fix ydb scan

---
 cloud/disk_manager/internal/pkg/cells/storage/storage.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/storage/storage.go b/cloud/disk_manager/internal/pkg/cells/storage/storage.go
index bf64291d6a0..60460dd1616 100644
--- a/cloud/disk_manager/internal/pkg/cells/storage/storage.go
+++ b/cloud/disk_manager/internal/pkg/cells/storage/storage.go
@@ -70,7 +70,7 @@ func scanClusterCapacity(
 	err = res.ScanNamed(
 		persistence.OptionalWithDefault("zone_id", &capacity.ZoneID),
 		persistence.OptionalWithDefault("cell_id", &capacity.CellID),
-		persistence.OptionalWithDefault("kind", kind),
+		persistence.OptionalWithDefault("kind", &kind),
 		persistence.OptionalWithDefault("total_bytes", &capacity.TotalBytes),
 		persistence.OptionalWithDefault("free_bytes", &capacity.FreeBytes),
 		persistence.OptionalWithDefault("created_at", &capacity.CreatedAt),

From 613fcceb8237054d50953624e4b4a22849cd91b2 Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Wed, 26 Nov 2025 21:59:22 +0000
Subject: [PATCH 03/11] fix storage tests

---
 .../pkg/cells/storage/storage_test.go         | 23 +++++++++++++++----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
index fe43b7df384..e9a27e6c1c7 100644
--- a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
@@ -82,6 +82,8 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 	storage := newStorage(t, ctx, db)
 
 	zeroTime := time.Time{}
+	firstTime := time.Now()
+	secondTime := firstTime.Add(time.Hour)
 
 	capacity1 := ClusterCapacity{
 		ZoneID:     "zone-a",
@@ -89,6 +91,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		Kind:       types.DiskKind_DISK_KIND_HDD,
 		TotalBytes: 1024,
 		FreeBytes:  1024,
+		CreatedAt:  firstTime,
 	}
 
 	cellCapacity1 := ClusterCapacity{
@@ -97,6 +100,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		Kind:       types.DiskKind_DISK_KIND_HDD,
 		TotalBytes: 1024,
 		FreeBytes:  1024,
+		CreatedAt:  firstTime,
 	}
 
 	cellCapacitySsd1 := ClusterCapacity{
@@ -105,6 +109,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		Kind:       types.DiskKind_DISK_KIND_SSD,
 		TotalBytes: 1024,
 		FreeBytes:  1024,
+		CreatedAt:  firstTime,
 	}
 
 	capacity2 := ClusterCapacity{
@@ -113,6 +118,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		Kind:       types.DiskKind_DISK_KIND_HDD,
 		TotalBytes: 2048,
 		FreeBytes:  2048,
+		CreatedAt:  secondTime,
 	}
 
 	cellCapacity2 := ClusterCapacity{
@@ -121,6 +127,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		Kind:       types.DiskKind_DISK_KIND_HDD,
 		TotalBytes: 2048,
 		FreeBytes:  2048,
+		CreatedAt:  secondTime,
 	}
 
 	cellCapacitySsd2 := ClusterCapacity{
@@ -129,6 +136,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		Kind:       types.DiskKind_DISK_KIND_SSD,
 		TotalBytes: 2048,
 		FreeBytes:  2048,
+		CreatedAt:  secondTime,
 	}
 
 	err = storage.UpdateClusterCapacities(
@@ -212,13 +220,15 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 
 	storage := newStorage(t, ctx, db)
 
-	deleteBefore := time.Now()
+	deleteOlderThan := time.Now()
+	createdAt := deleteOlderThan.Add(time.Hour)
 	capacity1 := ClusterCapacity{
 		ZoneID:     "zone-a",
 		CellID:     "zone-a",
 		Kind:       types.DiskKind_DISK_KIND_HDD,
 		TotalBytes: 1024,
 		FreeBytes:  1024,
+		CreatedAt:  createdAt,
 	}
 	capacity2 := ClusterCapacity{
 		ZoneID:     "zone-a",
@@ -226,9 +236,10 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 		Kind:       types.DiskKind_DISK_KIND_HDD,
 		TotalBytes: 1024,
 		FreeBytes:  1024,
+		CreatedAt:  createdAt,
 	}
 
-	err = storage.UpdateClusterCapacities(ctx, []ClusterCapacity{capacity1}, deleteBefore)
+	err = storage.UpdateClusterCapacities(ctx, []ClusterCapacity{capacity1}, deleteOlderThan)
 	require.NoError(t, err)
 
 	capacities, err := storage.GetRecentClusterCapacities(
@@ -239,7 +250,7 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 	require.NoError(t, err)
 	require.ElementsMatch(t, capacities, []ClusterCapacity{capacity1})
 
-	err = storage.UpdateClusterCapacities(ctx, []ClusterCapacity{capacity2}, deleteBefore)
+	err = storage.UpdateClusterCapacities(ctx, []ClusterCapacity{capacity2}, deleteOlderThan)
 	require.NoError(t, err)
 
 	capacities, err = storage.GetRecentClusterCapacities(
@@ -250,16 +261,18 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 	require.NoError(t, err)
 	require.ElementsMatch(t, capacities, []ClusterCapacity{capacity1, capacity2})
 
-	deleteBefore = time.Now()
+	deleteOlderThan = time.Now()
+	createdAt = deleteOlderThan.Add(time.Hour)
 	capacity3 := ClusterCapacity{
 		ZoneID:     "zone-a",
 		CellID:     "zone-a",
 		Kind:       types.DiskKind_DISK_KIND_HDD,
 		TotalBytes: 2048,
 		FreeBytes:  2048,
+		CreatedAt:  createdAt,
 	}
 
-	err = storage.UpdateClusterCapacities(ctx, []ClusterCapacity{capacity3}, deleteBefore)
+	err = storage.UpdateClusterCapacities(ctx, []ClusterCapacity{capacity3}, deleteOlderThan)
 	require.NoError(t, err)
 
 	capacities, err = storage.GetRecentClusterCapacities(

From 45c3f999b02da271a57af80b500391f28cb15e5a Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Wed, 26 Nov 2025 23:17:51 +0000
Subject: [PATCH 04/11] support time.Equal in tests

---
 .../pkg/cells/storage/storage_test.go         | 70 +++++++++++++------
 1 file changed, 48 insertions(+), 22 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
index e9a27e6c1c7..5005d4d1bfd 100644
--- a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"sort"
 	"testing"
 	"time"
 
@@ -71,6 +72,16 @@ func newStorage(
 
 ////////////////////////////////////////////////////////////////////////////////
 
+// Uses time.Equal() to compare time.Time fields correctly, since == may fail
+// due to different internal representations.
+func clusterCapacitiesAreEqual(a ClusterCapacity, b ClusterCapacity) bool {
+	return a.CellID == b.CellID && a.ZoneID == b.ZoneID && a.Kind == b.Kind &&
+		a.FreeBytes == b.FreeBytes && a.TotalBytes == b.TotalBytes &&
+		a.CreatedAt.Equal(b.CreatedAt)
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 	ctx, cancel := context.WithCancel(newContext())
 	defer cancel()
@@ -142,7 +153,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 	err = storage.UpdateClusterCapacities(
 		ctx,
 		[]ClusterCapacity{capacity1, cellCapacity1, cellCapacitySsd1},
-		zeroTime, // deleteBefore
+		zeroTime, // deleteOlderThan
 	)
 	require.NoError(t, err)
 
@@ -152,10 +163,12 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		types.DiskKind_DISK_KIND_HDD,
 	)
 	require.NoError(t, err)
-	require.ElementsMatch(t, []ClusterCapacity{
-		capacity1,
-		cellCapacity1,
-	}, capacities)
+	sort.Slice(capacities, func(i, j int) bool {
+		return capacities[i].CellID < capacities[j].CellID
+	})
+	require.Len(t, capacities, 2)
+	require.True(t, clusterCapacitiesAreEqual(capacities[0], capacity1))
+	require.True(t, clusterCapacitiesAreEqual(capacities[1], cellCapacity1))
 
 	capacities, err = storage.GetRecentClusterCapacities(
 		ctx,
@@ -163,12 +176,13 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		types.DiskKind_DISK_KIND_SSD,
 	)
 	require.NoError(t, err)
-	require.ElementsMatch(t, []ClusterCapacity{cellCapacitySsd1}, capacities)
+	require.Len(t, capacities, 1)
+	require.True(t, clusterCapacitiesAreEqual(cellCapacitySsd1, capacities[0]))
 
 	err = storage.UpdateClusterCapacities(
 		ctx,
 		[]ClusterCapacity{capacity2},
-		zeroTime, // deleteBefore
+		zeroTime, // deleteOlderThan
 	)
 	require.NoError(t, err)
 
@@ -178,15 +192,17 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		types.DiskKind_DISK_KIND_HDD,
 	)
 	require.NoError(t, err)
-	require.ElementsMatch(t, []ClusterCapacity{
-		capacity2,
-		cellCapacity1,
-	}, capacities)
+	sort.Slice(capacities, func(i, j int) bool {
+		return capacities[i].CellID < capacities[j].CellID
+	})
+	require.Len(t, capacities, 2)
+	require.True(t, clusterCapacitiesAreEqual(capacities[0], capacity2))
+	require.True(t, clusterCapacitiesAreEqual(capacities[1], cellCapacity1))
 
 	err = storage.UpdateClusterCapacities(
 		ctx,
 		[]ClusterCapacity{cellCapacity2, cellCapacitySsd2},
-		zeroTime, // deleteBefore
+		zeroTime, // deleteOlderThan
 	)
 	require.NoError(t, err)
 
@@ -196,10 +212,12 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		types.DiskKind_DISK_KIND_HDD,
 	)
 	require.NoError(t, err)
-	require.ElementsMatch(t, []ClusterCapacity{
-		capacity2,
-		cellCapacity2,
-	}, capacities)
+	sort.Slice(capacities, func(i, j int) bool {
+		return capacities[i].CellID < capacities[j].CellID
+	})
+	require.Len(t, capacities, 2)
+	require.True(t, clusterCapacitiesAreEqual(capacities[0], capacity2))
+	require.True(t, clusterCapacitiesAreEqual(capacities[1], cellCapacity2))
 
 	capacities, err = storage.GetRecentClusterCapacities(
 		ctx,
@@ -207,7 +225,8 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		types.DiskKind_DISK_KIND_SSD,
 	)
 	require.NoError(t, err)
-	require.ElementsMatch(t, []ClusterCapacity{cellCapacitySsd2}, capacities)
+	require.Len(t, capacities, 1)
+	require.True(t, clusterCapacitiesAreEqual(cellCapacitySsd1, capacities[0]))
 }
 
 func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
@@ -248,7 +267,8 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 		types.DiskKind_DISK_KIND_HDD,
 	)
 	require.NoError(t, err)
-	require.ElementsMatch(t, capacities, []ClusterCapacity{capacity1})
+	require.Len(t, capacities, 1)
+	require.True(t, clusterCapacitiesAreEqual(capacity1, capacities[0]))
 
 	err = storage.UpdateClusterCapacities(ctx, []ClusterCapacity{capacity2}, deleteOlderThan)
 	require.NoError(t, err)
@@ -259,9 +279,14 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 		types.DiskKind_DISK_KIND_HDD,
 	)
 	require.NoError(t, err)
-	require.ElementsMatch(t, capacities, []ClusterCapacity{capacity1, capacity2})
-
-	deleteOlderThan = time.Now()
+	sort.Slice(capacities, func(i, j int) bool {
+		return capacities[i].CellID < capacities[j].CellID
+	})
+	require.Len(t, capacities, 2)
+	require.True(t, clusterCapacitiesAreEqual(capacities[0], capacity1))
+	require.True(t, clusterCapacitiesAreEqual(capacities[1], capacity2))
+
+	deleteOlderThan = createdAt.Add(time.Hour)
 	createdAt = deleteOlderThan.Add(time.Hour)
 	capacity3 := ClusterCapacity{
 		ZoneID:     "zone-a",
@@ -281,5 +306,6 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 		types.DiskKind_DISK_KIND_HDD,
 	)
 	require.NoError(t, err)
-	require.ElementsMatch(t, capacities, []ClusterCapacity{capacity3})
+	require.Len(t, capacities, 1)
+	require.True(t, clusterCapacitiesAreEqual(capacity3, capacities[0]))
 }

From 7d3a70f4d2a81ea2f3172693c6a26b0abc25cafa Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 27 Nov 2025 00:05:28 +0000
Subject: [PATCH 05/11] fix capacity task test

---
 .../collect_cluster_capacity_task_test.go     | 23 ++++++----
 .../pkg/cells/storage/storage_test.go         | 42 ++++++++++++-------
 2 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
index 78ebeaa5f3b..770092b2381 100644
--- a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
@@ -21,17 +21,24 @@ import (
 
 func matchClusterCapacities(t *testing.T, want []cells_storage.ClusterCapacity) interface{} {
 	return mock.MatchedBy(func(actual []cells_storage.ClusterCapacity) bool {
-		require.Equal(t, len(want), len(actual))
+		if len(want) != len(actual) {
+			return false
+		}
+
 		for i := range want {
-			expectd := want[i]
+			expected := want[i]
 			got := actual[i]
 
-			require.Equal(t, expectd.ZoneID, got.ZoneID)
-			require.Equal(t, expectd.CellID, got.CellID)
-			require.Equal(t, expectd.Kind, got.Kind)
-			require.Equal(t, expectd.FreeBytes, got.FreeBytes)
-			require.Equal(t, expectd.TotalBytes, got.TotalBytes)
-			require.WithinDuration(t, time.Now(), got.CreatedAt, 15*time.Minute)
+			if expected.FreeBytes != got.FreeBytes ||
+				expected.TotalBytes != got.TotalBytes ||
+				expected.CellID != got.CellID || expected.Kind != got.Kind ||
+				expected.ZoneID != got.ZoneID {
+				return false
+			}
+
+			if expected.CreatedAt.Sub(got.CreatedAt).Abs() >= 15*time.Minute {
+				return false
+			}
 		}
 		return true
 	})
diff --git a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
index 5005d4d1bfd..e6841a9658c 100644
--- a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
@@ -74,10 +74,20 @@ func newStorage(
 
 // Uses time.Equal() to compare time.Time fields correctly, since == may fail
 // due to different internal representations.
-func clusterCapacitiesAreEqual(a ClusterCapacity, b ClusterCapacity) bool {
-	return a.CellID == b.CellID && a.ZoneID == b.ZoneID && a.Kind == b.Kind &&
-		a.FreeBytes == b.FreeBytes && a.TotalBytes == b.TotalBytes &&
-		a.CreatedAt.Equal(b.CreatedAt)
+func requireClusterCapacitiesEqual(
+	t *testing.T,
+	expected ClusterCapacity,
+	actual ClusterCapacity,
+) {
+
+	require.Equal(t, expected.CellID, actual.CellID)
+	require.Equal(t, expected.ZoneID, actual.ZoneID)
+	require.Equal(t, expected.Kind, actual.Kind)
+	require.Equal(t, expected.FreeBytes, actual.FreeBytes)
+	require.Equal(t, expected.TotalBytes, actual.TotalBytes)
+
+	fmt.Printf("exp: %v, act: %v\n", expected.CreatedAt, actual.CreatedAt)
+	require.True(t, expected.CreatedAt.Equal(actual.CreatedAt))
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -167,8 +177,8 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		return capacities[i].CellID < capacities[j].CellID
 	})
 	require.Len(t, capacities, 2)
-	require.True(t, clusterCapacitiesAreEqual(capacities[0], capacity1))
-	require.True(t, clusterCapacitiesAreEqual(capacities[1], cellCapacity1))
+	requireClusterCapacitiesEqual(t, capacities[0], capacity1)
+	requireClusterCapacitiesEqual(t, capacities[1], cellCapacity1)
 
 	capacities, err = storage.GetRecentClusterCapacities(
 		ctx,
@@ -177,7 +187,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 	)
 	require.NoError(t, err)
 	require.Len(t, capacities, 1)
-	require.True(t, clusterCapacitiesAreEqual(cellCapacitySsd1, capacities[0]))
+	requireClusterCapacitiesEqual(t, cellCapacitySsd1, capacities[0])
 
 	err = storage.UpdateClusterCapacities(
 		ctx,
@@ -196,8 +206,8 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		return capacities[i].CellID < capacities[j].CellID
 	})
 	require.Len(t, capacities, 2)
-	require.True(t, clusterCapacitiesAreEqual(capacities[0], capacity2))
-	require.True(t, clusterCapacitiesAreEqual(capacities[1], cellCapacity1))
+	requireClusterCapacitiesEqual(t, capacities[0], capacity2)
+	requireClusterCapacitiesEqual(t, capacities[1], cellCapacity1)
 
 	err = storage.UpdateClusterCapacities(
 		ctx,
@@ -216,8 +226,8 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		return capacities[i].CellID < capacities[j].CellID
 	})
 	require.Len(t, capacities, 2)
-	require.True(t, clusterCapacitiesAreEqual(capacities[0], capacity2))
-	require.True(t, clusterCapacitiesAreEqual(capacities[1], cellCapacity2))
+	requireClusterCapacitiesEqual(t, capacities[0], capacity2)
+	requireClusterCapacitiesEqual(t, capacities[1], cellCapacity2)
 
 	capacities, err = storage.GetRecentClusterCapacities(
 		ctx,
@@ -226,7 +236,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 	)
 	require.NoError(t, err)
 	require.Len(t, capacities, 1)
-	require.True(t, clusterCapacitiesAreEqual(cellCapacitySsd1, capacities[0]))
+	requireClusterCapacitiesEqual(t, cellCapacitySsd1, capacities[0])
 }
 
 func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
@@ -268,7 +278,7 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 	)
 	require.NoError(t, err)
 	require.Len(t, capacities, 1)
-	require.True(t, clusterCapacitiesAreEqual(capacity1, capacities[0]))
+	requireClusterCapacitiesEqual(t, capacity1, capacities[0])
 
 	err = storage.UpdateClusterCapacities(ctx, []ClusterCapacity{capacity2}, deleteOlderThan)
 	require.NoError(t, err)
@@ -283,8 +293,8 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 		return capacities[i].CellID < capacities[j].CellID
 	})
 	require.Len(t, capacities, 2)
-	require.True(t, clusterCapacitiesAreEqual(capacities[0], capacity1))
-	require.True(t, clusterCapacitiesAreEqual(capacities[1], capacity2))
+	requireClusterCapacitiesEqual(t, capacities[0], capacity1)
+	requireClusterCapacitiesEqual(t, capacities[1], capacity2)
 
 	deleteOlderThan = createdAt.Add(time.Hour)
 	createdAt = deleteOlderThan.Add(time.Hour)
@@ -307,5 +317,5 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 	)
 	require.NoError(t, err)
 	require.Len(t, capacities, 1)
-	require.True(t, clusterCapacitiesAreEqual(capacity3, capacities[0]))
+	requireClusterCapacitiesEqual(t, capacity3, capacities[0])
 }

From 32e1ccedc739b5142ae391c7fea02fc2a8f0eadb Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 27 Nov 2025 00:47:29 +0000
Subject: [PATCH 06/11] finally all tests are fixed

---
 .../collect_cluster_capacity_task_test.go     |  4 ---
 .../pkg/cells/storage/storage_test.go         | 32 +++++++++----------
 2 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
index 770092b2381..52fe0b29095 100644
--- a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
@@ -35,10 +35,6 @@ func matchClusterCapacities(t *testing.T, want []cells_storage.ClusterCapacity)
 				expected.ZoneID != got.ZoneID {
 				return false
 			}
-
-			if expected.CreatedAt.Sub(got.CreatedAt).Abs() >= 15*time.Minute {
-				return false
-			}
 		}
 		return true
 	})
diff --git a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
index e6841a9658c..b9e3a8a1b18 100644
--- a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
@@ -72,9 +72,9 @@ func newStorage(
 
 ////////////////////////////////////////////////////////////////////////////////
 
-// Uses time.Equal() to compare time.Time fields correctly, since == may fail
+// Uses WithinDuration to compare CreatedAt field correctly, since == may fail
 // due to different internal representations.
-func requireClusterCapacitiesEqual(
+func requireClusterCapacitiesAreEqual(
 	t *testing.T,
 	expected ClusterCapacity,
 	actual ClusterCapacity,
@@ -85,9 +85,7 @@ func requireClusterCapacitiesEqual(
 	require.Equal(t, expected.Kind, actual.Kind)
 	require.Equal(t, expected.FreeBytes, actual.FreeBytes)
 	require.Equal(t, expected.TotalBytes, actual.TotalBytes)
-
-	fmt.Printf("exp: %v, act: %v\n", expected.CreatedAt, actual.CreatedAt)
-	require.True(t, expected.CreatedAt.Equal(actual.CreatedAt))
+	require.WithinDuration(t, expected.CreatedAt, actual.CreatedAt, time.Microsecond)
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -177,8 +175,8 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		return capacities[i].CellID < capacities[j].CellID
 	})
 	require.Len(t, capacities, 2)
-	requireClusterCapacitiesEqual(t, capacities[0], capacity1)
-	requireClusterCapacitiesEqual(t, capacities[1], cellCapacity1)
+	requireClusterCapacitiesAreEqual(t, capacities[0], capacity1)
+	requireClusterCapacitiesAreEqual(t, capacities[1], cellCapacity1)
 
 	capacities, err = storage.GetRecentClusterCapacities(
 		ctx,
@@ -187,7 +185,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 	)
 	require.NoError(t, err)
 	require.Len(t, capacities, 1)
-	requireClusterCapacitiesEqual(t, cellCapacitySsd1, capacities[0])
+	requireClusterCapacitiesAreEqual(t, cellCapacitySsd1, capacities[0])
 
 	err = storage.UpdateClusterCapacities(
 		ctx,
@@ -206,8 +204,8 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		return capacities[i].CellID < capacities[j].CellID
 	})
 	require.Len(t, capacities, 2)
-	requireClusterCapacitiesEqual(t, capacities[0], capacity2)
-	requireClusterCapacitiesEqual(t, capacities[1], cellCapacity1)
+	requireClusterCapacitiesAreEqual(t, capacities[0], capacity2)
+	requireClusterCapacitiesAreEqual(t, capacities[1], cellCapacity1)
 
 	err = storage.UpdateClusterCapacities(
 		ctx,
@@ -226,8 +224,8 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 		return capacities[i].CellID < capacities[j].CellID
 	})
 	require.Len(t, capacities, 2)
-	requireClusterCapacitiesEqual(t, capacities[0], capacity2)
-	requireClusterCapacitiesEqual(t, capacities[1], cellCapacity2)
+	requireClusterCapacitiesAreEqual(t, capacities[0], capacity2)
+	requireClusterCapacitiesAreEqual(t, capacities[1], cellCapacity2)
 
 	capacities, err = storage.GetRecentClusterCapacities(
 		ctx,
@@ -236,7 +234,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 	)
 	require.NoError(t, err)
 	require.Len(t, capacities, 1)
-	requireClusterCapacitiesEqual(t, cellCapacitySsd1, capacities[0])
+	requireClusterCapacitiesAreEqual(t, cellCapacitySsd1, capacities[0])
 }
 
 func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
@@ -278,7 +276,7 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 	)
 	require.NoError(t, err)
 	require.Len(t, capacities, 1)
-	requireClusterCapacitiesEqual(t, capacity1, capacities[0])
+	requireClusterCapacitiesAreEqual(t, capacity1, capacities[0])
 
 	err = storage.UpdateClusterCapacities(ctx, []ClusterCapacity{capacity2}, deleteOlderThan)
 	require.NoError(t, err)
@@ -293,8 +291,8 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 		return capacities[i].CellID < capacities[j].CellID
 	})
 	require.Len(t, capacities, 2)
-	requireClusterCapacitiesEqual(t, capacities[0], capacity1)
-	requireClusterCapacitiesEqual(t, capacities[1], capacity2)
+	requireClusterCapacitiesAreEqual(t, capacities[0], capacity1)
+	requireClusterCapacitiesAreEqual(t, capacities[1], capacity2)
 
 	deleteOlderThan = createdAt.Add(time.Hour)
 	createdAt = deleteOlderThan.Add(time.Hour)
@@ -317,5 +315,5 @@ func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {
 	)
 	require.NoError(t, err)
 	require.Len(t, capacities, 1)
-	requireClusterCapacitiesEqual(t, capacity3, capacities[0])
+	requireClusterCapacitiesAreEqual(t, capacity3, capacities[0])
 }

From 07609d1964011d79b2b576e5199140dabed5a1cc Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 27 Nov 2025 08:33:37 +0000
Subject: [PATCH 07/11] set correct value in tests

---
 cloud/disk_manager/internal/pkg/cells/storage/storage_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
index b9e3a8a1b18..452b890f43e 100644
--- a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
@@ -234,7 +234,7 @@ func TestGetRecentClusterCapacitiesReturnsOnlyRecentValue(t *testing.T) {
 	)
 	require.NoError(t, err)
 	require.Len(t, capacities, 1)
-	requireClusterCapacitiesAreEqual(t, cellCapacitySsd1, capacities[0])
+	requireClusterCapacitiesAreEqual(t, cellCapacitySsd2, capacities[0])
 }
 
 func TestUpdateClusterCapacitiesDeletesRecordsBeforeTimestamp(t *testing.T) {

From cd382e12aacab517b1bbe3366d39d8045f4082e5 Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 27 Nov 2025 09:01:05 +0000
Subject: [PATCH 08/11] simplify matchClusterCapacities

---
 .../collect_cluster_capacity_task_test.go     | 21 ++++++++-----------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
index 52fe0b29095..24745e91a9c 100644
--- a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
@@ -19,23 +19,20 @@ import (
 
 ////////////////////////////////////////////////////////////////////////////////
 
-func matchClusterCapacities(t *testing.T, want []cells_storage.ClusterCapacity) interface{} {
+func matchClusterCapacities(t *testing.T, expected cells_storage.ClusterCapacity) interface{} {
 	return mock.MatchedBy(func(actual []cells_storage.ClusterCapacity) bool {
-		if len(want) != len(actual) {
+		if len(actual) != 1 {
 			return false
 		}
 
-		for i := range want {
-			expected := want[i]
-			got := actual[i]
-
-			if expected.FreeBytes != got.FreeBytes ||
-				expected.TotalBytes != got.TotalBytes ||
-				expected.CellID != got.CellID || expected.Kind != got.Kind ||
-				expected.ZoneID != got.ZoneID {
-				return false
-			}
+		got := actual[0]
+		if expected.FreeBytes != got.FreeBytes ||
+			expected.TotalBytes != got.TotalBytes ||
+			expected.CellID != got.CellID || expected.Kind != got.Kind ||
+			expected.ZoneID != got.ZoneID {
+			return false
 		}
+
 		return true
 	})
 }

From 420c108d36f74c67e02bf8a555c7e8a5f24d4639 Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 27 Nov 2025 09:05:00 +0000
Subject: [PATCH 09/11] return time check into matchClusterCapacities

---
 .../internal/pkg/cells/collect_cluster_capacity_task_test.go   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
index 24745e91a9c..30eba211ce2 100644
--- a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
@@ -29,7 +29,8 @@ func matchClusterCapacities(t *testing.T, expected cells_storage.ClusterCapacity
 		if expected.FreeBytes != got.FreeBytes ||
 			expected.TotalBytes != got.TotalBytes ||
 			expected.CellID != got.CellID || expected.Kind != got.Kind ||
-			expected.ZoneID != got.ZoneID {
+			expected.ZoneID != got.ZoneID ||
+			time.Now().Sub(got.CreatedAt).Abs() >= time.Minute {
 			return false
 		}
 

From 05f3b8c80fcf9accf6fc73ea97cdcc6a7b16af17 Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Thu, 27 Nov 2025 09:17:23 +0000
Subject: [PATCH 10/11] fix build

---
 .../collect_cluster_capacity_task_test.go     | 98 ++++++++-----------
 1 file changed, 42 insertions(+), 56 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
index 30eba211ce2..71e4dccc988 100644
--- a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
@@ -79,14 +79,12 @@ func TestCollectClusterCapacityTask(t *testing.T) {
 		mock.Anything, // ctx.
 		matchClusterCapacities(
 			t,
-			[]cells_storage.ClusterCapacity{
-				{
-					ZoneID:     "zone-a",
-					CellID:     "zone-a-cell1",
-					Kind:       types.DiskKind_DISK_KIND_SSD,
-					FreeBytes:  1024,
-					TotalBytes: 2048,
-				},
+			cells_storage.ClusterCapacity{
+				ZoneID:     "zone-a",
+				CellID:     "zone-a-cell1",
+				Kind:       types.DiskKind_DISK_KIND_SSD,
+				FreeBytes:  1024,
+				TotalBytes: 2048,
 			},
 		),
 		deleteOlderThanExpectation,
@@ -95,14 +93,12 @@ func TestCollectClusterCapacityTask(t *testing.T) {
 		mock.Anything, // ctx.
 		matchClusterCapacities(
 			t,
-			[]cells_storage.ClusterCapacity{
-				{
-					ZoneID:     "zone-a",
-					CellID:     "zone-a",
-					Kind:       types.DiskKind_DISK_KIND_SSD,
-					FreeBytes:  1024,
-					TotalBytes: 2048,
-				},
+			cells_storage.ClusterCapacity{
+				ZoneID:     "zone-a",
+				CellID:     "zone-a",
+				Kind:       types.DiskKind_DISK_KIND_SSD,
+				FreeBytes:  1024,
+				TotalBytes: 2048,
 			},
 		),
 		deleteOlderThanExpectation,
@@ -111,14 +107,12 @@ func TestCollectClusterCapacityTask(t *testing.T) {
 		mock.Anything, // ctx.
 		matchClusterCapacities(
 			t,
-			[]cells_storage.ClusterCapacity{
-				{
-					ZoneID:     "zone-b",
-					CellID:     "zone-b",
-					Kind:       types.DiskKind_DISK_KIND_SSD,
-					FreeBytes:  1024,
-					TotalBytes: 2048,
-				},
+			cells_storage.ClusterCapacity{
+				ZoneID:     "zone-b",
+				CellID:     "zone-b",
+				Kind:       types.DiskKind_DISK_KIND_SSD,
+				FreeBytes:  1024,
+				TotalBytes: 2048,
 			},
 		),
 		deleteOlderThanExpectation,
@@ -195,14 +189,12 @@ func TestCollectClusterCapacityFailureNbsReturnsError(t *testing.T) {
 		mock.Anything, // ctx.
 		matchClusterCapacities(
 			t,
-			[]cells_storage.ClusterCapacity{
-				{
-					ZoneID:     "zone-a",
-					CellID:     "zone-a",
-					Kind:       types.DiskKind_DISK_KIND_SSD,
-					FreeBytes:  1024,
-					TotalBytes: 2048,
-				},
+			cells_storage.ClusterCapacity{
+				ZoneID:     "zone-a",
+				CellID:     "zone-a",
+				Kind:       types.DiskKind_DISK_KIND_SSD,
+				FreeBytes:  1024,
+				TotalBytes: 2048,
 			},
 		),
 		mock.Anything, // deleteOlderThan.
@@ -267,14 +259,12 @@ func TestCollectClusterCapacityFailureStorageReturnsError(t *testing.T) {
 		mock.Anything, // ctx.
 		matchClusterCapacities(
 			t,
-			[]cells_storage.ClusterCapacity{
-				{
-					ZoneID:     "zone-a",
-					CellID:     "zone-a",
-					Kind:       types.DiskKind_DISK_KIND_SSD,
-					FreeBytes:  1024,
-					TotalBytes: 2048,
-				},
+			cells_storage.ClusterCapacity{
+				ZoneID:     "zone-a",
+				CellID:     "zone-a",
+				Kind:       types.DiskKind_DISK_KIND_SSD,
+				FreeBytes:  1024,
+				TotalBytes: 2048,
 			},
 		),
 		mock.Anything, // deleteOlderThan.
@@ -283,14 +273,12 @@ func TestCollectClusterCapacityFailureStorageReturnsError(t *testing.T) {
 		mock.Anything,
 		matchClusterCapacities(
 			t,
-			[]cells_storage.ClusterCapacity{
-				{
-					ZoneID:     "zone-a",
-					CellID:     "zone-a-cell1",
-					Kind:       types.DiskKind_DISK_KIND_SSD,
-					FreeBytes:  1024,
-					TotalBytes: 2048,
-				},
+			cells_storage.ClusterCapacity{
+				ZoneID:     "zone-a",
+				CellID:     "zone-a-cell1",
+				Kind:       types.DiskKind_DISK_KIND_SSD,
+				FreeBytes:  1024,
+				TotalBytes: 2048,
 			},
 		),
 		mock.Anything, // deleteOlderThan.
@@ -354,14 +342,12 @@ func TestCollectClusterCapacityOneCellHasAlreadyBeenProcessed(t *testing.T) {
 		mock.Anything, // ctx.
 		matchClusterCapacities(
 			t,
-			[]cells_storage.ClusterCapacity{
-				{
-					ZoneID:     "zone-a",
-					CellID:     "zone-a",
-					Kind:       types.DiskKind_DISK_KIND_SSD,
-					FreeBytes:  1024,
-					TotalBytes: 2048,
-				},
+			cells_storage.ClusterCapacity{
+				ZoneID:     "zone-a",
+				CellID:     "zone-a",
+				Kind:       types.DiskKind_DISK_KIND_SSD,
+				FreeBytes:  1024,
+				TotalBytes: 2048,
 			},
 		),
 		mock.Anything, // deleteOlderThan.

From a241e61e01434572d39e3d9c57d64d8666c758e3 Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Wed, 3 Dec 2025 11:33:25 +0000
Subject: [PATCH 11/11] change ClusterCapacity comparation logic

---
 .../cells/collect_cluster_capacity_task_test.go  | 16 ++++++++++------
 .../internal/pkg/cells/storage/storage_test.go   |  8 +++-----
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
index 71e4dccc988..f052686588a 100644
--- a/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/collect_cluster_capacity_task_test.go
@@ -26,15 +26,12 @@ func matchClusterCapacities(t *testing.T, expected cells_storage.ClusterCapacity
 		}
 
 		got := actual[0]
-		if expected.FreeBytes != got.FreeBytes ||
-			expected.TotalBytes != got.TotalBytes ||
-			expected.CellID != got.CellID || expected.Kind != got.Kind ||
-			expected.ZoneID != got.ZoneID ||
-			time.Now().Sub(got.CreatedAt).Abs() >= time.Minute {
+		if expected.CreatedAt.Sub(got.CreatedAt).Abs() >= time.Minute {
 			return false
 		}
 
-		return true
+		got.CreatedAt = expected.CreatedAt
+		return got == expected
 	})
 }
 
@@ -85,6 +82,7 @@ func TestCollectClusterCapacityTask(t *testing.T) {
 				Kind:       types.DiskKind_DISK_KIND_SSD,
 				FreeBytes:  1024,
 				TotalBytes: 2048,
+				CreatedAt:  time.Now(),
 			},
 		),
 		deleteOlderThanExpectation,
@@ -99,6 +97,7 @@ func TestCollectClusterCapacityTask(t *testing.T) {
 				Kind:       types.DiskKind_DISK_KIND_SSD,
 				FreeBytes:  1024,
 				TotalBytes: 2048,
+				CreatedAt:  time.Now(),
 			},
 		),
 		deleteOlderThanExpectation,
@@ -113,6 +112,7 @@ func TestCollectClusterCapacityTask(t *testing.T) {
 				Kind:       types.DiskKind_DISK_KIND_SSD,
 				FreeBytes:  1024,
 				TotalBytes: 2048,
+				CreatedAt:  time.Now(),
 			},
 		),
 		deleteOlderThanExpectation,
@@ -195,6 +195,7 @@ func TestCollectClusterCapacityFailureNbsReturnsError(t *testing.T) {
 				Kind:       types.DiskKind_DISK_KIND_SSD,
 				FreeBytes:  1024,
 				TotalBytes: 2048,
+				CreatedAt:  time.Now(),
 			},
 		),
 		mock.Anything, // deleteOlderThan.
@@ -265,6 +266,7 @@ func TestCollectClusterCapacityFailureStorageReturnsError(t *testing.T) {
 				Kind:       types.DiskKind_DISK_KIND_SSD,
 				FreeBytes:  1024,
 				TotalBytes: 2048,
+				CreatedAt:  time.Now(),
 			},
 		),
 		mock.Anything, // deleteOlderThan.
@@ -279,6 +281,7 @@ func TestCollectClusterCapacityFailureStorageReturnsError(t *testing.T) {
 				Kind:       types.DiskKind_DISK_KIND_SSD,
 				FreeBytes:  1024,
 				TotalBytes: 2048,
+				CreatedAt:  time.Now(),
 			},
 		),
 		mock.Anything, // deleteOlderThan.
@@ -348,6 +351,7 @@ func TestCollectClusterCapacityOneCellHasAlreadyBeenProcessed(t *testing.T) {
 				Kind:       types.DiskKind_DISK_KIND_SSD,
 				FreeBytes:  1024,
 				TotalBytes: 2048,
+				CreatedAt:  time.Now(),
 			},
 		),
 		mock.Anything, // deleteOlderThan.
diff --git a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
index 452b890f43e..c153df0e097 100644
--- a/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
+++ b/cloud/disk_manager/internal/pkg/cells/storage/storage_test.go
@@ -80,12 +80,10 @@ func requireClusterCapacitiesAreEqual(
 	actual ClusterCapacity,
 ) {
 
-	require.Equal(t, expected.CellID, actual.CellID)
-	require.Equal(t, expected.ZoneID, actual.ZoneID)
-	require.Equal(t, expected.Kind, actual.Kind)
-	require.Equal(t, expected.FreeBytes, actual.FreeBytes)
-	require.Equal(t, expected.TotalBytes, actual.TotalBytes)
 	require.WithinDuration(t, expected.CreatedAt, actual.CreatedAt, time.Microsecond)
+	actual.CreatedAt = expected.CreatedAt
+
+	require.Equal(t, expected, actual)
 }
 
 ////////////////////////////////////////////////////////////////////////////////