kmsvc-sdk/internal_test_fake_test.go

package kmsvc

import (
	"context"
	"net"
	"testing"

	kafkamgmtv1 "forgejo.riotpiao.com/rock/kmsvc-proto/gen/kafkamgmt/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/test/bufconn"
)

// fakeQueueService is a hand-rolled implementation of QueueServiceServer for
// tests. Each field is an optional hook; unset hooks return Unimplemented.
type fakeQueueService struct {
	kafkamgmtv1.UnimplementedQueueServiceServer

	sendMessage             func(context.Context, *kafkamgmtv1.SendMessageRequest) (*kafkamgmtv1.SendMessageResponse, error)
	sendMessageBatch        func(context.Context, *kafkamgmtv1.SendMessageBatchRequest) (*kafkamgmtv1.SendMessageBatchResponse, error)
	receiveMessage          func(context.Context, *kafkamgmtv1.ReceiveMessageRequest) (*kafkamgmtv1.ReceiveMessageResponse, error)
	deleteMessage           func(context.Context, *kafkamgmtv1.DeleteMessageRequest) (*kafkamgmtv1.DeleteMessageResponse, error)
	deleteMessageBatch      func(context.Context, *kafkamgmtv1.DeleteMessageBatchRequest) (*kafkamgmtv1.DeleteMessageBatchResponse, error)
	changeMessageVisibility func(context.Context, *kafkamgmtv1.ChangeMessageVisibilityRequest) (*kafkamgmtv1.ChangeMessageVisibilityResponse, error)

	// lastIncomingAuth captures the authorization header seen by the most
	// recent call, for interceptor assertions.
	lastIncomingAuth string
}

func (f *fakeQueueService) captureAuth(ctx context.Context) {
	if md, ok := metadata.FromIncomingContext(ctx); ok {
		if vals := md.Get("authorization"); len(vals) > 0 {
			f.lastIncomingAuth = vals[0]
		}
	}
}

func (f *fakeQueueService) SendMessage(ctx context.Context, req *kafkamgmtv1.SendMessageRequest) (*kafkamgmtv1.SendMessageResponse, error) {
	f.captureAuth(ctx)
	if f.sendMessage != nil {
		return f.sendMessage(ctx, req)
	}
	return f.UnimplementedQueueServiceServer.SendMessage(ctx, req)
}

func (f *fakeQueueService) SendMessageBatch(ctx context.Context, req *kafkamgmtv1.SendMessageBatchRequest) (*kafkamgmtv1.SendMessageBatchResponse, error) {
	f.captureAuth(ctx)
	if f.sendMessageBatch != nil {
		return f.sendMessageBatch(ctx, req)
	}
	return f.UnimplementedQueueServiceServer.SendMessageBatch(ctx, req)
}

func (f *fakeQueueService) ReceiveMessage(ctx context.Context, req *kafkamgmtv1.ReceiveMessageRequest) (*kafkamgmtv1.ReceiveMessageResponse, error) {
	f.captureAuth(ctx)
	if f.receiveMessage != nil {
		return f.receiveMessage(ctx, req)
	}
	return f.UnimplementedQueueServiceServer.ReceiveMessage(ctx, req)
}

func (f *fakeQueueService) DeleteMessage(ctx context.Context, req *kafkamgmtv1.DeleteMessageRequest) (*kafkamgmtv1.DeleteMessageResponse, error) {
	f.captureAuth(ctx)
	if f.deleteMessage != nil {
		return f.deleteMessage(ctx, req)
	}
	return f.UnimplementedQueueServiceServer.DeleteMessage(ctx, req)
}

func (f *fakeQueueService) DeleteMessageBatch(ctx context.Context, req *kafkamgmtv1.DeleteMessageBatchRequest) (*kafkamgmtv1.DeleteMessageBatchResponse, error) {
	f.captureAuth(ctx)
	if f.deleteMessageBatch != nil {
		return f.deleteMessageBatch(ctx, req)
	}
	return f.UnimplementedQueueServiceServer.DeleteMessageBatch(ctx, req)
}

func (f *fakeQueueService) ChangeMessageVisibility(ctx context.Context, req *kafkamgmtv1.ChangeMessageVisibilityRequest) (*kafkamgmtv1.ChangeMessageVisibilityResponse, error) {
	f.captureAuth(ctx)
	if f.changeMessageVisibility != nil {
		return f.changeMessageVisibility(ctx, req)
	}
	return f.UnimplementedQueueServiceServer.ChangeMessageVisibility(ctx, req)
}

// newTestClient starts an in-process bufconn server backed by fake, dials a
// Client against it, and registers cleanup with t.
func newTestClient(t *testing.T, fake *fakeQueueService, opts ...Option) *Client {
	t.Helper()

	lis := bufconn.Listen(1024 * 1024)
	srv := grpc.NewServer()
	kafkamgmtv1.RegisterQueueServiceServer(srv, fake)
	go func() {
		_ = srv.Serve(lis)
	}()
	t.Cleanup(srv.Stop)

	dialer := func(ctx context.Context, _ string) (net.Conn, error) {
		return lis.DialContext(ctx)
	}

	o := &options{dialTimeout: 0}
	for _, opt := range opts {
		opt(o)
	}

	dialOpts := []grpc.DialOption{
		grpc.WithContextDialer(dialer),
		grpc.WithTransportCredentials(insecure.NewCredentials()),
	}
	if o.tokenSource != nil {
		dialOpts = append(dialOpts,
			grpc.WithUnaryInterceptor(authUnaryInterceptor(o.tokenSource)),
			grpc.WithStreamInterceptor(authStreamInterceptor(o.tokenSource)),
		)
	}

	conn, err := grpc.NewClient("passthrough:///bufnet", dialOpts...)
	if err != nil {
		t.Fatalf("dial bufconn: %v", err)
	}
	t.Cleanup(func() { _ = conn.Close() })

	return newFromConn(conn)
}