@@ -5,6 +5,7 @@ import { EventEmitter } from 'events';
55import {
66 AbstractCursor ,
77 ChangeStream ,
8+ ClientEncryption ,
89 ClientSession ,
910 Collection ,
1011 type CommandFailedEvent ,
@@ -37,13 +38,9 @@ import {
3738} from '../../mongodb' ;
3839import { ejson , getEnvironmentalOptions } from '../../tools/utils' ;
3940import type { TestConfiguration } from '../runner/config' ;
41+ import { EntityEventRegistry } from './entity_event_registry' ;
4042import { trace } from './runner' ;
41- import type {
42- ClientEncryption ,
43- ClientEntity ,
44- EntityDescription ,
45- ExpectedLogMessage
46- } from './schema' ;
43+ import type { ClientEntity , EntityDescription , ExpectedLogMessage } from './schema' ;
4744import {
4845 createClientEncryption ,
4946 makeConnectionString ,
@@ -357,9 +354,10 @@ export type Entity =
357354 | AbstractCursor
358355 | UnifiedChangeStream
359356 | GridFSBucket
357+ | Document
360358 | ClientEncryption
361359 | TopologyDescription // From recordTopologyDescription operation
362- | Document ; // Results from operations
360+ | number ;
363361
364362export type EntityCtor =
365363 | typeof UnifiedMongoClient
@@ -370,7 +368,7 @@ export type EntityCtor =
370368 | typeof AbstractCursor
371369 | typeof GridFSBucket
372370 | typeof UnifiedThread
373- | ClientEncryption ;
371+ | typeof ClientEncryption ;
374372
375373export type EntityTypeId =
376374 | 'client'
@@ -381,18 +379,26 @@ export type EntityTypeId =
381379 | 'thread'
382380 | 'cursor'
383381 | 'stream'
384- | 'clientEncryption' ;
382+ | 'clientEncryption'
383+ | 'errors'
384+ | 'failures'
385+ | 'events'
386+ | 'iterations'
387+ | 'successes' ;
385388
386389const ENTITY_CTORS = new Map < EntityTypeId , EntityCtor > ( ) ;
387390ENTITY_CTORS . set ( 'client' , UnifiedMongoClient ) ;
388391ENTITY_CTORS . set ( 'db' , Db ) ;
392+ ENTITY_CTORS . set ( 'clientEncryption' , ClientEncryption ) ;
389393ENTITY_CTORS . set ( 'collection' , Collection ) ;
390394ENTITY_CTORS . set ( 'session' , ClientSession ) ;
391395ENTITY_CTORS . set ( 'bucket' , GridFSBucket ) ;
392396ENTITY_CTORS . set ( 'thread' , UnifiedThread ) ;
393397ENTITY_CTORS . set ( 'cursor' , AbstractCursor ) ;
394398ENTITY_CTORS . set ( 'stream' , ChangeStream ) ;
395399
400+ const NO_INSTANCE_CHECK = [ 'errors' , 'failures' , 'events' , 'successes' , 'iterations' ] ;
401+
396402export class EntitiesMap < E = Entity > extends Map < string , E > {
397403 failPoints : FailPointMap ;
398404
@@ -435,15 +441,20 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
435441 getEntity ( type : 'thread' , key : string , assertExists ?: boolean ) : UnifiedThread ;
436442 getEntity ( type : 'cursor' , key : string , assertExists ?: boolean ) : AbstractCursor ;
437443 getEntity ( type : 'stream' , key : string , assertExists ?: boolean ) : UnifiedChangeStream ;
444+ getEntity ( type : 'iterations' , key : string , assertExists ?: boolean ) : number ;
445+ getEntity ( type : 'successes' , key : string , assertExists ?: boolean ) : number ;
446+ getEntity ( type : 'errors' , key : string , assertExists ?: boolean ) : Document [ ] ;
447+ getEntity ( type : 'failures' , key : string , assertExists ?: boolean ) : Document [ ] ;
448+ getEntity ( type : 'events' , key : string , assertExists ?: boolean ) : Document [ ] ;
438449 getEntity ( type : 'clientEncryption' , key : string , assertExists ?: boolean ) : ClientEncryption ;
439450 getEntity ( type : EntityTypeId , key : string , assertExists = true ) : Entity | undefined {
440451 const entity = this . get ( key ) ;
441452 if ( ! entity ) {
442453 if ( assertExists ) throw new Error ( `Entity '${ key } ' does not exist` ) ;
443454 return ;
444455 }
445- if ( type === 'clientEncryption' ) {
446- // we do not have instanceof checking here since csfle might not be installed
456+ if ( NO_INSTANCE_CHECK . includes ( type ) ) {
457+ // Skip constructor checks for interfaces.
447458 return entity ;
448459 }
449460 const ctor = ENTITY_CTORS . get ( type ) ;
@@ -499,6 +510,7 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
499510 entity . client . uriOptions
500511 ) ;
501512 const client = new UnifiedMongoClient ( uri , entity . client ) ;
513+ new EntityEventRegistry ( client , entity . client , map ) . register ( ) ;
502514 try {
503515 await client . connect ( ) ;
504516 } catch ( error ) {
0 commit comments