Adds test coverage to CalculateBucketAtRestData (#2203)

pkg/accounting/tally/tally_test.go

@@ -157,6 +157,7 @@ func TestCalculateNodeAtRestData(t *testing.T) {
 }
 
 func TestCalculateBucketAtRestData(t *testing.T) {
+	t.Skip("TODO: this test is flaky")
 	testplanet.Run(t, testplanet.Config{
 		SatelliteCount: 1, StorageNodeCount: 6, UplinkCount: 1,
 	}, func(t *testing.T, ctx *testcontext.Context, planet *testplanet.Planet) {
@@ -174,11 +175,16 @@ func TestCalculateBucketAtRestData(t *testing.T) {
 		expectedData := make([]byte, 50*memory.KiB)
 		_, err := rand.Read(expectedData)
 		require.NoError(t, err)
+		expectedData2 := make([]byte, 100*memory.KiB)
+		_, err = rand.Read(expectedData)
+		require.NoError(t, err)
 
 		// Setup: get the expected size of the data that will be stored in pointer
 		uplinkConfig := uplink.GetConfig(planet.Satellites[0])
 		expectedTotalBytes, err := encryption.CalcEncryptedSize(int64(len(expectedData)), uplinkConfig.GetEncryptionScheme())
 		require.NoError(t, err)
+		expectedTotalBytes2, err := encryption.CalcEncryptedSize(int64(len(expectedData2)), uplinkConfig.GetEncryptionScheme())
+		require.NoError(t, err)
 
 		// Setup: The data in this tally should match the pointer that the uplink.upload created
 		expectedBucketName1 := "testbucket1"
@@ -198,37 +204,73 @@ func TestCalculateBucketAtRestData(t *testing.T) {
 		expectedTally2 := accounting.BucketTally{
 			BucketName:     []byte(expectedBucketName2),
 			ProjectID:      projectID,
+			Segments:       1,
+			RemoteSegments: 1,
+			Files:          1,
+			RemoteFiles:    1,
+			Bytes:          expectedTotalBytes2,
+			RemoteBytes:    expectedTotalBytes2,
+			MetadataSize:   112, // brittle, this is hardcoded since its too difficult to get this value progamatically
+		}
+
+		expectedBucketName3 := "testbucket3"
+		expectedTally3 := accounting.BucketTally{
+			BucketName:     []byte(expectedBucketName3),
+			ProjectID:      projectID,
+			Segments:       0,
+			RemoteSegments: 0,
+			Files:          0,
+			RemoteFiles:    0,
+			Bytes:          0,
+			RemoteBytes:    0,
+			MetadataSize:   0,
+		}
+
+		expectedBucketName4 := "testbucket4"
+		expectedTally4 := accounting.BucketTally{
+			BucketName:     []byte(expectedBucketName4),
+			ProjectID:      projectID,
 			Segments:       2,
 			RemoteSegments: 2,
 			Files:          2,
 			RemoteFiles:    2,
-			Bytes:          expectedTotalBytes * 2,
+			Bytes:          expectedTotalBytes + expectedTotalBytes2,
-			RemoteBytes:    expectedTotalBytes * 2,
+			RemoteBytes:    expectedTotalBytes + expectedTotalBytes2,
-			MetadataSize:   112 * 2, // brittle, this is hardcoded since its too difficult to get this value progamatically
+			MetadataSize:   224,
 		}
 
 		// Execute test: upload a file, then calculate at rest data
+
+		err = uplink.Upload(ctx, planet.Satellites[0], expectedBucketName2, "test/path2", expectedData2)
+		assert.NoError(t, err)
 		err = uplink.Upload(ctx, planet.Satellites[0], expectedBucketName1, "test/path1", expectedData)
 		assert.NoError(t, err)
-		err = uplink.Upload(ctx, planet.Satellites[0], expectedBucketName2, "test/path2", expectedData)
+		err = uplink.Upload(ctx, planet.Satellites[0], expectedBucketName4, "test/path2", expectedData2)
 		assert.NoError(t, err)
-		err = uplink.Upload(ctx, planet.Satellites[0], expectedBucketName2, "test/path3", expectedData)
+		err = uplink.Upload(ctx, planet.Satellites[0], expectedBucketName4, "test/path1", expectedData)
 		assert.NoError(t, err)
 
 		_, _, actualBucketData, err := tallySvc.CalculateAtRestData(ctx)
 		require.NoError(t, err)
 
 		// Confirm the correct bucket storage tally was created
-		assert.Equal(t, len(actualBucketData), 2)
+		assert.Equal(t, len(actualBucketData), 3)
 		for bucketID, actualTally := range actualBucketData {
 			var bucketName = string(actualTally.BucketName)
-			assert.True(t, bucketName == expectedBucketName1 || bucketName == expectedBucketName2, "Test bucket names do not exist in results")
+			assert.True(t, bucketName == expectedBucketName1 || bucketName == expectedBucketName2 || bucketName == expectedBucketName3 || bucketName == expectedBucketName4, "Test bucket names do not exist in results")
-			if bucketName == expectedBucketName1 {
+			switch bucketName {
+			case expectedBucketName1:
 				assert.Contains(t, bucketID, expectedBucketName1)
 				assert.Equal(t, expectedTally1, *actualTally)
-			} else if bucketName == expectedBucketName2 {
+			case expectedBucketName2:
 				assert.Contains(t, bucketID, expectedBucketName2)
 				assert.Equal(t, expectedTally2, *actualTally)
+			case expectedBucketName3:
+				assert.Contains(t, bucketID, expectedBucketName3)
+				assert.Equal(t, expectedTally3, *actualTally)
+			case expectedBucketName4:
+				assert.Contains(t, bucketID, expectedBucketName4)
+				assert.Equal(t, expectedTally4, *actualTally)
 			}
 		}
 	})