Ensure MongoClient instances are closed in Scala tests (#830)

Co-authored-by: Ross Lawley <ross.lawley@gmail.com>

JAVA-4412

Author: jyemin

driver-scala/src/it/scala/org/mongodb/scala/ClientSideEncryptionBypassAutoEncryptionSpec.scala

@@ -60,25 +60,30 @@ class ClientSideEncryptionBypassAutoEncryptionSpec extends RequiresMongoDBISpec
       .codecRegistry(DEFAULT_CODEC_REGISTRY)
       .build
 
-    val clientEncrypted = MongoClient(clientSettings)
+    withTempClient(
+      clientSettings,
+      clientEncrypted => {
 
-    val fieldValue = BsonString("123456789")
+        val fieldValue = BsonString("123456789")
 
-    val dataKeyId = clientEncryption.createDataKey("local", DataKeyOptions()).head().futureValue
+        val dataKeyId = clientEncryption.createDataKey("local", DataKeyOptions()).head().futureValue
 
-    val encryptedFieldValue = clientEncryption
-      .encrypt(fieldValue, EncryptOptions("AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic").keyId(dataKeyId))
-      .head()
-      .futureValue
+        val encryptedFieldValue = clientEncryption
+          .encrypt(fieldValue, EncryptOptions("AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic").keyId(dataKeyId))
+          .head()
+          .futureValue
 
-    val collection: MongoCollection[Document] =
-      clientEncrypted.getDatabase(databaseName).getCollection[Document]("test")
+        val collection: MongoCollection[Document] =
+          clientEncrypted.getDatabase(databaseName).getCollection[Document]("test")
 
-    collection.insertOne(Document("encryptedField" -> encryptedFieldValue)).futureValue
+        collection.insertOne(Document("encryptedField" -> encryptedFieldValue)).futureValue
 
-    val result = collection.find().first().head().futureValue
+        val result = collection.find().first().head().futureValue
 
-    result.get[BsonString]("encryptedField") should equal(Some(fieldValue))
+        result.get[BsonString]("encryptedField") should equal(Some(fieldValue))
+
+      }
+    )
   }
 
 }

driver-scala/src/it/scala/org/mongodb/scala/RequiresMongoDBISpec.scala

@@ -16,30 +16,24 @@
 
 package org.mongodb.scala
 
-import com.mongodb.ClusterFixture.getServerApi
 import com.mongodb.connection.ServerVersion
 import org.mongodb.scala.bson.BsonString
 import org.scalatest._
 
 import scala.collection.JavaConverters._
 import scala.concurrent.duration.{ Duration, _ }
 import scala.concurrent.{ Await, ExecutionContext }
-import scala.util.{ Properties, Try }
 
 trait RequiresMongoDBISpec extends BaseSpec with BeforeAndAfterAll {
 
   implicit val ec: ExecutionContext = ExecutionContext.Implicits.global
 
-  private val DEFAULT_URI: String = "mongodb://localhost:27017/"
-  private val MONGODB_URI_SYSTEM_PROPERTY_NAME: String = "org.mongodb.test.uri"
   val WAIT_DURATION: Duration = 60.seconds
   private val DB_PREFIX = "mongo-scala-"
   private var _currentTestName: Option[String] = None
-  private var mongoDBOnline: Boolean = false
 
   protected override def runTest(testName: String, args: Args): Status = {
     _currentTestName = Some(testName.split("should")(1))
-    mongoDBOnline = isMongoDBOnline()
     super.runTest(testName, args)
   }
 
@@ -53,48 +47,33 @@ trait RequiresMongoDBISpec extends BaseSpec with BeforeAndAfterAll {
    */
   def collectionName: String = _currentTestName.getOrElse(suiteName).filter(_.isLetterOrDigit)
 
-  val mongoClientURI: String = {
-    val uri = Properties.propOrElse(MONGODB_URI_SYSTEM_PROPERTY_NAME, DEFAULT_URI)
-    if (!uri.isBlank) uri else DEFAULT_URI
-  }
-  val connectionString: ConnectionString = ConnectionString(mongoClientURI)
+  def mongoClientSettingsBuilder: MongoClientSettings.Builder = TestMongoClientHelper.mongoClientSettingsBuilder
 
-  def mongoClientSettingsBuilder: MongoClientSettings.Builder = {
-    val builder = MongoClientSettings.builder().applyConnectionString(connectionString)
-    if (getServerApi != null) {
-      builder.serverApi(getServerApi)
-    }
-    builder
-  }
+  val mongoClientSettings: MongoClientSettings = TestMongoClientHelper.mongoClientSettings
 
-  val mongoClientSettings: MongoClientSettings = mongoClientSettingsBuilder.build()
+  def mongoClient(): MongoClient = TestMongoClientHelper.mongoClient
 
-  def mongoClient(): MongoClient = MongoClient(mongoClientSettings)
-
-  def isMongoDBOnline(): Boolean = {
-    Try(Await.result(MongoClient(mongoClientSettings).listDatabaseNames().toFuture(), WAIT_DURATION)).isSuccess
-  }
-
-  def hasSingleHost(): Boolean = {
-    new ConnectionString(mongoClientURI).getHosts.size() == 1
+  def checkMongoDB(): Unit = {
+    if (!TestMongoClientHelper.isMongoDBOnline) {
+      cancel("No Available Database")
+    }
   }
 
-  def checkMongoDB() {
-    if (!mongoDBOnline) {
-      cancel("No Available Database")
+  def withTempClient(mongoClientSettings: MongoClientSettings, testCode: MongoClient => Any): Unit = {
+    val client = MongoClient(mongoClientSettings)
+    try {
+      testCode(client)
+    } finally {
+      client.close()
     }
   }
 
   def withClient(testCode: MongoClient => Any): Unit = {
     checkMongoDB()
-    val client = mongoClient()
-    try testCode(client) // loan the client
-    finally {
-      client.close()
-    }
+    testCode(TestMongoClientHelper.mongoClient) // loan the client
   }
 
-  def withDatabase(dbName: String)(testCode: MongoDatabase => Any) {
+  def withDatabase(dbName: String)(testCode: MongoDatabase => Any): Unit = {
     withClient { client =>
       val databaseName = if (dbName.startsWith(DB_PREFIX)) dbName.take(63) else s"$DB_PREFIX$dbName".take(63) // scalastyle:ignore
       val mongoDatabase = client.getDatabase(databaseName)
@@ -108,7 +87,7 @@ trait RequiresMongoDBISpec extends BaseSpec with BeforeAndAfterAll {
 
   def withDatabase(testCode: MongoDatabase => Any): Unit = withDatabase(databaseName)(testCode: MongoDatabase => Any)
 
-  def withCollection(testCode: MongoCollection[Document] => Any) {
+  def withCollection(testCode: MongoCollection[Document] => Any): Unit = {
     withDatabase(databaseName) { mongoDatabase =>
       val mongoCollection = mongoDatabase.getCollection(collectionName)
       try testCode(mongoCollection) // "loan" the fixture to the test
@@ -119,19 +98,25 @@ trait RequiresMongoDBISpec extends BaseSpec with BeforeAndAfterAll {
     }
   }
 
-  lazy val isSharded: Boolean = if (!mongoDBOnline) {
+  lazy val isSharded: Boolean = if (!TestMongoClientHelper.isMongoDBOnline) {
     false
   } else {
     Await
-      .result(mongoClient().getDatabase("admin").runCommand(Document("isMaster" -> 1)).toFuture(), WAIT_DURATION)
+      .result(
+        mongoClient().getDatabase("admin").runCommand(Document("isMaster" -> 1)).toFuture(),
+        WAIT_DURATION
+      )
       .getOrElse("msg", BsonString(""))
       .asString()
       .getValue == "isdbgrid"
   }
 
   lazy val buildInfo: Document = {
-    if (mongoDBOnline) {
-      Await.result(mongoClient().getDatabase("admin").runCommand(Document("buildInfo" -> 1)).toFuture(), WAIT_DURATION)
+    if (TestMongoClientHelper.isMongoDBOnline) {
+      Await.result(
+        mongoClient().getDatabase("admin").runCommand(Document("buildInfo" -> 1)).toFuture(),
+        WAIT_DURATION
+      )
     } else {
       Document()
     }
@@ -158,26 +143,14 @@ trait RequiresMongoDBISpec extends BaseSpec with BeforeAndAfterAll {
   }
 
   override def beforeAll() {
-    if (mongoDBOnline) {
-      val client = mongoClient()
-      Await.result(client.getDatabase(databaseName).drop().toFuture(), WAIT_DURATION)
-      client.close()
+    if (TestMongoClientHelper.isMongoDBOnline) {
+      Await.result(TestMongoClientHelper.mongoClient.getDatabase(databaseName).drop().toFuture(), WAIT_DURATION)
     }
   }
 
   override def afterAll() {
-    if (mongoDBOnline) {
-      val client = mongoClient()
-      Await.result(client.getDatabase(databaseName).drop().toFuture(), WAIT_DURATION)
-      client.close()
-    }
-  }
-
-  Runtime.getRuntime.addShutdownHook(new ShutdownHook())
-
-  private[mongodb] class ShutdownHook extends Thread {
-    override def run() {
-      mongoClient().getDatabase(databaseName).drop()
+    if (TestMongoClientHelper.isMongoDBOnline) {
+      Await.result(TestMongoClientHelper.mongoClient.getDatabase(databaseName).drop().toFuture(), WAIT_DURATION)
     }
   }
 

driver-scala/src/it/scala/org/mongodb/scala/TestMongoClientHelper.scala

@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2008 - 2013 10gen, Inc. <http://10gen.com>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.mongodb.scala
+
+import com.mongodb.ClusterFixture.getServerApi
+import org.mongodb.scala.syncadapter.WAIT_DURATION
+
+import scala.concurrent.Await
+import scala.util.{ Properties, Try }
+
+object TestMongoClientHelper {
+  private val DEFAULT_URI: String = "mongodb://localhost:27017/"
+  private val MONGODB_URI_SYSTEM_PROPERTY_NAME: String = "org.mongodb.test.uri"
+
+  val mongoClientURI: String = {
+    val uri = Properties.propOrElse(MONGODB_URI_SYSTEM_PROPERTY_NAME, DEFAULT_URI)
+    if (!uri.isBlank) uri else DEFAULT_URI
+  }
+  val connectionString: ConnectionString = ConnectionString(mongoClientURI)
+
+  def mongoClientSettingsBuilder: MongoClientSettings.Builder = {
+    val builder = MongoClientSettings.builder().applyConnectionString(connectionString)
+    if (getServerApi != null) {
+      builder.serverApi(getServerApi)
+    }
+    builder
+  }
+
+  val mongoClientSettings: MongoClientSettings = mongoClientSettingsBuilder.build()
+  val mongoClient: MongoClient = MongoClient(mongoClientSettings)
+
+  def isMongoDBOnline: Boolean = {
+    Try(Await.result(TestMongoClientHelper.mongoClient.listDatabaseNames().toFuture(), WAIT_DURATION)).isSuccess
+  }
+
+  def hasSingleHost: Boolean = {
+    TestMongoClientHelper.connectionString.getHosts.size() == 1
+  }
+
+  Runtime.getRuntime.addShutdownHook(new ShutdownHook())
+
+  private[mongodb] class ShutdownHook extends Thread {
+    override def run() {
+      mongoClient.close()
+    }
+  }
+}

driver-scala/src/it/scala/org/mongodb/scala/documentation/DocumentationExampleSpec.scala

@@ -17,8 +17,8 @@
 package org.mongodb.scala.documentation
 
 import java.util.concurrent.atomic.AtomicBoolean
-
 import com.mongodb.client.model.changestream.{ ChangeStreamDocument, FullDocument }
+import org.mongodb.scala.TestMongoClientHelper.hasSingleHost
 import org.mongodb.scala._
 import org.mongodb.scala.bson.conversions.Bson
 import org.mongodb.scala.bson.{ BsonArray, BsonDocument, BsonNull, BsonString, BsonValue }
@@ -596,7 +596,7 @@ class DocumentationExampleSpec extends RequiresMongoDBISpec with FuturesSpec {
   }
 
   it should "be able to watch" in withCollection { collection =>
-    assume(serverVersionAtLeast(List(3, 6, 0)) && !hasSingleHost())
+    assume(serverVersionAtLeast(List(3, 6, 0)) && !hasSingleHost)
     val inventory: MongoCollection[Document] = collection
     val stop: AtomicBoolean = new AtomicBoolean(false)
     new Thread(new Runnable {

driver-scala/src/it/scala/org/mongodb/scala/documentation/DocumentationTransactionsExampleSpec.scala

@@ -16,6 +16,7 @@
 
 package org.mongodb.scala.documentation
 
+import org.mongodb.scala.TestMongoClientHelper.hasSingleHost
 import org.mongodb.scala._
 import org.mongodb.scala.model.{ Filters, Updates }
 import org.mongodb.scala.result.{ InsertOneResult, UpdateResult }
@@ -36,7 +37,7 @@ class DocumentationTransactionsExampleSpec extends RequiresMongoDBISpec {
   // end implicit functions
 
   "The Scala driver" should "be able to commit a transaction" in withClient { client =>
-    assume(serverVersionAtLeast(List(4, 0, 0)) && !hasSingleHost())
+    assume(serverVersionAtLeast(List(4, 0, 0)) && !hasSingleHost)
     client.getDatabase("hr").drop().execute()
     client.getDatabase("hr").createCollection("employees").execute()
     client.getDatabase("hr").createCollection("events").execute()

driver-scala/src/it/scala/org/mongodb/scala/gridfs/GridFSObservableSpec.scala

@@ -284,29 +284,30 @@ class GridFSObservableSpec extends RequiresMongoDBISpec with FuturesSpec with Be
   }
 
   it should "use the user provided codec registries for encoding / decoding data" in {
-    val client = MongoClient(
+    withTempClient(
       mongoClientSettingsBuilder
         .uuidRepresentation(UuidRepresentation.STANDARD)
-        .build()
+        .build(),
+      client => {
+        val database = client.getDatabase(databaseName)
+        val uuid = UUID.randomUUID()
+        val fileMeta = new org.bson.Document("uuid", uuid)
+        val bucket = GridFSBucket(database)
+
+        val fileId = bucket
+          .uploadFromObservable(
+            "myFile",
+            Observable(Seq(ByteBuffer.wrap(multiChunkString.getBytes()))),
+            new GridFSUploadOptions().metadata(fileMeta)
+          )
+          .head()
+          .futureValue
+
+        val fileAsDocument = filesCollection.find[BsonDocument]().head().futureValue
+        fileAsDocument.getDocument("metadata").getBinary("uuid").getType should equal(4.toByte)
+        fileAsDocument.getDocument("metadata").getBinary("uuid").asUuid() should equal(uuid)
+      }
     )
-
-    val database = client.getDatabase(databaseName)
-    val uuid = UUID.randomUUID()
-    val fileMeta = new org.bson.Document("uuid", uuid)
-    val bucket = GridFSBucket(database)
-
-    val fileId = bucket
-      .uploadFromObservable(
-        "myFile",
-        Observable(Seq(ByteBuffer.wrap(multiChunkString.getBytes()))),
-        new GridFSUploadOptions().metadata(fileMeta)
-      )
-      .head()
-      .futureValue
-
-    val fileAsDocument = filesCollection.find[BsonDocument]().head().futureValue
-    fileAsDocument.getDocument("metadata").getBinary("uuid").getType should equal(4.toByte)
-    fileAsDocument.getDocument("metadata").getBinary("uuid").asUuid() should equal(uuid)
   }
 
   it should "handle missing file name data when downloading" in {