Fix awaits { } lambda signature

lupuuss · Aug 8, 2024 · 64b7f8a · 64b7f8a
1 parent 7d90e2d
commit 64b7f8a
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/mokkery-coroutines/src/commonMain/kotlin/dev/mokkery/coroutines/answering/AwaitsApi.kt b/mokkery-coroutines/src/commonMain/kotlin/dev/mokkery/coroutines/answering/AwaitsApi.kt
@@ -1,5 +1,6 @@
 package dev.mokkery.coroutines.answering
 
+import dev.mokkery.answering.CallArgs
 import dev.mokkery.answering.SuspendAnsweringScope
 import dev.mokkery.answering.SuspendCallDefinitionScope
 import dev.mokkery.coroutines.internal.answering.AwaitAnswer
@@ -23,6 +24,10 @@ public infix fun <T> SuspendAnsweringScope<T>.awaits(deferred: Deferred<T>) {
 /**
  * Function call awaits on each call for a [Deferred] provided by [provider].
  */
-public infix fun <T> SuspendAnsweringScope<T>.awaits(provider: (SuspendCallDefinitionScope<T>) -> Deferred<T>) {
-    awaits(AwaitDeferred(description = { "{...}" }, deferred = { SuspendCallDefinitionScope<T>(it).let(provider) }))
+public infix fun <T> SuspendAnsweringScope<T>.awaits(provider: SuspendCallDefinitionScope<T>.(CallArgs) -> Deferred<T>) {
+    val awaitable = AwaitDeferred(
+        description = { "{...}" },
+        deferred = { provider(SuspendCallDefinitionScope<T>(it), CallArgs(it.args)) }
+    )
+    awaits(awaitable)
 }
diff --git a/test-mokkery/src/coroutinesTest/kotlin/dev/mokkery/test/CoroutinesAnsweringTest.kt b/test-mokkery/src/coroutinesTest/kotlin/dev/mokkery/test/CoroutinesAnsweringTest.kt
@@ -55,6 +55,13 @@ class CoroutinesAnsweringTest {
         assertEquals(1, mock.fetchWithDefault(0))
     }
 
+    @Test
+    fun testAwaitsProvidedDeferred() = runTest {
+        everySuspend { mock.fetchWithDefault(any()) } awaits { (i: Int) -> CompletableDeferred(i) }
+        assertEquals(0, mock.fetchWithDefault(0))
+        assertEquals(1, mock.fetchWithDefault(1))
+    }
+
     @Test
     fun testAwaitsAllDeferred() = runTest {
         val deferreds = listOf(CompletableDeferred("1"), CompletableDeferred("2"))