A Java dynamic proxy that calls either the real or mocked implementation of an interface

MockProxy.kt

import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import java.lang.reflect.InvocationHandler
import java.lang.reflect.InvocationTargetException
import java.lang.reflect.Method
import java.lang.reflect.Proxy
import kotlin.coroutines.Continuation
import kotlin.coroutines.cancellation.CancellationException
import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED

object MockProxy {

    inline fun <reified T : Any> create(
        actual: T,
        mock: T,
        controller: MockProxyController,
    ): T {
        return create(clazz = T::class.java, actual = actual, mock = mock, controller = controller)
    }

    /**
     * Returns a Java dynamic proxy instance for the specified interface.
     *
     * Methods invocations will be re-directed to the provided [mock] if
     * [MockProxyController.areMocksEnabled] returns true. If it returns
     * false (or a mocked method throws a [MockDisabledException]) the
     * [actual] implementation of the interface will be called.
     *
     * As a convenience, the return of mocked **suspend** functions can be
     * delayed. The duration of the delay can be controlled in
     * [MockProxyController.suspendFunctionsDelay]. A delay value of 0 or
     * less will result in no delay.
     *
     * A convenient way to implement a [mock], would be to implement the
     * interface and use the [actual] class as a delegate, so only methods
     * that actually need to be mocked can be overriden.
     *
     * Throwing a [MockDisabledException] can be useful if you want to keep
     * mocks for several methods, but need to enable only part of them at
     * the same time.
     *
     * @param clazz the interface for the proxy to implement
     * @param actual the actual, or real implementation of the interface
     * @param mock the implementation that contains the mocks
     * @param controller used to control the global behaviour or mocks
     * @return a proxy instance that implements the specified interface
     *
     * @see Proxy.newProxyInstance
     *
     * @author Roman Kamyshnikov
     */
    @Suppress("UNCHECKED_CAST")
    fun <T : Any> create(
        clazz: Class<T>,
        actual: T,
        mock: T,
        controller: MockProxyController,
    ): T {
        return Proxy.newProxyInstance(
            clazz.classLoader,
            arrayOf(clazz),
            MockProxyInvocationHandler(actual = actual, mock = mock, controller = controller)
        ) as T
    }

}

object MockDisabledException : Exception()

/**
 * Base interface to control the behaviour of a [MockProxy].
 *
 * - areMocksEnabled is used to decide whether to call the actual
 *   or mock implementation
 * - suspendFunctionsDelay sets the delay that the proxy should wait
 *   before returning the result of a mocked method.
 */
interface MockProxyController {
    val areMocksEnabled: Boolean
    val suspendFunctionsDelay: Long
}

class DefaultMockProxyController(
    override var areMocksEnabled: Boolean = false,
    override var suspendFunctionsDelay: Long = 3000L,
) : MockProxyController

internal class MockProxyInvocationHandler<T : Any>(
    private val actual: T,
    private val mock: T,
    private val controller: MockProxyController,
) : InvocationHandler {

    private val targetObject: T
        get() = if (controller.areMocksEnabled) mock else actual

    @Suppress("UNCHECKED_CAST")
    override fun invoke(proxy: Any, method: Method, args: Array<out Any>?): Any? {
        val safeArgs = args ?: emptyArray()
        val continuation = args?.firstNotNullOfOrNull { it as? Continuation<Any?> }

        return if (continuation == null) {
            invokeMethod(method, targetObject, safeArgs)
        } else {
            invokeSuspendingFunction(method, targetObject, safeArgs, continuation)
            COROUTINE_SUSPENDED
        }
    }

    private fun invokeMethod(method: Method, obj: Any, safeArgs: Array<out Any>): Any? {
        return try {
            method.invoke(obj, *safeArgs)
        } catch (exception: Exception) {
            when (val actualException = exception.unwrap()) {
                is MockDisabledException -> invokeMethod(method, actual, safeArgs)
                else -> throw actualException
            }
        }
    }

    private fun invokeSuspendingFunction(
        method: Method, obj: Any, safeArgs: Array<out Any>, continuation: Continuation<Any?>
    ) {
        val delayDuration = controller.suspendFunctionsDelay.takeIf { obj == mock } ?: 0L

        CoroutineScope(continuation.context).launch {
            try {
                coroutineScope {
                    launch {
                        val result = method.invoke(obj, *safeArgs)
                        if (result != COROUTINE_SUSPENDED) {
                            continuation.resumeWithAfterDelay(Result.success(result), delayDuration)
                        }
                    }
                }
            } catch (exception: Exception) {
                when (val actualException = exception.unwrap()) {
                    is MockDisabledException -> invokeSuspendingFunction(method, actual, safeArgs, continuation)
                    is CancellationException -> continuation.resumeWith(Result.failure(actualException))
                    else -> continuation.resumeWithAfterDelay(Result.failure(actualException), delayDuration)
                }
            }
        }
    }

    private fun Throwable.unwrap(): Throwable {
        return (this as? InvocationTargetException)?.cause ?: this
    }

    private suspend fun Continuation<Any?>.resumeWithAfterDelay(result: Result<Any?>, timeMillis: Long) {
        delay(timeMillis)
        resumeWith(result)
    }

}

MockProxyExample.kt

interface SomeApi {
    suspend fun foo(): String
    suspend fun bar(): String
    suspend fun baz(): String
}

class SomeApiImpl : SomeApi {
    override suspend fun foo() = "Actual foo"
    override suspend fun bar() = "Actual bar"
    override suspend fun baz() = "Actual baz"
}

class SomeApiMockController : MockProxyController {
    override val areMocksEnabled: Boolean
        get() = mockFoo || mockBaz
    override var suspendFunctionsDelay: Long = 3000L
    var mockFoo = true
    var mockBaz = false
}

class SomeApiMock(
    private val actual: SomeApi,
    private val controller: SomeApiMockController,
) : SomeApi by actual {

    private inline fun <T> getResult(
        isMockEnabled: Boolean = true,
        method: () -> T,
    ): T {
        if (isMockEnabled) return method() else throw MockDisabledException
    }

    override suspend fun foo() = getResult(controller.mockFoo) { "Mocked foo" }
    override suspend fun baz() = getResult(controller.mockBaz) { "Mocked baz" }
}

suspend fun example() {
    val actual = SomeApiImpl()
    val controller = SomeApiMockController()
    val mock = SomeApiMock(actual, controller)
    val proxy = MockProxy.create(actual, mock, controller)

    proxy.foo() // returns "Mocked foo"
    proxy.bar() // returns "Actual bar"
    proxy.baz() // returns "Actual baz", because of the thrown MockDisabledException
}