gitlab.com/gitlab-org/labkit@v1.21.0/correlation/grpc/server_interceptors_test.go

package grpccorrelation

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/labkit/correlation"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

var (
	_ grpc.ServerTransportStream = (*mockServerTransportStream)(nil)
	_ grpc.ServerStream          = (*mockServerStream)(nil)
)

type tcType struct {
	name               string
	md                 metadata.MD
	withoutPropagation bool

	expectRandom       bool
	expectedClientName string
}

func TestServerCorrelationInterceptors(t *testing.T) {
	tests := []tcType{
		{
			name: "default",
			md: metadata.Pairs(
				metadataCorrelatorKey,
				correlationID,
				metadataClientNameKey,
				clientName,
			),
			expectedClientName: clientName,
		},
		{
			name: "id present but not trusted",
			md: metadata.Pairs(
				metadataCorrelatorKey,
				correlationID,
			),
			withoutPropagation: true,
			expectRandom:       true,
		},
		{
			name: "id present, trusted but empty",
			md: metadata.Pairs(
				metadataCorrelatorKey,
				"",
			),
			withoutPropagation: true,
			expectRandom:       true,
		},
		{
			name:               "id absent and not trusted",
			md:                 metadata.Pairs(),
			withoutPropagation: true,
			expectRandom:       true,
		},
		{
			name:         "id absent and trusted",
			md:           metadata.Pairs(),
			expectRandom: true,
		},
		{
			name:         "no metadata",
			md:           nil,
			expectRandom: true,
		},
	}

	t.Run("unary", func(t *testing.T) {
		for _, tc := range tests {
			t.Run(tc.name, testUnaryServerCorrelationInterceptor(tc, false))
			t.Run(tc.name+" (reverse)", testUnaryServerCorrelationInterceptor(tc, true))
		}
	})
	t.Run("streaming", func(t *testing.T) {
		for _, tc := range tests {
			t.Run(tc.name, testStreamingServerCorrelationInterceptor(tc, false))
			t.Run(tc.name+" (reverse)", testStreamingServerCorrelationInterceptor(tc, true))
		}
	})
}

func testUnaryServerCorrelationInterceptor(tc tcType, reverseCorrelationID bool) func(*testing.T) {
	return func(t *testing.T) {
		t.Helper()

		sts := &mockServerTransportStream{}
		ctx := grpc.NewContextWithServerTransportStream(context.Background(), sts)
		if tc.md != nil {
			ctx = metadata.NewIncomingContext(ctx, tc.md)
		}
		interceptor := UnaryServerCorrelationInterceptor(constructServerOpts(tc, reverseCorrelationID)...)
		_, err := interceptor(
			ctx,
			nil,
			nil,
			func(ctx context.Context, req interface{}) (interface{}, error) {
				testServerCtx(ctx, t, tc, reverseCorrelationID, sts.header)
				return nil, nil
			},
		)
		require.NoError(t, err)
	}
}

func testStreamingServerCorrelationInterceptor(tc tcType, reverseCorrelationID bool) func(*testing.T) {
	return func(t *testing.T) {
		t.Helper()

		ctx := context.Background()
		if tc.md != nil {
			ctx = metadata.NewIncomingContext(ctx, tc.md)
		}
		ss := &mockServerStream{
			ctx: ctx,
		}
		interceptor := StreamServerCorrelationInterceptor(constructServerOpts(tc, reverseCorrelationID)...)
		err := interceptor(
			nil,
			ss,
			nil,
			func(srv interface{}, stream grpc.ServerStream) error {
				testServerCtx(stream.Context(), t, tc, reverseCorrelationID, ss.header)
				return nil
			},
		)
		require.NoError(t, err)
	}
}

func constructServerOpts(tc tcType, reverseCorrelationID bool) []ServerCorrelationInterceptorOption {
	var opts []ServerCorrelationInterceptorOption
	if tc.withoutPropagation {
		opts = append(opts, WithoutPropagation())
	}
	if reverseCorrelationID {
		opts = append(opts, WithReversePropagation())
	}
	return opts
}

func testServerCtx(ctx context.Context, t *testing.T, tc tcType, reverseCorrelationID bool, header metadata.MD) {
	t.Helper()

	actualID := correlation.ExtractFromContext(ctx)
	if tc.expectRandom {
		assert.NotEqual(t, correlationID, actualID)
		assert.NotEmpty(t, actualID)
	} else {
		assert.Equal(t, correlationID, actualID)
	}
	vals := header.Get(metadataCorrelatorKey)
	if reverseCorrelationID {
		assert.Equal(t, []string{actualID}, vals)
	} else {
		assert.Empty(t, vals)
	}
	assert.Equal(t, tc.expectedClientName, correlation.ExtractClientNameFromContext(ctx))
}

type mockServerTransportStream struct {
	header metadata.MD
}

func (s *mockServerTransportStream) Method() string {
	panic("implement me")
}

func (s *mockServerTransportStream) SetHeader(md metadata.MD) error {
	s.header = metadata.Join(s.header, md)
	return nil
}

func (s *mockServerTransportStream) SendHeader(md metadata.MD) error {
	panic("implement me")
}

func (s *mockServerTransportStream) SetTrailer(md metadata.MD) error {
	panic("implement me")
}

type mockServerStream struct {
	ctx    context.Context
	header metadata.MD
}

func (s *mockServerStream) SetHeader(md metadata.MD) error {
	s.header = metadata.Join(s.header, md)
	return nil
}

func (s *mockServerStream) SendHeader(md metadata.MD) error {
	panic("implement me")
}

func (s *mockServerStream) SetTrailer(md metadata.MD) {
	panic("implement me")
}

func (s *mockServerStream) Context() context.Context {
	return s.ctx
}

func (s *mockServerStream) SendMsg(m interface{}) error {
	panic("implement me")
}

func (s *mockServerStream) RecvMsg(m interface{}) error {
	panic("implement me")
}