@@ -84,6 +84,11 @@ type WriteConcernErrorData struct {
84
84
ErrInfo bson.Raw `bson:"errInfo,omitempty"`
85
85
}
86
86
87
+ type failPoint struct {
88
+ name string
89
+ client * mongo.Client
90
+ }
91
+
87
92
// T is a wrapper around testing.T.
88
93
type T struct {
89
94
// connsCheckedOut is the net number of connections checked out during test execution.
@@ -103,7 +108,7 @@ type T struct {
103
108
createdColls []* Collection // collections created in this test
104
109
proxyDialer * proxyDialer
105
110
dbName , collName string
106
- failPointNames [] string
111
+ failPoints [] failPoint
107
112
minServerVersion string
108
113
maxServerVersion string
109
114
validTopologies []TopologyKind
@@ -128,15 +133,16 @@ type T struct {
128
133
succeeded []* event.CommandSucceededEvent
129
134
failed []* event.CommandFailedEvent
130
135
131
- Client * mongo.Client
132
- fpClient * mongo.Client
133
- DB * mongo.Database
134
- Coll * mongo.Collection
136
+ Client * mongo.Client
137
+ fpClients map [ * mongo.Client ] bool
138
+ DB * mongo.Database
139
+ Coll * mongo.Collection
135
140
}
136
141
137
142
func newT (wrapped * testing.T , opts ... * Options ) * T {
138
143
t := & T {
139
- T : wrapped ,
144
+ T : wrapped ,
145
+ fpClients : make (map [* mongo.Client ]bool ),
140
146
}
141
147
for _ , opt := range opts {
142
148
for _ , optFn := range opt .optFuncs {
@@ -176,16 +182,10 @@ func New(wrapped *testing.T, opts ...*Options) *T {
176
182
177
183
t := newT (wrapped , opts ... )
178
184
179
- fpOpt := t .clientOpts
180
- if fpOpt != nil {
181
- fpOpt .AutoEncryptionOptions = nil
182
- }
183
- t .fpClient = t .createTestClient (fpOpt )
184
-
185
185
// only create a client if it needs to be shared in sub-tests
186
186
// otherwise, a new client will be created for each subtest
187
187
if t .shareClient != nil && * t .shareClient {
188
- t .Client = t . createTestClient (t . clientOpts )
188
+ t .createTestClient ()
189
189
}
190
190
191
191
wrapped .Cleanup (t .cleanup )
@@ -209,6 +209,12 @@ func (t *T) cleanup() {
209
209
// always disconnect the client regardless of clientType because Client.Disconnect will work against
210
210
// all deployments
211
211
_ = 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
212
218
}
213
219
214
220
// Run creates a new T instance for a sub-test and runs the given callback. It also creates a new collection using the
@@ -231,21 +237,14 @@ func (t *T) RunOpts(name string, opts *Options, callback func(mt *T)) {
231
237
sub .AddMockResponses (sub .mockResponses ... )
232
238
}
233
239
234
- if sub .fpClient == nil {
235
- clientOpts := sub .clientOpts
236
- if clientOpts != nil {
237
- clientOpts .AutoEncryptionOptions = nil
238
- }
239
- sub .fpClient = sub .createTestClient (clientOpts )
240
- }
241
240
// for shareClient, inherit the client from the parent
242
241
if sub .shareClient != nil && * sub .shareClient && sub .clientType == t .clientType {
243
242
sub .Client = t .Client
244
243
}
245
244
// only create a client if not already set
246
245
if sub .Client == nil {
247
246
if sub .createClient == nil || * sub .createClient {
248
- sub .Client = sub . createTestClient (sub . clientOpts )
247
+ sub .createTestClient ()
249
248
}
250
249
}
251
250
// create a collection for this test
@@ -270,7 +269,9 @@ func (t *T) RunOpts(name string, opts *Options, callback func(mt *T)) {
270
269
}
271
270
// only disconnect client if it's not being shared
272
271
if sub .shareClient == nil || ! * sub .shareClient {
273
- _ = sub .Client .Disconnect (context .Background ())
272
+ if _ , ok := sub .fpClients [sub .Client ]; ! ok {
273
+ _ = sub .Client .Disconnect (context .Background ())
274
+ }
274
275
}
275
276
assert .Equal (sub , 0 , sessions , "%v sessions checked out" , sessions )
276
277
assert .Equal (sub , 0 , conns , "%v connections checked out" , conns )
@@ -419,8 +420,10 @@ func (t *T) ResetClient(opts *options.ClientOptions) {
419
420
t .clientOpts = opts
420
421
}
421
422
422
- _ = t .Client .Disconnect (context .Background ())
423
- t .Client = t .createTestClient (t .clientOpts )
423
+ if _ , ok := t .fpClients [t .Client ]; ! ok {
424
+ _ = t .Client .Disconnect (context .Background ())
425
+ }
426
+ t .createTestClient ()
424
427
t .DB = t .Client .Database (t .dbName )
425
428
t .Coll = t .DB .Collection (t .collName , t .collOpts )
426
429
@@ -576,7 +579,8 @@ func (t *T) SetFailPoint(fp FailPoint) {
576
579
if err := SetFailPoint (fp , t .Client ); err != nil {
577
580
t .Fatal (err )
578
581
}
579
- t .failPointNames = append (t .failPointNames , fp .ConfigureFailPoint )
582
+ t .fpClients [t .Client ] = true
583
+ t .failPoints = append (t .failPoints , failPoint {fp .ConfigureFailPoint , t .Client })
580
584
}
581
585
582
586
// SetFailPointFromDocument sets the fail point represented by the given document for the client associated with T. This
@@ -588,30 +592,35 @@ func (t *T) SetFailPointFromDocument(fp bson.Raw) {
588
592
t .Fatal (err )
589
593
}
590
594
595
+ t .fpClients [t .Client ] = true
591
596
name := fp .Index (0 ).Value ().StringValue ()
592
- t .failPointNames = append (t .failPointNames , name )
597
+ t .failPoints = append (t .failPoints , failPoint { name , t . Client } )
593
598
}
594
599
595
600
// TrackFailPoint adds the given fail point to the list of fail points to be disabled when the current test finishes.
596
601
// This function does not create a fail point on the server.
597
- func (t * T ) TrackFailPoint (fpName string ) {
598
- t .failPointNames = append (t .failPointNames , fpName )
602
+ func (t * T ) TrackFailPoint (fpName string , client * mongo.Client ) {
603
+ t .fpClients [client ] = true
604
+ t .failPoints = append (t .failPoints , failPoint {fpName , client })
599
605
}
600
606
601
607
// ClearFailPoints disables all previously set failpoints for this test.
602
608
func (t * T ) ClearFailPoints () {
603
- db := t .Client .Database ("admin" )
604
- for _ , fp := range t .failPointNames {
609
+ for _ , fp := range t .failPoints {
605
610
cmd := bson.D {
606
- {"configureFailPoint" , fp },
611
+ {"configureFailPoint" , fp . name },
607
612
{"mode" , "off" },
608
613
}
609
- err := db .RunCommand (context .Background (), cmd ).Err ()
614
+ err := fp . client . Database ( "admin" ) .RunCommand (context .Background (), cmd ).Err ()
610
615
if err != nil {
611
- t .Fatalf ("error clearing fail point %s: %v" , fp , err )
616
+ t .Fatalf ("error clearing fail point %s: %v" , fp .name , err )
617
+ }
618
+ if fp .client != t .Client {
619
+ _ = fp .client .Disconnect (context .Background ())
620
+ t .fpClients [fp .client ] = false
612
621
}
613
622
}
614
- t .failPointNames = t .failPointNames [:0 ]
623
+ t .failPoints = t .failPoints [:0 ]
615
624
}
616
625
617
626
// CloneDatabase modifies the default database for this test to match the given options.
@@ -639,7 +648,8 @@ func sanitizeCollectionName(db string, coll string) string {
639
648
return coll
640
649
}
641
650
642
- func (t * T ) createTestClient (clientOpts * options.ClientOptions ) * mongo.Client {
651
+ func (t * T ) createTestClient () {
652
+ clientOpts := t .clientOpts
643
653
if clientOpts == nil {
644
654
// default opts
645
655
clientOpts = options .Client ().SetWriteConcern (MajorityWc ).SetReadPreference (PrimaryRp )
@@ -697,20 +707,17 @@ func (t *T) createTestClient(clientOpts *options.ClientOptions) *mongo.Client {
697
707
})
698
708
}
699
709
700
- var client * mongo.Client
701
- var err error
710
+ var uriOpts * options.ClientOptions
702
711
switch t .clientType {
703
712
case Pinned :
704
713
// pin to first mongos
705
714
pinnedHostList := []string {testContext .connString .Hosts [0 ]}
706
- uriOpts := options .Client ().ApplyURI (testContext .connString .Original ).SetHosts (pinnedHostList )
707
- client , err = mongo .NewClient (uriOpts , clientOpts )
715
+ uriOpts = options .Client ().ApplyURI (testContext .connString .Original ).SetHosts (pinnedHostList )
708
716
case Mock :
709
717
// clear pool monitor to avoid configuration error
710
718
clientOpts .PoolMonitor = nil
711
719
t .mockDeployment = newMockDeployment ()
712
720
clientOpts .Deployment = t .mockDeployment
713
- client , err = mongo .NewClient (clientOpts )
714
721
case Proxy :
715
722
t .proxyDialer = newProxyDialer ()
716
723
clientOpts .SetDialer (t .proxyDialer )
@@ -720,23 +727,23 @@ func (t *T) createTestClient(clientOpts *options.ClientOptions) *mongo.Client {
720
727
case Default :
721
728
// Use a different set of options to specify the URI because clientOpts may already have a URI or host seedlist
722
729
// specified.
723
- var uriOpts * options.ClientOptions
724
730
if clientOpts .Deployment == nil {
725
731
// Only specify URI if the deployment is not set to avoid setting topology/server options along with the
726
732
// deployment.
727
733
uriOpts = options .Client ().ApplyURI (testContext .connString .Original )
728
734
}
729
-
730
- // Pass in uriOpts first so clientOpts wins if there are any conflicting settings.
731
- client , err = mongo .NewClient (uriOpts , clientOpts )
732
735
}
736
+ t .clientOpts = options .MergeClientOptions (uriOpts , clientOpts )
737
+
738
+ var err error
739
+ // Pass in uriOpts first so clientOpts wins if there are any conflicting settings.
740
+ t .Client , err = mongo .NewClient (t .clientOpts )
733
741
if err != nil {
734
742
t .Fatalf ("error creating client: %v" , err )
735
743
}
736
- if err := client .Connect (context .Background ()); err != nil {
744
+ if err := t . Client .Connect (context .Background ()); err != nil {
737
745
t .Fatalf ("error connecting client: %v" , err )
738
746
}
739
- return client
740
747
}
741
748
742
749
func (t * T ) createTestCollection () {
0 commit comments