paul-samuels.com

Thoughts on software development languages, tools and techniques.

@[email protected]

© Paul Samuels 2024. All rights reserved.

Test state restoration helper

20 Feb 2021

I’ve noticed that a fairly common thing to do in unit tests, especially when isolation isn’t perfect, is to update a value during a test and restore it to its original value after. This is to avoid leaking state between tests.

Given some poorly isolated code a test case might look something like this:

final class AccountViewControllerTests: XCTestCase {
    var initialUser: User?
    
    override func setUp() {
        super.setUp()
        
        initialUser = Current.user        // Capture the current state
        Current.user = .authenticatedUser // Set new state for the test
    }
    
    override func tearDown() {
        Current.user = initialUser        // Restore the original state after the test
        
        super.tearDown()
    }
    
    func testLogoutButtonIsVisible() {
        let accountViewController = AccountViewController()
        
        accountViewController.loadViewIfNeeded()
        accountViewController.viewWillAppear(true)
        
        XCTAssertFalse(accountViewController.logoutButton.isHidden)
    }
    
    // More tests
}

This is a lot of busy work and it’s easy to mess up. Thankfully we can add a small extension to handle the caching dance for us:

extension XCTestCase {
    func testCaseSet<T: AnyObject, U>(_ newValue: U, for keyPath: ReferenceWritableKeyPath<T, U>, on subject: T) {
        let intitial = subject[keyPath: keyPath]
        
        subject[keyPath: keyPath] = newValue
        
        addTeardownBlock {
            subject[keyPath: keyPath] = intitial
        }
    }
}

Using the above snippet we can more succinctly state that we want to upate a value for the duration of the test and don’t need to worry about how to actually do it.

final class AccountViewControllerWithExtensionTests: XCTestCase {
    override func setUp() {
        super.setUp()
        
        testCaseSet(.authenticatedUser, for: \.user, on: Current)
    }
    
    func testLogoutButtonIsVisible() {
        let accountViewController = AccountViewController()
        
        accountViewController.loadViewIfNeeded()
        accountViewController.viewWillAppear(true)
        
        XCTAssertFalse(accountViewController.logoutButton.isHidden)
    }
    
    // More tests
}

Side note

The original example is long winded and could be made more concise by using addTeardownBlock, which would look like this

final class AccountViewControllerTests: XCTestCase {
    override func setUp() {
        super.setUp()
        
        let initialUser = Current.user    // Capture the current state
        Current.user = .authenticatedUser // Set new state for the test
        
        addTeardownBlock {
            Current.user = initialUser    // Restore the original state after the test
        }
    }
        
    func testLogoutButtonIsVisible() {
        let accountViewController = AccountViewController()
        
        accountViewController.loadViewIfNeeded()
        accountViewController.viewWillAppear(true)
        
        XCTAssertFalse(accountViewController.logoutButton.isHidden)
    }
    
    // More tests
}

Although this is definitely an improvement it still means the author has to know the pattern and not make a mistake. Having a higher level abstraction removes the chances of messing things up and hopefully reduces the burden on future readers of the code.

Related Posts