pool_test.go

Documentation: google.golang.org/api/transport/grpc

     1  // Copyright 2020 Google LLC.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package grpc
     6  
     7  import (
     8  	"context"
     9  	"net"
    10  	"testing"
    11  
    12  	"google.golang.org/api/option"
    13  	"google.golang.org/grpc"
    14  )
    15  
    16  func TestPool(t *testing.T) {
    17  	conn1 := &grpc.ClientConn{}
    18  	conn2 := &grpc.ClientConn{}
    19  
    20  	pool := &roundRobinConnPool{
    21  		conns: []*grpc.ClientConn{
    22  			conn1, conn2,
    23  		},
    24  	}
    25  
    26  	if got := pool.Conn(); got != conn2 {
    27  		t.Errorf("pool.Conn() #1 got %v; want conn2 (%v)", got, conn2)
    28  	}
    29  
    30  	if got := pool.Conn(); got != conn1 {
    31  		t.Errorf("pool.Conn() #2 got %v; want conn1 (%v)", got, conn1)
    32  	}
    33  
    34  	if got := pool.Conn(); got != conn2 {
    35  		t.Errorf("pool.Conn() #3 got %v; want conn2 (%v)", got, conn2)
    36  	}
    37  }
    38  
    39  func TestClose(t *testing.T) {
    40  	_, l := mockServer(t)
    41  
    42  	pool := &roundRobinConnPool{}
    43  	for i := 0; i < 4; i++ {
    44  		conn, err := grpc.Dial(l.Addr().String(), grpc.WithInsecure())
    45  		if err != nil {
    46  			t.Fatal(err)
    47  		}
    48  		pool.conns = append(pool.conns, conn)
    49  	}
    50  
    51  	if err := pool.Close(); err != nil {
    52  		t.Fatalf("pool.Close: %v", err)
    53  	}
    54  }
    55  
    56  // Regression test for https://github.com/googleapis/google-cloud-go/issues/1780
    57  func TestWithGRPCConnAndPoolSize(t *testing.T) {
    58  	_, l := mockServer(t)
    59  
    60  	conn, err := grpc.Dial(l.Addr().String(), grpc.WithInsecure())
    61  	if err != nil {
    62  		t.Fatal(err)
    63  	}
    64  
    65  	ctx := context.Background()
    66  	connPool, err := DialPool(ctx,
    67  		option.WithGRPCConn(conn),
    68  		option.WithGRPCConnectionPool(4),
    69  	)
    70  	if err != nil {
    71  		t.Fatal(err)
    72  	}
    73  
    74  	if err := connPool.Close(); err != nil {
    75  		t.Fatalf("pool.Close: %v", err)
    76  	}
    77  }
    78  
    79  func mockServer(t *testing.T) (*grpc.Server, net.Listener) {
    80  	t.Helper()
    81  
    82  	l, err := net.Listen("tcp", "localhost:0")
    83  	if err != nil {
    84  		t.Fatal(err)
    85  	}
    86  
    87  	s := grpc.NewServer()
    88  	go s.Serve(l)
    89  
    90  	return s, l
    91  }
    92
View as plain text