pomerium/pkg/grpc/databroker/databroker_test.go

package databroker

import (
	"context"
	"net"
	"strconv"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	status "google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/anypb"
	"google.golang.org/protobuf/types/known/wrapperspb"
)

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)
		})
	}
}

type mockServer struct {
	DataBrokerServiceServer

	getAll func(context.Context, *GetAllRequest) (*GetAllResponse, error)
}

func (m *mockServer) GetAll(ctx context.Context, req *GetAllRequest) (*GetAllResponse, error) {
	return m.getAll(ctx, req)
}

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

	t.Run("resource exhausted", func(t *testing.T) {
		li, err := net.Listen("tcp", "127.0.0.1:0")
		if !assert.NoError(t, err) {
			return
		}
		defer li.Close()

		m := &mockServer{
			getAll: func(ctx context.Context, req *GetAllRequest) (*GetAllResponse, error) {
				any, _ := anypb.New(wrapperspb.String("TEST"))
				var records []*Record
				for i := 0; i < 1000000; i++ {
					records = append(records, &Record{
						Type: req.GetType(),
						Data: any,
					})
				}
				return &GetAllResponse{Records: records}, 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)

		res, err := GetAllPages(ctx, c, &GetAllRequest{
			Type: "TEST",
		})
		assert.Error(t, err)
		assert.Equal(t, codes.ResourceExhausted, status.Code(err))
		assert.Nil(t, res)
	})
	t.Run("with paging", func(t *testing.T) {
		li, err := net.Listen("tcp", "127.0.0.1:0")
		if !assert.NoError(t, err) {
			return
		}
		defer li.Close()

		m := &mockServer{
			getAll: func(ctx context.Context, req *GetAllRequest) (*GetAllResponse, error) {
				pageToken, _ := strconv.Atoi(req.GetPageToken())

				any, _ := anypb.New(wrapperspb.String("TEST"))
				var records []*Record
				for i := pageToken; i < pageToken+10000 && i < 1000000; i++ {
					records = append(records, &Record{
						Type: req.GetType(),
						Data: any,
					})
				}
				if len(records) == 0 {
					return &GetAllResponse{}, nil
				}
				return &GetAllResponse{Records: records, NextPageToken: strconv.Itoa(pageToken + 10000)}, 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)

		res, err := GetAllPages(ctx, c, &GetAllRequest{
			Type: "TEST",
		})
		assert.NoError(t, err)
		assert.NotEqual(t, codes.ResourceExhausted, status.Code(err))
		assert.Len(t, res.GetRecords(), 1000000)
	})
}