Watch out for protocol extensions in your Swift API (unit tests trap).

We all love protocol extensions, one of the most powerful element of protocol oriented programming (POP) in Swift. Despite their unquestionable benefits, there are some rare cases where you should avoid them. In this article, let me demonstrate a potential trap that your API consumer may fall into when trying to unit test a code that depends on some protocol extension function.


Quick reminder: method dispatching

In Swift we have three kinds of method dispatching: static, vtable and message dispatching. If you are not familiar with this terms, let me recommend a great post from Riazlab’s. In short, dispatching methods use different techniques to choose a concrete implementation of you function to execute if the same signature is defined in several places (e.g. parent class or protocol extension). For instance, 1) having a class that inherits from NSObject will always use message dispatching and 2) value types (structs, enums) will always use static dispatch.


OK, let’s go back to our dangerous scenario, where as an API creators (and we all are API designers, do you remember presentation by John Sundell?) we follow best practices and provide protocol abstraction for our public API. For the sake of this article, let’s assume we wish to expose a logger class that may log verbose and error messages, as shown below:

At glance it looks OK, but what if we wish to add a convenient API functions to log particular level? Protocol extension looks as a perfect match:

So far so good. If we have some hypothetical API consumer where generic Logger abstraction is used in System class, it can directly call a method to log verbosely, like below:

Start function uses convenient API to log a verbose message

This works as expected, but let’s go over unit tests of a System class?

Testing API protocol extensions

Let us verify that our System actually logs a verbose message to a logger, every time we start it. Piece of cake. First we have to create a mock that conforms to a Logger protocol and keep track of all calls to verbose function. At the end, verify that verbose function has been called as expected:

It may be surprising for you, especially knowing that System actually logs a verbose message in a production code, but this test fails with a message:

XCTAssertEqual failed: (“[]”) is not equal to (“[“System started”]”)

Failing scenario

To understand what went wrong, we have to analyse dispatching method used to execute logger.verbose(message:) inside System.start(). System instance has a reference to a protocol Logger which provides function verbose in a protocol extension. By design Swift will always use static dispatch to call protocol extension’s implementation, no matter if actual implementation of a logger (LogMock in our case) implements verbose(_:String) or not.

logger.verbose always calls a function from an extension due to static dispatch

One way to workaround it to verifying if LogMock.log(_:message) has been called correctly. It may seem like a rational idea at first, but but there’s an inherent problem with this approach —in System class tests, we have just began testing real verbose(message:) implementation from the API, written by a third-party developer. Potential implementation change or bug in Logger.verbose(_:String) could influence our test result. Unit test, as even name suggests, should validate the single unit of a code (here System class), in a highly isolated context.

How to define API protocol correctly?

Solution to this problem is really simple. As API designer all you have to do is to include all the functions/variables that you want to expose in a protocol extension into a main protocol definition, like:

This changes a dispatching method of a verbose(message:) and error(message:) functions into vtable , which resolves an implementation in a runtime. As a result, our test scenario executes verbose(message:) from a LogMock implementation, as you would expect:

Calling chain when verbose function is defined also in a protocol definition

Customer flow

Above fix is dedicated to API creators, but even if you don’t have access to modify protocol declaration, there is a remedy for your troubles. You’ll need to depend on your custom protocol that inherits from an original one and includes all declarations that API designer exposes only from protocol extension. In the meantime, while waiting for API designer’s fix, you are not blocked anymore.

Summary

We talked about protocol extensions and saw that testing a code that depends on a protocol extensions could be tricky. Fortunately, there is a really simple remedy for that — just include functions/variables definitions from a public protocol extension in your protocol definition. So when you hear yourself saying, “Let’s just implement it in a protocol extensions”, you’ll know it’s time to backpedal and ensure that static dispatch would not thwart customer’s unit tests.

For all their advantages, dynamic nature of vtable dispatching has some performance overhead comparing static one, which does not have to perform a table lookup. However, difference isn’t noticeable in most cases so it sounds as a reasonable tradeoff to make your API testable.