Simplified mock library

mock.kt

package id.thony.mockbyreflection

import java.lang.reflect.Method
import java.lang.reflect.Proxy
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.locks.ReentrantLock

inline fun <reified T> mock(): T = mock(T::class.java)

fun <T> every(stubBlock: () -> T): MockScope<T> = MockScope(stubBlock)

@Suppress("UNCHECKED_CAST")
fun <T> mock(mock: Class<T>): T {
    return Proxy.newProxyInstance(
        mock.classLoader,
        arrayOf(mock),
        MockInvocationHandler()
    ) as T
}

val Lock = ReentrantLock()
val RecordingMock = AtomicBoolean(false)
val ReturnsDeque = ArrayDeque<Any>()
val MockSet = mutableSetOf<MockMethod>()

class NotRegisteredException(method: Method) : RuntimeException("$method is not mocked.")

class MockScope<T>(private val block: () -> T) {
    infix fun returns(what: T) {
        Lock.lock()
        try {
            RecordingMock.set(true)
            ReturnsDeque.addLast(what as Any)
            block.invoke()
        } finally {
            Lock.unlock()
        }
    }
}

MockInvocationHandler.kt

package id.thony.mockbyreflection

import java.lang.reflect.InvocationHandler
import java.lang.reflect.Method

class MockInvocationHandler : InvocationHandler {
    override fun invoke(proxy: Any?, method: Method?, args: Array<out Any>?): Any {
        method ?: throw IllegalStateException("method args is null")
        val argsList = args?.toList() ?: emptyList()

        if (RecordingMock.compareAndSet(true, false)) {
            val returnForThisInvoke = ReturnsDeque.removeFirst()
            val mockMethod = MockMethod(method, argsList, returnForThisInvoke)
            MockSet.add(mockMethod)
            return returnForThisInvoke
        }

        return MockSet.firstOrNull { it.method == method && it.args == argsList }?.returns
            ?: throw NotRegisteredException(method)
    }
}

MockMethod.kt

package id.thony.mockbyreflection

import java.lang.reflect.Method

data class MockMethod(
    val method: Method,
    val args: List<Any>,
    val returns: Any
)