Mocking composition of Traits with Self-Types in Scala

Scala’s Traits are a powerful modeling tool, they let us implement “composition over inheritance” to achieve more flexibility in object-oriented design. However, there are some situations where mocking such composed classes in unit tests is difficult if not impossible. Let’s see one of these situations: composition of traits with self-types.

Self-Types

By using self-types we can enforce a dependency constraint in a trait we created, indicating that any class that mixes our trait must also mix the trait that we specified as self-type. For instance:

trait OurTrait { this: ASelfTypeTrait =>
  def doThisAndThat() = …
}
// Wrong: must extend also ASelfTypeTrait
class OurCustomClass extends OurTrait {
...
}
// OK
class OurCustomClass extends OurTrait with ASelfTypeTrait {
}

OK, now let’s get into the case.

Our Case

Let’s assume that our codebase contains a trait like this:

trait BaseTrait {

  val optMultiplier: Option[Int]
  val multiplier = optMultiplier.getOrElse(0)

  def multiplyIt(v: Float): Float = v * multiplier
  def squareIt(v: Float): Float = v * v

}

and it also contains another trait that uses the above one as self-type:

trait IntermediateTrait { this: BaseTrait =>

  def oneTenth(v: Float): Float = v / 10

}

Our code base also contains two other traits, each one defining separate sets of methods, in order to achieve separation of concerns:

trait FirstUpperTrait extends IntermediateTrait { this: BaseTrait =>

  def squarePolinomial(v: Float): Float =
    squareIt(v) + multiplyIt(v)

}
trait SecondUpperTrait
  extends IntermediateTrait { this: BaseTrait =>

  def quarterIt(v: Float): Float = v / 4

}

Now we create a component as the composition of the previous traits:

class MyComponentImpl(
  override val optMultiplier: Option[Int] = Some(3)
) extends BaseTrait
  with IntermediateTrait
  with FirstUpperTrait
  with SecondUpperTrait

This component comprises all basic operations needed in the rest of our code and also ensures optMultiplier is initialized just once. The example seems trivial but imagine this is not just an Option[Int] but a database connection or some other expensive resource that needs to be initialized just once.

In the end, we have a consumer component (let’s call it UserComponent) that happens to depend on operations defined by FirstUpperTrait:

class UserComponent(val cmp: FirstUpperTrait) {

  private val Value = 4

  def process(): Float =
    cmp.squarePolinomial(Value) + cmp.oneTenth(Value)

We follow best practices so we will create a unit test for UserComponent. We will code it using ScalaTest:

class UserComponentTest extends AnyFlatSpec with MockFactory {

  behavior of "UserComponentTest"

  it should "process" in {
    val myCmp = mock[FirstUpperTrait]

    (myCmp.squarePolinomial _).expects(4).returns(28.0F)
    (myCmp.oneTenth _).expects(4).returns(0.4F)

    val userComponent = new UserComponent(myCmp)

    val result = userComponent.process()

    assertResult(result)(28.4F)
  }

}

Then we try to compile and launch it…Oops! What happened?

illegal inheritance;
 self-type delegation.test.FirstUpperTrait does not conform to delegation.test.FirstUpperTrait’s selftype delegation.test.FirstUpperTrait with delegation.test.BaseTrait
 val myCmp = mock[FirstUpperTrait]

That means the mock we try to create is a new class that extends FirstUpperTrait and, as such, it should also extend its self-type ( BaseTrait ). That makes our UserComponent basically non-testable.

Solution

First, we will create a new trait AbstractFirstUpperTrait that will declare al operations consumed by UserComponent:

trait AbstractFirstUpperTrait {

  def squarePolinomial(v: Float): Float
  def oneTenth(v: Float): Float
}

Second, we make FirstUpperTrait extend also this new trait:

trait FirstUpperTrait extends AbstractFirstUpperTrait
  with IntermediateTrait { this: BaseTrait =>

  def squarePolinomial(v: Float): Float =
    squareIt(v) + multiplyIt(v)

}

Third, we’ll change UserComponent so that it has a dependency on AbstractFirstUpperTrait instead of FirstUpperTrait:

class UserComponent(val cmp: AbstractFirstUpperTrait) {

  private val Value = 4

  def process(): Float =
    cmp.squarePolinomial(Value) + cmp.oneTenth(Value)

}

Then, as last step, we will change just one line of code n our unit test:

val myCmp = mock[AbstractFirstUpperTrait]

Now let’s try building and launching again:

Alright, it worked!

Now, does that mean UserComponent and MyComponentImpl will work together? Yes, it will. Let’s see the sample main program using both classes:

object Main extends App {

  class MyComponentImpl(override val optMultiplier: Option[Int] = Some(3))
    extends BaseTrait
    with IntermediateTrait
    with FirstUpperTrait
    with SecondUpperTrait

  val myCmp = new MyComponentImpl()
  val userCmp = new UserComponent(myCmp)

  println(userCmp.process())
}

If I run it (from inside my IDE for brevity):

Same result as in the unit test.

Conclusion

In this article we have covered a sample case in which a class can become non-testable because of the use of self-types to constrain dependencies, and also a solution model the dependencies so that it can become testable.