Unit Testing Weak References / Retain Cycles in Swift

Unit Testing Weak References / Retain Cycles in Swift

Code that relies on memory tricks to work can be complicated, but there's a solid way of unit testing these cases. I see people wonder about this every once in a while, so I think it's a good moment to share the trick I use to achieve this.

Any sort of deallocation unit testing can be achieved by using autoreleasepool. I already wrote an article about it in the past, so if you're unfamiliar with this feature, here's a good place to start before reading this article.

In short, while in normal development a property's release is done at the autoreleasepool of the thread's RunLoop, you can create your own pools to have finer-grained control. For tests, pools can be used to validate weak references and test code that runs in a class' deinit.

Here's a simple example where we can see it in action. Let's assume that we want to test that a class' deinit is being called:

func testClassDeinitIsCalled() {
    let semaphore = DispatchSemaphore(value: 0)
    final class TestClass {
        let semaphore: DispatchSemaphore
        init(semaphore: DispatchSemaphore) { self.semaphore = semaphore }
        deinit {
            semaphore.signal()
        }
    }

    var cls: TestClass? = TestClass(semaphore: semaphore)
    cls = nil

    semaphore.wait()
}

This test will never succeed. While setting the property to nil does result in a cls.release() call by the compiler, the resources won't be freed until the main thread's current loop ends, which is going to be blocked by the test's execution. You could technically make this work by initializing the class in a different thread, but there's a much easier way!

func testClassDeinitIsCalled() {
    let semaphore = DispatchSemaphore(value: 0)
    final class TestClass {
        let semaphore: DispatchSemaphore
        init(semaphore: DispatchSemaphore) { self.semaphore = semaphore }
        deinit {
            semaphore.signal()
        }
    }

    autoreleasepool {
        let cls = TestClass(semaphore: semaphore)
    }

    semaphore.wait()
}

The test now passes, because any resources created by cls will be freed when the pool ends, resulting in deinit being called.

My favorite use of this trick is to test that a particular component isn't causing a retain cycle. If we have two types that use each other, you can test for a retain cycle by initializing both and checking that eliminating one doesn't result in the other keeping it alive:

class TypeA {
    weak var typeB: TypeB?
}

class TypeB {
    var typeA: TypeA?
}

func testNoRetainCycle() {

    let typeA = TypeA()
    weak var _typeB: TypeB? = nil

    autoreleasepool {
        let typeB = TypeB()
        typeB.typeA = typeA
        typeA.typeB = typeB
        _typeB = typeB
    }

    XCTAssertNil(_typeB)
}

Try modifying this test yourself to see what happens! If you modify TypeA to cause a retain cycle, the test will fail as the test's weak reference will be unable to deallocate.

This trick can also be done in reverse. In this case, we're using a weak property to test that TypeA does keep TypeB alive:

func testKeepsValueAlive() {

    let typeA = TypeA()
    weak var _typeB: TypeB? = nil

    autoreleasepool {
        let typeB = TypeB()
        typeA.typeB = typeB
        _typeB = typeB
    }

    XCTAssertNotNil(_typeB)
}

If TypeA has a weak reference or disposes of TypeB after it's set, the test will fail.