docstore/all: Fix offset handling and extend test coverage (google#3409)

docstore/awsdynamodb/query.go

@@ -471,30 +471,36 @@ type documentIterator struct {
 
 func (it *documentIterator) Next(ctx context.Context, doc driver.Document) error {
 	// Skip the first 'n' documents where 'n' is the offset.
-	if it.offset > 0 && it.count < it.offset {
-		it.curr++
-		it.count++
-		return it.Next(ctx, doc)
+	for it.count < it.offset {
+		if err := it.next(ctx, doc, false); err != nil {
+			return err
+		}
 	}
+	return it.next(ctx, doc, true)
+}
+
+func (it *documentIterator) next(ctx context.Context, doc driver.Document, decode bool) error {
 	// Only start counting towards the limit after the offset has been reached.
-	if it.limit > 0 && it.count >= it.offset+it.limit || it.curr >= len(it.items) && it.last == nil {
+	if it.limit > 0 && it.count >= it.offset+it.limit {
 		return io.EOF
 	}
-	if it.curr >= len(it.items) {
+	// it.items can be empty after a call to it.qr.run, but unless it.last is nil there may be more items.
+	for it.curr >= len(it.items) {
 		// Make a new query request at the end of this page.
+		if it.last == nil {
+			return io.EOF
+		}
 		var err error
 		it.items, it.last, it.asFunc, err = it.qr.run(ctx, it.last)
 		if err != nil {
 			return err
 		}
 		it.curr = 0
 	}
-	// If there are no more items, return EOF.
-	if len(it.items) == 0 {
-		return io.EOF
-	}
-	if err := decodeDoc(&dyn.AttributeValue{M: it.items[it.curr]}, doc); err != nil {
-		return err
+	if decode {
+		if err := decodeDoc(&dyn.AttributeValue{M: it.items[it.curr]}, doc); err != nil {
+			return err
+		}
 	}
 	it.curr++
 	it.count++

docstore/awsdynamodb/query_test.go

@@ -15,13 +15,16 @@
 package awsdynamodb
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
+	dyn "github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"gocloud.dev/docstore/driver"
@@ -574,3 +577,103 @@ func TestCopyTopLevel(t *testing.T) {
 		}
 	}
 }
+
+func Test_documentIterator_Next(t *testing.T) {
+	type fields struct {
+		qr     *queryRunner
+		items  []map[string]*dyn.AttributeValue
+		curr   int
+		offset int
+		limit  int
+		count  int
+		last   map[string]*dyn.AttributeValue
+		asFunc func(i interface{}) bool
+	}
+	type args struct {
+		ctx context.Context
+		doc driver.Document
+	}
+	tests := []struct {
+		name    string
+		fields  fields
+		args    args
+		wantErr bool
+	}{
+		{
+			name: "nextWithNoDecodeError",
+			fields: fields{
+				qr: &queryRunner{},
+				items: []map[string]*dyn.AttributeValue{
+					{"key": {M: map[string]*dyn.AttributeValue{"key": {S: aws.String("value")}}}},
+				},
+				curr:   0,
+				offset: 0,
+				limit:  0,
+				count:  0,
+				last:   map[string]*dyn.AttributeValue{},
+			},
+			args: args{
+				ctx: context.Background(),
+				doc: drivertest.MustDocument(map[string]interface{}{}),
+			},
+			wantErr: false,
+		},
+		{
+			name: "nextWithDecodeError",
+			fields: fields{
+				qr: &queryRunner{},
+				items: []map[string]*dyn.AttributeValue{
+					{"key": {M: nil}}, // set M to nil to trigger decode error
+				},
+				curr:   0,
+				offset: 0,
+				limit:  0,
+				count:  0,
+				last:   map[string]*dyn.AttributeValue{},
+			},
+			args: args{
+				ctx: context.Background(),
+				doc: drivertest.MustDocument(map[string]interface{}{}),
+			},
+			wantErr: true,
+		},
+		{
+			name: "nextWhereCurrIsGreaterThanOrEqualToItemsAndLastIsNotNil",
+			fields: fields{
+				qr: &queryRunner{
+					scanIn: &dyn.ScanInput{},
+					// hack to return error from run
+					beforeRun: func(asFunc func(i interface{}) bool) error { return errors.New("invalid") },
+				},
+				items:  []map[string]*dyn.AttributeValue{{"key": {M: map[string]*dyn.AttributeValue{"key": {S: aws.String("value"), M: nil}}}}},
+				curr:   1,
+				offset: 0,
+				limit:  0,
+				count:  0,
+				last:   map[string]*dyn.AttributeValue{"key": {S: aws.String("value")}},
+			},
+			args: args{
+				ctx: context.Background(),
+				doc: drivertest.MustDocument(map[string]interface{}{}),
+			},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			it := &documentIterator{
+				qr:     tt.fields.qr,
+				items:  tt.fields.items,
+				curr:   tt.fields.curr,
+				offset: tt.fields.offset,
+				limit:  tt.fields.limit,
+				count:  tt.fields.count,
+				last:   tt.fields.last,
+				asFunc: tt.fields.asFunc,
+			}
+			if err := it.Next(tt.args.ctx, tt.args.doc); (err != nil) != tt.wantErr {
+				t.Errorf("documentIterator.Next() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}