updates

Author: qingyang-hu

mongo/integration/mtest/mongotest.go

@@ -133,14 +133,16 @@ type T struct {
 	succeeded   []*event.CommandSucceededEvent
 	failed      []*event.CommandFailedEvent
 
-	Client *mongo.Client
-	DB     *mongo.Database
-	Coll   *mongo.Collection
+	Client    *mongo.Client
+	fpClients map[*mongo.Client]bool
+	DB        *mongo.Database
+	Coll      *mongo.Collection
 }
 
 func newT(wrapped *testing.T, opts ...*Options) *T {
 	t := &T{
-		T: wrapped,
+		T:         wrapped,
+		fpClients: make(map[*mongo.Client]bool),
 	}
 	for _, opt := range opts {
 		for _, optFn := range opt.optFuncs {
@@ -207,6 +209,12 @@ func (t *T) cleanup() {
 	// always disconnect the client regardless of clientType because Client.Disconnect will work against
 	// all deployments
 	_ = t.Client.Disconnect(context.Background())
+	for client, v := range t.fpClients {
+		if v {
+			client.Disconnect(context.Background())
+		}
+	}
+	t.fpClients = make(map[*mongo.Client]bool)
 }
 
 // Run creates a new T instance for a sub-test and runs the given callback. It also creates a new collection using the
@@ -261,7 +269,9 @@ func (t *T) RunOpts(name string, opts *Options, callback func(mt *T)) {
 			}
 			// only disconnect client if it's not being shared
 			if sub.shareClient == nil || !*sub.shareClient {
-				_ = sub.Client.Disconnect(context.Background())
+				if v, ok := sub.fpClients[sub.Client]; !ok || !v {
+					_ = sub.Client.Disconnect(context.Background())
+				}
 			}
 			assert.Equal(sub, 0, sessions, "%v sessions checked out", sessions)
 			assert.Equal(sub, 0, conns, "%v connections checked out", conns)
@@ -410,7 +420,9 @@ func (t *T) ResetClient(opts *options.ClientOptions) {
 		t.clientOpts = opts
 	}
 
-	_ = t.Client.Disconnect(context.Background())
+	if v, ok := t.fpClients[t.Client]; !ok || !v {
+		_ = t.Client.Disconnect(context.Background())
+	}
 	t.createTestClient()
 	t.DB = t.Client.Database(t.dbName)
 	t.Coll = t.DB.Collection(t.collName, t.collOpts)
@@ -564,42 +576,31 @@ func (t *T) SetFailPoint(fp FailPoint) {
 		}
 	}
 
-	client, err := mongo.NewClient(t.clientOpts)
-	if err != nil {
-		t.Fatalf("error creating client: %v", err)
-	}
-	if err = client.Connect(context.Background()); err != nil {
-		t.Fatalf("error connecting client: %v", err)
-	}
-	if err = SetFailPoint(fp, client); err != nil {
+	if err := SetFailPoint(fp, t.Client); err != nil {
 		t.Fatal(err)
 	}
-	t.failPoints = append(t.failPoints, failPoint{fp.ConfigureFailPoint, client})
+	t.fpClients[t.Client] = true
+	t.failPoints = append(t.failPoints, failPoint{fp.ConfigureFailPoint, t.Client})
 }
 
 // SetFailPointFromDocument sets the fail point represented by the given document for the client associated with T. This
 // method assumes that the given document is in the form {configureFailPoint: <failPointName>, ...}. Commands to create
 // the failpoint will appear in command monitoring channels. The fail point will be automatically disabled after this
 // test has run.
 func (t *T) SetFailPointFromDocument(fp bson.Raw) {
-	client, err := mongo.NewClient(t.clientOpts)
-	if err != nil {
-		t.Fatalf("error creating client: %v", err)
-	}
-	if err = client.Connect(context.Background()); err != nil {
-		t.Fatalf("error connecting client: %v", err)
-	}
-	if err = SetRawFailPoint(fp, client); err != nil {
+	if err := SetRawFailPoint(fp, t.Client); err != nil {
 		t.Fatal(err)
 	}
 
+	t.fpClients[t.Client] = true
 	name := fp.Index(0).Value().StringValue()
-	t.failPoints = append(t.failPoints, failPoint{name, client})
+	t.failPoints = append(t.failPoints, failPoint{name, t.Client})
 }
 
 // TrackFailPoint adds the given fail point to the list of fail points to be disabled when the current test finishes.
 // This function does not create a fail point on the server.
 func (t *T) TrackFailPoint(fpName string, client *mongo.Client) {
+	t.fpClients[client] = true
 	t.failPoints = append(t.failPoints, failPoint{fpName, client})
 }
 
@@ -614,7 +615,10 @@ func (t *T) ClearFailPoints() {
 		if err != nil {
 			t.Fatalf("error clearing fail point %s: %v", fp.name, err)
 		}
-		_ = fp.client.Disconnect(context.Background())
+		if fp.client != t.Client {
+			_ = fp.client.Disconnect(context.Background())
+			t.fpClients[fp.client] = false
+		}
 	}
 	t.failPoints = t.failPoints[:0]
 }