test(NODE-3049): add astrolabe support to the UnifiedTestRunner (#3805)

Author: durran

package.json

@@ -143,6 +143,7 @@
     "check:unit": "mocha test/unit",
     "check:ts": "node ./node_modules/typescript/bin/tsc -v && node ./node_modules/typescript/bin/tsc --noEmit",
     "check:atlas": "mocha --config test/manual/mocharc.json test/manual/atlas_connectivity.test.js",
+    "check:drivers-atlas-testing": "mocha --config test/mocha_mongodb.json test/atlas/drivers_atlas_testing.test.ts",
     "check:adl": "mocha --config test/mocha_mongodb.json test/manual/atlas-data-lake-testing",
     "check:aws": "nyc mocha --config test/mocha_mongodb.json test/integration/auth/mongodb_aws.test.ts",
     "check:oidc": "mocha --config test/mocha_mongodb.json test/manual/mongodb_oidc.prose.test.ts",

test/atlas/drivers_atlas_testing.test.ts

@@ -0,0 +1,10 @@
+import { runUnifiedSuite } from '../tools/unified-spec-runner/runner';
+
+describe('Node Driver Atlas Testing', async function () {
+  // Astrolabe can, well, take some time. In some cases up to 800s to
+  // reconfigure clusters.
+  this.timeout(0);
+  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+  const spec = JSON.parse(process.env.WORKLOAD_SPECIFICATION!);
+  runUnifiedSuite([spec]);
+});

test/tools/reporter/mongodb_reporter.js

@@ -103,7 +103,7 @@ class MongoDBMochaReporter extends mocha.reporters.Spec {
       catchErr(test => this.testEnd(test))
     );
 
-    process.on('SIGINT', () => this.end(true));
+    process.prependListener('SIGINT', () => this.end(true));
   }
 
   start() {}
@@ -183,7 +183,11 @@ class MongoDBMochaReporter extends mocha.reporters.Spec {
     } catch (error) {
       console.error(chalk.red(`Failed to output xunit report! ${error}`));
     } finally {
-      if (ctrlC) process.exit(1);
+      // Dont exit the process on Astrolabe testing, let it interrupt and
+      // finish naturally.
+      if (!process.env.WORKLOAD_SPECIFICATION) {
+        if (ctrlC) process.exit(1);
+      }
     }
   }
 

test/tools/runner/config.ts

@@ -376,3 +376,18 @@ export class AtlasTestConfiguration extends TestConfiguration {
     return process.env.MONGODB_URI!;
   }
 }
+
+/**
+ * Test configuration specific to Astrolabe testing.
+ */
+export class AstrolabeTestConfiguration extends TestConfiguration {
+  override newClient(): MongoClient {
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    return new MongoClient(process.env.DRIVERS_ATLAS_TESTING_URI!);
+  }
+
+  override url(): string {
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    return process.env.DRIVERS_ATLAS_TESTING_URI!;
+  }
+}

test/tools/runner/hooks/configuration.js

@@ -7,7 +7,7 @@ require('source-map-support').install({
 const path = require('path');
 const fs = require('fs');
 const { MongoClient } = require('../../../mongodb');
-const { TestConfiguration } = require('../config');
+const { AstrolabeTestConfiguration, TestConfiguration } = require('../config');
 const { getEnvironmentalOptions } = require('../../utils');
 const mock = require('../../mongodb-mock/index');
 const { inspect } = require('util');
@@ -105,6 +105,10 @@ const skipBrokenAuthTestBeforeEachHook = function ({ skippedTests } = { skippedT
 };
 
 const testConfigBeforeHook = async function () {
+  if (process.env.DRIVERS_ATLAS_TESTING_URI) {
+    this.configuration = new AstrolabeTestConfiguration(process.env.DRIVERS_ATLAS_TESTING_URI, {});
+    return;
+  }
   // TODO(NODE-5035): Implement OIDC support. Creating the MongoClient will fail
   // with "MongoInvalidArgumentError: AuthMechanism 'MONGODB-OIDC' not supported"
   // as is expected until that ticket goes in. Then this condition gets removed.

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

@@ -0,0 +1,44 @@
+import { writeFile } from 'node:fs/promises';
+
+import * as path from 'path';
+
+import type { EntitiesMap } from './entities';
+import { trace } from './runner';
+
+/**
+ * Writes the entities saved from the loop operations run in the
+ * Astrolabe workload executor to the required files.
+ */
+export class AstrolabeResultsWriter {
+  constructor(private entities: EntitiesMap) {
+    this.entities = entities;
+  }
+
+  async write(): Promise<void> {
+    // Write the events.json to the execution directory.
+    const errors = this.entities.getEntity('errors', 'errors', false);
+    const failures = this.entities.getEntity('failures', 'failures', false);
+    const events = this.entities.getEntity('events', 'events', false);
+    const iterations = this.entities.getEntity('iterations', 'iterations', false);
+    const successes = this.entities.getEntity('successes', 'successes', false);
+
+    // Write the events.json to the execution directory.
+    trace('writing events.json');
+    await writeFile(
+      path.join(process.env.OUTPUT_DIRECTORY ?? '', 'events.json'),
+      JSON.stringify({ events: events ?? [], errors: errors ?? [], failures: failures ?? [] })
+    );
+
+    // Write the results.json to the execution directory.
+    trace('writing results.json');
+    await writeFile(
+      path.join(process.env.OUTPUT_DIRECTORY ?? '', 'results.json'),
+      JSON.stringify({
+        numErrors: errors?.length ?? 0,
+        numFailures: failures?.length ?? 0,
+        numSuccesses: successes ?? 0,
+        numIterations: iterations ?? 0
+      })
+    );
+  }
+}

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

@@ -5,6 +5,7 @@ import { EventEmitter } from 'events';
 import {
   AbstractCursor,
   ChangeStream,
+  ClientEncryption,
   ClientSession,
   Collection,
   type CommandFailedEvent,
@@ -37,13 +38,9 @@ import {
 } from '../../mongodb';
 import { ejson, getEnvironmentalOptions } from '../../tools/utils';
 import type { TestConfiguration } from '../runner/config';
+import { EntityEventRegistry } from './entity_event_registry';
 import { trace } from './runner';
-import type {
-  ClientEncryption,
-  ClientEntity,
-  EntityDescription,
-  ExpectedLogMessage
-} from './schema';
+import type { ClientEntity, EntityDescription, ExpectedLogMessage } from './schema';
 import {
   createClientEncryption,
   makeConnectionString,
@@ -357,9 +354,10 @@ export type Entity =
   | AbstractCursor
   | UnifiedChangeStream
   | GridFSBucket
+  | Document
   | ClientEncryption
   | TopologyDescription // From recordTopologyDescription operation
-  | Document; // Results from operations
+  | number;
 
 export type EntityCtor =
   | typeof UnifiedMongoClient
@@ -370,7 +368,7 @@ export type EntityCtor =
   | typeof AbstractCursor
   | typeof GridFSBucket
   | typeof UnifiedThread
-  | ClientEncryption;
+  | typeof ClientEncryption;
 
 export type EntityTypeId =
   | 'client'
@@ -381,18 +379,26 @@ export type EntityTypeId =
   | 'thread'
   | 'cursor'
   | 'stream'
-  | 'clientEncryption';
+  | 'clientEncryption'
+  | 'errors'
+  | 'failures'
+  | 'events'
+  | 'iterations'
+  | 'successes';
 
 const ENTITY_CTORS = new Map<EntityTypeId, EntityCtor>();
 ENTITY_CTORS.set('client', UnifiedMongoClient);
 ENTITY_CTORS.set('db', Db);
+ENTITY_CTORS.set('clientEncryption', ClientEncryption);
 ENTITY_CTORS.set('collection', Collection);
 ENTITY_CTORS.set('session', ClientSession);
 ENTITY_CTORS.set('bucket', GridFSBucket);
 ENTITY_CTORS.set('thread', UnifiedThread);
 ENTITY_CTORS.set('cursor', AbstractCursor);
 ENTITY_CTORS.set('stream', ChangeStream);
 
+const NO_INSTANCE_CHECK = ['errors', 'failures', 'events', 'successes', 'iterations'];
+
 export class EntitiesMap<E = Entity> extends Map<string, E> {
   failPoints: FailPointMap;
 
@@ -435,15 +441,20 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
   getEntity(type: 'thread', key: string, assertExists?: boolean): UnifiedThread;
   getEntity(type: 'cursor', key: string, assertExists?: boolean): AbstractCursor;
   getEntity(type: 'stream', key: string, assertExists?: boolean): UnifiedChangeStream;
+  getEntity(type: 'iterations', key: string, assertExists?: boolean): number;
+  getEntity(type: 'successes', key: string, assertExists?: boolean): number;
+  getEntity(type: 'errors', key: string, assertExists?: boolean): Document[];
+  getEntity(type: 'failures', key: string, assertExists?: boolean): Document[];
+  getEntity(type: 'events', key: string, assertExists?: boolean): Document[];
   getEntity(type: 'clientEncryption', key: string, assertExists?: boolean): ClientEncryption;
   getEntity(type: EntityTypeId, key: string, assertExists = true): Entity | undefined {
     const entity = this.get(key);
     if (!entity) {
       if (assertExists) throw new Error(`Entity '${key}' does not exist`);
       return;
     }
-    if (type === 'clientEncryption') {
-      // we do not have instanceof checking here since csfle might not be installed
+    if (NO_INSTANCE_CHECK.includes(type)) {
+      // Skip constructor checks for interfaces.
       return entity;
     }
     const ctor = ENTITY_CTORS.get(type);
@@ -499,6 +510,7 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
           entity.client.uriOptions
         );
         const client = new UnifiedMongoClient(uri, entity.client);
+        new EntityEventRegistry(client, entity.client, map).register();
         try {
           await client.connect();
         } catch (error) {

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

@@ -0,0 +1,76 @@
+import {
+  COMMAND_FAILED,
+  COMMAND_STARTED,
+  COMMAND_SUCCEEDED,
+  CONNECTION_CHECK_OUT_FAILED,
+  CONNECTION_CHECK_OUT_STARTED,
+  CONNECTION_CHECKED_IN,
+  CONNECTION_CHECKED_OUT,
+  CONNECTION_CLOSED,
+  CONNECTION_CREATED,
+  CONNECTION_POOL_CLEARED,
+  CONNECTION_POOL_CLOSED,
+  CONNECTION_POOL_CREATED,
+  CONNECTION_POOL_READY,
+  CONNECTION_READY
+} from '../../mongodb';
+import { type EntitiesMap, type UnifiedMongoClient } from './entities';
+import { type ClientEntity } from './schema';
+
+/**
+ * Maps the names of the events the unified runner passes and maps
+ * them to the names of the events emitted in the driver.
+ */
+const MAPPINGS = {
+  PoolCreatedEvent: CONNECTION_POOL_CREATED,
+  PoolReadyEvent: CONNECTION_POOL_READY,
+  PoolClearedEvent: CONNECTION_POOL_CLEARED,
+  PoolClosedEvent: CONNECTION_POOL_CLOSED,
+  ConnectionCreatedEvent: CONNECTION_CREATED,
+  ConnectionReadyEvent: CONNECTION_READY,
+  ConnectionClosedEvent: CONNECTION_CLOSED,
+  ConnectionCheckOutStartedEvent: CONNECTION_CHECK_OUT_STARTED,
+  ConnectionCheckOutFailedEvent: CONNECTION_CHECK_OUT_FAILED,
+  ConnectionCheckedOutEvent: CONNECTION_CHECKED_OUT,
+  ConnectionCheckedInEvent: CONNECTION_CHECKED_IN,
+  CommandStartedEvent: COMMAND_STARTED,
+  CommandSucceededEvent: COMMAND_SUCCEEDED,
+  CommandFailedEvent: COMMAND_FAILED
+};
+
+/**
+ * Registers events that need to be stored in the entities map, since
+ * the UnifiedMongoClient does not contain a cyclical dependency on the
+ * entities map itself.
+ */
+export class EntityEventRegistry {
+  constructor(
+    private client: UnifiedMongoClient,
+    private clientEntity: ClientEntity,
+    private entitiesMap: EntitiesMap
+  ) {
+    this.client = client;
+    this.clientEntity = clientEntity;
+    this.entitiesMap = entitiesMap;
+  }
+
+  /**
+   * Connect the event listeners on the client and the entities map.
+   */
+  register(): void {
+    if (this.clientEntity.storeEventsAsEntities) {
+      for (const { id, events } of this.clientEntity.storeEventsAsEntities) {
+        this.entitiesMap.set(id, []);
+        for (const eventName of events) {
+          // Need to map the event names to the Node event names.
+          this.client.on(MAPPINGS[eventName], () => {
+            this.entitiesMap.getEntity('events', id).push({
+              name: eventName,
+              observedAt: Date.now()
+            });
+          });
+        }
+      }
+    }
+  }
+}