github.com/blend/go-sdk@v1.20220411.3/grpcutil/check_connectivity_state_test.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package grpcutil
     9  
    10  import (
    11  	"context"
    12  	"testing"
    13  	"time"
    14  
    15  	"google.golang.org/grpc/connectivity"
    16  
    17  	"github.com/blend/go-sdk/assert"
    18  )
    19  
    20  type mockGetConnectionState connectivity.State
    21  
    22  func (m mockGetConnectionState) GetConnectionState() connectivity.State {
    23  	return connectivity.State(m)
    24  }
    25  
    26  type mockGetConnectionStateMany chan connectivity.State
    27  
    28  func (m mockGetConnectionStateMany) GetConnectionState() connectivity.State {
    29  	return connectivity.State(<-m)
    30  }
    31  
    32  func Test_CheckConnectivityState(t *testing.T) {
    33  	t.Parallel()
    34  	its := assert.New(t)
    35  
    36  	checker := CheckConnectivityState(
    37  		mockGetConnectionState(connectivity.Ready),
    38  	)
    39  	err := checker.Check(context.Background())
    40  	its.Nil(err)
    41  }
    42  
    43  func Test_CheckConnectivityState_failure(t *testing.T) {
    44  	t.Parallel()
    45  	its := assert.New(t)
    46  
    47  	checker := CheckConnectivityState(
    48  		mockGetConnectionState(connectivity.TransientFailure),
    49  		OptRetryCheckConnectivityStateRetryBackoff(time.Microsecond),
    50  		OptRetryCheckConnectivityStateRetryTimeout(time.Millisecond),
    51  	)
    52  	err := checker.Check(context.Background())
    53  	its.NotNil(err)
    54  }
    55  
    56  func Test_CheckConnectivityState_retry_success(t *testing.T) {
    57  	t.Parallel()
    58  	its := assert.New(t)
    59  
    60  	states := mockGetConnectionStateMany(make(chan connectivity.State, 5))
    61  	states <- connectivity.Connecting
    62  	states <- connectivity.Connecting
    63  	states <- connectivity.Connecting
    64  	states <- connectivity.Connecting
    65  	states <- connectivity.Ready
    66  
    67  	checker := CheckConnectivityState(
    68  		mockGetConnectionState(connectivity.Ready),
    69  		OptRetryCheckConnectivityStateRetryBackoff(time.Microsecond),
    70  		OptRetryCheckConnectivityStateRetryTimeout(time.Millisecond),
    71  	)
    72  	err := checker.Check(context.Background())
    73  	its.Nil(err)
    74  }