package grpccorrelation

import (
	"context"
	"testing"

	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/labkit/correlation"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

func TestUnaryServerCorrelationInterceptor(t *testing.T) {
	tests := []struct {
		name        string
		md          metadata.MD
		interceptor grpc.UnaryServerInterceptor
		test        func(*testing.T) grpc.UnaryHandler
	}{
		{
			name: "default",
			md: metadata.Pairs(
				metadataCorrelatorKey,
				correlationID,
				metadataClientNameKey,
				clientName,
			),
			interceptor: UnaryServerCorrelationInterceptor(),
			test: func(*testing.T) grpc.UnaryHandler {
				return func(ctx context.Context, req interface{}) (interface{}, error) {
					require.Equal(t, correlationID, correlation.ExtractFromContext(ctx))
					require.Equal(t, clientName, correlation.ExtractClientNameFromContext(ctx))
					return nil, nil
				}
			},
		},
		{
			name: "id present but not trusted",
			md: metadata.Pairs(
				metadataCorrelatorKey,
				correlationID,
			),
			interceptor: UnaryServerCorrelationInterceptor(WithoutPropagation()),
			test:        unaryExpectRandomID,
		},
		{
			name: "id present, trusted but empty",
			md: metadata.Pairs(
				metadataCorrelatorKey,
				"",
			),
			interceptor: UnaryServerCorrelationInterceptor(WithoutPropagation()),
			test:        unaryExpectRandomID,
		},
		{
			name:        "id absent and not trusted",
			md:          metadata.Pairs(),
			interceptor: UnaryServerCorrelationInterceptor(WithoutPropagation()),
			test:        unaryExpectRandomID,
		},
		{
			name:        "id absent and trusted",
			md:          metadata.Pairs(),
			interceptor: UnaryServerCorrelationInterceptor(),
			test:        unaryExpectRandomID,
		},
		{
			name:        "no metadata",
			md:          nil,
			interceptor: UnaryServerCorrelationInterceptor(),
			test:        unaryExpectRandomID,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			ctx := context.Background()
			if tc.md != nil {
				ctx = metadata.NewIncomingContext(context.Background(), tc.md)
			}
			_, err := tc.interceptor(
				ctx,
				nil,
				nil,
				tc.test(t),
			)
			require.NoError(t, err)
		})
	}
}

func TestStreamingServerCorrelationInterceptor(t *testing.T) {
	tests := []struct {
		name        string
		md          metadata.MD
		interceptor grpc.StreamServerInterceptor
		test        func(*testing.T) grpc.StreamHandler
	}{
		{
			name: "default",
			md: metadata.Pairs(
				metadataCorrelatorKey,
				correlationID,
				metadataClientNameKey,
				clientName,
			),
			interceptor: StreamServerCorrelationInterceptor(),
			test: func(*testing.T) grpc.StreamHandler {
				return func(srv interface{}, stream grpc.ServerStream) error {
					ctx := stream.Context()
					require.Equal(t, correlationID, correlation.ExtractFromContext(ctx))
					require.Equal(t, clientName, correlation.ExtractClientNameFromContext(ctx))
					return nil
				}
			},
		},
		{
			name: "id present but not trusted",
			md: metadata.Pairs(
				metadataCorrelatorKey,
				correlationID,
			),
			interceptor: StreamServerCorrelationInterceptor(WithoutPropagation()),
			test:        streamExpectRandomID,
		},
		{
			name: "id present, trusted but empty",
			md: metadata.Pairs(
				metadataCorrelatorKey,
				"",
			),
			interceptor: StreamServerCorrelationInterceptor(WithoutPropagation()),
			test:        streamExpectRandomID,
		},
		{
			name:        "id absent and not trusted",
			md:          metadata.Pairs(),
			interceptor: StreamServerCorrelationInterceptor(WithoutPropagation()),
			test:        streamExpectRandomID,
		},
		{
			name:        "id absent and trusted",
			md:          metadata.Pairs(),
			interceptor: StreamServerCorrelationInterceptor(),
			test:        streamExpectRandomID,
		},
		{
			name:        "no metadata",
			md:          nil,
			interceptor: StreamServerCorrelationInterceptor(),
			test:        streamExpectRandomID,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			ctx := context.Background()
			if tc.md != nil {
				ctx = metadata.NewIncomingContext(context.Background(), tc.md)
			}
			err := tc.interceptor(
				nil,
				&grpc_middleware.WrappedServerStream{WrappedContext: ctx},
				nil,
				tc.test(t),
			)
			require.NoError(t, err)
		})
	}
}

func unaryExpectRandomID(t *testing.T) grpc.UnaryHandler {
	return func(ctx context.Context, req interface{}) (interface{}, error) {
		actualID := correlation.ExtractFromContext(ctx)
		require.NotEqual(t, correlationID, actualID)
		require.NotEmpty(t, actualID)
		return nil, nil
	}
}

func streamExpectRandomID(t *testing.T) grpc.StreamHandler {
	return func(srv interface{}, stream grpc.ServerStream) error {
		actualID := correlation.ExtractFromContext(stream.Context())
		require.NotEqual(t, correlationID, actualID)
		require.NotEmpty(t, actualID)
		return nil
	}
}
