Allow SingleResult Decode and DecodeBytes methods to be called more than once.

GODRIVER-1004

Change-Id: Ia1a7471dda022414acccba2b81e13822a24285fb

Author: Divjot Arora

mongo/single_result.go

@@ -39,51 +39,49 @@ func (sr *SingleResult) Decode(v interface{}) error {
 	if sr.reg == nil {
 		return bson.ErrNilRegistry
 	}
-	switch {
-	case sr.rdr != nil:
-		if v == nil {
-			return nil
-		}
-		return bson.UnmarshalWithRegistry(sr.reg, sr.rdr, v)
-	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
-			}
-			return ErrNoDocuments
-		}
-		if v == nil {
-			return nil
-		}
-		return sr.cur.Decode(v)
-	}
 
-	return ErrNoDocuments
+	if sr.err = sr.setRdrContents(); sr.err != nil {
+		return sr.err
+	}
+	return bson.UnmarshalWithRegistry(sr.reg, sr.rdr, v)
 }
 
 // DecodeBytes will return a copy of the document as a bson.Raw. If there was an
 // error from the operation that created this SingleResult then the error
 // will be returned. If there were no returned documents, ErrNoDocuments is
 // returned.
 func (sr *SingleResult) DecodeBytes() (bson.Raw, error) {
+	if sr.err != nil {
+		return nil, sr.err
+	}
+
+	if sr.err = sr.setRdrContents(); sr.err != nil {
+		return nil, sr.err
+	}
+	return sr.rdr, nil
+}
+
+// setRdrContents will set the contents of rdr by iterating the underlying cursor if necessary.
+func (sr *SingleResult) setRdrContents() error {
 	switch {
 	case sr.err != nil:
-		return nil, sr.err
+		return sr.err
 	case sr.rdr != nil:
-		return sr.rdr, nil
+		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 nil, err
+				return err
 			}
-			return nil, ErrNoDocuments
+
+			return ErrNoDocuments
 		}
-		return sr.cur.Current, nil
+		sr.rdr = sr.cur.Current
+		return nil
 	}
 
-	return nil, ErrNoDocuments
+	return ErrNoDocuments
 }
 
 // Err will return the error from the operation that created this SingleResult.

mongo/single_result_test.go

@@ -0,0 +1,48 @@
+// Copyright (C) MongoDB, Inc. 2017-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 mongo
+
+import (
+	"bytes"
+	"testing"
+
+	"go.mongodb.org/mongo-driver/bson"
+)
+
+func TestSingleResult(t *testing.T) {
+	t.Run("TestDecode", func(t *testing.T) {
+		t.Run("DecodeTwice", func(t *testing.T) {
+			// Test that Decode and DecodeBytes can be called more than once
+			c, err := newCursor(newTestBatchCursor(1, 1), bson.DefaultRegistry)
+			if err != nil {
+				t.Fatalf("error creating cursor: %v", err)
+			}
+
+			sr := &SingleResult{cur: c, reg: bson.DefaultRegistry}
+			var firstDecode, secondDecode bson.Raw
+			if err = sr.Decode(&firstDecode); err != nil {
+				t.Fatalf("error on first Decode call: %v", err)
+			}
+			if err = sr.Decode(&secondDecode); err != nil {
+				t.Fatalf("error on second Decode call: %v", err)
+			}
+			decodeBytes, err := sr.DecodeBytes()
+			if err != nil {
+				t.Fatalf("error on DecodeBytes call: %v", err)
+			}
+
+			if !bytes.Equal(firstDecode, secondDecode) {
+				t.Fatalf("Decode contents do not match; first returned %v, second returned %v",
+					firstDecode, secondDecode)
+			}
+			if !bytes.Equal(firstDecode, decodeBytes) {
+				t.Fatalf("Decode and DecodeBytes contents do not match; Decode returned %v, DecodeBytes "+
+					"returned %v", firstDecode, decodeBytes)
+			}
+		})
+	})
+}