test: refactor unified event capture logic for ease of use

Author: dariakp

test/tools/unified-spec-runner/entities.ts

@@ -140,6 +140,17 @@ export class UnifiedMongoClient extends MongoClient {
     return this.ignoredEvents.includes(e.commandName);
   }
 
+  getCapturedEvents(eventType: string): CommandEvent[] | CmapEvent[] {
+    switch (eventType) {
+      case 'command':
+        return this.commandEvents;
+      case 'cmap':
+        return this.cmapEvents;
+      default:
+        throw new Error(`Unknown eventType: ${eventType}`);
+    }
+  }
+
   // NOTE: pushCommandEvent must be an arrow function
   pushCommandEvent: (e: CommandEvent) => void = e => {
     if (!this.isIgnored(e)) {
@@ -152,22 +163,14 @@ export class UnifiedMongoClient extends MongoClient {
     this.cmapEvents.push(e);
   };
 
-  stopCapturingEvents(pushFn: PushFunction): void {
-    const observedEvents = [...this.observedCommandEvents, ...this.observedCmapEvents];
-    for (const eventName of observedEvents) {
-      this.off(eventName, pushFn);
-    }
-  }
-
   /** Disables command monitoring for the client and returns a list of the captured events. */
-  stopCapturingCommandEvents(): CommandEvent[] {
-    this.stopCapturingEvents(this.pushCommandEvent);
-    return this.commandEvents;
-  }
-
-  stopCapturingCmapEvents(): CmapEvent[] {
-    this.stopCapturingEvents(this.pushCmapEvent);
-    return this.cmapEvents;
+  stopCapturingEvents(): void {
+    for (const eventName of this.observedCommandEvents) {
+      this.off(eventName, this.pushCommandEvent);
+    }
+    for (const eventName of this.observedCmapEvents) {
+      this.off(eventName, this.pushCmapEvent);
+    }
   }
 }
 

test/tools/unified-spec-runner/runner.ts

@@ -8,7 +8,7 @@ import { ReadPreference } from '../../../src/read_preference';
 import { TopologyType } from '../../../src/sdam/common';
 import { ns } from '../../../src/utils';
 import { ejson } from '../utils';
-import { CmapEvent, CommandEvent, EntitiesMap } from './entities';
+import { CmapEvent, CommandEvent, EntitiesMap, UnifiedMongoClient } from './entities';
 import { matchesEvents } from './match';
 import { executeOperationAndCheck } from './operations';
 import * as uni from './schema';
@@ -198,13 +198,12 @@ async function runUnifiedTest(
       }
     }
 
-    const clientCommandEvents = new Map<string, CommandEvent[]>();
-    const clientCmapEvents = new Map<string, CmapEvent[]>();
+    const clientList = new Map<string, UnifiedMongoClient>();
     // If any event listeners were enabled on any client entities,
     // the test runner MUST now disable those event listeners.
     for (const [id, client] of entities.mapOf('client')) {
-      clientCommandEvents.set(id, client.stopCapturingCommandEvents());
-      clientCmapEvents.set(id, client.stopCapturingCmapEvents());
+      client.stopCapturingEvents();
+      clientList.set(id, client);
     }
 
     if (test.expectEvents) {
@@ -213,11 +212,13 @@ async function runUnifiedTest(
         const eventType = expectedEventsForClient.eventType;
         // If no event type is provided it defaults to 'command', so just
         // check for 'cmap' here for now.
-        const actualEvents =
-          eventType === 'cmap' ? clientCmapEvents.get(clientId) : clientCommandEvents.get(clientId);
-
-        expect(actualEvents, `No client entity found with id ${clientId}`).to.exist;
-        matchesEvents(expectedEventsForClient, actualEvents!, entities);
+        const testClient = clientList.get(clientId);
+        expect(testClient, `No client entity found with id ${clientId}`).to.exist;
+        matchesEvents(
+          expectedEventsForClient,
+          testClient!.getCapturedEvents(eventType ?? 'command'),
+          entities
+        );
       }
     }