github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/grpc/xds/internal/balancer/priority/ignore_resolve_now_test.go (about)

     1  /*
     2   *
     3   * Copyright 2021 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  package priority
    20  
    21  import (
    22  	"context"
    23  	"testing"
    24  	"time"
    25  
    26  	"github.com/hxx258456/ccgo/grpc/balancer"
    27  	"github.com/hxx258456/ccgo/grpc/balancer/roundrobin"
    28  	"github.com/hxx258456/ccgo/grpc/internal/testutils"
    29  	"github.com/hxx258456/ccgo/grpc/resolver"
    30  )
    31  
    32  const resolveNowBalancerName = "test-resolve-now-balancer"
    33  
    34  var resolveNowBalancerCCCh = testutils.NewChannel()
    35  
    36  type resolveNowBalancerBuilder struct {
    37  	balancer.Builder
    38  }
    39  
    40  func (r *resolveNowBalancerBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
    41  	resolveNowBalancerCCCh.Send(cc)
    42  	return r.Builder.Build(cc, opts)
    43  }
    44  
    45  func (r *resolveNowBalancerBuilder) Name() string {
    46  	return resolveNowBalancerName
    47  }
    48  
    49  func init() {
    50  	balancer.Register(&resolveNowBalancerBuilder{
    51  		Builder: balancer.Get(roundrobin.Name),
    52  	})
    53  }
    54  
    55  func (s) TestIgnoreResolveNowBalancerBuilder(t *testing.T) {
    56  	resolveNowBB := balancer.Get(resolveNowBalancerName)
    57  	// Create a build wrapper, but will not ignore ResolveNow().
    58  	ignoreResolveNowBB := newIgnoreResolveNowBalancerBuilder(resolveNowBB, false)
    59  
    60  	cc := testutils.NewTestClientConn(t)
    61  	tb := ignoreResolveNowBB.Build(cc, balancer.BuildOptions{})
    62  	defer tb.Close()
    63  
    64  	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
    65  	defer cancel()
    66  	// This is the balancer.ClientConn that the inner resolverNowBalancer is
    67  	// built with.
    68  	balancerCCI, err := resolveNowBalancerCCCh.Receive(ctx)
    69  	if err != nil {
    70  		t.Fatalf("timeout waiting for ClientConn from balancer builder")
    71  	}
    72  	balancerCC := balancerCCI.(balancer.ClientConn)
    73  
    74  	// Call ResolveNow() on the CC, it should be forwarded.
    75  	balancerCC.ResolveNow(resolver.ResolveNowOptions{})
    76  	select {
    77  	case <-cc.ResolveNowCh:
    78  	case <-time.After(time.Second):
    79  		t.Fatalf("timeout waiting for ResolveNow()")
    80  	}
    81  
    82  	// Update ignoreResolveNow to true, call ResolveNow() on the CC, they should
    83  	// all be ignored.
    84  	ignoreResolveNowBB.updateIgnoreResolveNow(true)
    85  	for i := 0; i < 5; i++ {
    86  		balancerCC.ResolveNow(resolver.ResolveNowOptions{})
    87  	}
    88  	select {
    89  	case <-cc.ResolveNowCh:
    90  		t.Fatalf("got unexpected ResolveNow() call")
    91  	case <-time.After(time.Millisecond * 100):
    92  	}
    93  
    94  	// Update ignoreResolveNow to false, new ResolveNow() calls should be
    95  	// forwarded.
    96  	ignoreResolveNowBB.updateIgnoreResolveNow(false)
    97  	balancerCC.ResolveNow(resolver.ResolveNowOptions{})
    98  	select {
    99  	case <-cc.ResolveNowCh:
   100  	case <-time.After(time.Second):
   101  		t.Fatalf("timeout waiting for ResolveNow()")
   102  	}
   103  }