Revamp of ongoing stubbing for suspend functions

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

@@ -0,0 +1,213 @@
+/*
+ * The MIT License
+ *
+ * Copyright (c) 2018 Niek Haarman
+ * Copyright (c) 2007 Mockito contributors
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package org.mockito.kotlin
+
+import kotlinx.coroutines.delay
+import org.mockito.internal.stubbing.answers.Returns
+import org.mockito.kotlin.internal.SuspendableAnswer
+import org.mockito.kotlin.internal.lastInvocationMethodIsSuspend
+import org.mockito.stubbing.Answer
+import org.mockito.stubbing.OngoingStubbing
+import kotlin.reflect.KClass
+
+class CoroutinesOngoingStubbing<T>(val mockitoOngoingStubbing: OngoingStubbing<T>) {
+    /**
+     * Sets a return value to be returned when the method is called.
+     *
+     * Alias for [thenReturn].
+     */
+    infix fun doReturn(t: T): CoroutinesOngoingStubbing<T> {
+        return thenReturn(t)
+    }
+
+    /**
+     * Sets a return value to be returned when the method is called.
+     */
+    @Suppress("UNCHECKED_CAST")
+    infix fun thenReturn(t: T): CoroutinesOngoingStubbing<T> {
+        return thenAnswer(Returns(t))
+    }
+
+    /**
+     * Sets an answer for the suspendable function using a suspendable lambda.
+     *
+     * Alias for [thenAnswer].
+     */
+    infix fun doAnswer(answer: suspend (KInvocationOnMock) -> T?): CoroutinesOngoingStubbing<T> {
+        return thenAnswer(answer)
+    }
+
+    /**
+     * Sets an answer for the suspendable function using a suspendable lambda.
+     *
+     * Alias for [thenAnswer].
+     *
+     * This deprecated method was added for backwards compatibility.
+     */
+    @Deprecated(
+        " use doAnswer() or thenAnswer() instead.",
+        level = DeprecationLevel.WARNING
+    )
+    infix fun doSuspendableAnswer(answer: suspend (KInvocationOnMock) -> T?): CoroutinesOngoingStubbing<T> {
+        return thenAnswer(answer)
+    }
+
+    /**
+     * Sets an answer for the suspendable function using a suspendable lambda.
+     */
+    infix fun thenAnswer(answer: suspend (KInvocationOnMock) -> T?): CoroutinesOngoingStubbing<T> {
+        return thenAnswer(SuspendableAnswer(answer))
+    }
+
+    /**
+     * Sets an answer for the suspendable function by wrapping a non-suspendable Mockito Answer.
+     *
+     * Alias for [thenAnswer].
+     */
+    infix fun doAnswer(answer: Answer<*>): CoroutinesOngoingStubbing<T> {
+        return thenAnswer(answer)
+    }
+
+    /**
+     * Sets an answer for the suspendable function by wrapping a non-suspendable Mockito Answer.
+     */
+    infix fun thenAnswer(answer: Answer<*>): CoroutinesOngoingStubbing<T> {
+        return if (mockitoOngoingStubbing.lastInvocationMethodIsSuspend ?: false) {
+            answer.wrapAsSuspendableAnswer()
+        } else {
+            answer // to support stubbing a sync method call while using wheneverBlocking()/onBlocking()
+        }.let {
+            CoroutinesOngoingStubbing(mockitoOngoingStubbing.thenAnswer(it))
+        }
+    }
+
+    private fun Answer<*>.wrapAsSuspendableAnswer(): Answer<out Any?> =
+        (this as? SuspendableAnswer<*>) ?: SuspendableAnswer(
+            { invocation ->
+                suspendToEnforceProperValueBoxing()
+                this.answer(invocation)
+            }
+        )
+
+    private suspend fun suspendToEnforceProperValueBoxing() {
+        // delaying for 1 ms, forces a suspension to happen.
+        // This (somehow) ensures that value class instances will be properly boxed when the
+        // answer is yielded by the mock
+        delay(1)
+    }
+
+    /**
+     * Sets consecutive return values to be returned when the method is called.
+     *
+     * Alias for [OngoingStubbing.thenReturn].
+     */
+    fun doReturn(t: T, vararg ts: T): CoroutinesOngoingStubbing<T> {
+        return thenReturn(listOf(t, *ts))
+    }
+
+    /**
+     * Sets consecutive return values to be returned when the method is called.
+     */
+    fun doReturnConsecutively(vararg ts: T): CoroutinesOngoingStubbing<T> {
+        return doReturnConsecutively(listOf(*ts))
+    }
+
+    /**
+     * Sets consecutive return values to be returned when the method is called.
+     */
+    infix fun doReturnConsecutively(ts: List<T>): CoroutinesOngoingStubbing<T> {
+        return thenReturn(ts)
+    }
+
+    /**
+     * Sets Throwable objects to be thrown when the method is called.
+     *
+     * Alias for [OngoingStubbing.thenThrow].
+     */
+    infix fun doThrow(t: Throwable): CoroutinesOngoingStubbing<T> {
+        return CoroutinesOngoingStubbing(
+            mockitoOngoingStubbing.thenThrow(t)
+        )
+    }
+
+    /**
+     * Sets Throwable objects to be thrown when the method is called.
+     *
+     * Alias for [OngoingStubbing.doThrow].
+     */
+    fun doThrow(
+        t: Throwable,
+        vararg ts: Throwable
+    ): CoroutinesOngoingStubbing<T> {
+        return CoroutinesOngoingStubbing(
+            mockitoOngoingStubbing.thenThrow(t, *ts)
+        )
+    }
+
+    /**
+     * Sets a Throwable type to be thrown when the method is called.
+     */
+    infix fun doThrow(t: KClass<out Throwable>): CoroutinesOngoingStubbing<T> {
+        return CoroutinesOngoingStubbing(
+            mockitoOngoingStubbing.thenThrow(t.java)
+        )
+    }
+
+    /**
+     * Sets Throwable classes to be thrown when the method is called.
+     */
+    fun doThrow(
+        t: KClass<out Throwable>,
+        vararg ts: KClass<out Throwable>
+    ): CoroutinesOngoingStubbing<T> {
+        return CoroutinesOngoingStubbing(
+            mockitoOngoingStubbing.thenThrow(
+                t.java,
+                *ts.map { it.java }.toTypedArray()
+            )
+        )
+    }
+
+    private fun thenReturn(values: List<T>): CoroutinesOngoingStubbing<T> {
+        return thenAnswer(
+            values.map { value ->
+                SuspendableAnswer(
+                    {
+                        suspendToEnforceProperValueBoxing()
+                        value
+                    }
+                )
+            }
+        )
+    }
+
+    private fun thenAnswer(answers: List<SuspendableAnswer<T>>): CoroutinesOngoingStubbing<T> {
+        return CoroutinesOngoingStubbing(
+            answers
+                .fold(mockitoOngoingStubbing) { stubbing, answer -> stubbing.thenAnswer(answer) }
+        )
+    }
+}

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

@@ -28,12 +28,66 @@ 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 lastIndex  = invocationOnMock.arguments.size - 1
+        return invocationOnMock.getArgument(lastIndex)
+    }
+
+    /**
+     * 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()
 }