Automated detection of memory leaks in Unit Tests

LeakCheckTestCase.swift

class LeakCheckTestCase: XCTestCase {

    private var excludedProperties = [String: Any]()
    private var weakReferences = NSMapTable<NSString, AnyObject>.weakToWeakObjects()

    // 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
            }
        }

        // 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.excludedProperties[propertyName] == nil,
                    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 {

        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 }

            // 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!")
            }
        }

        // 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!")
            }
        }

        excludedProperties.removeAll()
        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 {

        // 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
    }
}