Unit test for unary tenancy handler

Signed-off-by: Ed Snible <[email protected]>
jaegertracing · Jul 11, 2022 · 8e7cbc7 · 8e7cbc7
1 parent 6541aa5
commit 8e7cbc7
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 2 deletions.
diff --git a/pkg/tenancy/grpc.go b/pkg/tenancy/grpc.go
@@ -81,7 +81,7 @@ func tenantFromMetadata(md metadata.MD, tenancyHeader string) (string, error) {
 	return tenants[0], nil
 }
 
-// NewGuardingStreamInterceptor blocks handling of streams whose tenancy header doesn't meet tenancy requirements.
+// NewGuardingUnaryInterceptor blocks handling of RPCs whose tenancy header doesn't meet tenancy requirements.
 // It also ensures the tenant is directly in the context, rather than context metadata.
 func NewGuardingUnaryInterceptor(tc *TenancyConfig) grpc.UnaryServerInterceptor {
 	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {

diff --git a/pkg/tenancy/grpc_test.go b/pkg/tenancy/grpc_test.go
@@ -24,7 +24,7 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
-func TestTenancyInterceptor(t *testing.T) {
+func TestStreamingTenancyInterceptor(t *testing.T) {
 	tests := []struct {
 		name          string
 		tenancyConfig *TenancyConfig
@@ -96,3 +96,74 @@ func TestTenancyInterceptor(t *testing.T) {
 		})
 	}
 }
+
+func TestUnaryTenancyInterceptor(t *testing.T) {
+	tests := []struct {
+		name          string
+		tenancyConfig *TenancyConfig
+		ctx           context.Context
+		errMsg        string
+	}{
+		{
+			name:          "missing tenant context",
+			tenancyConfig: NewTenancyConfig(&Options{Enabled: true}),
+			ctx:           context.Background(),
+			errMsg:        "rpc error: code = PermissionDenied desc = missing tenant header",
+		},
+		{
+			name:          "invalid tenant context",
+			tenancyConfig: NewTenancyConfig(&Options{Enabled: true, Tenants: []string{"megacorp"}}),
+			ctx:           WithTenant(context.Background(), "acme"),
+			errMsg:        "rpc error: code = PermissionDenied desc = unknown tenant header",
+		},
+		{
+			name:          "valid tenant context",
+			tenancyConfig: NewTenancyConfig(&Options{Enabled: true, Tenants: []string{"acme"}}),
+			ctx:           WithTenant(context.Background(), "acme"),
+			errMsg:        "",
+		},
+		{
+			name:          "invalid tenant header",
+			tenancyConfig: NewTenancyConfig(&Options{Enabled: true, Tenants: []string{"megacorp"}}),
+			ctx:           metadata.NewIncomingContext(context.Background(), map[string][]string{"x-tenant": {"acme"}}),
+			errMsg:        "rpc error: code = PermissionDenied desc = unknown tenant",
+		},
+		{
+			name:          "missing tenant header",
+			tenancyConfig: NewTenancyConfig(&Options{Enabled: true, Tenants: []string{"megacorp"}}),
+			ctx:           metadata.NewIncomingContext(context.Background(), map[string][]string{}),
+			errMsg:        "rpc error: code = PermissionDenied desc = missing tenant header",
+		},
+		{
+			name:          "valid tenant header",
+			tenancyConfig: NewTenancyConfig(&Options{Enabled: true, Tenants: []string{"acme"}}),
+			ctx:           metadata.NewIncomingContext(context.Background(), map[string][]string{"x-tenant": {"acme"}}),
+			errMsg:        "",
+		},
+		{
+			name:          "extra tenant header",
+			tenancyConfig: NewTenancyConfig(&Options{Enabled: true, Tenants: []string{"acme"}}),
+			ctx:           metadata.NewIncomingContext(context.Background(), map[string][]string{"x-tenant": {"acme", "megacorp"}}),
+			errMsg:        "rpc error: code = PermissionDenied desc = extra tenant header",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			interceptor := NewGuardingUnaryInterceptor(test.tenancyConfig)
+			usi := &grpc.UnaryServerInfo{}
+			iface := 0
+			handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+				// do nothing
+				return req, nil
+			}
+			_, err := interceptor(test.ctx, iface, usi, handler)
+			if test.errMsg == "" {
+				assert.NoError(t, err)
+			} else {
+				require.Error(t, err)
+				assert.Equal(t, test.errMsg, err.Error())
+			}
+		})
+	}
+}