Lift assertions out of withContext calls

michaelbull · michaelbull · commit 8b4fbb84b143 · 2025-09-14T19:08:19.000+01:00
diff --git a/src/main/kotlin/com/github/michaelbull/jdbc/Transaction.kt b/src/main/kotlin/com/github/michaelbull/jdbc/Transaction.kt
@@ -2,6 +2,7 @@ package com.github.michaelbull.jdbc
 
 import com.github.michaelbull.jdbc.context.CoroutineTransaction
 import com.github.michaelbull.jdbc.context.connection
+import com.github.michaelbull.jdbc.context.transaction
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.withContext
@@ -28,7 +29,7 @@ suspend inline fun <T> transaction(crossinline block: suspend CoroutineScope.()
     }
 
     val ctx = currentCoroutineContext()
-    val existingTransaction = ctx[CoroutineTransaction]
+    val existingTransaction = ctx.transaction
 
     return when {
         existingTransaction == null -> {
diff --git a/src/main/kotlin/com/github/michaelbull/jdbc/context/CoroutineTransaction.kt b/src/main/kotlin/com/github/michaelbull/jdbc/context/CoroutineTransaction.kt
@@ -3,9 +3,13 @@ package com.github.michaelbull.jdbc.context
 import kotlin.coroutines.AbstractCoroutineContextElement
 import kotlin.coroutines.CoroutineContext
 
+@PublishedApi
+internal val CoroutineContext.transaction: CoroutineTransaction?
+    get() = get(CoroutineTransaction)
+
 @PublishedApi
 internal class CoroutineTransaction(
-    private var completed: Boolean = false
+    private var completed: Boolean = false,
 ) : AbstractCoroutineContextElement(CoroutineTransaction) {
 
     companion object Key : CoroutineContext.Key<CoroutineTransaction>
diff --git a/src/test/kotlin/com/github/michaelbull/jdbc/ConnectionTest.kt b/src/test/kotlin/com/github/michaelbull/jdbc/ConnectionTest.kt
@@ -39,18 +39,18 @@ class ConnectionTest {
             every { connection } returns newConnection
         }
 
-        withContext(CoroutineDataSource(dataSource)) {
-            val actual = withConnection {
+        val actual = withContext(CoroutineDataSource(dataSource)) {
+            withConnection {
                 coroutineContext.connection
             }
-
-            assertEquals(newConnection, actual)
         }
+
+        assertEquals(newConnection, actual)
     }
 
     @Test
     fun `withConnection adds new connection to context if existing connection isClosed returns true`() = runTest {
-        val existingConnection = mockk<Connection> {
+        val closedConnection = mockk<Connection> {
             every { isClosed } returns true
         }
 
@@ -62,18 +62,18 @@ class ConnectionTest {
             every { connection } returns newConnection
         }
 
-        withContext(CoroutineDataSource(dataSource) + CoroutineConnection(existingConnection)) {
-            val actual = withConnection {
+        val actual = withContext(CoroutineDataSource(dataSource) + CoroutineConnection(closedConnection)) {
+            withConnection {
                 coroutineContext.connection
             }
-
-            assertEquals(newConnection, actual)
         }
+
+        assertEquals(newConnection, actual)
     }
 
     @Test
     fun `withConnection adds new connection to context if existing connection isClosed throws exception`() = runTest {
-        val existingConnection = mockk<Connection> {
+        val brokenConnection = mockk<Connection> {
             every { isClosed } throws SQLException()
         }
 
@@ -85,30 +85,30 @@ class ConnectionTest {
             every { connection } returns newConnection
         }
 
-        withContext(CoroutineDataSource(dataSource) + CoroutineConnection(existingConnection)) {
-            val actual = withConnection {
+        val actual = withContext(CoroutineDataSource(dataSource) + CoroutineConnection(brokenConnection)) {
+            withConnection {
                 coroutineContext.connection
             }
-
-            assertEquals(newConnection, actual)
         }
+
+        assertEquals(newConnection, actual)
     }
 
     @Test
-    fun `withConnection reuses existing connection in context if not closed`() = runTest {
-        val existing = mockk<Connection> {
+    fun `withConnection reuses open connection`() = runTest {
+        val openConnection = mockk<Connection> {
             every { isClosed } returns false
         }
 
         val dataSource = mockk<DataSource>()
 
-        withContext(CoroutineDataSource(dataSource) + CoroutineConnection(existing)) {
-            val actual = withConnection {
+        val actual = withContext(CoroutineDataSource(dataSource) + CoroutineConnection(openConnection)) {
+            withConnection {
                 coroutineContext.connection
             }
-
-            assertEquals(existing, actual)
         }
+
+        assertEquals(openConnection, actual)
     }
 
     @Test
@@ -153,20 +153,20 @@ class ConnectionTest {
 
     @Test
     fun `withConnection does not close connection if connection was not added to context`() = runTest {
-        val existing = mockk<Connection> {
+        val openConnection = mockk<Connection> {
             every { isClosed } returns false
         }
 
         val dataSource = mockk<DataSource> {
-            every { connection } returns existing
+            every { connection } returns openConnection
         }
 
-        withContext(CoroutineDataSource(dataSource) + CoroutineConnection(existing)) {
+        withContext(CoroutineDataSource(dataSource) + CoroutineConnection(openConnection)) {
             withConnection {
                 /* empty */
             }
         }
 
-        verify(exactly = 0) { existing.close() }
+        verify(exactly = 0) { openConnection.close() }
     }
 }
diff --git a/src/test/kotlin/com/github/michaelbull/jdbc/TransactionTest.kt b/src/test/kotlin/com/github/michaelbull/jdbc/TransactionTest.kt
@@ -2,6 +2,7 @@ package com.github.michaelbull.jdbc
 
 import com.github.michaelbull.jdbc.context.CoroutineConnection
 import com.github.michaelbull.jdbc.context.CoroutineTransaction
+import com.github.michaelbull.jdbc.context.transaction
 import io.mockk.every
 import io.mockk.just
 import io.mockk.mockk
@@ -66,13 +67,13 @@ class TransactionTest {
     fun `transaction reuses existing transaction in context if incomplete`() = runTest {
         val incompleteTransaction = CoroutineTransaction(completed = false)
 
-        withContext(incompleteTransaction) {
-            val actual = transaction {
-                coroutineContext[CoroutineTransaction]
+        val actual = withContext(incompleteTransaction) {
+            transaction {
+                coroutineContext.transaction
             }
-
-            assertEquals(incompleteTransaction, actual)
         }
+
+        assertEquals(incompleteTransaction, actual)
     }
 
     @Test
@@ -104,13 +105,13 @@ class TransactionTest {
             every { autoCommit } returns true
         }
 
-        withContext(CoroutineConnection(connection)) {
-            val transaction = transaction {
-                coroutineContext[CoroutineTransaction]
+        val actual = withContext(CoroutineConnection(connection)) {
+            transaction {
+                coroutineContext.transaction
             }
-
-            assertNotNull(transaction)
         }
+
+        assertNotNull(actual)
     }
 
     @Test
@@ -119,13 +120,13 @@ class TransactionTest {
             every { autoCommit } returns true
         }
 
-        withContext(CoroutineConnection(connection)) {
-            val transaction = runTransactionally {
-                coroutineContext[CoroutineTransaction]
+        val actual = withContext(CoroutineConnection(connection)) {
+            runTransactionally {
+                coroutineContext.transaction
             }
-
-            assertNotNull(transaction)
         }
+
+        assertNotNull(actual)
     }
 
     @Test
@@ -135,7 +136,9 @@ class TransactionTest {
         }
 
         withContext(CoroutineConnection(connection)) {
-            runTransactionally {}
+            runTransactionally {
+                /* empty */
+            }
         }
 
         verify(exactly = 1) { connection.commit() }
diff --git a/src/test/kotlin/com/github/michaelbull/jdbc/context/CoroutineConnectionTest.kt b/src/test/kotlin/com/github/michaelbull/jdbc/context/CoroutineConnectionTest.kt
@@ -21,9 +21,10 @@ class CoroutineConnectionTest {
     fun `connection returns connection if in context`() = runTest {
         val expected = mockk<Connection>()
 
-        withContext(CoroutineConnection(expected)) {
-            val actual = coroutineContext.connection
-            assertEquals(expected, actual)
+        val actual = withContext(CoroutineConnection(expected)) {
+            coroutineContext.connection
         }
+
+        assertEquals(expected, actual)
     }
 }
diff --git a/src/test/kotlin/com/github/michaelbull/jdbc/context/CoroutineDataSourceTest.kt b/src/test/kotlin/com/github/michaelbull/jdbc/context/CoroutineDataSourceTest.kt
@@ -21,9 +21,10 @@ class CoroutineDataSourceTest {
     fun `dataSource returns connection if in context`() = runTest {
         val expected = mockk<DataSource>()
 
-        withContext(CoroutineDataSource(expected)) {
-            val actual = coroutineContext.dataSource
-            assertEquals(expected, actual)
+        val actual = withContext(CoroutineDataSource(expected)) {
+            coroutineContext.dataSource
         }
+
+        assertEquals(expected, actual)
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -39,18 +39,18 @@ class ConnectionTest {`
`39`	`39`	`every { connection } returns newConnection`
`40`	`40`	`}`
`41`	`41`
`42`		`- withContext(CoroutineDataSource(dataSource)) {`
`43`		`- val actual = withConnection {`
	`42`	`+ val actual = withContext(CoroutineDataSource(dataSource)) {`
	`43`	`+ withConnection {`
`44`	`44`	`coroutineContext.connection`
`45`	`45`	`}`
`46`		`-`
`47`		`- assertEquals(newConnection, actual)`
`48`	`46`	`}`
	`47`	`+`
	`48`	`+ assertEquals(newConnection, actual)`
`49`	`49`	`}`
`50`	`50`
`51`	`51`	`@Test`
`52`	`52`	fun `withConnection adds new connection to context if existing connection isClosed returns true`() = runTest {
`53`		`- val existingConnection = mockk<Connection> {`
	`53`	`+ val closedConnection = mockk<Connection> {`
`54`	`54`	`every { isClosed } returns true`
`55`	`55`	`}`
`56`	`56`
`@@ -62,18 +62,18 @@ class ConnectionTest {`
`62`	`62`	`every { connection } returns newConnection`
`63`	`63`	`}`
`64`	`64`
`65`		`- withContext(CoroutineDataSource(dataSource) + CoroutineConnection(existingConnection)) {`
`66`		`- val actual = withConnection {`
	`65`	`+ val actual = withContext(CoroutineDataSource(dataSource) + CoroutineConnection(closedConnection)) {`
	`66`	`+ withConnection {`
`67`	`67`	`coroutineContext.connection`
`68`	`68`	`}`
`69`		`-`
`70`		`- assertEquals(newConnection, actual)`
`71`	`69`	`}`
	`70`	`+`
	`71`	`+ assertEquals(newConnection, actual)`
`72`	`72`	`}`
`73`	`73`
`74`	`74`	`@Test`
`75`	`75`	fun `withConnection adds new connection to context if existing connection isClosed throws exception`() = runTest {
`76`		`- val existingConnection = mockk<Connection> {`
	`76`	`+ val brokenConnection = mockk<Connection> {`
`77`	`77`	`every { isClosed } throws SQLException()`
`78`	`78`	`}`
`79`	`79`
`@@ -85,30 +85,30 @@ class ConnectionTest {`
`85`	`85`	`every { connection } returns newConnection`
`86`	`86`	`}`
`87`	`87`
`88`		`- withContext(CoroutineDataSource(dataSource) + CoroutineConnection(existingConnection)) {`
`89`		`- val actual = withConnection {`
	`88`	`+ val actual = withContext(CoroutineDataSource(dataSource) + CoroutineConnection(brokenConnection)) {`
	`89`	`+ withConnection {`
`90`	`90`	`coroutineContext.connection`
`91`	`91`	`}`
`92`		`-`
`93`		`- assertEquals(newConnection, actual)`
`94`	`92`	`}`
	`93`	`+`
	`94`	`+ assertEquals(newConnection, actual)`
`95`	`95`	`}`
`96`	`96`
`97`	`97`	`@Test`
`98`		- fun `withConnection reuses existing connection in context if not closed`() = runTest {
`99`		`- val existing = mockk<Connection> {`
	`98`	+ fun `withConnection reuses open connection`() = runTest {
	`99`	`+ val openConnection = mockk<Connection> {`
`100`	`100`	`every { isClosed } returns false`
`101`	`101`	`}`
`102`	`102`
`103`	`103`	`val dataSource = mockk<DataSource>()`
`104`	`104`
`105`		`- withContext(CoroutineDataSource(dataSource) + CoroutineConnection(existing)) {`
`106`		`- val actual = withConnection {`
	`105`	`+ val actual = withContext(CoroutineDataSource(dataSource) + CoroutineConnection(openConnection)) {`
	`106`	`+ withConnection {`
`107`	`107`	`coroutineContext.connection`
`108`	`108`	`}`
`109`		`-`
`110`		`- assertEquals(existing, actual)`
`111`	`109`	`}`
	`110`	`+`
	`111`	`+ assertEquals(openConnection, actual)`
`112`	`112`	`}`
`113`	`113`
`114`	`114`	`@Test`
`@@ -153,20 +153,20 @@ class ConnectionTest {`
`153`	`153`
`154`	`154`	`@Test`
`155`	`155`	fun `withConnection does not close connection if connection was not added to context`() = runTest {
`156`		`- val existing = mockk<Connection> {`
	`156`	`+ val openConnection = mockk<Connection> {`
`157`	`157`	`every { isClosed } returns false`
`158`	`158`	`}`
`159`	`159`
`160`	`160`	`val dataSource = mockk<DataSource> {`
`161`		`- every { connection } returns existing`
	`161`	`+ every { connection } returns openConnection`
`162`	`162`	`}`
`163`	`163`
`164`		`- withContext(CoroutineDataSource(dataSource) + CoroutineConnection(existing)) {`
	`164`	`+ withContext(CoroutineDataSource(dataSource) + CoroutineConnection(openConnection)) {`
`165`	`165`	`withConnection {`
`166`	`166`	`/* empty */`
`167`	`167`	`}`
`168`	`168`	`}`
`169`	`169`
`170`		`- verify(exactly = 0) { existing.close() }`
	`170`	`+ verify(exactly = 0) { openConnection.close() }`
`171`	`171`	`}`
`172`	`172`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,7 @@ package com.github.michaelbull.jdbc`
`2`	`2`
`3`	`3`	`import com.github.michaelbull.jdbc.context.CoroutineConnection`
`4`	`4`	`import com.github.michaelbull.jdbc.context.CoroutineTransaction`
	`5`	`+import com.github.michaelbull.jdbc.context.transaction`
`5`	`6`	`import io.mockk.every`
`6`	`7`	`import io.mockk.just`
`7`	`8`	`import io.mockk.mockk`
`@@ -66,13 +67,13 @@ class TransactionTest {`
`66`	`67`	fun `transaction reuses existing transaction in context if incomplete`() = runTest {
`67`	`68`	`val incompleteTransaction = CoroutineTransaction(completed = false)`
`68`	`69`
`69`		`- withContext(incompleteTransaction) {`
`70`		`- val actual = transaction {`
`71`		`- coroutineContext[CoroutineTransaction]`
	`70`	`+ val actual = withContext(incompleteTransaction) {`
	`71`	`+ transaction {`
	`72`	`+ coroutineContext.transaction`
`72`	`73`	`}`
`73`		`-`
`74`		`- assertEquals(incompleteTransaction, actual)`
`75`	`74`	`}`
	`75`	`+`
	`76`	`+ assertEquals(incompleteTransaction, actual)`
`76`	`77`	`}`
`77`	`78`
`78`	`79`	`@Test`
`@@ -104,13 +105,13 @@ class TransactionTest {`
`104`	`105`	`every { autoCommit } returns true`
`105`	`106`	`}`
`106`	`107`
`107`		`- withContext(CoroutineConnection(connection)) {`
`108`		`- val transaction = transaction {`
`109`		`- coroutineContext[CoroutineTransaction]`
	`108`	`+ val actual = withContext(CoroutineConnection(connection)) {`
	`109`	`+ transaction {`
	`110`	`+ coroutineContext.transaction`
`110`	`111`	`}`
`111`		`-`
`112`		`- assertNotNull(transaction)`
`113`	`112`	`}`
	`113`	`+`
	`114`	`+ assertNotNull(actual)`
`114`	`115`	`}`
`115`	`116`
`116`	`117`	`@Test`
`@@ -119,13 +120,13 @@ class TransactionTest {`
`119`	`120`	`every { autoCommit } returns true`
`120`	`121`	`}`
`121`	`122`
`122`		`- withContext(CoroutineConnection(connection)) {`
`123`		`- val transaction = runTransactionally {`
`124`		`- coroutineContext[CoroutineTransaction]`
	`123`	`+ val actual = withContext(CoroutineConnection(connection)) {`
	`124`	`+ runTransactionally {`
	`125`	`+ coroutineContext.transaction`
`125`	`126`	`}`
`126`		`-`
`127`		`- assertNotNull(transaction)`
`128`	`127`	`}`
	`128`	`+`
	`129`	`+ assertNotNull(actual)`
`129`	`130`	`}`
`130`	`131`
`131`	`132`	`@Test`
`@@ -135,7 +136,9 @@ class TransactionTest {`
`135`	`136`	`}`
`136`	`137`
`137`	`138`	`withContext(CoroutineConnection(connection)) {`
`138`		`- runTransactionally {}`
	`139`	`+ runTransactionally {`
	`140`	`+ /* empty */`
	`141`	`+ }`
`139`	`142`	`}`
`140`	`143`
`141`	`144`	`verify(exactly = 1) { connection.commit() }`
Original file line number	Diff line number	Diff line change
`@@ -21,9 +21,10 @@ class CoroutineConnectionTest {`
`21`	`21`	fun `connection returns connection if in context`() = runTest {
`22`	`22`	`val expected = mockk<Connection>()`
`23`	`23`
`24`		`- withContext(CoroutineConnection(expected)) {`
`25`		`- val actual = coroutineContext.connection`
`26`		`- assertEquals(expected, actual)`
	`24`	`+ val actual = withContext(CoroutineConnection(expected)) {`
	`25`	`+ coroutineContext.connection`
`27`	`26`	`}`
	`27`	`+`
	`28`	`+ assertEquals(expected, actual)`
`28`	`29`	`}`
`29`	`30`	`}`
Original file line number	Diff line number	Diff line change
`@@ -21,9 +21,10 @@ class CoroutineDataSourceTest {`
`21`	`21`	fun `dataSource returns connection if in context`() = runTest {
`22`	`22`	`val expected = mockk<DataSource>()`
`23`	`23`
`24`		`- withContext(CoroutineDataSource(expected)) {`
`25`		`- val actual = coroutineContext.dataSource`
`26`		`- assertEquals(expected, actual)`
	`24`	`+ val actual = withContext(CoroutineDataSource(expected)) {`
	`25`	`+ coroutineContext.dataSource`
`27`	`26`	`}`
	`27`	`+`
	`28`	`+ assertEquals(expected, actual)`
`28`	`29`	`}`
`29`	`30`	`}`