Search code examples
swiftunit-testingtraitsmixinsstub

How to stub Swift "Trait/Mixin" method in struct/class for testing


I have recently read about how to add "Traits/Mixins" to a struct/class in Swift by creating a protocol and extending that protocol with a default implementation. This is great as it allows me to add functionality to view the controller without having to add a bunch of helper objects to said view controller. My question is, how do I stub calls that are provided by these default implementations?

Here is a simple example:

protocol CodeCop {
  func shouldAllowExecution() -> Bool
}

extension CodeCop {
  func shouldAllowExecution() -> Bool {
    return arc4random_uniform(2) == 0
  }
}

struct Worker : CodeCop {
  func doSomeStuff() -> String {
    if shouldAllowExecution() {
       return "Cop allowed it"
     } else {
       return "Cop said no"
    }
  }
}

If I wanted to write two tests, one that verifies that the String "Cop allowed it" is returned by doStuff() when CodeCop does not allow execution, and another test that verifies that the String "Cop said no" is returned by doStuff() when CodeCop does not allow execution.


Solution

  • This is simple enough to do by writing an additional protocol in your test target, called CodeCopStub, that inherits from CodeCop:

    protocol CodeCopStub: CodeCop {
        // CodeCopStub declares a static value on the implementing type
        // that you can use to control what is returned by
        // `shouldAllowExecution()`.
        //
        // Note that this has to be static, because you can't add stored instance
        // variables in extensions.
        static var allowed: Bool { get }
    }
    

    Then extend CodeCopStub's shouldAllowExecution() method, inherited from CodeCop, to return a value depending on that new static variable allowed. This overrides the original CodeCop implementation for any type that implements CodeCopStub.

    extension CodeCopStub {
        func shouldAllowExecution() -> Bool {
            // We use `Self` here to refer to the implementing type (`Worker` in
            // this case).
            return Self.allowed
        }
    }
    

    All you have left to do at this point is to make Worker conform to CodeCopStub:

    extension Worker: CodeCopStub {
        // It doesn't matter what the initial value of this variable is, because
        // you're going to set it in every test, but it has to have one because
        // it's static.
        static var allowed: Bool = false
    }
    

    Your tests will then look something like this:

    func testAllowed() {
        // Create the worker.
        let worker = Worker()
        // Because `Worker` has been extended to conform to `CodeCopStub`, it will
        // have this static property. Set it to true to cause
        // `shouldAllowExecution()` to return `true`.
        Worker.allowed = true
    
        // Call the method and get the result.
        let actualResult = worker.doSomeStuff()
        // Make sure the result was correct.
        let expectedResult = "Cop allowed it"
        XCTAssertEqual(expectedResult, actualResult)
    }
    
    func testNotAllowed() {
        // Same stuff as last time...
        let worker = Worker()
        // ...but you tell it not to allow it.
        Worker.allowed = false
    
        let actualResult = worker.doSomeStuff()
        // This time, the expected result is different.
        let expectedResult = "Cop said no"
        XCTAssertEqual(expectedResult, actualResult)
    }
    

    Remember that all of this code should go in your test target, not your main target. By putting it in your test target, none of it will affect your original code, and no modification to the original is required.