From 4a103890d3db68328163a152e37dbcd2a416e97b Mon Sep 17 00:00:00 2001
From: Neal Beeken <neal.beeken@mongodb.com>
Date: Wed, 21 Jul 2021 19:48:49 -0400
Subject: [PATCH] fix(NODE-3442): AsyncIterator has incorrect return type
 (#2916)

---
 package.json                                  |  1 +
 src/cursor/abstract_cursor.ts                 |  9 +++++--
 test/functional/cursor_async_iterator.test.js | 24 +++++++++++++++----
 test/types/community/cursor.test-d.ts         |  2 +-
 test/unit/execute_legacy_operation.test.js    | 15 +++---------
 5 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/package.json b/package.json
index 6074d01760..8f8b19e87d 100644
--- a/package.json
+++ b/package.json
@@ -57,6 +57,7 @@
     "@types/whatwg-url": "^8.2.1",
     "@typescript-eslint/eslint-plugin": "^4.19.0",
     "@typescript-eslint/parser": "^4.19.0",
+    "bluebird": "^3.7.2",
     "chai": "^4.2.0",
     "chai-subset": "^1.6.0",
     "chalk": "^4.1.0",
diff --git a/src/cursor/abstract_cursor.ts b/src/cursor/abstract_cursor.ts
index e2addc9738..9593749866 100644
--- a/src/cursor/abstract_cursor.ts
+++ b/src/cursor/abstract_cursor.ts
@@ -221,9 +221,14 @@ export abstract class AbstractCursor<
     return this[kDocuments].splice(0, number ?? this[kDocuments].length);
   }
 
-  [Symbol.asyncIterator](): AsyncIterator<TSchema | null> {
+  [Symbol.asyncIterator](): AsyncIterator<TSchema, void> {
     return {
-      next: () => this.next<TSchema>().then(value => ({ value, done: value === null }))
+      next: () =>
+        this.next().then(value =>
+          value !== null && value !== undefined
+            ? { value, done: false }
+            : { value: undefined, done: true }
+        )
     };
   }
 
diff --git a/test/functional/cursor_async_iterator.test.js b/test/functional/cursor_async_iterator.test.js
index 1830902929..944235b92d 100644
--- a/test/functional/cursor_async_iterator.test.js
+++ b/test/functional/cursor_async_iterator.test.js
@@ -2,6 +2,7 @@
 
 const { expect } = require('chai');
 const Sinon = require('sinon');
+const { Promise: BluebirdPromise } = require('bluebird');
 
 describe('Cursor Async Iterator Tests', function () {
   context('default promise library', function () {
@@ -87,11 +88,12 @@ describe('Cursor Async Iterator Tests', function () {
   context('custom promise library', () => {
     let client, collection, promiseSpy;
     before(async function () {
-      class CustomPromise extends Promise {}
-      promiseSpy = Sinon.spy(CustomPromise.prototype, 'then');
-      client = this.configuration.newClient({}, { promiseLibrary: CustomPromise });
+      promiseSpy = Sinon.spy(BluebirdPromise.prototype, 'then');
+      client = this.configuration.newClient({}, { promiseLibrary: BluebirdPromise });
 
-      await client.connect();
+      const connectPromise = client.connect();
+      expect(connectPromise).to.be.instanceOf(BluebirdPromise);
+      await connectPromise;
       const docs = Array.from({ length: 1 }).map((_, index) => ({ foo: index, bar: 1 }));
 
       collection = client.db(this.configuration.db).collection('async_cursor_tests');
@@ -121,5 +123,19 @@ describe('Cursor Async Iterator Tests', function () {
       expect(countBeforeIteration).to.not.equal(promiseSpy.callCount);
       expect(promiseSpy.called).to.equal(true);
     });
+
+    it('should properly use custom promise manual iteration', async function () {
+      const cursor = collection.find();
+
+      const iterator = cursor[Symbol.asyncIterator]();
+      let isDone;
+      do {
+        const promiseFromIterator = iterator.next();
+        expect(promiseFromIterator).to.be.instanceOf(BluebirdPromise);
+        const { done, value } = await promiseFromIterator;
+        if (done) expect(value).to.be.a('undefined');
+        isDone = done;
+      } while (!isDone);
+    });
   });
 });
diff --git a/test/types/community/cursor.test-d.ts b/test/types/community/cursor.test-d.ts
index ba1b13d8f8..2ea825b388 100644
--- a/test/types/community/cursor.test-d.ts
+++ b/test/types/community/cursor.test-d.ts
@@ -94,7 +94,7 @@ expectType<{ name: string }[]>(
 
 void async function () {
   for await (const item of cursor) {
-    if (!item) break;
+    expectNotType<{ foo: number } | null>(item);
     expectType<number>(item.foo);
   }
 };
diff --git a/test/unit/execute_legacy_operation.test.js b/test/unit/execute_legacy_operation.test.js
index 4d7fd049e8..f85babb062 100644
--- a/test/unit/execute_legacy_operation.test.js
+++ b/test/unit/execute_legacy_operation.test.js
@@ -28,9 +28,8 @@ describe('executeLegacyOperation', function () {
     expect(caughtError).to.equal(expectedError);
   });
 
-  it('should reject promise with errors on throw errors, and rethrow error', function (done) {
+  it('should reject promise with errors on throw errors, and rethrow error', function () {
     const expectedError = new Error('THIS IS AN ERROR');
-    let callbackError;
 
     const topology = {
       logicalSessionTimeoutMinutes: null
@@ -39,18 +38,10 @@ describe('executeLegacyOperation', function () {
       throw expectedError;
     };
 
-    const callback = err => (callbackError = err);
     const options = { skipSessions: true };
 
-    executeLegacyOperation(topology, operation, [{}, null], options).then(null, callback);
-
-    setTimeout(() => {
-      try {
-        expect(callbackError).to.equal(expectedError);
-        done();
-      } catch (e) {
-        done(e);
-      }
+    return executeLegacyOperation(topology, operation, [{}, null], options).then(null, err => {
+      expect(err).to.equal(expectedError);
     });
   });
 });