Added test for DataSourceConnectionWarmup

cohort-core/src/main/kotlin/com/sksamuel/cohort/db/DataSourceConnectionWarmup.kt

@@ -0,0 +1,28 @@
+package com.sksamuel.cohort.db
+
+import com.sksamuel.cohort.Warmup
+import javax.sql.DataSource
+import kotlin.time.Duration
+import kotlin.time.Duration.Companion.seconds
+
+/**
+ * A Cohort [Warmup] that warms a [DataSource] by executing a query.
+ *
+ * Uses the JDBC4 method isValid(timeout) with the given [timeout] to check that the connection
+ * returned is open and usable.
+ */
+class DataSourceConnectionWarmup(
+   private val ds: DataSource,
+   private val query: String,
+   private val timeout: Duration = 1.seconds,
+) : Warmup {
+
+   override val name: String = "datasource_connection_warmup"
+
+   override suspend fun warm(iteration: Int) {
+      ds.connection.use { conn ->
+         conn.isValid(timeout.inWholeSeconds.toInt())
+         conn.createStatement().use { it.execute(query) }
+      }
+   }
+}

cohort-core/src/main/kotlin/com/sksamuel/cohort/db/DataSourceWarmup.kt

@@ -1,6 +1,5 @@
 package com.sksamuel.cohort.db
 
-import com.sksamuel.cohort.Warmup
 import com.sksamuel.cohort.WarmupHealthCheck
 import javax.sql.DataSource
 import kotlin.time.Duration
@@ -30,24 +29,3 @@ class DataSourceWarmup(
    }
 }
 
-/**
- * A Cohort [WarmupHealthCheck] that warms a [DataSource] by executing a query.
- *
- * Uses the JDBC4 method isValid(timeout) with the given [timeout] to check that the connection
- * returned is open and usable.
- */
-class DataSourceConnectionWarmup(
-   private val ds: DataSource,
-   private val query: String,
-   private val timeout: Duration = 1.seconds,
-) : Warmup {
-
-   override val name: String = "datasource_connection_warmup"
-
-   override suspend fun warm(iteration: Int) {
-      ds.connection.use { conn ->
-         conn.isValid(timeout.inWholeSeconds.toInt())
-         conn.createStatement().use { it.execute(query) }
-      }
-   }
-}

cohort-core/src/test/kotlin/com/sksamuel/cohort/db/DataSourceConnectionWarmupTest.kt

@@ -0,0 +1,23 @@
+package com.sksamuel.cohort.db
+
+import com.sksamuel.cohort.WarmupRegistry
+import com.sksamuel.cohort.WarmupState
+import io.kotest.core.spec.style.FunSpec
+import io.kotest.matchers.shouldBe
+import kotlinx.coroutines.delay
+import kotlin.time.Duration.Companion.seconds
+
+class DataSourceConnectionWarmupTest : FunSpec({
+
+   test("happy path") {
+      val ds = createHikariDS()
+      ds.connection.use { it.createStatement().executeUpdate("CREATE TABLE foo (id int)") }
+      val warmups = WarmupRegistry {
+         register(DataSourceConnectionWarmup(ds, "select * from foo"), 3.seconds)
+      }
+      warmups.state() shouldBe WarmupState.Running
+      delay(5.seconds)
+      warmups.state() shouldBe WarmupState.Completed
+   }
+
+})

cohort-core/src/test/kotlin/com/sksamuel/cohort/db/DatabaseConnectionHealthCheckTest.kt

@@ -8,6 +8,12 @@ import io.kotest.matchers.shouldBe
 
 class DatabaseConnectionHealthCheckTest : FunSpec({
 
+   test("should return healthy if query null and connection is open") {
+      DatabaseConnectionHealthCheck(
+         createHikariDS(),
+      ).check().status shouldBe HealthStatus.Healthy
+   }
+
    test("should return healthy for valid query") {
       val ds = createHikariDS()
       ds.connection.use { it.createStatement().executeUpdate("CREATE TABLE foo (id int)") }