Add suspendable answer (#357)

This code was inspired by:
- Java wrapper part: https://github.com/square/retrofit/blob/8c93b59dbc57841959f5237cb141ce0b3c18b778/retrofit/src/main/java/retrofit2/HttpServiceMethod.java#L168
- Kotlin wrapper part: https://github.com/square/retrofit/blob/8c93b59dbc57841959f5237cb141ce0b3c18b778/retrofit/src/main/java/retrofit2/KotlinExtensions.kt#L83-L98

Fixes #346

3 files changed, 115 insertions, 2 deletions

diff --git a/mockito-kotlin/build.gradle b/mockito-kotlin/build.gradle
index 14882a1..8385b69 100644
--- a/mockito-kotlin/build.gradle
+++ b/mockito-kotlin/build.gradle
@@ -31,7 +31,7 @@ dependencies {
     testCompile 'com.nhaarman:expect.kt:1.0.0'
 
     testCompile  "org.jetbrains.kotlin:kotlin-stdlib:$kotlin_version"
-    testCompile 'org.jetbrains.kotlinx:kotlinx-coroutines-core:1.0.0'
+    testCompile 'org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.0'
 
     testImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-android:1.0.0"
 }
diff --git a/mockito-kotlin/src/main/kotlin/org/mockito/kotlin/OngoingStubbing.kt b/mockito-kotlin/src/main/kotlin/org/mockito/kotlin/OngoingStubbing.kt
index 3d97ce1..78548c8 100644
--- a/mockito-kotlin/src/main/kotlin/org/mockito/kotlin/OngoingStubbing.kt
+++ b/mockito-kotlin/src/main/kotlin/org/mockito/kotlin/OngoingStubbing.kt
@@ -26,10 +26,12 @@
 package org.mockito.kotlin
 
 import org.mockito.Mockito
+import org.mockito.internal.invocation.InterceptedInvocation
 import org.mockito.invocation.InvocationOnMock
 import org.mockito.stubbing.Answer
 import org.mockito.stubbing.OngoingStubbing
-import kotlin.DeprecationLevel.ERROR
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn
 import kotlin.reflect.KClass
 
 
@@ -124,3 +126,15 @@ infix fun <T> OngoingStubbing<T>.doAnswer(answer: Answer<*>): OngoingStubbing<T>
 infix fun <T> OngoingStubbing<T>.doAnswer(answer: (InvocationOnMock) -> T?): OngoingStubbing<T> {
     return thenAnswer(answer)
 }
+
+infix fun <T> OngoingStubbing<T>.doSuspendableAnswer(answer: suspend (InvocationOnMock) -> T?): OngoingStubbing<T> {
+    return thenAnswer {
+        //all suspend functions/lambdas has Continuation as the last argument.
+        //InvocationOnMock does not see last argument
+        val rawInvocation = it as InterceptedInvocation
+        val continuation = rawInvocation.rawArguments.last() as Continuation<T?>
+
+        // https://youtrack.jetbrains.com/issue/KT-33766#focus=Comments-27-3707299.0-0
+        answer.startCoroutineUninterceptedOrReturn(it, continuation)
+    }
+}
diff --git a/mockito-kotlin/src/test/kotlin/test/CoroutinesTest.kt b/mockito-kotlin/src/test/kotlin/test/CoroutinesTest.kt
index 5ca6eb6..51084e0 100644
--- a/mockito-kotlin/src/test/kotlin/test/CoroutinesTest.kt
+++ b/mockito-kotlin/src/test/kotlin/test/CoroutinesTest.kt
@@ -7,8 +7,12 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.withContext
+import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.actor
+import org.junit.Assert.assertEquals
 import org.junit.Test
 import org.mockito.kotlin.*
+import java.util.*
 
 
 class CoroutinesTest {
@@ -157,11 +161,106 @@ class CoroutinesTest {
             verify(testSubject).suspending()
         }
     }
+
+    @Test
+    fun answerWithSuspendFunction() = runBlocking {
+        val fixture: SomeInterface = mock()
+
+        whenever(fixture.suspendingWithArg(any())).doSuspendableAnswer {
+            withContext(Dispatchers.Default) { it.getArgument<Int>(0) }
+        }
+
+        assertEquals(5, fixture.suspendingWithArg(5))
+    }
+
+    @Test
+    fun inplaceAnswerWithSuspendFunction() = runBlocking {
+        val fixture: SomeInterface = mock {
+            onBlocking { suspendingWithArg(any()) } doSuspendableAnswer {
+                withContext(Dispatchers.Default) { it.getArgument<Int>(0) }
+            }
+        }
+
+        assertEquals(5, fixture.suspendingWithArg(5))
+    }
+
+    @Test
+    fun callFromSuspendFunction() = runBlocking {
+        val fixture: SomeInterface = mock()
+
+        whenever(fixture.suspendingWithArg(any())).doSuspendableAnswer {
+            withContext(Dispatchers.Default) { it.getArgument<Int>(0) }
+        }
+
+        val result = async {
+            val answer = fixture.suspendingWithArg(5)
+
+            Result.success(answer)
+        }
+
+        assertEquals(5, result.await().getOrThrow())
+    }
+
+    @Test
+    fun callFromActor() = runBlocking {
+        val fixture: SomeInterface = mock()
+
+        whenever(fixture.suspendingWithArg(any())).doSuspendableAnswer {
+            withContext(Dispatchers.Default) { it.getArgument<Int>(0) }
+        }
+
+        val actor = actor<Optional<Int>> {
+            for (element in channel) {
+                fixture.suspendingWithArg(element.get())
+            }
+        }
+
+        actor.send(Optional.of(10))
+        actor.close()
+
+        verify(fixture).suspendingWithArg(10)
+
+        Unit
+    }
+
+    @Test
+    fun answerWithSuspendFunctionWithoutArgs() = runBlocking {
+        val fixture: SomeInterface = mock()
+
+        whenever(fixture.suspending()).doSuspendableAnswer {
+            withContext(Dispatchers.Default) { 42 }
+        }
+
+        assertEquals(42, fixture.suspending())
+    }
+
+    @Test
+    fun willAnswerWithControlledSuspend() = runBlocking {
+        val fixture: SomeInterface = mock()
+
+        val job = Job()
+
+        whenever(fixture.suspending()).doSuspendableAnswer {
+            job.join()
+            5
+        }
+
+        val asyncTask = async {
+            fixture.suspending()
+        }
+
+        job.complete()
+
+        withTimeout(100) {
+            assertEquals(5, asyncTask.await())
+        }
+    }
 }
 
 interface SomeInterface {
 
     suspend fun suspending(): Int
+    suspend fun suspendingWithArg(arg: Int): Int
     fun nonsuspending(): Int
 }