pool_test.go

Documentation: cloud.google.com/go/auth/grpctransport

     1  // Copyright 2023 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package grpctransport
    16  
    17  import (
    18  	"context"
    19  	"errors"
    20  	"net"
    21  	"testing"
    22  
    23  	"google.golang.org/grpc"
    24  	"google.golang.org/grpc/credentials/insecure"
    25  )
    26  
    27  func TestPool_RoundRobin(t *testing.T) {
    28  	conn1 := &grpc.ClientConn{}
    29  	conn2 := &grpc.ClientConn{}
    30  
    31  	pool := &roundRobinConnPool{
    32  		conns: []*grpc.ClientConn{
    33  			conn1, conn2,
    34  		},
    35  	}
    36  
    37  	if got := pool.Connection(); got != conn2 {
    38  		t.Errorf("pool.Conn() #1 = %v, want conn2 (%v)", got, conn2)
    39  	}
    40  	if got := pool.Connection(); got != conn1 {
    41  		t.Errorf("pool.Conn() #2 = %v, want conn1 (%v)", got, conn1)
    42  	}
    43  	if got := pool.Connection(); got != conn2 {
    44  		t.Errorf("pool.Conn() #3 = %v, want conn2 (%v)", got, conn2)
    45  	}
    46  	if got := pool.Len(); got != 2 {
    47  		t.Errorf("pool.Len() = %v, want %v", got, 2)
    48  	}
    49  }
    50  
    51  func TestPool_SingleConn(t *testing.T) {
    52  	conn1 := &grpc.ClientConn{}
    53  	pool := &singleConnPool{conn1}
    54  
    55  	if got := pool.Connection(); got != conn1 {
    56  		t.Errorf("pool.Conn() #1 = %v, want conn2 (%v)", got, conn1)
    57  	}
    58  	if got := pool.Connection(); got != conn1 {
    59  		t.Errorf("pool.Conn() #2 = %v, want conn1 (%v)", got, conn1)
    60  	}
    61  	if got := pool.Len(); got != 1 {
    62  		t.Errorf("pool.Len() = %v, want %v", got, 1)
    63  	}
    64  }
    65  
    66  func TestClose(t *testing.T) {
    67  	_, l := mockServer(t)
    68  
    69  	pool := &roundRobinConnPool{}
    70  	for i := 0; i < 4; i++ {
    71  		conn, err := grpc.Dial(l.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
    72  		if err != nil {
    73  			t.Fatal(err)
    74  		}
    75  		pool.conns = append(pool.conns, conn)
    76  	}
    77  
    78  	if err := pool.Close(); err != nil {
    79  		t.Fatalf("pool.Close: %v", err)
    80  	}
    81  }
    82  
    83  func TestWithEndpointAndPoolSize(t *testing.T) {
    84  	_, l := mockServer(t)
    85  	ctx := context.Background()
    86  	connPool, err := Dial(ctx, false, &Options{
    87  		Endpoint: l.Addr().String(),
    88  		PoolSize: 4,
    89  	})
    90  	if err != nil {
    91  		t.Fatal(err)
    92  	}
    93  
    94  	if err := connPool.Close(); err != nil {
    95  		t.Fatalf("pool.Close: %v", err)
    96  	}
    97  }
    98  
    99  func TestMultiError(t *testing.T) {
   100  	tests := []struct {
   101  		name string
   102  		errs multiError
   103  		want string
   104  	}{
   105  		{
   106  			name: "0 errors",
   107  			want: "(0 errors)",
   108  		},
   109  		{
   110  			name: "1 errors",
   111  			errs: []error{errors.New("the full error message")},
   112  			want: "the full error message",
   113  		},
   114  		{
   115  			name: "2 errors",
   116  			errs: []error{errors.New("foo"), errors.New("bar")},
   117  			want: "foo (and 1 other error)",
   118  		},
   119  		{
   120  			name: "3 errors",
   121  			errs: []error{errors.New("foo"), errors.New("bar"), errors.New("baz")},
   122  			want: "foo (and 2 other errors)",
   123  		},
   124  	}
   125  	for _, tt := range tests {
   126  		t.Run(tt.name, func(t *testing.T) {
   127  			got := tt.errs.Error()
   128  			if got != tt.want {
   129  				t.Fatalf("got %v, want %v", got, tt.want)
   130  			}
   131  		})
   132  	}
   133  }
   134  
   135  func mockServer(t *testing.T) (*grpc.Server, net.Listener) {
   136  	t.Helper()
   137  
   138  	l, err := net.Listen("tcp", "localhost:0")
   139  	if err != nil {
   140  		t.Fatal(err)
   141  	}
   142  
   143  	s := grpc.NewServer()
   144  	go s.Serve(l)
   145  
   146  	return s, l
   147  }
   148
View as plain text