77package topology
88
99import (
10+ "context"
11+ "errors"
1012 "testing"
1113
1214 "github.com/mongodb/mongo-go-driver/bson"
15+ "github.com/mongodb/mongo-go-driver/core/connection"
16+ "github.com/mongodb/mongo-go-driver/core/description"
17+ "github.com/mongodb/mongo-go-driver/core/wiremessage"
18+ "github.com/mongodb/mongo-go-driver/internal"
1319 "github.com/stretchr/testify/assert"
1420)
1521
@@ -29,3 +35,151 @@ func TestCursorNextDoesNotPanicIfContextisNil(t *testing.T) {
2935 })
3036 assert .True (t , iterNext )
3137}
38+
39+ func TestCursorLoopsUntilDocAvailable (t * testing.T ) {
40+ // Next should loop until at least one doc is available
41+ // Here, the mock pool and connection implementations (below) write
42+ // empty batch responses a few times before returning a non-empty batch
43+
44+ s := createDefaultConnectedServer (t , false )
45+ c := cursor {
46+ id : 1 ,
47+ batch : bson .NewArray (),
48+ server : s ,
49+ }
50+
51+ assert .True (t , c .Next (nil ))
52+ }
53+
54+ func TestCursorReturnsFalseOnContextCancellation (t * testing.T ) {
55+ // Next should return false if an error occurs
56+ // here the error is the Context being cancelled
57+
58+ s := createDefaultConnectedServer (t , false )
59+ c := cursor {
60+ id : 1 ,
61+ batch : bson .NewArray (),
62+ server : s ,
63+ }
64+
65+ ctx , cancel := context .WithCancel (context .Background ())
66+
67+ cancel ()
68+
69+ assert .False (t , c .Next (ctx ))
70+ }
71+
72+ func TestCursorNextReturnsFalseIfErrorOccurred (t * testing.T ) {
73+ // Next should return false if an error occurs
74+ // here the error is an invalid namespace (""."")
75+
76+ s := createDefaultConnectedServer (t , true )
77+ c := cursor {
78+ id : 1 ,
79+ batch : bson .NewArray (),
80+ server : s ,
81+ }
82+ assert .False (t , c .Next (nil ))
83+ }
84+
85+ func TestCursorNextReturnsFalseIfResIdZeroAndNoMoreDocs (t * testing.T ) {
86+ // Next should return false if the cursor id is 0 and there are no documents in the next batch
87+
88+ c := cursor {id : 0 , batch : bson .NewArray ()}
89+ assert .False (t , c .Next (nil ))
90+ }
91+
92+ func createDefaultConnectedServer (t * testing.T , willErr bool ) * Server {
93+ s , err := ConnectServer (nil , "127.0.0.1" )
94+ s .pool = & mockPool {t : t , willErr : willErr }
95+ if err != nil {
96+ assert .Fail (t , "Server creation failed" )
97+ }
98+
99+ return s
100+ }
101+
102+ func createOKBatchReplyDoc (id int64 , batchDocs * bson.Array ) * bson.Document {
103+ return bson .NewDocument (
104+ bson .EC .Int32 ("ok" , 1 ),
105+ bson .EC .SubDocument (
106+ "cursor" ,
107+ bson .NewDocument (
108+ bson .EC .Int64 ("id" , id ),
109+ bson .EC .Array ("nextBatch" , batchDocs ))))
110+ }
111+
112+ // Mock Pool implementation
113+ type mockPool struct {
114+ t * testing.T
115+ willErr bool
116+ writes int // the number of wire messages written so far
117+ }
118+
119+ func (m * mockPool ) Get (ctx context.Context ) (connection.Connection , * description.Server , error ) {
120+ m .writes ++
121+ return & mockConnection {willErr : m .willErr , writes : m .writes }, nil , nil
122+ }
123+
124+ func (* mockPool ) Connect (ctx context.Context ) error {
125+ return nil
126+ }
127+
128+ func (* mockPool ) Disconnect (ctx context.Context ) error {
129+ return nil
130+ }
131+
132+ func (* mockPool ) Drain () error {
133+ return nil
134+ }
135+
136+ // Mock Connection implementation that
137+ type mockConnection struct {
138+ t * testing.T
139+ willErr bool
140+ writes int // the number of wire messages written so far
141+ }
142+
143+ // this mock will not actually write anything
144+ func (* mockConnection ) WriteWireMessage (ctx context.Context , wm wiremessage.WireMessage ) error {
145+ select {
146+ case <- ctx .Done ():
147+ return errors .New ("intentional mock error" )
148+ default :
149+ return nil
150+ }
151+ }
152+
153+ // mock a read by returning an empty cursor result until
154+ func (m * mockConnection ) ReadWireMessage (ctx context.Context ) (wiremessage.WireMessage , error ) {
155+ if m .writes < 4 {
156+ // write empty batch
157+ d := createOKBatchReplyDoc (2 , bson .NewArray ())
158+
159+ return internal .MakeReply (m .t , d ), nil
160+ } else if m .willErr {
161+ // write error
162+ return nil , errors .New ("intentional mock error" )
163+ } else {
164+ // write non-empty batch
165+ d := createOKBatchReplyDoc (2 , bson .NewArray (bson .VC .String ("a" )))
166+
167+ return internal .MakeReply (m .t , d ), nil
168+ }
169+ }
170+
171+ func (* mockConnection ) Close () error {
172+ return nil
173+ }
174+
175+ func (* mockConnection ) Expired () bool {
176+ return false
177+ }
178+
179+ func (* mockConnection ) Alive () bool {
180+ return true
181+ }
182+
183+ func (* mockConnection ) ID () string {
184+ return ""
185+ }
0 commit comments