@@ -5,6 +5,7 @@ import { EventEmitter } from 'events';
5
5
import {
6
6
AbstractCursor ,
7
7
ChangeStream ,
8
+ ClientEncryption ,
8
9
ClientSession ,
9
10
Collection ,
10
11
type CommandFailedEvent ,
@@ -37,13 +38,9 @@ import {
37
38
} from '../../mongodb' ;
38
39
import { ejson , getEnvironmentalOptions } from '../../tools/utils' ;
39
40
import type { TestConfiguration } from '../runner/config' ;
41
+ import { EntityEventRegistry } from './entity_event_registry' ;
40
42
import { 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' ;
47
44
import {
48
45
createClientEncryption ,
49
46
makeConnectionString ,
@@ -357,9 +354,10 @@ export type Entity =
357
354
| AbstractCursor
358
355
| UnifiedChangeStream
359
356
| GridFSBucket
357
+ | Document
360
358
| ClientEncryption
361
359
| TopologyDescription // From recordTopologyDescription operation
362
- | Document ; // Results from operations
360
+ | number ;
363
361
364
362
export type EntityCtor =
365
363
| typeof UnifiedMongoClient
@@ -370,7 +368,7 @@ export type EntityCtor =
370
368
| typeof AbstractCursor
371
369
| typeof GridFSBucket
372
370
| typeof UnifiedThread
373
- | ClientEncryption ;
371
+ | typeof ClientEncryption ;
374
372
375
373
export type EntityTypeId =
376
374
| 'client'
@@ -381,18 +379,26 @@ export type EntityTypeId =
381
379
| 'thread'
382
380
| 'cursor'
383
381
| 'stream'
384
- | 'clientEncryption' ;
382
+ | 'clientEncryption'
383
+ | 'errors'
384
+ | 'failures'
385
+ | 'events'
386
+ | 'iterations'
387
+ | 'successes' ;
385
388
386
389
const ENTITY_CTORS = new Map < EntityTypeId , EntityCtor > ( ) ;
387
390
ENTITY_CTORS . set ( 'client' , UnifiedMongoClient ) ;
388
391
ENTITY_CTORS . set ( 'db' , Db ) ;
392
+ ENTITY_CTORS . set ( 'clientEncryption' , ClientEncryption ) ;
389
393
ENTITY_CTORS . set ( 'collection' , Collection ) ;
390
394
ENTITY_CTORS . set ( 'session' , ClientSession ) ;
391
395
ENTITY_CTORS . set ( 'bucket' , GridFSBucket ) ;
392
396
ENTITY_CTORS . set ( 'thread' , UnifiedThread ) ;
393
397
ENTITY_CTORS . set ( 'cursor' , AbstractCursor ) ;
394
398
ENTITY_CTORS . set ( 'stream' , ChangeStream ) ;
395
399
400
+ const NO_INSTANCE_CHECK = [ 'errors' , 'failures' , 'events' , 'successes' , 'iterations' ] ;
401
+
396
402
export class EntitiesMap < E = Entity > extends Map < string , E > {
397
403
failPoints : FailPointMap ;
398
404
@@ -435,15 +441,20 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
435
441
getEntity ( type : 'thread' , key : string , assertExists ?: boolean ) : UnifiedThread ;
436
442
getEntity ( type : 'cursor' , key : string , assertExists ?: boolean ) : AbstractCursor ;
437
443
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 [ ] ;
438
449
getEntity ( type : 'clientEncryption' , key : string , assertExists ?: boolean ) : ClientEncryption ;
439
450
getEntity ( type : EntityTypeId , key : string , assertExists = true ) : Entity | undefined {
440
451
const entity = this . get ( key ) ;
441
452
if ( ! entity ) {
442
453
if ( assertExists ) throw new Error ( `Entity '${ key } ' does not exist` ) ;
443
454
return ;
444
455
}
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.
447
458
return entity ;
448
459
}
449
460
const ctor = ENTITY_CTORS . get ( type ) ;
@@ -499,6 +510,7 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
499
510
entity . client . uriOptions
500
511
) ;
501
512
const client = new UnifiedMongoClient ( uri , entity . client ) ;
513
+ new EntityEventRegistry ( client , entity . client , map ) . register ( ) ;
502
514
try {
503
515
await client . connect ( ) ;
504
516
} catch ( error ) {
0 commit comments