google.golang.org/grpc@v1.62.1/internal/testutils/state.go (about)

     1  /*
     2   *
     3   * Copyright 2023 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 testutils
    20  
    21  import (
    22  	"context"
    23  	"testing"
    24  
    25  	"google.golang.org/grpc/connectivity"
    26  )
    27  
    28  // A StateChanger reports state changes, e.g. a grpc.ClientConn.
    29  type StateChanger interface {
    30  	// Connect begins connecting the StateChanger.
    31  	Connect()
    32  	// GetState returns the current state of the StateChanger.
    33  	GetState() connectivity.State
    34  	// WaitForStateChange returns true when the state becomes s, or returns
    35  	// false if ctx is canceled first.
    36  	WaitForStateChange(ctx context.Context, s connectivity.State) bool
    37  }
    38  
    39  // StayConnected makes sc stay connected by repeatedly calling sc.Connect()
    40  // until the state becomes Shutdown or until ithe context expires.
    41  func StayConnected(ctx context.Context, sc StateChanger) {
    42  	for {
    43  		state := sc.GetState()
    44  		switch state {
    45  		case connectivity.Idle:
    46  			sc.Connect()
    47  		case connectivity.Shutdown:
    48  			return
    49  		}
    50  		if !sc.WaitForStateChange(ctx, state) {
    51  			return
    52  		}
    53  	}
    54  }
    55  
    56  // AwaitState waits for sc to enter stateWant or fatal errors if it doesn't
    57  // happen before ctx expires.
    58  func AwaitState(ctx context.Context, t *testing.T, sc StateChanger, stateWant connectivity.State) {
    59  	t.Helper()
    60  	for state := sc.GetState(); state != stateWant; state = sc.GetState() {
    61  		if !sc.WaitForStateChange(ctx, state) {
    62  			t.Fatalf("Timed out waiting for state change.  got %v; want %v", state, stateWant)
    63  		}
    64  	}
    65  }
    66  
    67  // AwaitNotState waits for sc to leave stateDoNotWant or fatal errors if it
    68  // doesn't happen before ctx expires.
    69  func AwaitNotState(ctx context.Context, t *testing.T, sc StateChanger, stateDoNotWant connectivity.State) {
    70  	t.Helper()
    71  	for state := sc.GetState(); state == stateDoNotWant; state = sc.GetState() {
    72  		if !sc.WaitForStateChange(ctx, state) {
    73  			t.Fatalf("Timed out waiting for state change.  got %v; want NOT %v", state, stateDoNotWant)
    74  		}
    75  	}
    76  }
    77  
    78  // AwaitNoStateChange expects ctx to be canceled before sc's state leaves
    79  // currState, and fatal errors otherwise.
    80  func AwaitNoStateChange(ctx context.Context, t *testing.T, sc StateChanger, currState connectivity.State) {
    81  	t.Helper()
    82  	if sc.WaitForStateChange(ctx, currState) {
    83  		t.Fatalf("State changed from %q to %q when no state change was expected", currState, sc.GetState())
    84  	}
    85  }