@@ -303,15 +303,26 @@ operations.set('insertMany', async ({ entities, operation }) => {
303
303
} ) ;
304
304
305
305
operations . set ( 'iterateUntilDocumentOrError' , async ( { entities, operation } ) => {
306
- const cursorOrStream =
307
- entities . getEntity ( 'stream' , operation . object , false ) ??
308
- entities . getEntity ( 'cursor' , operation . object , false ) ;
306
+ function getChangeStream ( ) : UnifiedChangeStream | null {
307
+ try {
308
+ const changeStream = entities . getEntity ( 'stream' , operation . object ) ;
309
+ return changeStream ;
310
+ } catch ( e ) {
311
+ return null ;
312
+ }
313
+ }
309
314
310
- if ( cursorOrStream == null ) {
311
- throw new Error ( `Unable to get entity - expect stream or cursor but received ${ operation } ` ) ;
315
+ const changeStream = getChangeStream ( ) ;
316
+ if ( changeStream == null ) {
317
+ // iterateUntilDocumentOrError is used for changes streams and regular cursors.
318
+ // we have no other way to distinguish which scenario we are testing when we run an
319
+ // iterateUntilDocumentOrError operation, so we first try to get the changeStream and
320
+ // if that fails, we know we need to get a cursor
321
+ const cursor = entities . getEntity ( 'cursor' , operation . object ) ;
322
+ return await cursor . next ( ) ;
312
323
}
313
324
314
- return await cursorOrStream . next ( ) ;
325
+ return changeStream . next ( ) ;
315
326
} ) ;
316
327
317
328
operations . set ( 'listCollections' , async ( { entities, operation } ) => {
0 commit comments