@@ -45,6 +45,18 @@ func compareErrors(err1, err2 error) bool {
45
45
return true
46
46
}
47
47
48
+ var _ description.ServerSelector = & mockServerSelector {}
49
+
50
+ type mockServerSelector struct {
51
+ selectServerCalls atomic.Int64
52
+ }
53
+
54
+ func (m * mockServerSelector ) SelectServer (description.Topology , []description.Server ) ([]description.Server , error ) {
55
+ m .selectServerCalls .Add (1 )
56
+
57
+ return nil , nil
58
+ }
59
+
48
60
func TestServerSelection (t * testing.T ) {
49
61
var selectFirst serverselector.Func = func (_ description.Topology , candidates []description.Server ) ([]description.Server , error ) {
50
62
if len (candidates ) == 0 {
@@ -263,6 +275,30 @@ func TestServerSelection(t *testing.T) {
263
275
_ , err = topo .SelectServer (context .Background (), & serverselector.Write {})
264
276
assert .Equal (t , ErrSubscribeAfterClosed , err , "expected error %v, got %v" , ErrSubscribeAfterClosed , err )
265
277
})
278
+ t .Run ("if no servers are suitable, block on topology updates" , func (t * testing.T ) {
279
+ // Create a connected Topology with no selectable servers.
280
+ topo , err := New (nil )
281
+ require .NoError (t , err )
282
+ atomic .StoreInt64 (& topo .state , topologyConnected )
283
+
284
+ mss := & mockServerSelector {}
285
+
286
+ ctx , cancel := context .WithTimeout (context .Background (), 50 * time .Millisecond )
287
+ defer cancel ()
288
+
289
+ _ , err = topo .SelectServer (ctx , mss )
290
+ assert .ErrorIs (t , err , context .DeadlineExceeded , "expected context deadline exceeded error" )
291
+
292
+ // Expect SelectServer to be called twice: once for the fast path and
293
+ // once to select from the topology updates subscription.
294
+ //
295
+ // Note: The second call is due to Topology.Subscript pre-populating the
296
+ // channel with the current topology. It's not clear what the purpose of
297
+ // that behavior is. The main goal of this assertion is to make sure the
298
+ // subscription path blocks on updates and doesn't turn into a busy
299
+ // wait.
300
+ assert .Equal (t , int64 (2 ), mss .selectServerCalls .Load (), "expected SelectServer to be called once" )
301
+ })
266
302
}
267
303
268
304
func TestSessionTimeout (t * testing.T ) {
0 commit comments