add test, fix bug

internal/zero/controller/usagereporter/usagereporter_test.go

@@ -1,14 +1,125 @@
 package usagereporter
 
 import (
+	"context"
+	"errors"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
+	"google.golang.org/protobuf/types/known/timestamppb"
 
+	"github.com/pomerium/pomerium/internal/databroker"
+	"github.com/pomerium/pomerium/internal/testutil"
+	databrokerpb "github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/grpc/session"
+	"github.com/pomerium/pomerium/pkg/grpc/user"
 	"github.com/pomerium/pomerium/pkg/zero/cluster"
 )
 
+type mockAPI struct {
+	reportUsage func(ctx context.Context, req cluster.ReportUsageRequest) error
+}
+
+func (m mockAPI) ReportUsage(ctx context.Context, req cluster.ReportUsageRequest) error {
+	return m.reportUsage(ctx, req)
+}
+
+func TestUsageReporter(t *testing.T) {
+	t.Parallel()
+
+	ctx, clearTimeout := context.WithTimeout(context.Background(), time.Second*10)
+	t.Cleanup(clearTimeout)
+
+	ctx, cancel := context.WithCancel(ctx)
+	t.Cleanup(cancel)
+
+	cc := testutil.NewGRPCServer(t, func(srv *grpc.Server) {
+		databrokerpb.RegisterDataBrokerServiceServer(srv, databroker.New())
+	})
+	t.Cleanup(func() { cc.Close() })
+
+	tm1 := time.Date(2024, time.September, 11, 11, 56, 0, 0, time.UTC)
+
+	requests := make(chan cluster.ReportUsageRequest, 1)
+
+	client := databrokerpb.NewDataBrokerServiceClient(cc)
+	ur := New(mockAPI{
+		reportUsage: func(ctx context.Context, req cluster.ReportUsageRequest) error {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case requests <- req:
+			}
+			return nil
+		},
+	}, "bQjwPpxcwJRbvsSMFgbZFkXmxFJ", time.Millisecond*100)
+
+	eg, ctx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		return ur.Run(ctx, client)
+	})
+	eg.Go(func() error {
+		_, err := databrokerpb.Put(ctx, client,
+			&session.Session{
+				Id:       "S1a",
+				UserId:   "U1",
+				IssuedAt: timestamppb.New(tm1),
+			},
+			&session.Session{
+				Id:       "S1b",
+				UserId:   "U1",
+				IssuedAt: timestamppb.New(tm1),
+			})
+		if err != nil {
+			return err
+		}
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case req := <-requests:
+			assert.Equal(t, cluster.ReportUsageRequest{
+				Users: []cluster.ReportUsageUser{{
+					LastSignedInAt: tm1,
+					PseudonymousId: "095xqqsjEEgYf5Yf+TAjWjooMQyh6jSV5SCPGe9eqvg=",
+				}},
+			}, req, "should send a single usage record")
+		}
+
+		_, err = databrokerpb.Put(ctx, client,
+			&user.User{
+				Id:    "U1",
+				Email: "u1@example.com",
+			})
+		if err != nil {
+			return err
+		}
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case req := <-requests:
+			assert.Equal(t, cluster.ReportUsageRequest{
+				Users: []cluster.ReportUsageUser{{
+					LastSignedInAt:    tm1,
+					PseudonymousEmail: "iq8/fj+uZaKitkWY12JIQgKJ5KIP+E0Cmy/HpxpdBXY=",
+					PseudonymousId:    "095xqqsjEEgYf5Yf+TAjWjooMQyh6jSV5SCPGe9eqvg=",
+				}},
+			}, req, "should send another usage record with the email set")
+		}
+
+		cancel()
+		return nil
+	})
+	err := eg.Wait()
+	if err != nil && !errors.Is(ctx.Err(), context.Canceled) {
+		assert.NoError(t, err)
+	}
+}
+
 func Test_coalesce(t *testing.T) {
 	t.Parallel()
 
@@ -40,3 +151,14 @@ func Test_convertUsageReporterRecords(t *testing.T) {
 		lastSignedInAt: tm1,
 	}}), "should leave empty email")
 }
+
+func Test_latest(t *testing.T) {
+	t.Parallel()
+
+	tm1 := time.Date(2024, time.September, 11, 11, 56, 0, 0, time.UTC)
+	tm2 := time.Date(2024, time.September, 12, 11, 56, 0, 0, time.UTC)
+
+	assert.Equal(t, tm2, latest(tm1, tm2))
+	assert.Equal(t, tm2, latest(tm2, tm1), "should ignore ordering")
+	assert.Equal(t, tm1, latest(tm1, time.Time{}), "should handle zero time")
+}