package kmsvc import ( "context" "net" "testing" kafkamgmtv1 "forgejo.riotpiao.homelab.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) }