Automatic Stubbing of Network Requests to Deflakify Automation Testing

Apple introduced automated UI testing in Xcode 7. This was a great addition for developers because this native support promised, among other things, an improvement in the flakiness notoriously associated with automation tests. As many of us developers have experienced, tests can sometimes fail even when there has been no modification to the test or underlying feature code.

Flakiness can often come from external factors. One such source of flakiness is live data, e.g. data loaded from a web service API. Data from the web service often drives the state of the UI. A test might assume the UI is in a certain state, e.g. the like button shows as “liked” or that a playlist has 10 songs. But when the data from the web service changes, the test fails.

Such occurrences are relatively common, and there are many other scenarios, such as network connectivity issues, in which these false negatives can occur. In these situations, the functionality of the app is not broken; it’s just that our assumptions about the data are no longer valid. What’s more, debugging the cause of these failures can be a time-consuming affair.

In this blog post, I’ll discuss how we at SoundCloud implemented a simple system for automatically recording and replaying the network responses that occur during the course of an Xcode UI automation test in order to remove any flakiness that may be introduced from the existence of live data coming from a web service.

It should be noted that third-party options such as DVR provide similar functionality. However, given that automatic stubbing of network requests can be achieved with a handful of classes and extensions, this post aims to showcase a lightweight approach that doesn’t require adding yet another dependency to your project.

An example project showcasing the approach taken can be found here. Open ArtistsFeature.swift for instructions.

Overview

The first step is to create an AutomationTestURLSession object that is capable of stubbing the app’s network responses with the contents of a JSON file, since this is the data format we expect from our web service.

The next step is to inject this session whenever tests are running.

Finally, we will implement a system for automatically recording and replaying network responses in order to create an easy-to-use API.

Network stubbing overview

The figure above shows how an AutomationTestURLSession is injected into the app in order to stub network responses using JSON files.

Session Stubbing

Let’s say we have a Client object our app uses to make all network requests. A simplified version might look like this:

class Client {
    let session: URLSession

    init(session: URLSession) {
        self.session = session
    }

    func makeRequest(_ request: URLRequest, completion: @escaping (Data?, Error?) -> Void) -> URLSessionDataTask {
        let task = session.dataTask(with: request) { data, response, error in
            // ...
        }
        task.resume()
        return task
    }
}

In order to replace/stub network responses with the contents of JSON files on disk, what we do is replace the app’s URLSession with our own AutomationTestsURLSession when tests are running. The AutomationTestsURLSession will then be responsible for stubbing the responses. To facilitate this, we start with a couple of protocols that our Client’s URLSession will conform to:

protocol URLSessionManaging: class {
    func dataTaskFromRequest(_ request: URLRequest, completionHandler: @escaping (Data?, URLResponse?, Error?) -> Void) -> URLSessionDataTasking
}
protocol URLSessionDataTasking: class {
    func resume()
    func suspend()
    func cancel()
}

With these protocols in place, we can easily make URLSession and URLSessionDataTask conform to it:

extension URLSession: URLSessionManaging {
    func dataTaskFromRequest(_ request: URLRequest, completionHandler: @escaping (Data?, URLResponse?, Error?) -> Void) -> URLSessionDataTasking {
        return dataTask(with: request, completionHandler: completionHandler) as URLSessionDataTasking
    }
}

extension URLSessionDataTask: URLSessionDataTasking {}

Now that we have URLSession and URLSessionDataTask conforming to our protocols, we can depend on the protocols from the Client instead of the concrete object types:

class Client {
    let session: URLSessionManaging

    init(session: URLSessionManaging = URLSession.shared) {
        self.session = session
    }

    func makeRequest(_ request: URLRequest, completion: @escaping (Result) -> Void) -> URLSessionDataTasking {
        let task = session.dataTaskFromRequest(request) { data, response, error in
			//...
        }
        task.resume()
        return task
    }
}

Now we can create our AutomationTestsURLSession to stub network responses with the contents of JSON files by making it conform to the same URLSessionManaging protocol:

class AutomationTestsURLSession: URLSessionManaging {

    func dataTaskFromRequest(_ request: URLRequest, completionHandler: @escaping (Data?, URLResponse?, Error?) -> Void) -> URLSessionDataTasking {
        return AutomationTestsDataTask(url: request.url!, completion: completionHandler)
    }
}

In the code above, AutomationTestsURLSession returns an instance of AutomationTestsDataTask, which conforms to our URLSessionDataTasking protocol instead of the standard URLSessionDataTask.

typealias DataCompletion = (Data?, URLResponse?, Error?) -> Void

class AutomationTestsDataTask: URLSessionDataTasking {
    private let request: URLRequest
    private let completion: DataCompletion

    init(request: URLRequest, completion: @escaping DataCompletion) {
        self.request = request
        self.completion = completion
    }

    func resume() {
        if let json = ProcessInfo.processInfo.environment[url.absoluteString] {
            let response = HTTPURLResponse(url: url, statusCode: 200, httpVersion: nil, headerFields: nil)
            let data = json.data(using: .utf8)
            completion(data, response, nil)
        }
    }

    func cancel() { }
    func suspend() { }
}

The resume method of our AutomationTestsDataTask is where the stubbing actually occurs.

In order to explain the first line of the resume method, it’s useful to know that Xcode automation tests run to the app in a separate process, which limits the ways in which the app and the tests can share data. One common way to pass data between the tests and the app is by using the launchEnvironment property on XCUIApplication. This is basically a dictionary of type [String: String]. Keys and values you set in tests will be available to the app at runtime via the environment property of ProcessInfo.processInfo.

Here we check to see if there is some JSON in ProcessInfo’s environment dictionary for the data task’s request, as identified by its URL. If we have some JSON for a given request, then we convert it to data and call completion immediately. This results in the stubbed data being returned to our application instead of the live data from the API that would usually be returned!

Now all we need is a way of detecting that automation tests are running so that we can swap out the session object the Client uses.

Session Injection during Automation Tests

At SoundCloud, we create a separate app delegate object to use when running automation tests, e.g. AutomationTestsApplicationDelegate. This can be a good place to set up application state that your tests rely on and/or clean up that state between test runs.

We can hook into this system to detect whether or not tests are running. First we need to set an environment variable whenever tests are running. We can use XCUIApplication’s launchEnvironment property to set a flag. Say we’re writing a test for a sign-in flow in our app. We could alter the implementation of setUp, like so:

class SignInFeature: XCTestCase {

    private var app: XCUIApplication!

    override func setUp() {
        super.setUp()
        app = XCUIApplication()
        app.launchEnvironment = ["IsAutomationTestsRunning": "YES"]
        app.launch()
    }

    func test_signIn() {

        // ... test code

    }
}

We can dynamically set an app delegate class by deleting the standard @UIApplicationDelegate annotation that is added to the default app delegate class that Xcode creates and adding a main.swift file instead:

import UIKit

let isAcceptanceTestsRunning = ProcessInfo.processInfo.environment["IsAutomationTestsRunning"] != nil
let appDelegateClass: AnyClass = isAcceptanceTestsRunning ? AutomationTestsApplicationDelegate.self : AppDelegate.self

let args = UnsafeMutableRawPointer(CommandLine.unsafeArgv)
.bindMemory(to: UnsafeMutablePointer<Int8>.self, capacity: Int(CommandLine.argc))
UIApplicationMain(CommandLine.argc, args, nil, NSStringFromClass(appDelegateClass))

The main thing we do here is use the isAutomationTestsRunning flag in order to decide which app delegate object to use.

We can use this to create a simple helper method on UIApplication to tell if tests are running:

extension UIApplication {

    static var isAutomationTest: Bool {
        guard let delegate = UIApplication.shared.delegate, delegate is AutomationTestsApplicationDelegate else {
            return false
        }
        return true
    }
}

Then all we need to do is create the correct session depending upon whether or not tests are running and use it wherever we usually create the client:

func makeClient() -> Client {
    let session: URLSessionManaging = UIApplication.isAutomationTest ? AutomationTestsURLSession() : URLSession.shared
    let client = Client(session: session)
    return client
}

With this, we already have the core of what we need to stub requests in our tests. However, using it as is would place a certain burden on developers; in order to stub requests, they first need to know what requests are made during a test, something which could potentially be a lot. Then they would need to manually make the requests, e.g. in Postman; save the responses to a file; and manually specify which stub should replace which request. The implementation currently lacks the API that would enable this final requirement, but since I’ve promised you automatic recording/replay functionality, I’ll skip straight to the good stuff!

Recording Stubs

What we’re aiming for is a simple API that takes the burden off developers and enables them to do network stubbing quickly. Our goal is to be able to call one API for recording, run our test, and if it succeeds, then switch to another API for replaying the recorded stubs — something like the record/replay functionality that is built into Xcode automation tests themselves.

To achieve this, we rely on an important mapping throughout: URLRequest -> Stubbed Response. We will basically use the URLRequest to construct a path on disk to the stubbed response. We can then use this path to save responses to disk as JSON files when we’re in “record mode” and load them back when we’re in “replay mode.”

Recording of responses during a test

The gif above shows stubs being recorded for the test_singIn test.

Environment Setup

Before we get started with the automation and application side setup, it’s worth spending a little time on the part that binds the two together. We will need to pass data between the two processes, and because the environment and launchEnvironment dictionaries of ProcessInfo and XCUIApplication, respectively, are just “stringly typed,” we can add some syntactic sugar on top to make the whole thing a bit more type-safe:

public enum EnvironmentKey: String {
    case stubbedTestName
    case networkStubsDir
    case recordMode
}

public enum RecordMode: String {
    case recording
    case replaying
}

final class Environment {

    private let processInfo = ProcessInfo.processInfo

    subscript(key: EnvironmentKey) -> String? {
        return processInfo.environment[key.rawValue]
    }
}

extension XCUIApplication {
    func setEnvironmentValue(_ environmentValue: String, forKey key: EnvironmentKey) {
        var launchEnv = launchEnvironment
        launchEnv[key.rawValue] = environmentValue
        launchEnvironment = launchEnv
    }
}

Here we basically define an EnvironmentKey enum that allows us to wrap the environment and launchEnvironment dictionaries of ProcessInfo and XCUIApplication in order to be able to pass data between the app and tests in a type-safe way. Along the way, we’ll discover what the various environment keys are used for, but most should be fairly self-explanatory.

Automation Tests Setup

We need to pass the name of the running test through to the app so that our AutomationTestsURLSession can construct the directory in which to record/replay stubs. We can benefit from the fact that each test in a given test case is essentially a new instance of the test case. This is how XCTest helps ensure that individual tests don’t interfere with each other. In fact, for an example test case called, say, SignInFeature, if you print self.description during a test, you will see something similar to the following:

-[SignInFeature test_signIn]

The output above reflects the test case name (SignInFeature) and the actual running test (test_signIn). This is perfect for our needs because it provides us with an easy way to associate a test name with the requests that are made during the test:

class SignInFeature: XCTestCase {

    private var app: XCUIApplication!

    override func setUp() {
        super.setUp()
        app = XCUIApplication()
        app.launchEnvironment = ["IsAutomationTestsRunning": "YES"]
        app.startRecordingStubs(for: self)
        app.launch()
    }

    override func tearDown() {
        app = nil
        super.tearDown()
    }


    func test_signIn() {
	    //...
    }
}

We implement startRecordingStubs as an extension on XCUIApplication, passing self so that we can use self.description to get a name for the test. Although relying on the output from self.description may seem brittle, its convenience was too good for this developer to ignore:

extension XCUIApplication {

    func startRecordingStubs(for testCase: XCTestCase, file: StaticString = #file, line: UInt = #line) {
        setupFileStubs(for: testCase, recordingMode: .recording, file: file, line: line)
    }

    private func setupFileStubs(for testCase: XCTestCase, recordingMode: RecordMode, file: StaticString = #file, line: UInt = #line) {
        let testName = String(describing: testCase)
        setEnvironmentValue(recordingMode.rawValue, forKey: .recordMode)
        setEnvironmentValue(testName, forKey: .stubbedTestName)
        if let networkStubsDir = networkStubsDirURL(file: file, line: line) {
            setEnvironmentValue(networkStubsDir.path, forKey: .networkStubsDir)
        }
    }
}

The code above basically sets test_signIn into “record mode” so that the app knows to start recording network requests during its execution. We’ll later use stubbedTestName and networkStubsDir EnvironmentKey cases from the app to construct the path to the stubs on disk. The .networkStubsDir key will point to the root directory where all stubs are stored. The .stubbedTestName key will point to the name of the currently running test, used to create the path to the directory for all stubbed requests for this test.

For the implementation of networkStubsDirURL, we take advantage of the default #file argument passed into every method, which gives the absolute path to the current file on disk. Using this, we can construct a new path for where we will record our network stubs, relative to some directory that we know will always be there, like AutomationTests/ in the example below:

private func networkStubsDirURL(file: StaticString, line: UInt = #line) -> URL? {
    let filePath = file.description
    guard let range = filePath.range(of: "AutomationTests/", options: [.literal]) else {
        XCTFail("AutomationTests directory does not exist!", file: file, line: line)
        return nil
    }
    let testsDirIndex = filePath.index(before: range.upperBound)
    let automationTestsPath = filePath[...testsDirIndex]
    return URL(fileURLWithPath: String(automationTestsPath)).appendingPathComponent("NetworkStubs", isDirectory: true)
}

App-Side Setup

In order to record the network responses in the app when the tests are running, we need to add a bit of functionality to our AutomationTestsURLSession:

class AutomationTestsURLSession: NSObject, URLSessionManaging {

    private let environment = Environment()
    private let _session = URLSession(configuration: .default)
    private let testStubsRecorder = TestStubsRecorder()

    func dataTaskFromRequest(_ request: URLRequest, completionHandler: @escaping CompletionBlock) -> URLSessionDataTasking {

        if environment.isRecordingStubs {
            let completion = testStubsRecorder.recordResult(of: request, with: completionHandler)
            return _session.dataTask(with: request, completionHandler: completion)
        }

        if shouldStub(request) {
            return AutomationTestsDataTask(request: request, completion: completionHandler)
        }

        return _session.dataTask(with: request, completionHandler: completionHandler)
    }


    private func shouldStub(_ request: URLRequest) -> Bool {
        return NetworkResponseStubStorage().stubExists(for: request)
    }

}

The logic of dataTaskFromRequest(_:completionHandler:) should be easy enough to read:

  • If we’re recording, record the stub and return the standard URLSessionDataTask to the application.
  • Otherwise, if we should stub the response, return an instance of our AutomationTestsDataTask, which will do the stubbing.
  • Otherwise, bypass the AutomationTestsURLSession by immediately returning a standard URLSessionDataTask to the application.

We’ve added three new private properties. The first is an instance of our Environment class. It has a new property, isRecordingStubs, which informs the AutomationTestsURLSession of whether or not it should record the responses from the network. This is based on the call we made back in setUp of our test case, i.e. app.startRecordingStubs(for: self):

var isRecordingStubs: Bool {
	guard let recordModelValue = self[.recordMode] else { return false }
	return recordModelValue == RecordMode.recording.rawValue
}

The second property, _session, is an internal URLSession property. We use this when we want the network response to complete as normal, i.e. in points 1 and 3 above.

The final property, testStubsRecorder, is an instance of a new class, TestStubsRecorder, which actually saves the responses to disk:

typealias NetworkResponse = (data: Data?, response: URLResponse?, error: Error?)
typealias CompletionBlock = (Data?, URLResponse?, Error?) -> Void

final class TestStubsRecorder {

    private let environment = Environment()
    private let fileManager = FileManager.default

    func recordResult(of request: URLRequest, with completionHandler: @escaping CompletionBlock) -> CompletionBlock {

        func completionWrapper(networkResponse: NetworkResponse) {
            record(networkResponse: networkResponse, for: request)
            completionHandler(networkResponse.data, networkResponse.response, networkResponse.error)
        }

        return completionWrapper
    }

    private func record(networkResponse: NetworkResponse, for request: URLRequest) {

        // handle errors

        guard let testName = environment[.stubbedTestName] else { return }

        guard let data = networkResponse.data else { return }

        // create dir for test case if it doesn't already exist
        createTestCaseDirIfRequired(forTestName: testName)

        // create and return path where stub JSON file will be saved (inside the test case dir)
        let stubPath = makeTestStubPath(forTestName: testName, request: request)

        fileManager.createFile(atPath: stubPath.path, contents: data, attributes: nil)
    }
}

The main method we’re interested in here is recordResult(of:with:), which our AutomationTestsURLSession calls whenever “record mode” is enabled. It uses a local function which wraps the passed completion block. This is the completion block that will return network responses to the application. We actually return this wrapped function and pass it as the completion parameter of the internal URLSession’s URLSessionDataTask that our AutomationTestsURLSession holds so that when the data task’s completion block gets called, our recordResult(of:with:) method runs first, thereby writing the results to disk. This method then internally calls the URLSessionDataTask’s completion block so that the data can also be returned to the app as normal. The implementations of createTestCaseDirIfRequired and makeTestStubPath are omitted for brevity.

One thing not explained so far is the AutomationTestsURLSession’s shouldStub method. This makes use of a NetworkResponseStubStorage object, which is a simple object that makes use of the Environment object and FileManager in order to construct the file paths to the stubbed responses on disk. If shouldStub returns true, we return an instance of our AutomationTestsDataTask, which actually does the stubbing:

final class NetworkResponseStubStorage {

    private let environment = Environment()
    private let fileManager = FileManager.default

    func stub(for request: URLRequest) -> Data? {
        // return stub
    }

    func stubExists(for request: URLRequest) -> Bool {
        // check if stub exists
    }
}

We can take advantage of the NetworkResponseStubStorage object just introduced to update the resume method of AutomationTestsDataTask:

func resume() {
    if let jsonData = NetworkResponseStubStorage().stub(for: request),
        let url = request.url {
        let response = HTTPURLResponse(url: url, statusCode: 200, httpVersion: nil, headerFields: nil)
        completion(jsonData, response, nil)
    }
}

You might notice that, for the sake of brevity, we are hardcoding the statusCode parameter of the response object. If your app explicitly handles these status codes, you may need to extend the solution to also record/replay the status code for a given request.

The only remaining thing to do is actually return the stubbed response during test runs. In order to do this, we switch the callback in the setUp method of SignInFeature from app.startRecordingStubs(for: self) to app.replayStubs(for: self). The implementation is a simple one-liner added to the extension on XCUIApplication where we implemented startRecordingStubs(for:):

func replayStubs(for testCase: XCTestCase, file: StaticString = #file, line: UInt = #line) {
    setupFileStubs(for: testCase, recordingMode: .replaying, file: file, line: line)
}

So there we have it. With a relatively small amount of code, we’ve been able to design a powerful feature that should go some way to alleviating flakiness in automated tests in a way that is easy to use and saves development time.

Conclusion

In this post, we first demonstrated how it’s possible to stub the responses to network requests made during an automation test in order to improve the stability of our tests with just a small amount of code. We then showed how the recording and replaying of these stubs can be automated using a simple API in order to remove the burden from the developer and enable this stubbing to be done quickly.

References