Managing memory and avoiding leaks is a crucial part of building any kind of program. Thankfully Swift makes this relatively easy in most situations - thanks to Automatic Reference Counting (or ARC for short). However, there are still times when it can be quite easy to make a mistake that causes memory to be leaked.

Memory leaks in Swift are often the product of a retain cycle, when one object will hold a strong reference to an object that also strongly references the original object. So A retains B and B retains A. These kinds of issues can sometimes be tricky to debug and lead to hard to reproduce crashes.

This week, let's take a look at how we can set up unit tests to both help us identify memory leaks, and also make it easier to avoid common mistakes that could end up causing leaks in the future.

Delegates

The delegate pattern is very common on Apple's platforms. It's also a pretty nice pattern - it's simple and enables two objects to have a loosely coupled relationship but still communicate with each other in a nice way.

Let's say that we have a view controller that displays a list of the user's friends in our app. To enable the view controller to send events back to its owner, we add a delegate API to it, like this:

class FriendListViewController: UIViewController {
   var delegate: FriendListViewControllerDelegate?
}

At first glance, it may seem like nothing is wrong with the above code, but once you look a bit closer you'll probably spot the problem - the delegate is being strongly referenced. This will most likely end up causing a retain cycle, since the owner of a FriendListViewController instance will probably also be its delegate - causing both of the objects to strongly reference each other.

To ensure that we don't make these kind of mistakes, either now or in the future, we can setup a unit test to ensure that FriendListViewController doesn't retain its delegate:

class FriendListViewControllerTests: XCTestCase {
    func testDelegateNotRetained() {
        let controller = FriendListViewController()

        // Create a strong reference to a delegate object, then
        // assign it as the view controller's delegate
        var delegate = FriendListViewControllerDelegateMock()
        controller.delegate = delegate

        // Re-assign the strong reference to a new object, which
        // should cause the original object to be released, thus
        // setting the view controller's delegate to nil
        delegate = FriendListViewControllerDelegateMock()
        XCTAssertNil(controller.delegate)
    }
}

The above test will initially fail, which is great - because it will later let us assure that we have indeed fixed the problem. In this case, the fix is easy - we'll simply make the delegate weak:

class FriendListViewController: UIViewController {
    weak var delegate: FriendListViewControllerDelegate?
}

If we now run our test again, it'll turn green - and we have both fixed a memory leak in our app, and also added a safeguard so that this bug doesn't come back again in the future (which is super useful when doing tasks like refactoring).

Another tip is to use SwiftLint, which will warn you in case you define a delegate property without making it weak 👍

Observers

The observer pattern is also quite common when enabling an object to easily notify multiple other objects of various events. Just like with delegates, we most likely don't want to retain observers, since they often will keep a strong reference to the object they are observing.

Let's say we have a UserManager, which actually does retain all of its observers, by storing them in an array:

class UserManager {
    private var observers = [UserManagerObserver]()

    func addObserver(_ observer: UserManagerObserver) {
        observers.append(observer)
    }
}

Just like when implementing the delegate pattern, it can be quite easy to accidentally store observers like we do above - which will most likely end up causing memory leaks in our app.

Thankfully, the problem is also easily reproduced in a test case:

class UserManagerTests: XCTestCase {
    func testObserversNotRetained() {
        let manager = UserManager()

        // Create both a strong and a weak local reference to an
        // observer, which we then add to our UserManager
        var observer = UserManagerObserverMock()
        weak var weakObserver = observer
        manager.addObserver(observer)

        // If we re-assign the strong reference to a new object,
        // we expect the weak reference to become nil, since
        // observers shouldn't be retained
        observer = UserManagerObserverMock()
        XCTAssertNil(weakObserver)
    }
}

Again the above test will fail - which is awesome, since we have a solid way to reproduce the problem. To fix it, we're going to need a small wrapper type to be able to weakly reference our observers (since arrays always strongly reference their elements):

private extension UserManager {
    struct ObserverWrapper {
        weak var observer: UserManagerObserver?
    }
}

Then, we'll simply update UserManager to instead store wrapper values instead of directly storing observer objects, which will make our test pass.

class UserManager {
    private var observers = [ObserverWrapper]()

    func addObserver(_ observer: UserManagerObserver) {
        let wrapper = ObserverWrapper(observer: observer)
        observers.append(wrapper)
    }
}

The above technique is super useful in general when you need to weakly store objects in any kind of collection. Important to keep in mind though, is to clean up any wrappers which objects have been released. A nice way to do that is to filter out all expired wrappers when we are already iterating over them, for example when we are notifying them of an event, like this:

private func notifyObserversOfUserChange() {
    observers = observers.filter { wrapper in
        guard let observer = wrapper.observer else {
            return false
        }

        observer.userManager(self, userDidChange: user)
        return true
    }
}

Closures

Finally, let's take a look at how unit tests can help us identify and prevent memory leaks when implementing a closure-based API. Closures can generally be a quite common source of memory related bugs and leaks, since they per default strongly capture all of the objects that are used within them (for more on this topic, check out "Capturing objects in Swift closures").

Let's say that we're building an ImageLoader, that lets us load remote images over the network, and run a completion handler once done:

class ImageLoader {
    func loadImage(from url: URL,
                   completionHandler: @escaping (UIImage) -> Void) {
        ...
    }
}

One common mistake when implementing something like the above is to accidentally keep completion handlers around even after their operation has been finished. In order to enable cancellation or batching, we might be storing completion handlers in some form of collection, and forgetting to remove handlers can end up causing an overuse of memory and leaks.

So how can we use a unit test to make sure that completion handlers are indeed removed as soon as they have been run? In the other situations we were using weak references to the objects that were used as delegates or observers, but we can't hold weak references to closures 🤔.

What we can do instead, is to use object capturing to associate an object with a closure, and then use that object to verify that the closure was released:

class ImageLoaderTests: XCTestCase {
    func testCompletionHandlersRemoved() {
        // Setup an image loader with a mocked network manager
        let networkManager = NetworkManagerMock()
        let loader = ImageLoader(networkManager: networkManager)

        // Mock a response for a given URL
        let url = URL(fileURLWithPath: "image")
        let data = UIImagePNGRepresentation(UIImage())
        let response = networkManager.mockResponse(for: url, with: data)

        // Create an object (it can be of any type), and hold both
        // a strong and a weak reference to it
        var object = NSObject()
        weak var weakObject = object

        loader.loadImage(from: url) { [object] image in
            // Capture the object in the closure (note that we need to use
            // a capture list like [object] above in order for the object
            // to be captured by reference instead of by pointer value)
            _ = object
        }

        // Send the response, which should cause the above closure to be run
        // and then removed & released
        response.send()

        // When we re-assign our local strong reference to a new object the
        // weak reference should become nil, since the closure should have been
        // run and removed at this point
        object = NSObject()
        XCTAssertNil(weakObject)
    }
}

We now have a guarantee that our image loader does not keep old closures around and have removed another potential source of memory leaks 🎉.

Conclusion

While using unit tests this way can seem a bit overkill at first, it can be a really nice tool to apply in situations when you're either trying to reproduce an existing memory leak or when you want to add an extra level of protection against accidentally introducing one in the future.

Even though tests like these won't completely protect you against memory leaks, they can potentially reduce the amount of time you'll have to spend hunting down the cause of memory related issues in the future.

What do you think? Have you been using unit testing this way before or is it something that you'll try out? Let me know, along with any questions, comments or feedback that you have - on Twitter @johnsundell.

Thanks for reading! 🚀