fix: Various issues found in testing Ruby client

tests/mutaterows_test.go

@@ -336,17 +336,18 @@ func TestMutateRows_Retry_ExponentialBackoff(t *testing.T) {
 
 	// 4b. Log the retry delays
 	origReq := <-recorder
-	firstRetry := <-recorder
-	secondRetry := <-recorder
-	thirdRetry := <-recorder
 
-	firstDelay := int(firstRetry.ts.UnixMilli() - origReq.ts.UnixMilli())
-	secondDelay := int(secondRetry.ts.UnixMilli() - firstRetry.ts.UnixMilli())
-	thirdDelay := int(thirdRetry.ts.UnixMilli() - secondRetry.ts.UnixMilli())
-
-	// Different clients may have different behaviors, we log the delays for informational purpose.
-	// Example: For the first retry delay, C++ client uses 100ms but Java client uses 10ms.
-	t.Logf("The three retry delays are: %dms, %dms, %dms", firstDelay, secondDelay, thirdDelay)
+	for n := 1; n < numRPCs; n += 1 {
+		select {
+		case retry := <- recorder:
+			delay := int(retry.ts.UnixMilli() - origReq.ts.UnixMilli())
+			// Different clients may have different behaviors, we log the delays for informational purpose.
+			// Example: For the first retry delay, C++ client uses 100ms but Java client uses 10ms.
+			t.Logf("Retry #%d delay: %dms", n, delay)
+		case <-time.After(500 * time.Millisecond):
+			t.Logf("Retry #%d: Timeout waiting for retry (expecting %d retries)", n, numRPCs - 1)
+		}
+	}
 }
 
 // TestMutateRows_Generic_MultiStreams tests that client can have multiple concurrent streams.
@@ -496,14 +497,19 @@ func TestMutateRows_Retry_WithRoutingCookie(t *testing.T) {
 	// 4b. Verify routing cookie is seen
 	// Ignore the first metadata which won't have the routing cookie
 	var _ = <-mdRecorder
-	// second metadata which comes from the retry attempt should have a routing cookie field
-	md1 := <-mdRecorder
-	val := md1["x-goog-cbt-cookie-test"]
-	assert.NotEmpty(t, val)
-	if len(val) == 0 {
-		return
-	}
-	assert.Equal(t, cookie, val[0])
+
+	select {
+	case md1 := <-mdRecorder:
+		// second metadata which comes from the retry attempt should have a routing cookie field
+		val := md1["x-goog-cbt-cookie-test"]
+		assert.NotEmpty(t, val)
+		if len(val) == 0 {
+			return
+		}
+		assert.Equal(t, cookie, val[0])
+	case <- time.After(100 * time.Millisecond):
+		t.Error("Timeout waiting for requests on recorder channel")
+	}
 }
 
 // TestMutateRows_Retry_WithRetryInfo tests that client is handling RetryInfo correctly.

tests/readrow_test.go

@@ -397,5 +397,7 @@ func TestReadRow_Retry_WithRetryInfo(t *testing.T) {
 	firstReqTs := firstReq.ts.Unix()
 	retryReqTs := retryReq.ts.Unix()
 
-	assert.True(t, retryReqTs-firstReqTs >= 2)
+	delta := retryReqTs - firstReqTs
+
+	assert.True(t, delta >= 2, fmt.Sprintf("Expected retry to happen after 2ms, got %dms", delta))
 }

tests/readrows_test.go

@@ -136,10 +136,8 @@ func TestReadRows_ReverseScans_FeatureFlag_Enabled(t *testing.T) {
 	md := <-mdRecords
 
 	ff, err := getClientFeatureFlags(md)
-
 	assert.Nil(t, err, "failed to decode client feature flags")
-
-	assert.True(t, ff.ReverseScans, "client does must enable ReverseScans feature flag")
+	assert.True(t, ff != nil && ff.ReverseScans, "client must enable ReverseScans feature flag")
 }
 
 // TestReadRows_NoRetry_OutOfOrderError_Reverse tests that client will fail on receiving out of order row keys for reverse scans.
@@ -892,8 +890,13 @@ func TestReadRows_Retry_WithRoutingCookie_MultipleErrorResponses(t *testing.T) {
 	// 4a. Verify that the read succeeds
 	checkResultOkStatus(t, res)
 	assert.Equal(t, 2, len(res.GetRows()))
-	assert.Equal(t, "row-01", string(res.Rows[0].Key))
-	assert.Equal(t, "row-05", string(res.Rows[1].Key))
+	if len(res.GetRows()) > 0 {
+		assert.Equal(t, "row-01", string(res.Rows[0].Key))
+	}
+
+	if len(res.GetRows()) > 1 {
+		assert.Equal(t, "row-05", string(res.Rows[1].Key))
+	}
 
 	// 4b. Verify routing cookie is seen
 	// Ignore the first metadata which won't have the routing cookie

tests/samplerowkeys_test.go

@@ -313,9 +313,13 @@ func TestSampleRowKeys_Retry_WithRetryInfo(t *testing.T) {
 
 	// 4b. Verify retry backoff time is correct
 	firstReq := <-recorder
-	retryReq := <-recorder
 	firstReqTs := firstReq.ts.Unix()
-	retryReqTs := retryReq.ts.Unix()
 
-	assert.True(t, retryReqTs-firstReqTs >= 2)
+	select {
+	case retryReq := <-recorder:
+		retryReqTs := retryReq.ts.Unix()
+		assert.True(t, retryReqTs-firstReqTs >= 2)
+	case <-time.After(2 * time.Second):
+		t.Error("Timeout waiting for retry request")
+	}
 }