129 lines
4.6 KiB
Go
129 lines
4.6 KiB
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)
|
|
}
|