In the previous post I showcased how we could use our test suite to help us find potential memory leaks. This novel approach to Unit Tests was good at catching instances where we forgot to cleanup our properties and SUTs during the tearDown, but let’s not kid ourselves, short of adding assert clauses at the end of each test, we wouldn’t be detecting sneakier leaks. What’s more, that solution didn’t handle more complex testing scenarios with non-class entities and constants.

Fret not. What if I told you we could narrow those gaps? (leak pun intended 🚱)

# Recap of Part 1

This was the previous naive approach:

  • We create an utility class which every unit test inherits from.
  • This parent class checks which SUTs are leaking during super.tearDown() using reflection, after the SUTs have been reset with = nil in self.tearDown().
  • A test will fail if a SUT property is not nil during NaiveLeakCheckTestCase.tearDown(), meaning that a reference is being retained somewhere.
class NaiveLeakCheckTestCase: XCTestCase {

    override func tearDown() {

        defer { super.tearDown() }

        let mirror = Mirror(reflecting: self)

        mirror.children.forEach { label, value in
            guard let propertyName = label else { return }

            // NOTE: If the value is something, e.g. not `nil` (a.k.a. not `Optional<Any>.none`),
            // then we have a potential leak.
            if case Optional<Any>.some(_) = value {
                XCTFail("\"\(propertyName)\" is a potential memory leak!")
            }
        }
    }
}

class SampleNonLeakingTestCase: NaiveLeakCheckTestCase {
    // ... SUT declaration, setup, etc.

    override func tearDown() {

        // Make sure that our memory leak inspection is the last thing to happen
        defer { super.tearDown() }

        // reset SUTs
        parentSUT = nil
        childSUT = nil
    }

    // ... Actual SUT tests
}

# Plenty of room for improvement

At first glance the previous approach seems to do everything we set out to do, but that’s only because we’ve been been blinded by the simplicity of the example scenario that I had used. If the subject under test (SUT) is leaking within itself then sut = nil will still leave a lingering reference but XCTAssertNil(sut) will still be true, resulting in a false-negative for our leak checker and a false sense of security that everything is ok with out SUT. What if we also want to use a struct, enum, or anything other than a class in our tests? What about using constants that don’t need resetting in between tests? How do lazy properties affect ou approach? All of these would present themselves as false-positives, throwing a wrench in our plans for automated memory leak detection. There’s nothing worse than a failing test that you just don’t quite know how to fix. So let’s try and address these new constraints:

  • We want to automate the asserts for lingering references at the end of a test.
  • We want to be able to distinguish between SUT properties and constants.
  • We want to be able to use non-reference types such as structs and enums.
  • We want to be able to use lazy-loaded properties.

# Making NaiveLeakCheckTestCase more readable

Before we start addressing our new constraints, NaiveLeakCheckTestCase could be more readable, so for NaiveReadableLeakCheckTestCase we are going to move that confusing mess that is if case Optional<Any>.some(_) = wrappedValue into a more pleasantly named isSomeValue() function. We’re also going to start overriding tearDownWithError() instead of tearDown() so as to guarantee that this is the last thing our test case classes do. (If you want more information about the lifecycle of XCTestCase in swift, you can check out the official documentation)

class NaiveReadableLeakCheckTestCase: XCTestCase {

    // MARK: - TearDown

    override func tearDownWithError() throws {

        let mirror = Mirror(reflecting: self)
        mirror.children.forEach { label, wrappedValue in

            guard let propertyName = label else { return }

            // NOTE: If the value is something, e.g. not `nil`, then we either are just not resetting the
            // value to `nil` during tearDown, or we have a potential leak.
            if isSomeValue(wrappedValue) {
                XCTFail("📄♻️ \"\(propertyName)\" is missing from tearDown, or is a potential memory leak!")
            }
        }

        try super.tearDownWithError()
    }

    // MARK: - Utils

    private func isSomeValue(_ wrappedValue: Any) -> Bool {

        // NOTE: This weird syntax is due to the fact that non-optional `Any` can be `nil`, even
        // though it does not conform to Optional. So `<some nil Any> == nil` returns true or
        // won't even compile, depending on how it's done.
        //
        // This case comparison using the optional protocol, however, does work.
        //
        // I've also seen it done using `String(describing: <some Any>) == "nil"`,
        // but it feels more hacky.

        // NOTE: swiftlint attempts to fix `Optional<Any>.none` as `Any?.none`
        // but the compiler doesn't like that, so we disable the rule just this once.

        // swiftlint:disable:next syntactic_sugar
        if case Optional<Any>.some(_) = wrappedValue {
            return true
        }
        return false
    }
}

This little detour is going to pay out soon enough, I promise.

# Adding weak references and non-reference types to the mix

Let’s look at the following test scenario which uses the NaiveReadableLeakCheckTestCase for a single object that we know for a fact is leaking internally.

import XCTest

fileprivate class LeakingPerson {

    var name: String
    lazy var greeting: () -> String = { // there's a missing `[unowned self] in` here, causing the leak
        return "Hello \(self.name)!"
    }

    init(name: String) {
        self.name = name
    }
}

class SomeLeakingClassTestCase: NaiveReadableLeakCheckTestCase {

    private var leakingSUT: LeakingPerson!

    override func setUp() {

        super.setUp()

        leakingSUT = LeakingPerson(name: "aclima")
    }

    override func tearDown() {

        defer { super.tearDown() }

        leakingSUT = nil
    }

    // MARK: - Tests

    func test_WithLeakingSUT_ShouldNotBeNil() {

        XCTAssertNotNil(leakingSUT)
        XCTAssertEqual(leakingSUT.greeting(), "Hello aclima!") // this causes the leak
    }

Using our naive approach, this test would still be successful, when in reality it shouldn’t. If we want to be sure that leakingSUT is leaking we need to create an auxiliary weak reference to it, and inspect that at the end of our test. There’s no need to assign nil to it, since we want to check that leakingSUT is being properly deallocated during super.tearDown().

import XCTest

fileprivate class LeakingPerson {

    var name: String
    lazy var greeting: () -> String = { // there's a missing `[unowned self] in` here, causing the leak
        return "Hello \(self.name)!"
    }

    init(name: String) {
        self.name = name
    }
}

class SomeLeakingClassTestCase: NaiveReadableLeakCheckTestCase {

    private var leakingSUT: LeakingPerson!
    private weak var auxWeakReference: LeakingPerson!

    override func setUp() {

        super.setUp()

        leakingSUT = LeakingPerson(name: "aclima")
        auxWeakReference = leakingSUT // weak reference to the SUT
    }

    override func tearDown() {

        defer { super.tearDown() }

        leakingSUT = nil

        // NOTE: We don't perform `auxWeakReference = nil` so that `super.tearDown()` will
        // assert that auxWeakReference isn't `nil`, meaning that leakingSUT is leaking somewhere!
    }

    // MARK: - Tests

    func test_WithLeakingSUT_ShouldNotBeNil() {

        XCTAssertNotNil(leakingSUT)
        XCTAssertEqual(leakingSUT.greeting(), "Hello aclima!") // this causes the leak
        XCTAssertEqual(leakingSUT, auxWeakReference)
    }

This looks great! Until we realize that doing this by hand for every class in our tests is going to troublesome and error prone… Let’s automate it instead!

We are going to need:

  • a collection to hold our weak references, which should be created before the test
  • a way of assigning the weak references at some point before the test ends and super.tearDownWithError() checks for leaks
  • to check the weak references during super.tearDownWithError()

Starting with the first requirement, since in swift we can’t create a collection of weak references with something like var weakReferences = [weak SomeObject](), we might try and roll out our own solution, something along the lines of

class WeakReference {

    weak var objectValue: AnyObject?

    init(objectValue: AnyObject) {
        self.objectValue = objectValue
    }
}

var weakReferences = [WeakReference]()

But as it turns out, there’s a simple but not widely known collection that ensures weak references. I’m talking about NSMapTable, which we can use as simply as

var weakReferences = NSMapTable<NSString, AnyObject>.weakToWeakObjects()

Next up, we want a way of assigning our weak variables after a test ends and before super.tearDownWithError() kick in. Again, there’s a tool for that, it’s the addTearDownBlock closure. Once more, we can use reflection to inspect our classes’ properties and assign our class instance references to our weak variables! But in doing so, we stumble upon another of our previously stated constraints: How do we distinguish a “reference” from a “non-reference” type property?

This can also be accomplished with reflection! We can use a Mirror to inspect the properties themselves and their displayStyle like so

Mirror(reflecting: property).displayStyle == .class

Due to the way that reflection works in swift, we’re also going to need a new utility method for unwrapping our properties (which might be optionals, a.k.a nil) safely.

private func unwrap<T>(_ any: T) -> Any {

    let mirror = Mirror(reflecting: any)
    guard mirror.displayStyle == .optional, let first = mirror.children.first else {
        return any
    }
    return unwrap(first.value)
}

Putting it all together, during setUpWithError() we will create a tearDownBlock where we use reflection to go through our classes’ properties and, if they hold a non-nil value and are a class, we add them to our collection of weak references. Later on, during tearDownWithError() we iterate over our collection fo weak references and check if the value is non-nil, in which case we are certain that the object is either leaking internally, or being retained somewhere else. This can be translated to code as follows.

class WeakRefLeakCheckTestCase: XCTestCase {

    private var weakReferences = NSMapTable<NSString, AnyObject>.weakToWeakObjects()

    // MARK: - SetUp

    override func setUpWithError() throws {

        try super.setUpWithError()

        // NOTE: Just before the test ends and executes the tearDown, create a weak reference for each object.
        // After the tearDown, check if those properties are `nil`. If they aren't, we have a leak candidate.
        // Ignore any non-reference type (e.g. structs and enums) during this stage.
        addTeardownBlock { [unowned self] in

            let mirror = Mirror(reflecting: self)
            mirror.children.forEach { label, wrappedValue in

                guard
                    let propertyName = label,
                    self.isSomeValue(wrappedValue)
                else { return }

                // NOTE: We need to unwrap the optional value to check its underlying type.
                let unwrappedValue = self.unwrap(wrappedValue)
                if Mirror(reflecting: unwrappedValue).displayStyle == .class {
                    self.weakReferences.setObject(unwrappedValue as AnyObject, forKey: propertyName as NSString)
                }
            }
        }
    }

    // MARK: - TearDown

    override func tearDownWithError() throws {

        // ...

        // check the weak reference we created before the tearDown
        weakReferences.keyEnumerator().allObjects.forEach { key in
            if let objectName = key as? NSString, weakReferences.object(forKey: objectName) != nil {
                XCTFail("🧮🚰 \"\(objectName)\" is a potential memory leak!")
            }
        }

        weakReferences.removeAllObjects()

        try super.tearDownWithError()
    }

    // MARK: - Utils

    private func unwrap<T>(_ any: T) -> Any {

        let mirror = Mirror(reflecting: any)
        guard mirror.displayStyle == .optional, let first = mirror.children.first else {
            return any
        }
        return unwrap(first.value)
    }

    private func isSomeValue(_ wrappedValue: Any) -> Bool {
        // ...
     }
}

# Distinguishing between constants and SUTs

Let’s look at the following test scenario which uses the NaiveReadableLeakCheckTestCase and a bunch of non-reference types (Int struct, custom struct, and custom enum) as well as constants.

import XCTest

fileprivate struct SomeStruct {
    let someValue: Int
}

fileprivate enum SomeEnum {
    case someCase
}

class NonClassWithTearDownTestCase: NaiveReadableLeakCheckTestCase {

    private let someConstantStruct = SomeStruct(someValue: 666)
    private var someStruct: SomeStruct!

    private let someConstantEnum: SomeEnum = .someCase
    private var someEnum: SomeEnum!

    private let someConstantInt: Int = 1337
    private var someInt: Int!

    override func setUp() {

        super.setUp()

        someStruct = SomeStruct(someValue: 666)
        someEnum = SomeEnum.someCase
        someInt = 1337
    }

    override func tearDown() {

        defer { super.tearDown() }

        someStruct = nil
        someEnum = nil
        someInt = nil
    }

    func test_WithVariousNonReferenceTypes_ShouldBeEqual() {

        XCTAssertEqual(someConstantEnum, .someCase)
        XCTAssertEqual(someEnum, .someCase)

        XCTAssertEqual(someConstantStruct.someValue, 666)
        XCTAssertEqual(someStruct.someValue, 666)

        XCTAssertEqual(someConstantInt, 1337)
        XCTAssertEqual(someInt, 1337)
    }
}

If you’ve been paying attention, then you already know that this TestCase would break our naive implementation with false-positives for the constant properties someConstantEnum, someConstantStruct, and someConstantInt. This is due to the fact that these properties never get reset during teardDown(). But the thing is, these ones aren’t supposed to be nil, they’re constants after all! What we need is some way to mark them as exceptions and skip them once its time to check for memory leaks.

Similarly to how we addressed the collection of weak references in a previous step, we should get by with another collection that keeps track of which properties are constants. And to check if a property is a constant all we have to do is use reflection and inspect it before we assign values during setUp(). If it has a value beforehand we’re just go ahead and assume that it’s a constant. This is another great opportunity to look at the lifecycle of a test case, and super.setUpWithError() is guaranteed to be executed before setUp(), so we just need to override it in our ConstantsLeakCheckTestCase.

To sum it up:

  • Create a collection to exclude properties from our memory leak checking endeavors.
  • Add to it properties that already have assigned values during setUpWithError().
  • Ignore those properties when checking for memory leaks during tearDownWithError().

All of this can be easily translated into code as such:

class ConstantsLeakCheckTestCase: XCTestCase {

    private var excludedProperties = [String: Any]()

    // MARK: - SetUp

    override func setUpWithError() throws {

        try super.setUpWithError()

        // NOTE: Before running the actual test, register which properties already have values
        // assigned to them. These should be constants which we don't need/want to setup every time
        // and can be excluded from the memory leak search.
        let mirror = Mirror(reflecting: self)
        mirror.children.forEach { label, wrappedValue in

            guard let propertyName = label else { return }

            if isSomeValue(wrappedValue) {
                excludedProperties[propertyName] = wrappedValue
            }
        }
    }

    // MARK: - TearDown

    override func tearDownWithError() throws {

        let mirror = Mirror(reflecting: self)
        mirror.children.forEach { label, wrappedValue in

            guard let propertyName = label else { return }

            // Ignore excluded properties
            guard excludedProperties[propertyName] == nil else { return }

            // NOTE: If the value is something, e.g. not `nil`, then we either are just not resetting the
            // value to `nil` during tearDown, or we have a potential leak.
            if isSomeValue(wrappedValue) {
                XCTFail("📄♻️ \"\(propertyName)\" is missing from tearDown, or is a potential memory leak!")
            }
        }

        excludedProperties.removeAll()

        try super.tearDownWithError()
    }

    // MARK: - Utils ...
}

# Dealing with lazy-loaded properties

Lastly, we should address the elephant in the room: lazy-loaded variables. These variables are only computed when they are needed and they can rely on closures to execute instructions at that point. Here’s a couple of lazy variable examples:

class SomeClass {

    lazy var someLazyConstant: Int = 1337

    lazy var someLazyVariable: Int = {
        someLazyConstant * 10
    }()

    lazy var someLazyVariableClosure: () -> Int = { [unowned self] in
        return self.someLazyConstant * 100
    }
}

The trouble with trying to evaluate lazy variables while checking for memory leaks is that, they don’t really hold a value until you use them, and we can’t just “open” them up and see if executing them might lead to a memory leak, because in doing so they can alter the test results. Thus, my proposed solution when tackling lazy variables and checking for memory leaks is to just ignore them and let them do their thing. If they do end up causing a memory leak, then hopefully our automated weak references should be able to catch it.

In order to determine if a property is a lazy variable, we can use reflection yet again, this time taking into consideration the internal representation that swift uses to distinguish lazy variables from other data types. It turns out that we just need to check the properties internal name and if it matches the $__lazy_storage_$_{property_name} form (official source code here), we have a lazy variable on our hands and we can just ignore it. We can achieve this with a simple guard clause during tearDownWithError().

class LazyVarLeakCheckTestCase: XCTestCase {

    // MARK: - SetUp ...

    // MARK: - TearDown

    override func tearDownWithError() throws {

        let mirror = Mirror(reflecting: self)
        mirror.children.forEach { label, wrappedValue in

            guard let propertyName = label else { return }

            // Lazy-loaded properties can present themselves as false-positives, so we filter them
            guard !propertyName.hasPrefix("$__lazy_storage_$_") else { return }

            // NOTE: If the value is something, e.g. not `nil`, then we either are just not resetting the
            // value to `nil` during tearDown, or we have a potential leak.
            if isSomeValue(wrappedValue) {
                XCTFail("📄♻️ \"\(propertyName)\" is missing from tearDown, or is a potential memory leak!")
            }
        }

        try super.tearDownWithError()
    }

    // MARK: - Utils ...
}

# Putting it all together

If you read all the way through this post, congratulations and thank you for sticking around, I put a lot of time and effort into breaking down all the improvements in a somewhat coherent and understandable way (at least I hope so 🤞). If you just want the copy-pastable snippet for the final form of LeakCheckTestCase, you’ve come to the right place. Here it is:

# Final thoughts

A big thanks goes out to my colleague Filipe for idea-bouncing with me on how to automate memory leak detection within our project for the better part of the last two months. It was quite an insightful side-quest born from the absolute horror of trying to use Instrument’s Leaks in a large app with far too many use cases to count.

I’m sure there’s still other ways of improving this approach, and I’m especially not satisfied with the way we handle lazy variables: we’re just skipping them altogether and hoping that our weak references will be enough to catch any leaks afterwards 🙈. There’s bound to be other testing use cases which I just haven’t encountered and that will either result in more false-positives (incorrectly reported leaks) or false-negatives (undetected leaks that get through), so if you know of any I’d love to hear more!



# Sources