mockito-kotlin 4.x + kotlinx.coroutines 1.6.x で doReturnから始めたい + delayをサボりたい人向けのKotlin extensions

MockitoSuspendableExtension.kt

package co.uzzu.coroutines.testing.mockito

import kotlin.coroutines.Continuation
import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.delay
import kotlinx.coroutines.test.TestScope
import org.mockito.Mockito
import org.mockito.internal.invocation.InterceptedInvocation
import org.mockito.invocation.InvocationOnMock
import org.mockito.kotlin.doSuspendableAnswer
import org.mockito.stubbing.Answer
import org.mockito.stubbing.BaseStubber
import org.mockito.stubbing.OngoingStubbing
import org.mockito.stubbing.Stubber

@OptIn(ExperimentalCoroutinesApi::class)
fun <T> TestScope.doSuspendableReturn(t: T): Stubber {
    return doSuspendableAnswer {
        delay(1L)
        t
    }
}

@OptIn(ExperimentalCoroutinesApi::class)
fun TestScope.doSuspendableThrow(throwable: Throwable): Stubber {
    return doSuspendableAnswer {
        delay(1L)
        throw throwable
    }
}

@OptIn(ExperimentalCoroutinesApi::class)
fun <T> TestScope.doSuspendableAnswer(answer: suspend (InvocationOnMock) -> T?): Stubber {
    return Mockito.doAnswer(SuspendableAnswer(answer))
}

fun <T> BaseStubber.doSuspendableReturn(t: T): Stubber {
    return doSuspendableAnswer {
        delay(1L)
        t
    }
}

fun <T> BaseStubber.doSuspendableThrow(throwable: Throwable): Stubber {
    return doSuspendableAnswer {
        delay(1L)
        throw throwable
    }
}

fun <T> BaseStubber.doSuspendableAnswer(answer: suspend (InvocationOnMock) -> T?): Stubber {
    return doAnswer(SuspendableAnswer(answer))
}

infix fun <T> OngoingStubbing<T>.doSuspendableReturn(t: T): OngoingStubbing<T> {
    return doSuspendableAnswer {
        delay(1L)
        t
    }
}

infix fun <T> OngoingStubbing<T>.doSuspendableThrow(throwable: Throwable): OngoingStubbing<T> {
    return doSuspendableAnswer {
        delay(1L)
        throw throwable
    }
}

/**
 * @see [https://github.com/mockito/mockito-kotlin/blob/4.0.0/mockito-kotlin/src/main/kotlin/org/mockito/kotlin/internal/SuspendableAnswer.kt]
 */
private class SuspendableAnswer<T>(
    private val body: suspend (InvocationOnMock) -> T?,
) : Answer<T> {

    @Suppress("UNCHECKED_CAST")
    override fun answer(invocation: InvocationOnMock?): T {
        //all suspend functions/lambdas has Continuation as the last argument.
        //InvocationOnMock does not see last argument
        val rawInvocation = invocation as InterceptedInvocation
        val continuation = rawInvocation.rawArguments.last() as Continuation<T?>

        // https://youtrack.jetbrains.com/issue/KT-33766#focus=Comments-27-3707299.0-0
        return body.startCoroutineUninterceptedOrReturn(invocation, continuation) as T
    }
}