Effective extension function refactoring in Kotlin: companion object use case

In Kotlin we have a great extension functions functionality which enables us to write more expressive and concise code. Under the hood, they are simple static methods that risk damaging our codebases if we use them incorrectly. In this article, I describe how to deal with extension functions that have transformed over time from lightweight additions to code into monsters with domain-coupled logic.

Starting point

Let’s imagine we have the following function:

Whatever way it is implemented it makes no difference: either we can use top level functions or methods of a singleton. The following refactoring approach can be applied to them all.

As you can see, these methods check if biometrics are available on the current device. Pretty simple and straightforward logic to be implemented as an extension method.

Testability

Such extension methods (also top-level functions) are basically @JVMStatic methods of a particular utility class.

Here is a Java equivalent of the method:

Generally speaking, we are using an old-fashioned singleton class (defined in a scope of class via static modifier) that we can access from any place in our code to reuse the logic behind it.

So, what is the main issue with singletons?

Testability.

Imagine the following case:

How testable is this code?

Not very, because even if you provide Context mock in the test, you still won’t be able to use it efficiently because it is used indirectly inside BiometricManager. You need to know the implementation details of BiometricManager and how exactly it uses Context in order to find out how to set up a mock properly.

This issue might be solved by using Robolectric or running this test on a device, but do we actually need to do this? Robolectic and device tests take much longer to run.

Logic complication

Also, what might happen to utility functions?

They can get really complex, but we would be too late to notice.

In the light of the previous example, imagine we have a new requirement — everywhere we check hardware availability of biometrics, we also need to check AB-test value to enable this functionality.

Firstly, the code that I am going to show you next is not code that should end up in production but rather it is merely an example.

Secondly, sooner or later you will encounter this issue in the production code anyway. None of us is perfect.

So, we end up with quite an ugly extension function blending domain and data logic, which might prove hard to refactor because it is used in dozens of places.

Big codebase problem

What is the problem that arises when implementing a new class that handles the logic and injects it correctly into the constructor on every class that uses it?

Pull request size.

In big codebases, the following function can be used in 10–100 different places of the codebase. Every place requires changes in the corresponding class constructor. If your app is a multi-module app with direct explicit dependencies, you may also need to explicitly declare a new class as a dependency of every module where it is used.

With all these changes your pull request might end up being huge and, as such, harder to review. It will also be easier to overlook an error. By using the approach below we can implement every step as a separate pull request.

Replace with singleton

The first step is to admit the problem of singleton usage. We need to replace the implicit singleton with an explicit one.

Interface companion object magic

Now we have a class that we can work with.

Because testability is a requirement, we will want to replace direct usages of BiometricsUtils class with an interface in the future.

For now, the interface should look like this:

But in the future, we will want to remove Context and AbTestStore from the method parameters list because they are the same between different method invocations — Context is an application context and AbTestStore is a local singleton in our dependency graph.

For now, we should stick to the first variant and migrate to the second one in the additional step at the end.

So, at this point we have the interface definition and singleton class.

How can we actually link them together in a manner that will mean we don’t have to change the usages of this method?

A companion object is the answer.

Companion objects were introduced early in the development of Kotlin, even prior to the release of 1.0. At that stage, the advantages of companion objects over regular objects and usual Java static methods were not obvious. This was especially because the word Companion had to be used every time you wanted to access it.

Companion usage requirement was removed later on and that was a blessing. Now we can access a companion object in a Java static functions manner.

Furthermore, Kotlin compiler is also smart enough to distinguish method calls between interface and its companion object.

Also, companion objects in Kotlin have the ability to extend classes and interfaces because it is a regular singleton class.

This leads us to interface’s companion object that implements enclosing interface.

As you can see, we can invoke the abstract function bar on interface Foo by delegating it to the companion object of Foo.

Let’s use this technique to refactor our code.

After this change, we can still use BiometricsUtils.canUseBiometrics(applicationContext, abTestStore) without any issues. This gets us one step closer to finishing our refactoring.

Inject interface with a default value

Because we have the interface now, we actually have something that can be provided as a constructor parameter.

We are using BiometricsUtils.Companion as a default value for biometricsUtils parameter. It means that you do not need to change the existing code that creates this class.

But this change is also quite important in a different way. It finally allows us to test ScreenViewModel with JVM tests.

BiometricsUtils is an interface, and we can provide a mock (or stub) in the test.

Because we are only using applicationContext and abTestStore as parameters of canUseBiometrics, we can provide empty mocks. We no longer have to mock any other method of these classes as we did previously. The logic behind canUseBiometrics method can be unit tested separately, so we can be sure that it works as expected.

Remove default value

Now we can remove the default value of biometricsUtils parameter and provide the actual value via our DI system.

Move towards to better interface

Now we can move biometricsUtils parameters to the constructor of a class and update our usages one by one. Let’s deprecate our currently existing function and introduce a new one. Also, because we already removed all usages of BiometricsUtils.Companion in our codebase, we can remove it here as well.

Then we add a new implementation of BiometricsUtils.

So, we can provide the new class via the DI system.

And finally, we can remove all usages of the deprecated method and update tests that use it.

Conclusion

The result was that we were able to refactor extension method usage step by step in a compatible manner, introducing our changes little by little, avoiding disruption to our colleagues’ work and minimising the number of merge conflicts. Interface’s companion object is a powerful feature that allows the use of singletons that can be injected easily into constructor and replaced with mocks.