test(storage): update tests for Writer.ProgressFunc (#12784)

tritone · web-flow · commit 392940575591 · 2025-08-27T18:10:40.000-04:00
* fix(storage): call progress func only on flush

Writer.ProgressFunc should be called only on the flush of a chunk
buffer, not on takeover or finalize/close.

Fixes a behavior change introduced by the gRPC writer refactor.

* update tests

* update integration test
diff --git a/storage/client_test.go b/storage/client_test.go
@@ -956,35 +956,62 @@ func TestOpenWriterEmulated(t *testing.T) {
 	})
 }
 
-func TestWriterOneshotNoProgressReportEmulated(t *testing.T) {
+// Test that Writer.ProgressFunc is called at the expected intervals, showing
+// consistency across HTTP and gRPC.
+func TestWriterProgressFuncEmulated(t *testing.T) {
 	transportClientTest(context.Background(), t, func(t *testing.T, ctx context.Context, project, bucket string, client storageClient) {
 		_, err := client.CreateBucket(ctx, project, bucket, &BucketAttrs{
 			Name: bucket,
 		}, nil)
 		if err != nil {
 			t.Fatalf("client.CreateBucket: %v", err)
 		}
-		// Oneshot uploads can be forced with either a chunksize of 0, or a
-		// chunksize larger than the data size.
-		for _, chunksize := range []int{0, 16 * MiB} {
-			t.Run(fmt.Sprintf("data size %d chunksize %d", 3*MiB, chunksize), func(t *testing.T) {
-				prefix := time.Now().Nanosecond()
-				objName := fmt.Sprintf("%d-object-%d", prefix, time.Now().Nanosecond())
-
+		cases := []struct {
+			name          string
+			chunkSize     int
+			expectedCalls []int64
+		}{
+			{
+				name:          "one shot, zero chunkSize",
+				chunkSize:     0,
+				expectedCalls: []int64{},
+			},
+			{
+				name:          "one shot, chunk larger than data",
+				chunkSize:     16 * MiB,
+				expectedCalls: []int64{},
+			},
+			{
+				name:          "resumable, obj size evenly divisible by chunkSize",
+				chunkSize:     1 * MiB,
+				expectedCalls: []int64{MiB, 2 * MiB, 3 * MiB},
+			},
+			{
+				name:          "resumable, obj size not divisible by chunkSize",
+				chunkSize:     2 * MiB,
+				expectedCalls: []int64{2 * MiB, 3 * MiB},
+			},
+		}
+		for _, c := range cases {
+			t.Run(c.name, func(t *testing.T) {
 				vc := &Client{tc: client}
-				obj := vc.Bucket(bucket).Object(objName)
+				obj := vc.Bucket(bucket).Object(c.name)
+
 				w := obj.NewWriter(ctx)
-				w.ChunkSize = chunksize
-				progressCalls := 0
-				w.ProgressFunc = func(int64) { progressCalls++ }
+				w.ChunkSize = c.chunkSize
+				var gotCalls []int64
+				w.ProgressFunc = func(off int64) {
+					gotCalls = append(gotCalls, off)
+				}
+
 				if _, err := w.Write(randomBytes3MiB); err != nil {
-					t.Fatalf("writer.Write: %v", err)
+					t.Fatalf("writing data: %v", err)
 				}
 				if err := w.Close(); err != nil {
-					t.Fatalf("writer.Close: %v", err)
+					t.Fatalf("closing writer: %v", err)
 				}
-				if progressCalls != 0 {
-					t.Errorf("ProgressFunc was called %d times, expected 0", progressCalls)
+				if slices.Compare(gotCalls, c.expectedCalls) != 0 {
+					t.Errorf("progressFunc: got calls at %v, want %v", gotCalls, c.expectedCalls)
 				}
 				attrs, err := obj.Attrs(ctx)
 				if err != nil {
@@ -996,6 +1023,7 @@ func TestWriterOneshotNoProgressReportEmulated(t *testing.T) {
 			})
 		}
 	})
+
 }
 
 func TestOpenAppendableWriterEmulated(t *testing.T) {
diff --git a/storage/integration_test.go b/storage/integration_test.go
@@ -3414,7 +3414,7 @@ func TestIntegration_WriterAppendTakeover(t *testing.T) {
 				opts: &AppendableWriterOpts{
 					ChunkSize: 4 * MiB,
 				},
-				checkProgressOffsets: []int64{7 * MiB, 8 * MiB},
+				checkProgressOffsets: []int64{3 * MiB, 7 * MiB, 8 * MiB, 9 * MiB},
 			},
 			{
 				name:                "middle chunk takeover, small flush",
@@ -3456,7 +3456,7 @@ func TestIntegration_WriterAppendTakeover(t *testing.T) {
 				opts: &AppendableWriterOpts{
 					ChunkSize: 4 * MiB,
 				},
-				checkProgressOffsets: []int64{4 * MiB, 8 * MiB},
+				checkProgressOffsets: []int64{0, 4 * MiB, 8 * MiB, 9 * MiB},
 			},
 			{
 				name:           "last byte takeover",
@@ -3574,12 +3574,11 @@ func TestIntegration_WriterAppendTakeover(t *testing.T) {
 					t.Errorf("got object finalized at %v, want unfinalized", attrs.Finalized)
 				}
 				// Check ProgressFunc was called if applicable
-				// Skip this until progressFunc bug is fixed.
-				// if tc.checkProgressOffsets != nil {
-				// 	if !slices.Equal(gotOffsets, tc.checkProgressOffsets) {
-				// 		t.Errorf("progressFunc calls: got %v, want %v", gotOffsets, tc.checkProgressOffsets)
-				// 	}
-				// }
+				if tc.checkProgressOffsets != nil {
+					if !slices.Equal(gotOffsets, tc.checkProgressOffsets) {
+						t.Errorf("progressFunc calls: got %v, want %v", gotOffsets, tc.checkProgressOffsets)
+					}
+				}
 				if w2.Attrs() == nil {
 					t.Fatalf("takeover writer attrs: expected attrs, got nil")
 				}
