github.com/pachyderm/pachyderm@v1.13.4/src/client/client_test.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/pachyderm/pachyderm/src/client/pfs"
     9  	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
    10  	"google.golang.org/grpc"
    11  )
    12  
    13  type countingInterceptor struct {
    14  	s int
    15  	u int
    16  }
    17  
    18  func (i *countingInterceptor) unary() grpc.UnaryClientInterceptor {
    19  	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, call grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    20  		i.u++
    21  		return call(ctx, method, req, reply, cc, opts...)
    22  	}
    23  }
    24  
    25  func (i *countingInterceptor) stream() grpc.StreamClientInterceptor {
    26  	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
    27  		i.s++
    28  		return streamer(ctx, desc, cc, method, opts...)
    29  	}
    30  }
    31  
    32  func TestInterceptors(t *testing.T) {
    33  	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
    34  	defer cancel()
    35  	server, err := grpcutil.NewServer(ctx, false)
    36  	if err != nil {
    37  		t.Fatalf("server: %v", err)
    38  	}
    39  	defer server.Wait()
    40  
    41  	listener, err := server.ListenTCP("localhost", 0)
    42  	if err != nil {
    43  		t.Fatalf("listener: %v", err)
    44  	}
    45  	defer listener.Close()
    46  	pfs.RegisterAPIServer(server.Server, new(pfs.UnimplementedAPIServer))
    47  
    48  	interceptor := new(countingInterceptor)
    49  	c, err := NewFromAddress(listener.Addr().String(), WithAdditionalUnaryClientInterceptors(interceptor.unary()), WithAdditionalStreamClientInterceptors(interceptor.stream()))
    50  	if err != nil {
    51  		t.Fatalf("create client: %v", err)
    52  	}
    53  	defer c.Close()
    54  
    55  	// Unary call.
    56  	if err := c.CreateRepo("foo"); err == nil {
    57  		t.Fatal("create repo: expected error")
    58  	}
    59  	if got, want := interceptor.u, 1; got != want {
    60  		t.Errorf("unary call count:\n  got: %v\n want: %v", got, want)
    61  	}
    62  	if got, want := interceptor.s, 0; got != want {
    63  		t.Errorf("stream call count:\n  got: %v\n want: %v", got, want)
    64  	}
    65  
    66  	// Stream call.
    67  	if err := c.Fsck(true, func(*pfs.FsckResponse) error { return nil }); err == nil {
    68  		t.Fatal("fsck: expected error")
    69  	}
    70  	if got, want := interceptor.u, 1; got != want {
    71  		t.Errorf("unary call count:\n  got: %v\n want: %v", got, want)
    72  	}
    73  	if got, want := interceptor.s, 1; got != want {
    74  		t.Errorf("stream call count:\n  got: %v\n want: %v", got, want)
    75  	}
    76  }