From 1e58a4ce36a78fc92d59083a63ee36114204baf5 Mon Sep 17 00:00:00 2001
From: Durran Jordan <durran@gmail.com>
Date: Tue, 30 May 2023 15:45:06 -0400
Subject: [PATCH] test(NODE-5265): fix flaky operation count test (#3688)

---
 src/sdam/server.ts                            | 20 ++++++++++++++++---
 .../server-selection/operation_count.test.ts  | 17 ++++++++--------
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/src/sdam/server.ts b/src/sdam/server.ts
index b2a2f900015..de10935bfd6 100644
--- a/src/sdam/server.ts
+++ b/src/sdam/server.ts
@@ -343,13 +343,13 @@ export class Server extends TypedEventEmitter<ServerEvents> {
       return;
     }
 
-    this.s.operationCount += 1;
+    this.incrementOperationCount();
 
     this.pool.withConnection(
       conn,
       (err, conn, cb) => {
         if (err || !conn) {
-          this.s.operationCount -= 1;
+          this.decrementOperationCount();
           if (!err) {
             return cb(new MongoRuntimeError('Failed to create connection without error'));
           }
@@ -364,7 +364,7 @@ export class Server extends TypedEventEmitter<ServerEvents> {
           cmd,
           finalOptions,
           makeOperationHandler(this, conn, cmd, finalOptions, (error, response) => {
-            this.s.operationCount -= 1;
+            this.decrementOperationCount();
             cb(error, response);
           })
         );
@@ -420,6 +420,20 @@ export class Server extends TypedEventEmitter<ServerEvents> {
       }
     }
   }
+
+  /**
+   * Decrement the operation count, returning the new count.
+   */
+  private decrementOperationCount(): number {
+    return (this.s.operationCount -= 1);
+  }
+
+  /**
+   * Increment the operation count, returning the new count.
+   */
+  private incrementOperationCount(): number {
+    return (this.s.operationCount += 1);
+  }
 }
 
 function calculateRoundTripTime(oldRtt: number, duration: number): number {
diff --git a/test/integration/server-selection/operation_count.test.ts b/test/integration/server-selection/operation_count.test.ts
index 3791d1d9f54..27ced7f5d25 100644
--- a/test/integration/server-selection/operation_count.test.ts
+++ b/test/integration/server-selection/operation_count.test.ts
@@ -2,7 +2,7 @@ import { expect } from 'chai';
 import * as sinon from 'sinon';
 
 import { AbstractCursor, Collection, ConnectionPool, MongoClient } from '../../mongodb';
-import { FailPoint, sleep } from '../../tools/utils';
+import { FailPoint } from '../../tools/utils';
 
 const testMetadata: MongoDBMetadataUI = {
   requires: {
@@ -118,20 +118,21 @@ describe('Server Operation Count Tests', function () {
       const server = Array.from(client.topology.s.servers.values())[0];
       expect(server.s.operationCount).to.equal(0);
       const commandSpy = sinon.spy(server, 'command');
+      const incrementSpy = sinon.spy(server, 'incrementOperationCount');
+      const decrementSpy = sinon.spy(server, 'decrementOperationCount');
 
       const operationPromises = Array.from({ length: 10 }, () =>
         collection.insertOne({ count: 1 })
       );
 
-      // operation count is incremented after connection checkout, which happens asynchronously (even though there are plenty of connections in the pool).
-      // we sleep to give the event loop a turn so that all the commands check out a connection before asserting the operation count
-      await sleep(1);
-
-      expect(server.s.operationCount).to.equal(10);
-
-      await Promise.all(operationPromises);
+      await Promise.allSettled(operationPromises);
 
       expect(commandSpy.called).to.be.true;
+      // This test is flaky when sleeping and asserting the operation count after the sleep but before the
+      // promise execution, so we assert instead that the count was incremented 10 times and decremented 10
+      // times - the total number of operations.
+      expect(incrementSpy.callCount).to.equal(10);
+      expect(decrementSpy.callCount).to.equal(10);
       expect(server.s.operationCount).to.equal(0);
     });
   });