test(datastore): Retry aggregation query testcase on failure (#8487)

* test(datastore): remove explicit sleep
googleapis · Aug 29, 2023 · b5000f6 · b5000f6
1 parent 38a040e
commit b5000f6
Showing 1 changed file with 13 additions and 11 deletions.
diff --git a/datastore/integration_test.go b/datastore/integration_test.go
@@ -874,18 +874,20 @@ func TestIntegration_AggregationQueries(t *testing.T) {
 	}
 
 	for _, testCase := range testCases {
-		gotAggResult, gotErr := client.RunAggregationQuery(ctx, testCase.aggQuery)
-		gotFailure := gotErr != nil
+		testutil.Retry(t, 10, time.Second, func(r *testutil.R) {
+			gotAggResult, gotErr := client.RunAggregationQuery(ctx, testCase.aggQuery)
+			gotFailure := gotErr != nil
 
-		if gotFailure != testCase.wantFailure ||
-			(gotErr != nil && !strings.Contains(gotErr.Error(), testCase.wantErrMsg)) {
-			t.Errorf("%q: Mismatch in error got: %v, want: %q", testCase.desc, gotErr, testCase.wantErrMsg)
-			continue
-		}
-		if !reflect.DeepEqual(gotAggResult, testCase.wantAggResult) {
-			t.Errorf("%q: Mismatch in aggregation result got: %v, want: %v", testCase.desc, gotAggResult, testCase.wantAggResult)
-			continue
-		}
+			if gotFailure != testCase.wantFailure ||
+				(gotErr != nil && !strings.Contains(gotErr.Error(), testCase.wantErrMsg)) {
+				r.Errorf("%q: Mismatch in error got: %v, want: %q", testCase.desc, gotErr, testCase.wantErrMsg)
+				return
+			}
+			if gotErr == nil && !reflect.DeepEqual(gotAggResult, testCase.wantAggResult) {
+				r.Errorf("%q: Mismatch in aggregation result got: %v, want: %v", testCase.desc, gotAggResult, testCase.wantAggResult)
+				return
+			}
+		})
 	}
 
 }