GODRIVER-2161 Add NewSingleResultFromDocument and NewCursorFromDocuments functions (#806)

Author: benjirewis

mongo/cursor.go

@@ -15,6 +15,7 @@ import (
 
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/bsoncodec"
+	"go.mongodb.org/mongo-driver/x/bsonx"
 	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 	"go.mongodb.org/mongo-driver/x/mongo/driver"
 	"go.mongodb.org/mongo-driver/x/mongo/driver/session"
@@ -67,6 +68,47 @@ func newEmptyCursor() *Cursor {
 	return &Cursor{bc: driver.NewEmptyBatchCursor()}
 }
 
+// NewCursorFromDocuments creates a new Cursor pre-loaded with the provided documents, error and registry. If no registry is provided,
+// bson.DefaultRegistry will be used.
+//
+// The documents parameter must be a slice of documents. The slice may be nil or empty, but all elements must be non-nil.
+func NewCursorFromDocuments(documents []interface{}, err error, registry *bsoncodec.Registry) (*Cursor, error) {
+	if registry == nil {
+		registry = bson.DefaultRegistry
+	}
+
+	// Convert documents slice to a sequence-style byte array.
+	var docsBytes []byte
+	for _, doc := range documents {
+		switch t := doc.(type) {
+		case nil:
+			return nil, ErrNilDocument
+		case bsonx.Doc:
+			doc = t.Copy()
+		case []byte:
+			// Slight optimization so we'll just use MarshalBSON and not go through the codec machinery.
+			doc = bson.Raw(t)
+		}
+		var marshalErr error
+		docsBytes, marshalErr = bson.MarshalAppendWithRegistry(registry, docsBytes, doc)
+		if marshalErr != nil {
+			return nil, marshalErr
+		}
+	}
+
+	c := &Cursor{
+		bc:       driver.NewBatchCursorFromDocuments(docsBytes),
+		registry: registry,
+		err:      err,
+	}
+
+	// Initialize batch and batchLength here. The underlying batch cursor will be preloaded with the
+	// provided contents, and thus already has a batch before calls to Next/TryNext.
+	c.batch = c.bc.Batch()
+	c.batchLength = c.bc.Batch().DocumentCount()
+	return c, nil
+}
+
 // ID returns the ID of this cursor, or 0 if the cursor has been closed or exhausted.
 func (c *Cursor) ID() int64 { return c.bc.ID() }
 

mongo/cursor_test.go

@@ -8,6 +8,7 @@ package mongo
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"go.mongodb.org/mongo-driver/bson"
@@ -179,3 +180,57 @@ func TestCursor(t *testing.T) {
 		})
 	})
 }
+
+func TestNewCursorFromDocuments(t *testing.T) {
+	// Mock documents returned by Find in a Cursor.
+	t.Run("mock Find", func(t *testing.T) {
+		findResult := []interface{}{
+			bson.D{{"_id", 0}, {"foo", "bar"}},
+			bson.D{{"_id", 1}, {"baz", "qux"}},
+			bson.D{{"_id", 2}, {"quux", "quuz"}},
+		}
+		cur, err := NewCursorFromDocuments(findResult, nil, nil)
+		assert.Nil(t, err, "NewCursorFromDocuments error: %v", err)
+
+		// Assert that decoded documents are as expected.
+		var i int
+		for cur.Next(context.Background()) {
+			docBytes, err := bson.Marshal(findResult[i])
+			assert.Nil(t, err, "Marshal error: %v", err)
+			expectedDecoded := bson.Raw(docBytes)
+
+			var decoded bson.Raw
+			err = cur.Decode(&decoded)
+			assert.Nil(t, err, "Decode error: %v", err)
+			assert.Equal(t, expectedDecoded, decoded,
+				"expected decoded document %v of Cursor to be %v, got %v",
+				i, expectedDecoded, decoded)
+			i++
+		}
+		assert.Equal(t, 3, i, "expected 3 calls to cur.Next, got %v", i)
+
+		// Check for error on Cursor.
+		assert.Nil(t, cur.Err(), "Cursor error: %v", cur.Err())
+
+		// Assert that a call to cur.Close will not fail.
+		err = cur.Close(context.Background())
+		assert.Nil(t, err, "Close error: %v", err)
+	})
+
+	// Mock an error in a Cursor.
+	t.Run("mock Find with error", func(t *testing.T) {
+		mockErr := fmt.Errorf("mock error")
+		findResult := []interface{}{bson.D{{"_id", 0}, {"foo", "bar"}}}
+		cur, err := NewCursorFromDocuments(findResult, mockErr, nil)
+		assert.Nil(t, err, "NewCursorFromDocuments error: %v", err)
+
+		// Assert that a call to Next will return false because of existing error.
+		next := cur.Next(context.Background())
+		assert.False(t, next, "expected call to Next to return false, got true")
+
+		// Check for error on Cursor.
+		assert.NotNil(t, cur.Err(), "expected Cursor error, got nil")
+		assert.Equal(t, mockErr, cur.Err(), "expected Cursor error %v, got %v",
+			mockErr, cur.Err())
+	})
+}

mongo/integration/mock_find_test.go

@@ -0,0 +1,126 @@
+// Copyright (C) MongoDB, Inc. 2021-present.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+package integration
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/bsoncodec"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+// finder is an object that implements FindOne and Find.
+type finder interface {
+	FindOne(ctx context.Context, filter interface{}, opts ...*options.FindOneOptions) *mongo.SingleResult
+	Find(context.Context, interface{}, ...*options.FindOptions) (*mongo.Cursor, error)
+}
+
+// mockFinder implements finder.
+type mockFinder struct {
+	docs     []interface{}
+	err      error
+	registry *bsoncodec.Registry
+}
+
+// FindOne mocks a findOne operation using NewSingleResultFromDocument.
+func (mf *mockFinder) FindOne(_ context.Context, _ interface{}, _ ...*options.FindOneOptions) *mongo.SingleResult {
+	return mongo.NewSingleResultFromDocument(mf.docs[0], mf.err, mf.registry)
+}
+
+// Find mocks a find operation using NewCursorFromDocuments.
+func (mf *mockFinder) Find(context.Context, interface{}, ...*options.FindOptions) (*mongo.Cursor, error) {
+	return mongo.NewCursorFromDocuments(mf.docs, mf.err, mf.registry)
+}
+
+// ShopItem is an item with an associated ID and price.
+type ShopItem struct {
+	ID    int     `bson:"id"`
+	Price float64 `bson:"price"`
+}
+
+// getItem is an example function using the interface finder to test the mocking of SingleResult.
+func getItem(f finder, id int) (*ShopItem, error) {
+	res := f.FindOne(context.Background(), bson.D{{"id", id}})
+	var item ShopItem
+	err := res.Decode(&item)
+	return &item, err
+}
+
+// getItems is an example function using the interface finder to test the mocking of Cursor.
+func getItems(f finder) ([]ShopItem, error) {
+	cur, err := f.Find(context.Background(), bson.D{})
+	if err != nil {
+		return nil, err
+	}
+
+	var items []ShopItem
+	err = cur.All(context.Background(), &items)
+	return items, err
+}
+
+func TestMockFind(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().CreateClient(false))
+	defer mt.Close()
+
+	insertItems := []interface{}{
+		ShopItem{ID: 0, Price: 1.5},
+		ShopItem{ID: 1, Price: 5.7},
+		ShopItem{ID: 2, Price: 0.25},
+	}
+	insertItem := []interface{}{ShopItem{ID: 1, Price: 5.7}}
+
+	mt.Run("mongo.Collection can be passed as interface", func(mt *mtest.T) {
+		// Actually insert documents to collection.
+		_, err := mt.Coll.InsertMany(mtest.Background, insertItems)
+		assert.Nil(mt, err, "InsertMany error: %v", err)
+
+		// Assert that FindOne behaves as expected.
+		shopItem, err := getItem(mt.Coll, 1)
+		assert.Nil(mt, err, "getItem error: %v", err)
+		assert.Equal(mt, 1, shopItem.ID)
+		assert.Equal(mt, 5.7, shopItem.Price)
+
+		// Assert that Find behaves as expected.
+		shopItems, err := getItems(mt.Coll)
+		assert.Nil(mt, err, "getItems error: %v", err)
+		for i, shopItem := range shopItems {
+			expectedItem := insertItems[i].(ShopItem)
+			assert.Equal(mt, expectedItem.ID, shopItem.ID)
+			assert.Equal(mt, expectedItem.Price, shopItem.Price)
+		}
+	})
+
+	mt.Run("FindOne can be mocked", func(mt *mtest.T) {
+		// Mock a FindOne result with mockFinder.
+		mf := &mockFinder{docs: insertItem, err: nil, registry: nil}
+
+		// Assert that FindOne behaves as expected.
+		shopItem, err := getItem(mf, 1)
+		assert.Nil(mt, err, "getItem error: %v", err)
+		assert.Equal(mt, 1, shopItem.ID)
+		assert.Equal(mt, 5.7, shopItem.Price)
+	})
+
+	mt.Run("Find can be mocked", func(mt *mtest.T) {
+		// Mock a Find result with mockFinder.
+		mf := &mockFinder{docs: insertItems, err: nil, registry: nil}
+
+		// Assert that Find behaves as expected.
+		shopItems, err := getItems(mf)
+		assert.Nil(mt, err, "getItems error: %v", err)
+		for i, shopItem := range shopItems {
+			expectedItem := insertItems[i].(ShopItem)
+			assert.Equal(mt, expectedItem.ID, shopItem.ID)
+			assert.Equal(mt, expectedItem.Price, shopItem.Price)
+		}
+	})
+}

mongo/single_result.go

@@ -28,6 +28,31 @@ type SingleResult struct {
 	reg *bsoncodec.Registry
 }
 
+// NewSingleResultFromDocument creates a SingleResult with the provided error, registry, and an underlying Cursor pre-loaded with
+// the provided document, error and registry. If no registry is provided, bson.DefaultRegistry will be used. If an error distinct
+// from the one provided occurs during creation of the SingleResult, that error will be stored on the returned SingleResult.
+//
+// The document parameter must be a non-nil document.
+func NewSingleResultFromDocument(document interface{}, err error, registry *bsoncodec.Registry) *SingleResult {
+	if document == nil {
+		return &SingleResult{err: ErrNilDocument}
+	}
+	if registry == nil {
+		registry = bson.DefaultRegistry
+	}
+
+	cur, createErr := NewCursorFromDocuments([]interface{}{document}, err, registry)
+	if createErr != nil {
+		return &SingleResult{err: createErr}
+	}
+
+	return &SingleResult{
+		cur: cur,
+		err: err,
+		reg: registry,
+	}
+}
+
 // Decode will unmarshal the document represented by this SingleResult into v. If there was an error from the operation
 // that created this SingleResult, that error will be returned. If the operation returned no documents, Decode will
 // return ErrNoDocuments.
@@ -71,6 +96,7 @@ func (sr *SingleResult) setRdrContents() error {
 		return nil
 	case sr.cur != nil:
 		defer sr.cur.Close(context.TODO())
+
 		if !sr.cur.Next(context.TODO()) {
 			if err := sr.cur.Err(); err != nil {
 				return err

mongo/single_result_test.go

@@ -7,7 +7,9 @@
 package mongo
 
 import (
+	"context"
 	"errors"
+	"fmt"
 	"testing"
 
 	"go.mongodb.org/mongo-driver/bson"
@@ -49,3 +51,58 @@ func TestSingleResult(t *testing.T) {
 		assert.Equal(t, ErrNoDocuments, sr.Err(), "expected error %v, got %v", ErrNoDocuments, sr.Err())
 	})
 }
+
+func TestNewSingleResultFromDocument(t *testing.T) {
+	// Mock a document returned by FindOne in SingleResult.
+	t.Run("mock FindOne", func(t *testing.T) {
+		findOneResult := bson.D{{"_id", 2}, {"foo", "bar"}}
+		res := NewSingleResultFromDocument(findOneResult, nil, nil)
+
+		// Assert that first, decoded document is as expected.
+		findOneResultBytes, err := bson.Marshal(findOneResult)
+		assert.Nil(t, err, "Marshal error: %v", err)
+		expectedDecoded := bson.Raw(findOneResultBytes)
+		decoded, err := res.DecodeBytes()
+		assert.Nil(t, err, "DecodeBytes error: %v", err)
+		assert.Equal(t, expectedDecoded, decoded,
+			"expected decoded SingleResult to be %v, got %v", expectedDecoded, decoded)
+
+		// Assert that RDR contents are set correctly after Decode.
+		assert.NotNil(t, res.rdr, "expected non-nil rdr contents")
+		assert.Equal(t, expectedDecoded, res.rdr,
+			"expected RDR contents to be %v, got %v", expectedDecoded, res.rdr)
+
+		// Assert that a call to cur.Next will return false, as there was only one document in
+		// the slice passed to NewSingleResultFromDocument.
+		next := res.cur.Next(context.Background())
+		assert.False(t, next, "expected call to Next to return false, got true")
+
+		// Check for error on SingleResult.
+		assert.Nil(t, res.Err(), "SingleResult error: %v", res.Err())
+
+		// Assert that a call to cur.Close will not fail.
+		err = res.cur.Close(context.Background())
+		assert.Nil(t, err, "Close error: %v", err)
+	})
+
+	// Mock an error in SingleResult.
+	t.Run("mock FindOne with error", func(t *testing.T) {
+		mockErr := fmt.Errorf("mock error")
+		res := NewSingleResultFromDocument(bson.D{}, mockErr, nil)
+
+		// Assert that decoding returns the mocked error.
+		_, err := res.DecodeBytes()
+		assert.NotNil(t, err, "expected DecodeBytes error, got nil")
+		assert.Equal(t, mockErr, err, "expected error %v, got %v", mockErr, err)
+
+		// Check for error on SingleResult.
+		assert.NotNil(t, res.Err(), "expected SingleResult error, got nil")
+		assert.Equal(t, mockErr, res.Err(), "expected SingleResult error %v, got %v",
+			mockErr, res.Err())
+
+		// Assert that error is propagated to underlying cursor.
+		assert.NotNil(t, res.cur.err, "expected underlying cursor, got nil")
+		assert.Equal(t, mockErr, res.cur.err, "expected underlying cursor %v, got %v",
+			mockErr, res.cur.err)
+	})
+}

x/mongo/driver/batch_cursor.go

@@ -184,6 +184,21 @@ func NewEmptyBatchCursor() *BatchCursor {
 	return &BatchCursor{currentBatch: new(bsoncore.DocumentSequence)}
 }
 
+// NewBatchCursorFromDocuments returns a batch cursor with current batch set to a sequence-style
+// DocumentSequence containing the provided documents.
+func NewBatchCursorFromDocuments(documents []byte) *BatchCursor {
+	return &BatchCursor{
+		currentBatch: &bsoncore.DocumentSequence{
+			Data:  documents,
+			Style: bsoncore.SequenceStyle,
+		},
+		// BatchCursors created with this function have no associated ID nor server, so no getMore
+		// calls will be made.
+		id:     0,
+		server: nil,
+	}
+}
+
 // ID returns the cursor ID for this batch cursor.
 func (bc *BatchCursor) ID() int64 {
 	return bc.id