broad-announce/internal/grpcclient/grpcclient_test.go

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
}