Unit testing retain cycles

We all write retain cycles from time to time regardless of experience. Retain cycles are not always obvious to spot and can result in hours of debugging. There are of course great tools like the memory graph debugger available but debugging retain cycles can still be a painful and time consuming task.

The key thing to fixing retain cycles is detecting them. This post looks at some code you can incorporate into your unit tests to help with the discovery of retain cycles.

The basic idea

The ownership rules are fairly simple with ARC. An object will only be kept alive for as long as there is at least one strong reference to it.

class NoisyDeinit {
    deinit {
        print("I'm in deinit")
    }
}

var example: NoisyDeinit? = .init()

// Setting `example` (the only reference to the `NoisyDeinit` instance) to `nil`
// causes the instance to be deallocated and it's `deinit` will be invoked.
example = nil

Equally we know that a weak reference to an object will be nil‘d out when the last strong reference is released.

var example: NoisyDeinit? = .init()
weak var weakReference = example

assert(weakReference != nil)

// Setting `example` (the only reference to the `NoisyDeinit` instance) to `nil`
// also causes `weakReference` to be set to `nil`.
example = nil
assert(weakReference == nil)

Knowing the above we can write our tests in such a way that we have both a strong and weak reference to our object under test. After we have finished exercising our object we can set the strong reference to nil and then verify that the weak reference is also nil. If the weak reference is not nil at this point then we have to figure out what is causing the object to stay alive (this could be a retain cycle).

Let’s see how this would look. Here is a unit test without cycle checking:

final class SampleTests: XCTestCase {    
    func testGreeting() {
        let sut = Greeter()
        XCTAssertEqual(sut.greet("Paul"), "Hello Paul")
    }
}

In order to add this new checking we need to add three lines per object we want to check and make our original reference both var and Optional:

With this done the code looks like the below:

1
2
3
4
5
6
7
8
9
10
11
final class SampleTests: XCTestCase {
    func testGreeting() {
        var sut: Greeter? = .init()
        weak var weakSut = sut

        XCTAssertEqual(sut?.greet("Paul"), "Hello Paul")

        sut = nil
        XCTAssertNil(weakSut)
    }
}

1. A new weak var to hold the object who’s lifecycle we want to verify (line 4)
2. nil‘ing out the strong reference (line 8)
3. The assertion that the new variable does become nil (line 9)

Can we simplify this?

Adding 3 lines per object is a little tedious and error prone. For example you may accidentally forget any one of these steps and the validation will no longer work.

We can write a couple of helper functions that we can add as an extension on XCTestCase that allow us to get this down to just one line per object who’s lifecycle we want to validate.

First let’s add a function that allows us to validate that an object is correctly deallocated after we execute an arbitrary block of caller provided code. This will be useful for scenarios where you have an instance property that is holding onto your object.

1
2
3
4
5
6
7
8
9
10
11
12
extension XCTestCase {
    func assertNil(_ subject: AnyObject?, after: @escaping () -> Void, file: StaticString = #file, line: UInt = #line) {
        guard let value = subject else {
            return XCTFail("Argument must not be nil", file: file, line: line)
        }

        addTeardownBlock { [weak value] in
            after()
            XCTAssert(value == nil, "Expected subject to be nil after test! Retain cycle?", file: file, line: line)
        }
    }
}

• Lines 3-5 perform a little bit of validation. It’s programmer error to set the assertion up when the object is already nil so we guard against that scenario
• Lines 7-9 are enqueuing a closure to be invoked after the test has been run
  • Line 7 is where our weak reference to our object is created
  • Line 8 is where we execute our arbitrary closure
  • Line 9 is where we perform the assertion that our weak reference is nil‘d out

When using our helper function our unit test above becomes:

final class SampleTests: XCTestCase {
    var sut: Greeter!

    override func setUp() {
        super.setUp()
        sut = Greeter()
        assertNil(sut, after: { self.sut = nil })
    }

    func testGreeting() {
        XCTAssertEqual(sut.greet("Paul"), "Hello Paul")
    }
}

In scenarios where we don’t have an instance property holding onto our object we can provide a simpler function. We’ll write it so that it just calls through to our helper above:

extension XCTestCase {
    func assertNilAfterTest(_ subject: AnyObject?, file: StaticString = #file, line: UInt = #line) {
        assertNil(subject, after: {}, file: file, line: line)
    }
}

The above works because if there is nothing holding onto our subject outside the scope of the test function it should be naturally cleaned up by the fact that the only strong reference has gone out of scope. This allows for an even simpler test:

final class SampleTests: XCTestCase {
    func testGreeting() {
        let sut = Greeter()
        assertNilAfterTest(sut)
        XCTAssertEqual(sut.greet("Paul"), "Hello Paul")
    }
}

Conclusion

The two helper functions above make for a simple API that should hopefully be useful for helping detect those painful retain cycles before they become a real problem. The code is hopefully simple enough to understand and doesn’t require modifying existing tests too heavily (no subclassing etc).

For usage I am tempted in my own projects to litter this cycle checking throughout most tests and not make contrived tests that just create an object and check it gets deallocated. By putting this logic on most tests I can get a level of comfort that I am not creating retain cycles whilst exercising the various functions of an object.