@@ -84,6 +84,11 @@ type WriteConcernErrorData struct {
8484 ErrInfo bson.Raw `bson:"errInfo,omitempty"`
8585}
8686
87+ type failPoint struct {
88+ name string
89+ client * mongo.Client
90+ }
91+
8792// T is a wrapper around testing.T.
8893type T struct {
8994 // connsCheckedOut is the net number of connections checked out during test execution.
@@ -103,7 +108,7 @@ type T struct {
103108 createdColls []* Collection // collections created in this test
104109 proxyDialer * proxyDialer
105110 dbName , collName string
106- failPointNames [] string
111+ failPoints [] failPoint
107112 minServerVersion string
108113 maxServerVersion string
109114 validTopologies []TopologyKind
@@ -128,14 +133,16 @@ type T struct {
128133 succeeded []* event.CommandSucceededEvent
129134 failed []* event.CommandFailedEvent
130135
131- Client * mongo.Client
132- DB * mongo.Database
133- Coll * mongo.Collection
136+ Client * mongo.Client
137+ fpClients map [* mongo.Client ]bool
138+ DB * mongo.Database
139+ Coll * mongo.Collection
134140}
135141
136142func newT (wrapped * testing.T , opts ... * Options ) * T {
137143 t := & T {
138- T : wrapped ,
144+ T : wrapped ,
145+ fpClients : make (map [* mongo.Client ]bool ),
139146 }
140147 for _ , opt := range opts {
141148 for _ , optFn := range opt .optFuncs {
@@ -202,6 +209,12 @@ func (t *T) cleanup() {
202209 // always disconnect the client regardless of clientType because Client.Disconnect will work against
203210 // all deployments
204211 _ = t .Client .Disconnect (context .Background ())
212+ for client , v := range t .fpClients {
213+ if v {
214+ client .Disconnect (context .Background ())
215+ }
216+ }
217+ t .fpClients = nil
205218}
206219
207220// Run creates a new T instance for a sub-test and runs the given callback. It also creates a new collection using the
@@ -254,9 +267,11 @@ func (t *T) RunOpts(name string, opts *Options, callback func(mt *T)) {
254267 sub .ClearFailPoints ()
255268 sub .ClearCollections ()
256269 }
257- // only disconnect client if it's not being shared
270+ // only disconnect client if it's not being shared and not used by fail points.
258271 if sub .shareClient == nil || ! * sub .shareClient {
259- _ = sub .Client .Disconnect (context .Background ())
272+ if _ , ok := sub .fpClients [sub .Client ]; ! ok {
273+ _ = sub .Client .Disconnect (context .Background ())
274+ }
260275 }
261276 assert .Equal (sub , 0 , sessions , "%v sessions checked out" , sessions )
262277 assert .Equal (sub , 0 , conns , "%v connections checked out" , conns )
@@ -405,7 +420,10 @@ func (t *T) ResetClient(opts *options.ClientOptions) {
405420 t .clientOpts = opts
406421 }
407422
408- _ = t .Client .Disconnect (context .Background ())
423+ // Disconnect client if it is not being used by fail points.
424+ if _ , ok := t .fpClients [t .Client ]; ! ok {
425+ _ = t .Client .Disconnect (context .Background ())
426+ }
409427 t .createTestClient ()
410428 t .DB = t .Client .Database (t .dbName )
411429 t .Coll = t .DB .Collection (t .collName , t .collOpts )
@@ -562,7 +580,8 @@ func (t *T) SetFailPoint(fp FailPoint) {
562580 if err := SetFailPoint (fp , t .Client ); err != nil {
563581 t .Fatal (err )
564582 }
565- t .failPointNames = append (t .failPointNames , fp .ConfigureFailPoint )
583+ t .fpClients [t .Client ] = true
584+ t .failPoints = append (t .failPoints , failPoint {fp .ConfigureFailPoint , t .Client })
566585}
567586
568587// SetFailPointFromDocument sets the fail point represented by the given document for the client associated with T. This
@@ -574,30 +593,35 @@ func (t *T) SetFailPointFromDocument(fp bson.Raw) {
574593 t .Fatal (err )
575594 }
576595
596+ t .fpClients [t .Client ] = true
577597 name := fp .Index (0 ).Value ().StringValue ()
578- t .failPointNames = append (t .failPointNames , name )
598+ t .failPoints = append (t .failPoints , failPoint { name , t . Client } )
579599}
580600
581601// TrackFailPoint adds the given fail point to the list of fail points to be disabled when the current test finishes.
582602// This function does not create a fail point on the server.
583- func (t * T ) TrackFailPoint (fpName string ) {
584- t .failPointNames = append (t .failPointNames , fpName )
603+ func (t * T ) TrackFailPoint (fpName string , client * mongo.Client ) {
604+ t .fpClients [client ] = true
605+ t .failPoints = append (t .failPoints , failPoint {fpName , client })
585606}
586607
587608// ClearFailPoints disables all previously set failpoints for this test.
588609func (t * T ) ClearFailPoints () {
589- db := t .Client .Database ("admin" )
590- for _ , fp := range t .failPointNames {
610+ for _ , fp := range t .failPoints {
591611 cmd := bson.D {
592- {"configureFailPoint" , fp },
612+ {"configureFailPoint" , fp . name },
593613 {"mode" , "off" },
594614 }
595- err := db .RunCommand (context .Background (), cmd ).Err ()
615+ err := fp . client . Database ( "admin" ) .RunCommand (context .Background (), cmd ).Err ()
596616 if err != nil {
597- t .Fatalf ("error clearing fail point %s: %v" , fp , err )
617+ t .Fatalf ("error clearing fail point %s: %v" , fp .name , err )
618+ }
619+ if fp .client != t .Client {
620+ _ = fp .client .Disconnect (context .Background ())
621+ t .fpClients [fp .client ] = false
598622 }
599623 }
600- t .failPointNames = t .failPointNames [:0 ]
624+ t .failPoints = t .failPoints [:0 ]
601625}
602626
603627// CloneDatabase modifies the default database for this test to match the given options.
@@ -684,19 +708,17 @@ func (t *T) createTestClient() {
684708 })
685709 }
686710
687- var err error
711+ var uriOpts * options. ClientOptions
688712 switch t .clientType {
689713 case Pinned :
690714 // pin to first mongos
691715 pinnedHostList := []string {testContext .connString .Hosts [0 ]}
692- uriOpts := options .Client ().ApplyURI (testContext .connString .Original ).SetHosts (pinnedHostList )
693- t .Client , err = mongo .NewClient (uriOpts , clientOpts )
716+ uriOpts = options .Client ().ApplyURI (testContext .connString .Original ).SetHosts (pinnedHostList )
694717 case Mock :
695718 // clear pool monitor to avoid configuration error
696719 clientOpts .PoolMonitor = nil
697720 t .mockDeployment = newMockDeployment ()
698721 clientOpts .Deployment = t .mockDeployment
699- t .Client , err = mongo .NewClient (clientOpts )
700722 case Proxy :
701723 t .proxyDialer = newProxyDialer ()
702724 clientOpts .SetDialer (t .proxyDialer )
@@ -706,16 +728,17 @@ func (t *T) createTestClient() {
706728 case Default :
707729 // Use a different set of options to specify the URI because clientOpts may already have a URI or host seedlist
708730 // specified.
709- var uriOpts * options.ClientOptions
710731 if clientOpts .Deployment == nil {
711732 // Only specify URI if the deployment is not set to avoid setting topology/server options along with the
712733 // deployment.
713734 uriOpts = options .Client ().ApplyURI (testContext .connString .Original )
714735 }
715-
716- // Pass in uriOpts first so clientOpts wins if there are any conflicting settings.
717- t .Client , err = mongo .NewClient (uriOpts , clientOpts )
718736 }
737+ t .clientOpts = options .MergeClientOptions (uriOpts , clientOpts )
738+
739+ var err error
740+ // Pass in uriOpts first so clientOpts wins if there are any conflicting settings.
741+ t .Client , err = mongo .NewClient (t .clientOpts )
719742 if err != nil {
720743 t .Fatalf ("error creating client: %v" , err )
721744 }
0 commit comments