Find potential memory leaks with automated tests

Memory handling and types of references

Photo by Jorge Ramirez on Unsplash

Working in our iOS apps can be very challenging in terms of memory handling, this means we need to be very careful with the object instances allocated in memory and guarantee their correct deallocation.

As you may know instances are deallocated from memory by the ARC or Automatic Reference Counting and there are two types of references on iOS: Strong and Weak.

Strong references are the default type of reference, let’s see a simple object declaration. Imagine that we have a class called Person and we declare an object as follows:

var person = Person()

In the previous declaration we are creating a strong reference to a Person object. Each time we a create strong reference we increase the object retain count by 1. What does that mean? Why is this useful? It’s useful because the ARC won’t remove an object from memory until its retain count is equal to 0.

Weak references on the other hand don’t increase the object retain count. Let’s see an example:

var person = Person()
weak var anotherPerson: Person? = person

In the previous example we created a weak reference to the same person object and now we have the person strong reference and the anotherPerson weak reference. This means that even when we have two references to the same person object its retain count remains in 1 because one of them is a weak reference. Other important aspect of weak references is that they are optionals because we may access them in a time where they are already deallocated from memory.

Find retain cycles and memory leaks with testing

A retain cycle occurs when two objects have strong reference to each other. The result is that both objects retain themselves in memory because their retain count will never be 0. As developers we need to break this cycles so that the ARC will be able to free the space in memory occupied by those objects.

Imagine that we have two objects one of type HTTPClient which is a collaborator of an UserProfileLoader object.

class UserProfileLoader {
   private let client: HTTPClient
   private let url: URL
   init(url: URL, client: HTTPClient) {
      self.url = url
      self.client = client
   }
   enum Error: Swift.Error {
      case invalidProfileData
      case connectivity
   }
   enum Result: Equatable {
      case success(UserProfile)
      case failure(Error)
   }
   func loadProfile(completion: @escaping (Result) -> Void) {
      client.get(from: url) { data in
         let result = self.map(data: data)
         completion(result)
      }
   }
   func map(data: HTTPClientResult) -> Result {
      return .success(UserProfile(id: "an-user-id"))
   }
}

enum HTTPClientResult {
   case success(Data)
   case failure(Error)
}
protocol HTTPClient {
   func get(from url: URL, 
            completion: @escaping (HTTPClientResult) -> Void)
}
class HTTPClientSpy: HTTPClient {
   var completion: ((HTTPClientResult) -> Void)? = nil
   func get(from url: URL, 
            completion: @escaping (HTTPClientResult) -> Void) {
       
       self.completion = completion
   }
}

struct UserProfile: Equatable {
   var id: String
}

In the previous example our UserProfileLoader has a method called loadProfile that invokes the get method in the composed client of type HTTPClient. As you can see the implementation of get in our HTTPClientSpy stores a closure in a property called completion, this makes our closure accesible from our testings.

The above implementation contains a retain cycle for sure, but why and where? is not that obvious at a first glance isn’t it?, let’s take another look. The UserProfileLoader has a strong reference to the HTTPClient object, where? on the client property, the HTTPClient implementation on the other hand has a strong reference to the UserProfileLoader as well, where? there isn’t any UserProfileLoader stored in our HTTPClientSpy you may say, but there’s one for sure. Let’s take another look to our get method, this method is storing the received closure in the completion property, this closure has a strong reference to our UserProfileLoader object, where? in the map method invocation. As you can see we added a strong reference with the explicit use of self.

These kinds of mistakes are usually found in a lot of codebases, sometimes they are not that obvious and we are not exempt of falling into them. We can reduce that chances by using automated tests, but how? We must be sure that SUT or System Under Test and its collaborators are properly deallocated.

Let’s start by writing out first test, this test will cover the successful path. In this example our SUT will be our UserProfileLoader instance.

func test_load_withSuccesfulClientResponseDeliversSuccesLoadingResult() {
   let url = URL(string: "https://a-url.com")!
   let client = HTTPClientSpy()
   let sut = UserProfileLoader(url: url, client: client)
}

We have our givens, now we have to store our loading result in some place to make some assertions, as you remember the loadProfile method will complete with a closure of type (Result) -> Void. Let’s store that result in our test.

func test_load_withSuccesfulClientResponseDeliversSuccesLoadingResult() {
   let url = URL(string: "https://a-url.com")!
   let client = HTTPClientSpy()
   let sut = UserProfileLoader(url: url, client: client)
   var capturedResults = [UserProfileLoader.Result]()
   sut.loadProfile { capturedResults.append($0) }
}

In our previous test the load closure will never be executed, why? Because our loadProfile is calling the implementation of get, do you remember what does the implementation of get do? In our HTTPClientSpy instance the get method stores the received closure in a property. We must call that stored completion block in the test to get the result.

func test_load_withSuccesfulClientResponseDeliversSuccesLoadingResult() {
   let url = URL(string: "https://a-url.com")!
   let client = HTTPClientSpy()
   let sut = UserProfileLoader(url: url, client: client)
   var capturedResults = [UserProfileLoader.Result]()
   sut.loadProfile { capturedResults.append($0) }
   client.completion?(.success(Data()))
}

Now we can start to make some assertions with the capturedResultsarray. Our loadProfile method will always complete with a successful result containing and instance of UserProfile. That’s what we must expect in our test then.

class MemoryLeaksTests: XCTestCase {
   func test_load_withSuccesfulCallDeliversSuccesLoadingResult() {
      let url = URL(string: "https://a-url.com")!
      let client = HTTPClientSpy()
      let sut = UserProfileLoader(url: url, client: client)
      let expectedResult: UserProfileLoader.Result =
          .success(UserProfile(id: "an-user-id"))
      var capturedResults = [UserProfileLoader.Result]()
      sut.loadProfile { capturedResults.append($0) }
      client.completion?(.success(Data()))
      XCTAssertEqual(capturedResults, [expectedResult])
   }
}

Great! We have a passing test! Does that test ensures the correct deallocation of our SUT and its collaborators? Not yet. Let’s add some assertions to be sure that our test are deallocated from memory.

Our MemoryLeaksTests class extends from the XCTestCase class, the XCTestCase class provides us with a method called addTearDownBlock, this method receives a block of type () -> Void that will be executed each time a test function ends. We could add a tear down block right after the creation of our SUT and its collaborators.

class MemoryLeaksTests: XCTestCase {
   func test_load_withSuccesfulCallDeliversSuccesLoadingResult() {
      let url = URL(string: "https://a-url.com")!
      let client = HTTPClientSpy()
      let sut = UserProfileLoader(url: url, client: client)
      addTeardownBlock { [weak sut, weak client] in
         XCTAssertNil(client, "Instance has not been deallocated")
         XCTAssertNil(sut, "Instance has not been deallocated")
      }
      let expectedResult: UserProfileLoader.Result =
          .success(UserProfile(id: "an-user-id"))
      var capturedResults = [UserProfileLoader.Result]()
      sut.loadProfile { capturedResults.append($0) }
      client.completion?(.success(Data()))
      XCTAssertEqual(capturedResults, [expectedResult])
   }
}

Now we have our blocks in place. Notice that in our blocks we added weak references to the sut and client instances to avoid increasing their retain counts even more. Now our green test is not green anymore because the teardown blocks assertions are failing.

To remove the retain cycle we could avoid the explicit use of self in our loadProfile method. We have some options:

1. Use a weak reference in our closure and add guard let block to unwrap the result;
2. Turn our map method into a static method; or
3. Move the map method to a mapper class static method

For this example I’ll use the second option just to break the retain cycle and to have a passing test.

func loadProfile(completion: @escaping (Result) -> Void) {
   client.get(from: url) { data in
      completion(UserProfileLoader.map(data: data))
   }
}
static func map(data: HTTPClientResult) -> Result {
   return .success(UserProfile(id: "an-user-id"))
}

Now our test is passing but… it looks a little bit messy so let’s refactor our test to clean it a little bit. Imagine that we want to add more test functions(trust me we’ll want it 🙂), we would be creating our SUT and its collaborators and also adding the tear down blocks for each one of them, meaning that we would be repeating ourselves.

To refactor this test we will create some factory methods. The first one will be our makeSUT method, this method will return the SUT and its collaborators.

class MemoryLeaksTests: XCTestCase {
   func test_load_withSuccesfulCallDeliversSuccesLoadingResult() {
      let (sut, client) = 
          makeSUT(url: URL(string: "https://a-url.com")!)
      let expectedResult: UserProfileLoader.Result =
          .success(UserProfile(id: "an-user-id"))
      var capturedResults = [UserProfileLoader.Result]()
      sut.loadProfile { capturedResults.append($0) }
      client.completion?(.success(Data()))
      XCTAssertEqual(capturedResults, [expectedResult])
   }
   // MARK: - Helpers
   func makeSUT(url: URL,
                file: StaticString = #filePath,
                line: UInt = #line) -> (sut: UserProfileLoader, client: HTTPClientSpy) {
      let client = HTTPClientSpy()
      let sut = UserProfileLoader(url: url, client: client)
      addTeardownBlock { [weak sut, weak client] in
         XCTAssertNil(client, 
                      "Instance has not been deallocated", 
                      file: file, 
                      line: line)
         XCTAssertNil(sut, 
                      "Instance has not been deallocated", 
                      file: file, 
                      line: line)
      }
      return (sut, client)
   }
}

Now we have a cleaner test! Our test is more readable and we have another great advantage! Which one apart from cleaner test? you may say, If you take a look to the makeSUT factory method, this method is adding our tear down blocks, that means that every test that creates the SUT using this method will be automatically validating that the SUT and its collaborators are correctly deallocated from memory once the test is completed. If the SUT and its collaborators were not deallocated from memory the test will show the error on the makeSUT call because we are passing the file and line parameters to the XCTAssertNil.

Following this approach will give us the confidence to refactor and validating that we are doing a correct memory management.