github.com/hashicorp/go-plugin@v1.6.0/grpc_client_test.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package plugin
     5  
     6  import (
     7  	"context"
     8  	"reflect"
     9  	"testing"
    10  
    11  	grpctest "github.com/hashicorp/go-plugin/test/grpc"
    12  	"github.com/jhump/protoreflect/grpcreflect"
    13  	"google.golang.org/grpc"
    14  	reflectpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
    15  )
    16  
    17  func TestGRPC_App(t *testing.T) {
    18  	t.Run("default", func(t *testing.T) {
    19  		testGRPCClientApp(t, false)
    20  	})
    21  	t.Run("mux", func(t *testing.T) {
    22  		testGRPCClientApp(t, true)
    23  	})
    24  }
    25  
    26  func testGRPCClientApp(t *testing.T, multiplex bool) {
    27  	client, server := TestPluginGRPCConn(t, multiplex, map[string]Plugin{
    28  		"test": new(testGRPCInterfacePlugin),
    29  	})
    30  	defer client.Close()
    31  	defer server.Stop()
    32  
    33  	raw, err := client.Dispense("test")
    34  	if err != nil {
    35  		t.Fatalf("err: %s", err)
    36  	}
    37  
    38  	impl, ok := raw.(testInterface)
    39  	if !ok {
    40  		t.Fatalf("bad: %#v", raw)
    41  	}
    42  
    43  	result := impl.Double(21)
    44  	if result != 42 {
    45  		t.Fatalf("bad: %#v", result)
    46  	}
    47  
    48  	err = impl.Bidirectional()
    49  	if err != nil {
    50  		t.Fatal(err)
    51  	}
    52  }
    53  
    54  func TestGRPCConn_BidirectionalPing(t *testing.T) {
    55  	conn, _ := TestGRPCConn(t, func(s *grpc.Server) {
    56  		grpctest.RegisterPingPongServer(s, &pingPongServer{})
    57  	})
    58  	defer conn.Close()
    59  	pingPongClient := grpctest.NewPingPongClient(conn)
    60  
    61  	pResp, err := pingPongClient.Ping(context.Background(), &grpctest.PingRequest{})
    62  	if err != nil {
    63  		t.Fatal(err)
    64  	}
    65  	if pResp.Msg != "pong" {
    66  		t.Fatal("Bad PingPong")
    67  	}
    68  }
    69  
    70  func TestGRPCC_Stream(t *testing.T) {
    71  	t.Run("default", func(t *testing.T) {
    72  		testGRPCStream(t, false)
    73  	})
    74  	t.Run("mux", func(t *testing.T) {
    75  		testGRPCStream(t, true)
    76  	})
    77  }
    78  
    79  func testGRPCStream(t *testing.T, multiplex bool) {
    80  	client, server := TestPluginGRPCConn(t, multiplex, map[string]Plugin{
    81  		"test": new(testGRPCInterfacePlugin),
    82  	})
    83  	defer client.Close()
    84  	defer server.Stop()
    85  
    86  	raw, err := client.Dispense("test")
    87  	if err != nil {
    88  		t.Fatalf("err: %s", err)
    89  	}
    90  
    91  	impl, ok := raw.(testStreamer)
    92  	if !ok {
    93  		t.Fatalf("bad: %#v", raw)
    94  	}
    95  
    96  	expected := []int32{21, 22, 23, 24, 25, 26}
    97  	result, err := impl.Stream(21, 27)
    98  	if err != nil {
    99  		t.Fatal(err)
   100  	}
   101  
   102  	if !reflect.DeepEqual(result, expected) {
   103  		t.Fatalf("expected: %v\ngot: %v", expected, result)
   104  	}
   105  }
   106  
   107  func TestGRPC_Ping(t *testing.T) {
   108  	t.Run("default", func(t *testing.T) {
   109  		testGRPCClientPing(t, false)
   110  	})
   111  	t.Run("mux", func(t *testing.T) {
   112  		testGRPCClientPing(t, true)
   113  	})
   114  }
   115  
   116  func testGRPCClientPing(t *testing.T, multiplex bool) {
   117  	client, server := TestPluginGRPCConn(t, multiplex, map[string]Plugin{
   118  		"test": new(testGRPCInterfacePlugin),
   119  	})
   120  	defer client.Close()
   121  	defer server.Stop()
   122  
   123  	// Run a couple pings
   124  	if err := client.Ping(); err != nil {
   125  		t.Fatalf("err: %s", err)
   126  	}
   127  	if err := client.Ping(); err != nil {
   128  		t.Fatalf("err: %s", err)
   129  	}
   130  
   131  	// Close the remote end
   132  	server.server.Stop()
   133  
   134  	// Test ping fails
   135  	if err := client.Ping(); err == nil {
   136  		t.Fatal("should error")
   137  	}
   138  }
   139  
   140  func TestGRPC_Reflection(t *testing.T) {
   141  	t.Run("default", func(t *testing.T) {
   142  		testGRPCClientReflection(t, false)
   143  	})
   144  	t.Run("mux", func(t *testing.T) {
   145  		testGRPCClientReflection(t, true)
   146  	})
   147  }
   148  
   149  func testGRPCClientReflection(t *testing.T, multiplex bool) {
   150  	ctx := context.Background()
   151  
   152  	client, server := TestPluginGRPCConn(t, multiplex, map[string]Plugin{
   153  		"test": new(testGRPCInterfacePlugin),
   154  	})
   155  	defer client.Close()
   156  	defer server.Stop()
   157  
   158  	refClient := grpcreflect.NewClient(ctx, reflectpb.NewServerReflectionClient(client.Conn))
   159  
   160  	svcs, err := refClient.ListServices()
   161  	if err != nil {
   162  		t.Fatalf("err: %s", err)
   163  	}
   164  
   165  	// TODO: maybe only assert some specific services here to make test more resilient
   166  	expectedSvcs := []string{"grpc.health.v1.Health", "grpc.reflection.v1alpha.ServerReflection", "grpctest.Test", "plugin.GRPCBroker", "plugin.GRPCController", "plugin.GRPCStdio"}
   167  
   168  	if !reflect.DeepEqual(svcs, expectedSvcs) {
   169  		t.Fatalf("expected: %v\ngot: %v", expectedSvcs, svcs)
   170  	}
   171  
   172  	healthDesc, err := refClient.ResolveService("grpc.health.v1.Health")
   173  	if err != nil {
   174  		t.Fatalf("err: %s", err)
   175  	}
   176  
   177  	methods := healthDesc.GetMethods()
   178  	var methodNames []string
   179  	for _, m := range methods {
   180  		methodNames = append(methodNames, m.GetName())
   181  	}
   182  
   183  	expectedMethodNames := []string{"Check", "Watch"}
   184  
   185  	if !reflect.DeepEqual(methodNames, expectedMethodNames) {
   186  		t.Fatalf("expected: %v\ngot: %v", expectedMethodNames, methodNames)
   187  	}
   188  }