@@ -30,6 +30,7 @@ import (
3030 "google.golang.org/grpc/connectivity"
3131 "google.golang.org/grpc/credentials/insecure"
3232 "google.golang.org/grpc/internal/channelz"
33+ "google.golang.org/grpc/internal/envconfig"
3334 "google.golang.org/grpc/internal/grpcrand"
3435 "google.golang.org/grpc/internal/stubserver"
3536 "google.golang.org/grpc/internal/testutils"
@@ -381,7 +382,10 @@ func (s) TestPickFirst_StickyTransientFailure(t *testing.T) {
381382 wg .Wait ()
382383}
383384
385+ // Tests the PF LB policy with shuffling enabled.
384386func (s ) TestPickFirst_ShuffleAddressList (t * testing.T ) {
387+ defer func (old bool ) { envconfig .PickFirstLBConfig = old }(envconfig .PickFirstLBConfig )
388+ envconfig .PickFirstLBConfig = true
385389 const serviceConfig = `{"loadBalancingConfig": [{"pick_first":{ "shuffleAddressList": true }}]}`
386390
387391 // Install a shuffler that always reverses two entries.
@@ -390,6 +394,7 @@ func (s) TestPickFirst_ShuffleAddressList(t *testing.T) {
390394 grpcrand .Shuffle = func (n int , f func (int , int )) {
391395 if n != 2 {
392396 t .Errorf ("Shuffle called with n=%v; want 2" , n )
397+ return
393398 }
394399 f (0 , 1 ) // reverse the two addresses
395400 }
@@ -431,3 +436,42 @@ func (s) TestPickFirst_ShuffleAddressList(t *testing.T) {
431436 t .Fatal (err )
432437 }
433438}
439+
440+ // Tests the PF LB policy with the environment variable support of address list
441+ // shuffling disabled.
442+ func (s ) TestPickFirst_ShuffleAddressListDisabled (t * testing.T ) {
443+ defer func (old bool ) { envconfig .PickFirstLBConfig = old }(envconfig .PickFirstLBConfig )
444+ envconfig .PickFirstLBConfig = false
445+ const serviceConfig = `{"loadBalancingConfig": [{"pick_first":{ "shuffleAddressList": true }}]}`
446+
447+ // Install a shuffler that always reverses two entries.
448+ origShuf := grpcrand .Shuffle
449+ defer func () { grpcrand .Shuffle = origShuf }()
450+ grpcrand .Shuffle = func (n int , f func (int , int )) {
451+ if n != 2 {
452+ t .Errorf ("Shuffle called with n=%v; want 2" , n )
453+ return
454+ }
455+ f (0 , 1 ) // reverse the two addresses
456+ }
457+
458+ // Set up our backends.
459+ cc , r , backends := setupPickFirst (t , 2 )
460+ addrs := stubBackendsToResolverAddrs (backends )
461+
462+ ctx , cancel := context .WithTimeout (context .Background (), defaultTestTimeout )
463+ defer cancel ()
464+
465+ // Send a config with shuffling enabled. This will reverse the addresses,
466+ // so we should connect to backend 1 if shuffling is supported. However
467+ // with it disabled at the start of the test, we will connect to backend 0
468+ // instead.
469+ shufState := resolver.State {
470+ ServiceConfig : parseServiceConfig (t , r , serviceConfig ),
471+ Addresses : []resolver.Address {addrs [0 ], addrs [1 ]},
472+ }
473+ r .UpdateState (shufState )
474+ if err := pickfirst .CheckRPCsToBackend (ctx , cc , addrs [0 ]); err != nil {
475+ t .Fatal (err )
476+ }
477+ }
0 commit comments