|
@@ -0,0 +1,266 @@
|
|
|
|
|
+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
|
|
|
|
|
+}
|