add UIViewController lifecycle test and dispose used viewController (pointfreeco#308)

* add UIViewController lifecycle test and dispose used viewController

* add #if os(iOS)

* Apply suggestions from code review

Co-authored-by: Stephen Celis <stephen.celis@gmail.com>

Co-authored-by: Stephen Celis <stephen.celis@gmail.com>
Co-authored-by: Stephen Celis <stephen@stephencelis.com>

Author: sidepelican

Sources/SnapshotTesting/Common/View.swift

@@ -670,7 +670,7 @@ func prepareView(
   traits: UITraitCollection,
   view: UIView,
   viewController: UIViewController
-  ) {
+  ) -> () -> Void {
   let size = config.size ?? viewController.view.frame.size
   view.frame.size = size
   if view != viewController.view {
@@ -691,14 +691,16 @@ func prepareView(
       viewController: viewController
     )
   }
-  add(traits: traits, viewController: viewController, to: window)
+  let dispose = add(traits: traits, viewController: viewController, to: window)
 
   if size.width == 0 || size.height == 0 {
     // Try to call sizeToFit() if the view still has invalid size
     view.sizeToFit()
     view.setNeedsLayout()
     view.layoutIfNeeded()
   }
+
+  return dispose
 }
 
 func snapshotView(
@@ -710,7 +712,7 @@ func snapshotView(
   )
   -> Async<UIImage> {
     let initialFrame = view.frame
-    prepareView(
+    let dispose = prepareView(
       config: config,
       drawHierarchyInKeyWindow: drawHierarchyInKeyWindow,
       traits: traits,
@@ -719,7 +721,8 @@ func snapshotView(
     )
     // NB: Avoid safe area influence.
     if config.safeArea == .zero { view.frame.origin = .init(x: offscreen, y: offscreen) }
-    return view.snapshot ?? Async { callback in
+
+    return (view.snapshot ?? Async { callback in
       addImagesForRenderedViews(view).sequence().run { views in
         callback(
           renderer(bounds: view.bounds, for: traits).image { ctx in
@@ -733,7 +736,7 @@ func snapshotView(
         views.forEach { $0.removeFromSuperview() }
         view.frame = initialFrame
       }
-    }
+    }).map { dispose(); return $0 }
 }
 
 private let offscreen: CGFloat = 10_000
@@ -748,7 +751,7 @@ func renderer(bounds: CGRect, for traits: UITraitCollection) -> UIGraphicsImageR
   return renderer
 }
 
-private func add(traits: UITraitCollection, viewController: UIViewController, to window: UIWindow) {
+private func add(traits: UITraitCollection, viewController: UIViewController, to window: UIWindow) -> () -> Void {
   let rootViewController: UIViewController
   if viewController != window.rootViewController {
     rootViewController = UIViewController()
@@ -786,6 +789,12 @@ private func add(traits: UITraitCollection, viewController: UIViewController, to
 
   viewController.view.setNeedsLayout()
   viewController.view.layoutIfNeeded()
+
+  return {
+    rootViewController.beginAppearanceTransition(false, animated: false)
+    rootViewController.endAppearanceTransition()
+    window.rootViewController = nil
+  }
 }
 
 private final class Window: UIWindow {

Sources/SnapshotTesting/Snapshotting/UIView.swift

@@ -47,13 +47,14 @@ extension Snapshotting where Value == UIView, Format == String {
     )
     -> Snapshotting<UIView, String> {
       return SimplySnapshotting.lines.pullback { view in
-        prepareView(
+        let dispose = prepareView(
           config: .init(safeArea: .zero, size: size ?? view.frame.size, traits: traits),
           drawHierarchyInKeyWindow: false,
           traits: .init(),
           view: view,
           viewController: .init()
         )
+        defer { dispose() }
         return purgePointers(
           view.perform(Selector(("recursiveDescription"))).retain().takeUnretainedValue()
             as! String

Sources/SnapshotTesting/Snapshotting/UIViewController.swift

@@ -64,13 +64,14 @@ extension Snapshotting where Value == UIViewController, Format == String {
   /// A snapshot strategy for comparing view controllers based on their embedded controller hierarchy.
   public static var hierarchy: Snapshotting {
     return Snapshotting<String, String>.lines.pullback { viewController in
-      prepareView(
+      let dispose = prepareView(
         config: .init(),
         drawHierarchyInKeyWindow: false,
         traits: .init(),
         view: viewController.view,
         viewController: viewController
       )
+      defer { dispose() }
       return purgePointers(
         viewController.perform(Selector(("_printHierarchy"))).retain().takeUnretainedValue() as! String
       )
@@ -95,13 +96,14 @@ extension Snapshotting where Value == UIViewController, Format == String {
     )
     -> Snapshotting<UIViewController, String> {
       return SimplySnapshotting.lines.pullback { viewController in
-        prepareView(
+        let dispose = prepareView(
           config: .init(safeArea: config.safeArea, size: size ?? config.size, traits: config.traits),
           drawHierarchyInKeyWindow: false,
           traits: traits,
           view: viewController.view,
           viewController: viewController
         )
+        defer { dispose() }
         return purgePointers(
           viewController.view.perform(Selector(("recursiveDescription"))).retain().takeUnretainedValue()
             as! String

Tests/SnapshotTestingTests/SnapshotTestingTests.swift

@@ -789,6 +789,82 @@ final class SnapshotTestingTests: XCTestCase {
     #endif
   }
 
+  func testUIViewControllerLifeCycle() {
+    #if os(iOS)
+    class ViewController: UIViewController {
+      let viewDidLoadExpectation: XCTestExpectation
+      let viewWillAppearExpectation: XCTestExpectation
+      let viewDidAppearExpectation: XCTestExpectation
+      let viewWillDisappearExpectation: XCTestExpectation
+      let viewDidDisappearExpectation: XCTestExpectation
+      init(viewDidLoadExpectation: XCTestExpectation,
+           viewWillAppearExpectation: XCTestExpectation,
+           viewDidAppearExpectation: XCTestExpectation,
+           viewWillDisappearExpectation: XCTestExpectation,
+           viewDidDisappearExpectation: XCTestExpectation){
+        self.viewDidLoadExpectation = viewDidLoadExpectation
+        self.viewWillAppearExpectation = viewWillAppearExpectation
+        self.viewDidAppearExpectation = viewDidAppearExpectation
+        self.viewWillDisappearExpectation = viewWillDisappearExpectation
+        self.viewDidDisappearExpectation = viewDidDisappearExpectation
+        super.init(nibName: nil, bundle: nil)
+      }
+      required init?(coder: NSCoder) {
+        fatalError("init(coder:) has not been implemented")
+      }
+      override func viewDidLoad() {
+        super.viewDidLoad()
+        viewDidLoadExpectation.fulfill()
+      }
+      override func viewWillAppear(_ animated: Bool) {
+        super.viewWillAppear(animated)
+        viewWillAppearExpectation.fulfill()
+      }
+      override func viewDidAppear(_ animated: Bool) {
+        super.viewDidAppear(animated)
+        viewDidAppearExpectation.fulfill()
+      }
+      override func viewWillDisappear(_ animated: Bool) {
+        super.viewWillDisappear(animated)
+        viewWillDisappearExpectation.fulfill()
+      }
+      override func viewDidDisappear(_ animated: Bool) {
+        super.viewDidDisappear(animated)
+        viewDidDisappearExpectation.fulfill()
+      }
+    }
+
+    let viewDidLoadExpectation = expectation(description: "viewDidLoad")
+    let viewWillAppearExpectation = expectation(description: "viewWillAppear")
+    let viewDidAppearExpectation = expectation(description: "viewDidAppear")
+    let viewWillDisappearExpectation = expectation(description: "viewWillDisappear")
+    let viewDidDisappearExpectation = expectation(description: "viewDidDisappear")
+    viewWillAppearExpectation.expectedFulfillmentCount = 4
+    viewDidAppearExpectation.expectedFulfillmentCount = 4
+    viewWillDisappearExpectation.expectedFulfillmentCount = 4
+    viewDidDisappearExpectation.expectedFulfillmentCount = 4
+
+    let viewController = ViewController(
+      viewDidLoadExpectation: viewDidLoadExpectation,
+      viewWillAppearExpectation: viewWillAppearExpectation,
+      viewDidAppearExpectation: viewDidAppearExpectation,
+      viewWillDisappearExpectation: viewWillDisappearExpectation,
+      viewDidDisappearExpectation: viewDidDisappearExpectation
+    )
+
+    assertSnapshot(matching: viewController, as: .image)
+    assertSnapshot(matching: viewController, as: .image)
+
+    wait(for: [
+      viewDidLoadExpectation,
+      viewWillAppearExpectation,
+      viewDidAppearExpectation,
+      viewWillDisappearExpectation,
+      viewDidDisappearExpectation,
+    ], timeout: 1.0, enforceOrder: true)
+    #endif
+  }
+
   func testCALayer() {
     #if os(iOS)
     let layer = CALayer()