WIP: Fixing test failures

Author: W-A-James

src/sdam/server.ts

@@ -30,7 +30,6 @@ import {
   MongoInvalidArgumentError,
   MongoNetworkError,
   MongoNetworkTimeoutError,
-  MongoRuntimeError,
   MongoServerClosedError,
   type MongoServerError,
   needsRetryableWriteLabel
@@ -315,56 +314,53 @@ export class Server extends TypedEventEmitter<ServerEvents> {
     //       require special logic to decrement it again, or would double increment (the load
     //       balanced code makes a recursive call).  Instead, we increment the count after this
     //       check.
-
     if (this.loadBalanced && session && conn == null && isPinnableCommand(cmd, session)) {
       conn = await this.pool.checkOut();
       session.pin(conn);
       return this.command(ns, cmd, finalOptions);
     }
 
-    const runCommand = async (conn: Connection) => {
-      this.incrementOperationCount();
-      try {
-        return await conn.command(ns, cmd, finalOptions);
-      } catch (e) {
-        throw this.decorateAndHandleError(conn, cmd, finalOptions, e);
-      } finally {
-        this.decrementOperationCount();
-      }
-    };
-
-    const shouldReauth = (e: AnyError): boolean => {
-      return e instanceof MongoError && e.code === MONGODB_ERROR_CODES.Reauthenticate;
-    };
-
-    // FIXME: This is where withConnection logic should go
+    this.incrementOperationCount();
     if (!conn) {
       try {
         conn = await this.pool.checkOut();
-      } catch (e) {
-        if (!e) throw new MongoRuntimeError('Failed to create connection without error');
-        if (!(e instanceof PoolClearedError)) this.handleError(e);
+      } catch (checkoutError) {
+        if (!(checkoutError instanceof PoolClearedError)) this.handleError(checkoutError);
 
-        throw e;
+        throw checkoutError;
       }
     }
 
     // Reauth flow
-    let rv;
     try {
-      rv = await runCommand(conn);
-    } catch (e) {
-      if (shouldReauth(e)) {
+      return await this.executeCommand(conn, ns, cmd, finalOptions);
+    } catch (operationError) {
+      if (
+        operationError instanceof MongoError &&
+        operationError.code === MONGODB_ERROR_CODES.Reauthenticate
+      ) {
         conn = await this.pool.reauthenticateAsync(conn);
-        rv = await runCommand(conn);
+        return await this.executeCommand(conn, ns, cmd, finalOptions);
       } else {
-        throw e;
+        throw operationError;
       }
     } finally {
       this.pool.checkIn(conn);
+      this.decrementOperationCount();
     }
+  }
 
-    return rv;
+  private async executeCommand(
+    conn: Connection,
+    ns: MongoDBNamespace,
+    cmd: Document,
+    options?: CommandOptions
+  ): Promise<Document> {
+    try {
+      return await conn.command(ns, cmd, options);
+    } catch (commandOptions) {
+      throw this.decorateAndHandleError(conn, cmd, options, commandOptions);
+    }
   }
 
   /**
@@ -415,6 +411,10 @@ export class Server extends TypedEventEmitter<ServerEvents> {
     }
   }
 
+  /**
+   * Ensure that error is properly decorated and internal state is updated before throwing
+   * @internal
+   */
   private decorateAndHandleError(
     connection: Connection,
     cmd: Document,

test/integration/retryable-writes/retryable_writes.spec.prose.test.ts

@@ -278,15 +278,19 @@ describe('Retryable Writes Spec Prose', () => {
         const serverCommandStub = sinon.stub(Server.prototype, 'command');
         serverCommandStub
           .onCall(0)
-          .yieldsRight(
-            new MongoWriteConcernError({ errorLabels: ['RetryableWriteError'], code: 91 }, {})
+          .returns(
+            Promise.reject(
+              new MongoWriteConcernError({ errorLabels: ['RetryableWriteError'], code: 91 }, {})
+            )
           );
         serverCommandStub
           .onCall(1)
-          .yieldsRight(
-            new MongoWriteConcernError(
-              { errorLabels: ['RetryableWriteError', 'NoWritesPerformed'], errorCode: 10107 },
-              {}
+          .returns(
+            Promise.reject(
+              new MongoWriteConcernError(
+                { errorLabels: ['RetryableWriteError', 'NoWritesPerformed'], errorCode: 10107 },
+                {}
+              )
             )
           );
 

test/integration/server-selection/operation_count.test.ts

@@ -103,9 +103,9 @@ describe('Server Operation Count Tests', function () {
         const server = Array.from(client.topology.s.servers.values())[0];
         expect(server.s.operationCount).to.equal(0);
 
-        sinon.stub(ConnectionPool.prototype, 'checkOut').callsFake(function (cb) {
-          cb(new Error('unable to checkout connection'), undefined);
-        });
+        sinon
+          .stub(ConnectionPool.prototype, 'checkOut')
+          .returns(Promise.reject(new Error('unable to checkout connection')));
         const commandSpy = sinon.spy(server, 'command');
 
         const error = await collection.findOne({ count: 1 }).catch(e => e);

test/tools/spec-runner/context.js

@@ -143,23 +143,17 @@ class TestRunnerContext {
     return Promise.all(cleanupPromises).then(cleanup, cleanup);
   }
 
-  targetedFailPoint(options) {
+  async targetedFailPoint(options) {
     const session = options.session;
     const failPoint = options.failPoint;
     expect(session.isPinned).to.be.true;
 
-    return new Promise((resolve, reject) => {
-      const serverOrConnection = session.loadBalanced
-        ? session.pinnedConnection
-        : session.transaction.server;
+    const serverOrConnection = session.loadBalanced
+      ? session.pinnedConnection
+      : session.transaction.server;
 
-      serverOrConnection.command(ns('admin.$cmd'), failPoint, undefined, err => {
-        if (err) return reject(err);
-
-        this.appliedFailPoints.push(failPoint);
-        resolve();
-      });
-    });
+    await serverOrConnection.command(ns('admin.$cmd'), failPoint, {});
+    this.appliedFailPoints.push(failPoint);
   }
 
   enableFailPoint(failPoint) {

test/unit/operations/get_more.test.ts

@@ -202,7 +202,7 @@ describe('GetMoreOperation', function () {
         new ServerDescription('a:1'),
         {} as any
       );
-      sinon.stub(server, 'command').yieldsRight();
+      sinon.stub(server, 'command');
 
       it('should throw if the cursorId is undefined', async () => {
         const getMoreOperation = new GetMoreOperation(

test/unit/operations/kill_cursors.test.ts

@@ -91,7 +91,7 @@ describe('class KillCursorsOperation', () => {
         server,
         options
       ) as any;
-      const stub = sinon.stub(server, 'command').yieldsRight();
+      const stub = sinon.stub(server, 'command').returns(Promise.resolve({}));
       await killCursorsOperation.execute(server, undefined);
       expect(stub).to.have.been.calledOnceWith(namespace, {
         killCursors: namespace.collection,