Stub.scala

package nope.utils

trait Stub {

  class StubbedNullaryMethod[R] {
    @volatile
    protected var _function: () => R = () => throw new RuntimeException("Unset.")

    def setFunction(f: () => R) {
      _function = f
    }

    def setValue(value: R) {
      _function = () => value
    }

    def get: R = _function()
  }

  class PartialFunctionInterceptor[A, R] extends PartialFunction[A, R] {
    @volatile
    private var _value: PartialFunction[A, R] = PartialFunction.empty

    @volatile
    private var _invocations: Seq[A] = Vector.empty

    def invocations = _invocations

    def apply(arg: A): R = {
      _invocations +:= arg
      _value.apply(arg)
    }

    def addCases(newCases: PartialFunction[A, R]) {
      _value = _value orElse newCases
    }

    override def isDefinedAt(arg: A): Boolean = _value.isDefinedAt(arg)
  }

  abstract class StubbedMethod[A, R] {
    protected val underlyingImplementation = new PartialFunctionInterceptor[A, R]

    def invocationCount(f: A => Boolean = Function.const(true)) = {
      underlyingImplementation.invocations.count(f)
    }

    def setExpectations(pf: PartialFunction[A, R]) {
      underlyingImplementation.addCases(pf)
    }
  }

  final class StubbedMethod2[A1, A2, R] extends StubbedMethod[(A1, A2), R] {
    def apply(arg1: A1, arg2: A2): R = {
      underlyingImplementation((arg1, arg2))
    }
  }

  final class StubbedMethod3[A1, A2, A3, R] extends StubbedMethod[(A1, A2, A3), R] {
    def apply(arg1: A1, arg2: A2, arg3: A3): R = {
      underlyingImplementation((arg1, arg2, arg3))
    }
  }
  final class StubbedMethod4[A1, A2, A3, A4, R] extends StubbedMethod[(A1, A2, A3, A4), R] {
    def apply(arg1: A1, arg2: A2, arg3: A3, arg4: A4): R = {
      underlyingImplementation((arg1, arg2, arg3, arg4))
    }
  }
}

argh == arg17 in Base18 👅