Improving ongoing stubbing support for suspend functions, by introducing CoroutineAwareAnswer

Author: m-koops

mockito-kotlin/src/main/kotlin/org/mockito/kotlin/BDDMockito.kt

@@ -31,14 +31,14 @@ import org.mockito.BDDMockito
 import org.mockito.BDDMockito.BDDMyOngoingStubbing
 import org.mockito.BDDMockito.Then
 import org.mockito.invocation.InvocationOnMock
-import org.mockito.kotlin.internal.SuspendableAnswer
+import org.mockito.kotlin.internal.CoroutineAwareAnswer
 import org.mockito.stubbing.Answer
 import kotlin.reflect.KClass
 
 /**
  * Alias for [BDDMockito.given].
  */
-fun <T> given(methodCall: T): BDDMockito.BDDMyOngoingStubbing<T> {
+fun <T> given(methodCall: T): BDDMyOngoingStubbing<T> {
     return BDDMockito.given(methodCall)
 }
 
@@ -55,14 +55,14 @@ fun <T> given(methodCall: () -> T): BDDMyOngoingStubbing<T> {
  * Warning: Only last method call can be stubbed in the function.
  * other method calls are ignored!
  */
-fun <T> givenBlocking(methodCall: suspend CoroutineScope.() -> T): BDDMockito.BDDMyOngoingStubbing<T> {
+fun <T> givenBlocking(methodCall: suspend CoroutineScope.() -> T): BDDMyOngoingStubbing<T> {
     return runBlocking { BDDMockito.given(methodCall()) }
 }
 
 /**
  * Alias for [BDDMockito.then].
  */
-fun <T> then(mock: T): BDDMockito.Then<T> {
+fun <T> then(mock: T): Then<T> {
     return BDDMockito.then(mock)
 }
 
@@ -77,35 +77,35 @@ fun <T, R> Then<T>.shouldBlocking(f: suspend T.() -> R): R {
 /**
  * Alias for [BDDMyOngoingStubbing.will]
  * */
-infix fun <T> BDDMyOngoingStubbing<T>.will(value: Answer<T>): BDDMockito.BDDMyOngoingStubbing<T> {
+infix fun <T> BDDMyOngoingStubbing<T>.will(value: Answer<T>): BDDMyOngoingStubbing<T> {
     return will(value)
 }
 
 /**
- * Alias for [BBDMyOngoingStubbing.willAnswer], accepting a lambda.
+ * Alias for [BDDMyOngoingStubbing.willAnswer], accepting a lambda.
  */
-infix fun <T> BDDMyOngoingStubbing<T>.willAnswer(value: (InvocationOnMock) -> T?): BDDMockito.BDDMyOngoingStubbing<T> {
-    return willAnswer { value(it) }
+infix fun <T> BDDMyOngoingStubbing<T>.willAnswer(answer: (InvocationOnMock) -> T?): BDDMyOngoingStubbing<T> {
+    return willAnswer { answer(it) }
 }
 
 /**
- * Alias for [BBDMyOngoingStubbing.willAnswer], accepting a suspend lambda.
+ * Alias for [BDDMyOngoingStubbing.willAnswer], accepting a suspend lambda.
  */
-infix fun <T> BDDMyOngoingStubbing<T>.willSuspendableAnswer(value: suspend (InvocationOnMock) -> T?): BDDMockito.BDDMyOngoingStubbing<T> {
-    return willAnswer(SuspendableAnswer(value))
+infix fun <T> BDDMyOngoingStubbing<T>.willSuspendableAnswer(answer: suspend (KInvocationOnMock) -> T?): BDDMyOngoingStubbing<T> {
+    return willAnswer(CoroutineAwareAnswer(answer))
 }
 
 /**
- * Alias for [BBDMyOngoingStubbing.willReturn].
+ * Alias for [BDDMyOngoingStubbing.willReturn].
  */
-infix fun <T> BDDMyOngoingStubbing<T>.willReturn(value: () -> T): BDDMockito.BDDMyOngoingStubbing<T> {
+infix fun <T> BDDMyOngoingStubbing<T>.willReturn(value: () -> T): BDDMyOngoingStubbing<T> {
     return willReturn(value())
 }
 
 /**
- * Alias for [BBDMyOngoingStubbing.willThrow].
+ * Alias for [BDDMyOngoingStubbing.willThrow].
  */
-infix fun <T> BDDMyOngoingStubbing<T>.willThrow(value: () -> Throwable): BDDMockito.BDDMyOngoingStubbing<T> {
+infix fun <T> BDDMyOngoingStubbing<T>.willThrow(value: () -> Throwable): BDDMyOngoingStubbing<T> {
     return willThrow(value())
 }
 

mockito-kotlin/src/main/kotlin/org/mockito/kotlin/KInvocationOnMock.kt

@@ -28,12 +28,67 @@ package org.mockito.kotlin
 import org.mockito.invocation.InvocationOnMock
 
 class KInvocationOnMock(
-    private val invocationOnMock: InvocationOnMock
+    val invocationOnMock: InvocationOnMock
 ) : InvocationOnMock by invocationOnMock {
 
     operator fun <T> component1(): T = invocationOnMock.getArgument(0)
     operator fun <T> component2(): T = invocationOnMock.getArgument(1)
     operator fun <T> component3(): T = invocationOnMock.getArgument(2)
     operator fun <T> component4(): T = invocationOnMock.getArgument(3)
     operator fun <T> component5(): T = invocationOnMock.getArgument(4)
+
+    /**
+     * The first argument.
+     * @throws IndexOutOfBoundsException if the argument is not available.
+     */
+    inline fun <reified T> first(): T = invocationOnMock.getArgument(0)
+
+    /**
+     * The second argument.
+     * @throws IndexOutOfBoundsException if the argument is not available.
+     */
+    inline fun <reified T> second(): T = invocationOnMock.getArgument(1)
+
+    /**
+     * The third argument.
+     * @throws IndexOutOfBoundsException if the argument is not available.
+     */
+    inline fun <reified T> third(): T = invocationOnMock.getArgument(2)
+
+    /**
+     * The fourth argument.
+     * @throws IndexOutOfBoundsException if the argument is not available.
+     */
+    inline fun <reified T> fourth(): T = invocationOnMock.getArgument(3)
+
+    /**
+     * The fifth argument.
+     * @throws IndexOutOfBoundsException if the argument is not available.
+     */
+    inline fun <reified T> fifth(): T = invocationOnMock.getArgument(4)
+
+    /**
+     * The last argument.
+     * @throws IndexOutOfBoundsException if the argument is not available.
+     */
+    inline fun <reified T> last(): T {
+        val size  = invocationOnMock.arguments.size
+        require(size >= 1) { "The invocation was expected to have at least 1 argument but got none." }
+        return invocationOnMock.getArgument(size - 1)
+    }
+
+    /**
+     * The single argument.
+     * @throws IndexOutOfBoundsException if the argument is not available.
+     */
+    inline fun <reified T> single(): T {
+        val size  = invocationOnMock.arguments.size
+        require(size == 1) { "The invocation was expected to have exactly 1 argument but got $size." }
+        return first()
+    }
+
+    /**
+     * The all arguments.
+     */
+    fun all(): List<Any> = invocationOnMock.arguments.toList()
 }

mockito-kotlin/src/main/kotlin/org/mockito/kotlin/KStubbing.kt

@@ -25,21 +25,25 @@
 
 package org.mockito.kotlin
 
-import org.mockito.kotlin.internal.createInstance
 import kotlinx.coroutines.runBlocking
 import org.mockito.Mockito
 import org.mockito.exceptions.misusing.NotAMockException
+import org.mockito.internal.progress.ThreadSafeMockingProgress.mockingProgress
+import org.mockito.kotlin.internal.createInstance
 import org.mockito.stubbing.OngoingStubbing
 import org.mockito.stubbing.Stubber
 import kotlin.reflect.KClass
 
-inline fun <T : Any> stubbing(
-    mock: T,
-    stubbing: KStubbing<T>.(T) -> Unit
-) {
+/**
+ * Apply stubbing on the given mock. The stubbed behavior of the mock can then be specified in a supplied lambda.
+ */
+inline fun <T : Any> stubbing(mock: T, stubbing: KStubbing<T>.(T) -> Unit) {
     KStubbing(mock).stubbing(mock)
 }
 
+/**
+ * Apply stubbing on the given mock. The stubbed behavior of the mock can then be specified in a supplied lambda.
+ */
 inline fun <T : Any> T.stub(stubbing: KStubbing<T>.(T) -> Unit): T {
     return apply { KStubbing(this).stubbing(this) }
 }
@@ -49,44 +53,85 @@ class KStubbing<out T : Any>(val mock: T) {
         if (!mockingDetails(mock).isMock) throw NotAMockException("Stubbing target is not a mock!")
     }
 
-    fun <R> on(methodCall: R): OngoingStubbing<R> = Mockito.`when`(methodCall)
+    /**
+     * Enables stubbing methods. Use it when you want the mock to return particular value when particular method is called.
+     *
+     * Simply put: `When the x method is called then return y`
+     */
+    fun <R> on(methodCall: R): OngoingStubbing<R> = mockitoWhen(methodCall)
+
+    /**
+     * Enables stubbing methods. Use it when you want the mock to return particular value when particular method is called.
+     *
+     * Simply put: `When the x method is called then return y`
+     */
+    fun <R> on(methodCall: T.() -> R): OngoingStubbing<R> {
+        return try {
+            mockitoWhen(mock.methodCall())
+        } catch (e: NullPointerException) {
+            throw MockitoKotlinException(
+                  "NullPointerException thrown when stubbing.\nThis may be due to two reasons:\n\t- The method you're trying to stub threw an NPE: look at the stack trace below;\n\t- You're trying to stub a generic method: try `onGeneric` instead.",
+                  e
+            )
+        }
+    }
+
+    /**
+     * Enables stubbing suspend functions. Use it when you want the mock to return particular value when particular suspend function is called.
+     *
+     * Simply put: `When the x suspend function is called then return y`
+     */
+    fun <R> onBlocking(suspendFunctionCall: suspend T.() -> R): OngoingStubbing<R> {
+        return runBlocking { mockitoWhen(mock.suspendFunctionCall()) }
+    }
+
+    /**
+     * Enables stubbing generic methods with return type [R]. Use it when you want the mock to return particular value when particular method is called.
+     *
+     * Simply put: `When the x method is called then return y of type R`
+     */
+    inline fun <reified R : Any> onGeneric(noinline methodCall: T.() -> R?): OngoingStubbing<R?> {
+        return onGeneric(methodCall, R::class)
+    }
 
-    fun <R : Any> onGeneric(methodCall: T.() -> R?, c: KClass<R>): OngoingStubbing<R> {
+
+    /**
+     * Enables stubbing generic methods with return type [R]. Use it when you want the mock to return particular value when particular method is called.
+     *
+     * Simply put: `When the x method is called then return y of type R`
+     */
+    fun <R : Any> onGeneric(methodCall: T.() -> R?, c: KClass<R>): OngoingStubbing<R?> {
         val r = try {
             mock.methodCall()
-        } catch (e: NullPointerException) {
+        } catch (_: NullPointerException) {
             // An NPE may be thrown by the Kotlin type system when the MockMethodInterceptor returns a
             // null value for a non-nullable generic type.
             // We catch this NPE to return a valid instance.
             // The Mockito state has already been modified at this point to reflect
             // the wanted changes.
             createInstance(c)
         }
-        return Mockito.`when`(r)
+        return mockitoWhen(r)
     }
 
-    inline fun <reified R : Any> onGeneric(noinline methodCall: T.() -> R?): OngoingStubbing<R> {
-        return onGeneric(methodCall, R::class)
+    /**
+     * Completes stubbing a method, by addressing the method call to apply a given stubbed [org.mockito.stubbing.Answer] on.
+     * Use it when you want the mock to return particular value when particular method is called.
+     *
+     * Simply put: `Return y when the x method is called`
+     */
+    fun Stubber.on(methodCall: T.() -> Unit) {
+        this.`when`(mock).methodCall()
     }
 
-    fun <R> on(methodCall: T.() -> R): OngoingStubbing<R> {
-        return try {
-            Mockito.`when`(mock.methodCall())
-        } catch (e: NullPointerException) {
-            throw MockitoKotlinException(
-                  "NullPointerException thrown when stubbing.\nThis may be due to two reasons:\n\t- The method you're trying to stub threw an NPE: look at the stack trace below;\n\t- You're trying to stub a generic method: try `onGeneric` instead.",
-                  e
-            )
-        }
-    }
+    private fun <R> mockitoWhen(methodCall: R): OngoingStubbing<R> {
+        val ongoingStubbing = Mockito.`when`(methodCall)
 
-    fun <T : Any, R> KStubbing<T>.onBlocking(
-        m: suspend T.() -> R
-    ): OngoingStubbing<R> {
-        return runBlocking { Mockito.`when`(mock.m()) }
-    }
+        if (ongoingStubbing.getMock<T>() != mock) {
+            mockingProgress().reset()
+            throw IllegalArgumentException("Stubbing of another mock is not allowed")
+        }
 
-    fun Stubber.on(methodCall: T.() -> Unit) {
-        this.`when`(mock).methodCall()
+        return ongoingStubbing
     }
 }