fix tests

Author: rt-joe

main.go

@@ -133,7 +133,7 @@ func main() {
 
 	sugar.Debugf("Sending logs to timeseries database %s in AWS Region %s", cfg.databaseName, cfg.awsRegion)
 
-	timeStreamAdapter := newTimeStreamAdapter(sugar, cfg, nil, nil)
+	timeStreamAdapter := newTimeStreamAdapter(sugar, cfg, newTimestreamQueryPaginator, nil, nil)
 	if err := serve(ctx, sugar, cfg.listenAddr, timeStreamAdapter); err != nil {
 		sugar.Errorw("Failed to listen", "addr", cfg.listenAddr, "err", err)
 		os.Exit(1)

timestream.go

@@ -51,12 +51,20 @@ type TimestreamWriteApi interface {
 	WriteRecords(ctx context.Context, params *timestreamwrite.WriteRecordsInput, optFns ...func(*timestreamwrite.Options)) (*timestreamwrite.WriteRecordsOutput, error)
 }
 
+type NewQueryPaginator func(client TimestreamQueryApi, params *timestreamquery.QueryInput, optFns ...func(*timestreamquery.QueryPaginatorOptions)) PaginatorApi
+
+type PaginatorApi interface {
+	HasMorePages() bool
+	NextPage(ctx context.Context, optFns ...func(*timestreamquery.Options)) (*timestreamquery.QueryOutput, error)
+}
+
 type TimeStreamAdapter struct {
 	databaseName string
 	logger       *zap.SugaredLogger
 	tableName    string
 	TimestreamQueryApi
 	TimestreamWriteApi
+	NewQueryPaginator
 }
 
 type queryTask struct {
@@ -69,7 +77,11 @@ type writeTask struct {
 	dimensions  []writetypes.Dimension
 }
 
-func newTimeStreamAdapter(logger *zap.SugaredLogger, cfg *adapterCfg, writeSvc *timestreamwrite.Client, querySvc *timestreamquery.Client) TimeStreamAdapter {
+func newTimestreamQueryPaginator(client TimestreamQueryApi, params *timestreamquery.QueryInput, optFns ...func(*timestreamquery.QueryPaginatorOptions)) PaginatorApi {
+	return timestreamquery.NewQueryPaginator(client, params, optFns...)
+}
+
+func newTimeStreamAdapter(logger *zap.SugaredLogger, cfg *adapterCfg, newQueryPaginator NewQueryPaginator, writeSvc *timestreamwrite.Client, querySvc *timestreamquery.Client) TimeStreamAdapter {
 	tr := &http.Transport{
 		ResponseHeaderTimeout: 20 * time.Second,
 		// Using DefaultTransport values for other parameters: https://golang.org/pkg/net/http/#RoundTripper
@@ -105,6 +117,7 @@ func newTimeStreamAdapter(logger *zap.SugaredLogger, cfg *adapterCfg, writeSvc *
 	return TimeStreamAdapter{
 		TimestreamQueryApi: querySvc,
 		TimestreamWriteApi: writeSvc,
+		NewQueryPaginator:  newQueryPaginator,
 		databaseName:       cfg.databaseName,
 		logger:             logger,
 		tableName:          cfg.tableName,
@@ -229,7 +242,7 @@ func (t TimeStreamAdapter) runReadRequestQuery(ctx context.Context, q *prompb.Qu
 	}
 
 	var timeSeries []*prompb.TimeSeries
-	paginator := timestreamquery.NewQueryPaginator(t, &timestreamquery.QueryInput{
+	paginator := t.NewQueryPaginator(t, &timestreamquery.QueryInput{
 		QueryString: aws.String(task.query),
 	})
 

timestream_test.go

@@ -39,6 +39,7 @@ var (
 	timeStreamAdapter = &TimeStreamAdapter{
 		TimestreamQueryApi: &TimeStreamQueryMock{},
 		TimestreamWriteApi: &TimeStreamWriterMock{},
+		NewQueryPaginator:  NewPaginatorMock,
 		databaseName:       "mockDatabase",
 		logger:             zap.NewNop().Sugar(),
 		tableName:          "mockTable",
@@ -176,6 +177,31 @@ type TimeStreamWriterMock struct{}
 
 type TimeStreamQueryMock struct{}
 
+type PaginatorMock struct {
+	callCount int
+}
+
+func NewPaginatorMock(client TimestreamQueryApi, params *timestreamquery.QueryInput, optFns ...func(*timestreamquery.QueryPaginatorOptions)) PaginatorApi {
+	return &PaginatorMock{}
+}
+
+func (p *PaginatorMock) HasMorePages() bool {
+	p.callCount++
+	// Return true for the first two calls, then false.
+	return p.callCount <= 2
+}
+
+func (p *PaginatorMock) NextPage(ctx context.Context, optFns ...func(*timestreamquery.Options)) (*timestreamquery.QueryOutput, error) {
+	switch p.callCount {
+	case 1:
+		return queryOutput1, nil
+	case 2:
+		return queryOutput2, nil
+	default:
+		return nil, errors.New("no more pages")
+	}
+}
+
 func (t *TimeStreamWriterMock) WriteRecords(ctx context.Context, input *timestreamwrite.WriteRecordsInput, optFns ...func(*timestreamwrite.Options)) (*timestreamwrite.WriteRecordsOutput, error) {
 	for _, i := range input.Records {
 		if *i.MeasureName == "sample_name_error" {
@@ -195,13 +221,6 @@ func (t *TimeStreamQueryMock) Query(ctx context.Context, input *timestreamquery.
 	return nil, nil
 }
 
-func (t *TimeStreamQueryMock) QueryPages(ctx context.Context, input *timestreamquery.QueryInput, handler func(*timestreamquery.QueryOutput, bool) bool, optFns ...func(*timestreamquery.Options)) error {
-	handler(queryOutput0, false)
-	handler(queryOutput1, false)
-	handler(queryOutput2, true)
-	return nil
-}
-
 func TestTimeSteamAdapter_readLabels(t *testing.T) {
 	type args struct {
 		labels []prompb.Label