Add suspend answer; Fix #346

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

Author: neworld

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>.willAnswer(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)
+    }
+}

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,107 @@ class CoroutinesTest {
             verify(testSubject).suspending()
         }
     }
+
+    @Test
+    fun answerWithSuspendFunction() = runBlocking {
+        val fixture: SomeInterface = mock()
+
+        whenever(fixture.suspendingWithArg(any())).willAnswer {
+            withContext(Dispatchers.Default) { it.getArgument<Int>(0) }
+        }
+
+        assertEquals(5, fixture.suspendingWithArg(5))
+    }
+
+    @Test
+    fun inplaceAnswerWithSuspendFunction() = runBlocking {
+        val fixture: SomeInterface = mock {
+            onBlocking { suspendingWithArg(any()) } willAnswer {
+                withContext(Dispatchers.Default) { it.getArgument<Int>(0) }
+            }
+        }
+
+        assertEquals(5, fixture.suspendingWithArg(5))
+    }
+
+    @Test
+    fun callFromSuspendFunction() = runBlocking {
+        val fixture: SomeInterface = mock()
+
+        whenever(fixture.suspendingWithArg(any())).willAnswer {
+            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())).willAnswer {
+            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()).willAnswer {
+            withContext(Dispatchers.Default) { 42 }
+        }
+
+        assertEquals(42, fixture.suspending())
+    }
+
+    @Test
+    fun willAnswerWithControlledSuspend() = runBlocking {
+        val fixture: SomeInterface = mock()
+
+        val job = Job()
+
+        whenever(fixture.suspending()).willAnswer {
+            job.join()
+            5
+        }
+
+        val asyncTask = async {
+            fixture.suspending()
+        }
+
+        //TODO: on coroutines >1.3 it should be job.complete()
+        job.cancel()
+
+        withTimeout(100) {
+            assertEquals(5, asyncTask.await())
+        }
+    }
 }
 
 interface SomeInterface {
 
     suspend fun suspending(): Int
+    suspend fun suspendingWithArg(arg: Int): Int
     fun nonsuspending(): Int
 }