test(spanner): fix failed TestRsdBlockingStates test (#2597)

googleapis · Jul 14, 2020 · d1ef23f · d1ef23f
1 parent 1257173
commit d1ef23f
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 26 deletions.
diff --git a/spanner/internal/testutil/inmem_spanner_server.go b/spanner/internal/testutil/inmem_spanner_server.go
@@ -117,33 +117,35 @@ func (s *StatementResult) ToPartialResultSets(resumeToken []byte) (result []*spa
 	}
 
 	totalRows := uint64(len(s.ResultSet.Rows))
-	for {
-		rowCount := min(totalRows-startIndex, uint64(MaxRowsPerPartialResultSet))
-		rows := s.ResultSet.Rows[startIndex : startIndex+rowCount]
-		values := make([]*structpb.Value,
-			len(rows)*len(s.ResultSet.Metadata.RowType.Fields))
-		var idx int
-		for _, row := range rows {
-			for colIdx := range s.ResultSet.Metadata.RowType.Fields {
-				values[idx] = row.Values[colIdx]
-				idx++
+	if totalRows > 0 {
+		for {
+			rowCount := min(totalRows-startIndex, uint64(MaxRowsPerPartialResultSet))
+			rows := s.ResultSet.Rows[startIndex : startIndex+rowCount]
+			values := make([]*structpb.Value,
+				len(rows)*len(s.ResultSet.Metadata.RowType.Fields))
+			var idx int
+			for _, row := range rows {
+				for colIdx := range s.ResultSet.Metadata.RowType.Fields {
+					values[idx] = row.Values[colIdx]
+					idx++
+				}
+			}
+			var rt []byte
+			if len(s.ResumeTokens) == 0 {
+				rt = EncodeResumeToken(startIndex + rowCount)
+			} else {
+				rt = s.ResumeTokens[startIndex]
+			}
+			result = append(result, &spannerpb.PartialResultSet{
+				Metadata:    s.ResultSet.Metadata,
+				Values:      values,
+				ResumeToken: rt,
+			})
+
+			startIndex += rowCount
+			if startIndex == totalRows {
+				break
 			}
-		}
-		var rt []byte
-		if len(s.ResumeTokens) == 0 {
-			rt = EncodeResumeToken(startIndex + rowCount)
-		} else {
-			rt = s.ResumeTokens[startIndex]
-		}
-		result = append(result, &spannerpb.PartialResultSet{
-			Metadata:    s.ResultSet.Metadata,
-			Values:      values,
-			ResumeToken: rt,
-		})
-
-		startIndex += rowCount
-		if startIndex == totalRows {
-			break
 		}
 	}
 	return result, nil

diff --git a/spanner/read_test.go b/spanner/read_test.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -1140,14 +1141,17 @@ func TestRsdBlockingStates(t *testing.T) {
 			if err != nil {
 				t.Fatalf("failed to set up a result for a statement: %v", err)
 			}
+			var mutex = &sync.Mutex{}
 			var rs []*sppb.PartialResultSet
 			go func() {
 				for {
 					if !r.next() {
 						// Note that r.Next also exits on context cancel/timeout.
 						return
 					}
+					mutex.Lock()
 					rs = append(rs, r.get())
+					mutex.Unlock()
 				}
 			}()
 			// Verify that resumableStreamDecoder reaches expected state.
@@ -1160,6 +1164,8 @@ func TestRsdBlockingStates(t *testing.T) {
 				}
 				// Check if resumableStreamDecoder returns expected array of
 				// PartialResultSets.
+				mutex.Lock()
+				defer mutex.Unlock()
 				if !testEqual(rs, test.want) {
 					t.Fatalf("received PartialResultSets: \n%v\n, want \n%v\n", rs, test.want)
 				}