pomerium/pkg/grpc/databroker/databroker_test.go

package databroker

import (
	"context"
	"fmt"
	"net"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/structpb"

	"github.com/pomerium/pomerium/internal/testutil"
	"github.com/pomerium/pomerium/pkg/protoutil"
)

func TestApplyOffsetAndLimit(t *testing.T) {
	cases := []struct {
		name          string
		records       []*Record
		offset, limit int
		expect        []*Record
	}{
		{
			name:    "empty",
			records: nil,
			offset:  10,
			limit:   5,
			expect:  nil,
		},
		{
			name:    "less than limit",
			records: []*Record{{Id: "A"}, {Id: "B"}, {Id: "C"}, {Id: "D"}},
			offset:  1,
			limit:   10,
			expect:  []*Record{{Id: "B"}, {Id: "C"}, {Id: "D"}},
		},
		{
			name:    "more than limit",
			records: []*Record{{Id: "A"}, {Id: "B"}, {Id: "C"}, {Id: "D"}, {Id: "E"}, {Id: "F"}, {Id: "G"}, {Id: "H"}},
			offset:  3,
			limit:   2,
			expect:  []*Record{{Id: "D"}, {Id: "E"}},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			actual, cnt := ApplyOffsetAndLimit(tc.records, tc.offset, tc.limit)
			assert.Equal(t, len(tc.records), cnt)
			assert.Equal(t, tc.expect, actual)
		})
	}
}

func TestInitialSync(t *testing.T) {
	ctx, clearTimeout := context.WithTimeout(context.Background(), time.Second*10)
	defer clearTimeout()

	li, err := net.Listen("tcp", "127.0.0.1:0")
	if !assert.NoError(t, err) {
		return
	}
	defer li.Close()

	r1 := new(Record)
	r2 := new(Record)

	m := &mockServer{
		syncLatest: func(_ *SyncLatestRequest, stream DataBrokerService_SyncLatestServer) error {
			stream.Send(&SyncLatestResponse{
				Response: &SyncLatestResponse_Record{
					Record: r1,
				},
			})
			stream.Send(&SyncLatestResponse{
				Response: &SyncLatestResponse_Record{
					Record: r2,
				},
			})
			stream.Send(&SyncLatestResponse{
				Response: &SyncLatestResponse_Versions{
					Versions: &Versions{
						LatestRecordVersion: 2,
						ServerVersion:       1,
					},
				},
			})
			return nil
		},
	}

	srv := grpc.NewServer()
	RegisterDataBrokerServiceServer(srv, m)
	go srv.Serve(li)

	cc, err := grpc.Dial(li.Addr().String(), grpc.WithInsecure())
	if !assert.NoError(t, err) {
		return
	}
	defer cc.Close()

	c := NewDataBrokerServiceClient(cc)

	records, recordVersion, serverVersion, err := InitialSync(ctx, c, new(SyncLatestRequest))
	assert.NoError(t, err)
	assert.Equal(t, uint64(2), recordVersion)
	assert.Equal(t, uint64(1), serverVersion)
	testutil.AssertProtoEqual(t, []*Record{r1, r2}, records)
}

func TestOptimumPutRequestsFromRecords(t *testing.T) {
	var records []*Record
	for i := 0; i < 10_000; i++ {
		s := structpb.NewStructValue(&structpb.Struct{
			Fields: map[string]*structpb.Value{
				"long_string": structpb.NewStringValue(strings.Repeat("x", 987)),
			},
		})
		records = append(records, &Record{
			Id:   fmt.Sprintf("%d", i),
			Data: protoutil.NewAny(s),
		})
	}
	requests := OptimumPutRequestsFromRecords(records)
	for _, request := range requests {
		assert.LessOrEqual(t, proto.Size(request), maxMessageSize)
		assert.GreaterOrEqual(t, proto.Size(request), maxMessageSize/2)
	}
}

type mockServer struct {
	DataBrokerServiceServer

	syncLatest func(empty *SyncLatestRequest, server DataBrokerService_SyncLatestServer) error
}

func (m *mockServer) SyncLatest(req *SyncLatestRequest, stream DataBrokerService_SyncLatestServer) error {
	return m.syncLatest(req, stream)
}