Extensions to mockk to allow atomic changes to stubbed answers

MutableAnswers.kt

package ext.mockk

import io.mockk.Call
import io.mockk.MockKAnswerScope
import io.mockk.MockKStubScope
import java.util.concurrent.atomic.AtomicReference

infix fun <T, B> MockKStubScope<T, B>.returnsMutably(initialValue: T): MutableAnswer<T, B> =
  answersMutably { initialValue }

infix fun <T, B> MockKStubScope<T, B>.throwsMutably(throwable: Throwable): MutableAnswer<T, B> =
  answersMutably { throw throwable }

infix fun <T, B> MockKStubScope<T, B>.answersMutably(
  initialAnswer: MockKAnswerScope<T, B>.(Call) -> T,
): MutableAnswer<T, B> {
  val mutableAnswer = MutableAnswer(initialAnswer)
  this answers { call ->
    val answer = mutableAnswer.answer
    answer(call)
  }
  return mutableAnswer
}

class MutableAnswer<T, B> private constructor(
  private val atomicAnswer: AtomicReference<MockKAnswerScope<T, B>.(Call) -> T>,
) {

  internal constructor(initialAnswer: MockKAnswerScope<T, B>.(Call) -> T) : this(
    AtomicReference(initialAnswer),
  )

  infix fun nowReturns(value: T): Unit = nowAnswers { value }

  infix fun nowThrows(throwable: Throwable): Unit = nowAnswers { throw throwable }

  infix fun nowAnswers(newAnswer: MockKAnswerScope<T, B>.(Call) -> T) {
    atomicAnswer.set(newAnswer)
  }

  internal val answer: MockKAnswerScope<T, B>.(Call) -> T
    get() = atomicAnswer.get()
}

Used as so:

val charSequence = mockk<CharSequence>()
val charSequenceLengthStub = every { charSequence.length } returnsMutably 1
// do some work that asynchronously invokes charSequence.length, perhaps polling in a loop
charSequenceLengthStub nowReturns 2
// check that the asynchronous work reacts to the change in the stub