feat: Add retry info conformance test

tests/mock_helper.go

@@ -24,10 +24,12 @@ import (
 
 	"github.com/golang/protobuf/ptypes/wrappers"
 	btpb "google.golang.org/genproto/googleapis/bigtable/v2"
+	"google.golang.org/genproto/googleapis/rpc/errdetails"
 	"google.golang.org/genproto/googleapis/rpc/status"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	gs "google.golang.org/grpc/status"
+	drpb "google.golang.org/protobuf/types/known/durationpb"
 )
 
 var rowKeyPrefixRegex = regexp.MustCompile("^op[0-9]+-")
@@ -177,6 +179,16 @@ func mockReadRowsFnWithMetadata(recorder chan<- *readRowsReqRecord, mdRecorder c
 					trailer := metadata.Pairs("x-goog-cbt-cookie-test", action.routingCookie)
 					srv.SetTrailer(trailer)
 				}
+				// TODO check for feature flag
+				if action.retryInfo != "" {
+					st := gs.New(action.rpcError, "ReadRows failed")
+					delay, _ := time.ParseDuration(action.retryInfo)
+					retryInfo := &errdetails.RetryInfo{
+						RetryDelay: drpb.New(delay),
+					}
+					st, err = st.WithDetails(retryInfo)
+					return st.Err()
+				}
 				return gs.Error(action.rpcError, "ReadRows failed")
 			}
 

tests/readrows_test.go

@@ -841,4 +841,50 @@ func TestReadRows_Retry_WithRoutingCookie(t *testing.T) {
 	var _ = <-recorder
 	retryReq := <-recorder
 	assert.True(t, cmp.Equal(retryReq.req.GetRows().GetRowRanges()[0].StartKey, &btpb.RowRange_StartKeyOpen{StartKeyOpen: []byte("row-01")}))
+}
+
+// TestReadRows_Retry_WithRetryInfo tests that RetryInfo is handled correctly by the client.
+func TestReadRows_Retry_WithRetryInfo(t *testing.T) {
+	// 1. Instantiate the mock server
+	sequence := []*readRowsAction{
+		&readRowsAction{
+			chunks: []chunkData{
+				dummyChunkData("row-01", "v1", Commit)}},
+		&readRowsAction{rpcError: codes.Unavailable, retryInfo: "2s"}, // Error with retry info
+		&readRowsAction{
+			chunks: []chunkData{
+				dummyChunkData("row-05", "v5", Commit)}},
+	}
+	server := initMockServer(t)
+
+	recorder := make(chan *readRowsReqRecord, 2)
+	server.ReadRowsFn = mockReadRowsFn(recorder, sequence)
+
+	// 2. Build the request to test proxy
+	req := testproxypb.ReadRowsRequest{
+		ClientId: t.Name(),
+		Request: &btpb.ReadRowsRequest{
+			TableName: buildTableName("table"),
+		},
+	}
+
+	// 3. Perform the operation via test proxy
+	res := doReadRowsOp(t, server, &req, nil)
+
+	// 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))
+
+	// 4b. Verify retry request is correct
+	firstReq := <-recorder
+	retryReq := <-recorder
+	assert.True(t, cmp.Equal(retryReq.req.GetRows().GetRowRanges()[0].StartKey, &btpb.RowRange_StartKeyOpen{StartKeyOpen: []byte("row-01")}))
+
+	// 4c. Verify retry backoff time is correct
+	firstReqTs := firstReq.ts.Unix()
+	retryReqTs := retryReq.ts.Unix()
+
+	assert.True(t, retryReqTs - firstReqTs >= 2)
 }

tests/test_datatype.go

@@ -93,12 +93,15 @@ type chunkData struct {
 //     Effect: server will return an error after delay. chunks that are specified in the same action will be ignored.
 //  5. readRowsAction{rpcError: error, routingCookie: cookie}
 //     Effect: server will return an error with the routing cookie. Retry attempt header should have this cookie.
-//  6. To have a response stream with/without errors, a sequence of actions should be constructed.
+//  6. readRowsAction{rpcError: error, retryInfo: delay}
+//     Effect: server will return an error with RetryInfo which has the specific delay.
+//  7. To have a response stream with/without errors, a sequence of actions should be constructed.
 type readRowsAction struct {
 	chunks        []chunkData
 	rpcError      codes.Code
 	delayStr      string  // "" means zero delay; follow https://pkg.go.dev/time#ParseDuration otherwise
 	routingCookie string
+	retryInfo     string  // "" means no RetryInfo will be attached in the error status
 }
 func (a *readRowsAction) Validate() {
 	for _, chunk := range a.chunks {