@@ -44,6 +44,7 @@ import (
44
44
"golang.org/x/sync/errgroup"
45
45
"google.golang.org/api/iterator"
46
46
"google.golang.org/api/option"
47
+ "google.golang.org/genproto/googleapis/rpc/errdetails"
47
48
"google.golang.org/grpc/codes"
48
49
"google.golang.org/grpc/encoding/gzip"
49
50
"google.golang.org/grpc/metadata"
@@ -659,12 +660,29 @@ func TestClient_Single(t *testing.T) {
659
660
660
661
func TestClient_Single_Unavailable (t * testing.T ) {
661
662
t .Parallel ()
662
- err := testSingleQuery (t , status . Error (codes .Unavailable , "Temporary unavailable" ))
663
+ err := testSingleQuery (t , serverErrorWithMinimalRetryDelay (codes .Unavailable , "Temporary unavailable" ))
663
664
if err != nil {
664
665
t .Fatal (err )
665
666
}
666
667
}
667
668
669
+ func TestClient_Single_ResourceExhausted (t * testing.T ) {
670
+ t .Parallel ()
671
+ err := testSingleQuery (t , serverErrorWithMinimalRetryDelay (codes .ResourceExhausted , "Temporary server overload" ))
672
+ if err != nil {
673
+ t .Fatal (err )
674
+ }
675
+ }
676
+
677
+ func serverErrorWithMinimalRetryDelay (code codes.Code , msg string ) error {
678
+ st := status .New (code , msg )
679
+ retry := & errdetails.RetryInfo {
680
+ RetryDelay : durationpb .New (time .Nanosecond ),
681
+ }
682
+ st , _ = st .WithDetails (retry )
683
+ return st .Err ()
684
+ }
685
+
668
686
func TestClient_Single_InvalidArgument (t * testing.T ) {
669
687
t .Parallel ()
670
688
err := testSingleQuery (t , status .Error (codes .InvalidArgument , "Invalid argument" ))
@@ -1174,8 +1192,18 @@ func checkReqsForTransactionOptions(t *testing.T, server InMemSpannerServer, txo
1174
1192
1175
1193
func testSingleQuery (t * testing.T , serverError error ) error {
1176
1194
ctx := context .Background ()
1177
- server , client , teardown := setupMockedTestServer ( t )
1195
+ server , client , teardown := setupMockedTestServerWithConfig ( t , ClientConfig { SessionPoolConfig : SessionPoolConfig { MinOpened : 1 }} )
1178
1196
defer teardown ()
1197
+ // Wait until all sessions have been created, so we know that those requests will not interfere with the test.
1198
+ sp := client .idleSessions
1199
+ waitFor (t , func () error {
1200
+ sp .mu .Lock ()
1201
+ defer sp .mu .Unlock ()
1202
+ if uint64 (sp .idleList .Len ()) != sp .MinOpened {
1203
+ return fmt .Errorf ("num open sessions mismatch.\n Got: %d\n Want: %d" , sp .numOpened , sp .MinOpened )
1204
+ }
1205
+ return nil
1206
+ })
1179
1207
1180
1208
if serverError != nil {
1181
1209
server .TestSpanner .SetError (serverError )
0 commit comments