package grpcclient import ( "context" "io" "net" "sync/atomic" "testing" "time" pbv1 "git3.techno-world.net/lrosales/broad-announce/gen/go/broadannounce/v1" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/metadata" "google.golang.org/grpc/test/bufconn" ) const bufnetSize = 1 << 20 // 1 MB // mockIngestServer implements pbv1.IngestServer for testing. type mockIngestServer struct { pbv1.UnimplementedIngestServer alertsReceived atomic.Int64 rateLimitAfter int64 // after N alerts, start rate-limiting authKey string authFail atomic.Bool } func (m *mockIngestServer) StreamAlerts(stream pbv1.Ingest_StreamAlertsServer) error { md, ok := metadata.FromIncomingContext(stream.Context()) if !ok { return io.EOF } if m.authFail.Load() { return io.EOF } if len(m.authKey) > 0 { vals := md.Get("authorization") if len(vals) == 0 || vals[0] != "Bearer "+m.authKey { return io.EOF } } for { alert, err := stream.Recv() if err == io.EOF { return nil } if err != nil { return err } m.alertsReceived.Add(1) ack := &pbv1.Ack{ AlertId: alert.DedupeKey, DedupeKey: alert.DedupeKey, AcceptedAtMs: time.Now().UnixMilli(), } if m.rateLimitAfter > 0 && m.alertsReceived.Load() > m.rateLimitAfter { ack.Result = &pbv1.Ack_Error{ Error: &pbv1.Error{ Code: pbv1.Error_RATE_LIMITED, Message: "rate limited by test server", RetryAfterMs: 10, }, } } else { ack.Result = &pbv1.Ack_Ok{ Ok: &pbv1.Ok{}, } } if err := stream.Send(ack); err != nil { return err } } } func newTestServer(t *testing.T) (*grpc.Server, *bufconn.Listener) { lis := bufconn.Listen(bufnetSize) srv := grpc.NewServer() pbv1.RegisterIngestServer(srv, &mockIngestServer{authKey: "acme-001:prom:s3cret"}) go srv.Serve(lis) return srv, lis } func dialBufconn(ctx context.Context, t *testing.T, lis *bufconn.Listener) *grpc.ClientConn { conn, err := grpc.NewClient( "passthrough://bufconn", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return lis.Dial() }), grpc.WithTransportCredentials(insecure.NewCredentials()), ) if err != nil { t.Fatalf("grpc.NewClient: %v", err) } return conn } func TestClient_New(t *testing.T) { c, err := New("localhost:9090", WithAPIKey("acme-001:prom:s3cret"), WithMaxRetries(5), WithInsecure(), ) if err != nil { t.Fatalf("New() error = %v", err) } if c == nil { t.Fatal("New() returned nil client") } if c.maxRetries != 5 { t.Errorf("maxRetries = %d, want 5", c.maxRetries) } c.Close() } func TestClient_OptionDefaults(t *testing.T) { c, err := New("localhost:9090", WithInsecure()) if err != nil { t.Fatalf("New() error = %v", err) } if c.maxRetries != 3 { t.Errorf("default maxRetries = %d, want 3", c.maxRetries) } c.Close() } func TestStream_SendRecv(t *testing.T) { ctx := context.Background() srv, lis := newTestServer(t) defer srv.Stop() conn := dialBufconn(ctx, t, lis) defer conn.Close() client := pbv1.NewIngestClient(conn) md := metadata.Pairs("authorization", "Bearer acme-001:prom:s3cret") sctx := metadata.NewOutgoingContext(ctx, md) stream, err := client.StreamAlerts(sctx) if err != nil { t.Fatalf("StreamAlerts() error = %v", err) } alert := &pbv1.Alert{ CompanyId: "acme-001", SourceId: "prom", Severity: "critical", Title: "test alert", DedupeKey: "dk-001", ClientTsMs: time.Now().UnixMilli(), } if err := stream.Send(alert); err != nil { t.Fatalf("stream.Send() error = %v", err) } ack, err := stream.Recv() if err != nil { t.Fatalf("stream.Recv() error = %v", err) } if ack.DedupeKey != "dk-001" { t.Errorf("ack.DedupeKey = %q, want %q", ack.DedupeKey, "dk-001") } stream.CloseSend() } func TestStream_CloseSend(t *testing.T) { ctx := context.Background() srv, lis := newTestServer(t) defer srv.Stop() conn := dialBufconn(ctx, t, lis) defer conn.Close() client := pbv1.NewIngestClient(conn) md := metadata.Pairs("authorization", "Bearer acme-001:prom:s3cret") sctx := metadata.NewOutgoingContext(ctx, md) stream, err := client.StreamAlerts(sctx) if err != nil { t.Fatalf("StreamAlerts() error = %v", err) } if err := stream.CloseSend(); err != nil { t.Errorf("CloseSend() error = %v", err) } _, err = stream.Recv() if err != io.EOF { t.Errorf("Recv() after CloseSend = %v, want io.EOF", err) } } func TestRetryable(t *testing.T) { tests := []struct { code codes.Code want bool }{ {codes.Unavailable, true}, {codes.ResourceExhausted, true}, {codes.Internal, true}, {codes.OK, false}, {codes.InvalidArgument, false}, {codes.NotFound, false}, {codes.Unauthenticated, false}, } for _, tt := range tests { t.Run(tt.code.String(), func(t *testing.T) { if got := retryable(tt.code); got != tt.want { t.Errorf("retryable(%v) = %v, want %v", tt.code, got, tt.want) } }) } } func TestStream_SendRecv_MultipleAlerts(t *testing.T) { ctx := context.Background() srv, lis := newTestServer(t) defer srv.Stop() conn := dialBufconn(ctx, t, lis) defer conn.Close() client := pbv1.NewIngestClient(conn) md := metadata.Pairs("authorization", "Bearer acme-001:prom:s3cret") sctx := metadata.NewOutgoingContext(ctx, md) stream, err := client.StreamAlerts(sctx) if err != nil { t.Fatalf("StreamAlerts() error = %v", err) } const n = 5 for i := 0; i < n; i++ { alert := &pbv1.Alert{ CompanyId: "acme-001", SourceId: "prom", Severity: "info", Title: "test alert", DedupeKey: "dk-multi-" + string(rune('0'+i)), ClientTsMs: time.Now().UnixMilli(), } if err := stream.Send(alert); err != nil { t.Fatalf("stream.Send() error = %v", err) } ack, err := stream.Recv() if err != nil { t.Fatalf("stream.Recv() error = %v", err) } if ack.DedupeKey != alert.DedupeKey { t.Errorf("ack.DedupeKey = %q, want %q", ack.DedupeKey, alert.DedupeKey) } } stream.CloseSend() } // Verify generated types implement the expected interfaces. func TestProtoInterfaces(t *testing.T) { var _ pbv1.IngestClient = nil var _ pbv1.IngestServer = nil }