Easily test undo/redo in an AppKit app.

UndoAssertManager.swift

//
//  UndoAssertManager.swift
//  RetconTests
//
//  Created by Nathan Manceaux-Panot on 2022-10-23.
//

import Foundation
import AppKit

/// Used for testing undo/redo functionality. Records asserts for undo groups, then automatically performs undos and redos, running the asserts.
///
/// To use:
/// 1. In the test class's `setUp()` method, create a new `UndoAssertManager`, passing it the `UndoManager` your tested code will be talking to. Allow the `UndoAssertManager` configure the `UndoManager` (it sets `groupsByEvent` to false and opens an initial group).
/// 2. In your tests cases, call `endUndoGroupWithAsserts(_:)` once after you've set up the initial state, passing the initial state assert; and then whenever an undo group should end, passing it the assert for the state at that point.
/// 3. At the end of the test, call `testUndoRedo()` to execute the automatic undo/redo tests, running the recorded assertions.
/// 4. Optionally: in `tearDown()`, call `checkTestDidRun()`, which will fail if undo asserts have been recorded but `testRedoUndo()` hasn't been called.
class UndoAssertManager {
	typealias Assert = () -> Void
	
	private var undoManager: UndoManager
	
	/// An assert for the initial state, then one assert per undo point
	private var assertStack: [Assert] = []
	private var didRunTest = false
	
	/// Creates an `UndoAssertManager` that controls the specified `UndoManager`.
	///
	/// - Parameters:
	/// 	- undoManager: The `UndoManager` to use.
	/// 	- configureUndoManager:  If true, the undo manager is set to not group by event, and a new group is open.
	init(undoManager: UndoManager, configureUndoManager: Bool) {
		self.undoManager = undoManager
		
		if configureUndoManager {
			undoManager.groupsByEvent = false
			undoManager.beginUndoGrouping()
		}
	}
	
	/// Ends the current undo grouping, registers the assert to run whenever this undo state is reached, and runs the assert immediately.
	///
	/// - Parameters:
	/// 	- onlyRunOnReplay: If true, the assert is _not_ immediately executed. It will only be run when testing undo/redo.
	/// 	- assert: The block to run, which contains the asserts for this undo state.
	func endUndoGroupWithAsserts(onlyRunOnReplay: Bool = false, _ assert: @escaping Assert) {
		// Close the group
		// Even if this is the initial state assert, in case setting up the initial state ran undoable operations
		undoManager.endUndoGrouping()
		undoManager.beginUndoGrouping()
		
		// Register assert
		assertStack.append(assert)
		
		// Run assert immediately
		if !onlyRunOnReplay {
			assert()
		}
	}
	
	/// Undoes to the start, then redoes to the end, running registered asserts accordingly.
	func testUndoRedo() {
		assert(assertStack.count > 0, "No undo assert has been registered")
		assert(assertStack.count > 1, "Only the initial state assert has been registered")
		
		let lastAssertIndex = assertStack.count - 1
		var currentUndoPosition = lastAssertIndex
		
		// Close and undo last group, which should be empty
		undoManager.endUndoGrouping()
		undoManager.undo()
		
		// Test undo
		repeat {
			currentUndoPosition -= 1
			print("Undoing ← (now at \(currentUndoPosition)/\(lastAssertIndex))")
			
			undoManager.undo()
			assertStack[currentUndoPosition]()
		} while currentUndoPosition > 0
		
		// Test redo
		repeat {
			currentUndoPosition += 1
			print("Redoing → (now at \(currentUndoPosition)/\(lastAssertIndex))")
			
			undoManager.redo()
			assertStack[currentUndoPosition]()
		} while currentUndoPosition < lastAssertIndex
		
		// Remember the test was run
		didRunTest = true
	}
	
	/// If asserts were registered, checks that `testUndoRedo()` was called.
	func checkTestDidRun() {
		if assertStack.count > 0 && !didRunTest {
			assertionFailure("Did not run asserts (\(assertStack.count))")
		}
	}
}

Usage example:

func testLowercase() throws {
	let document = TextDocument()
	let undoAssertManager = UndoAssertManager(
		undoManager: document.undoManager!,
		configureUndoManager: true
	)
	
	// Setup
	document.text = "Hello, World!"
	
	undoAssertManager.endUndoGroupWithAsserts {
		XCTAssertEqual(document.text, "Hello, World!")
	}
	
	// Lowercase
	document.lowercaseText()
	
	undoAssertManager.endUndoGroupWithAsserts {
		XCTAssertEqual(document.text, "hello, world!")
	}
	
	// Test undo/redo
	undoAssertManager.testUndoRedo()
}

In this sample, the test first runs normally, executing both asserts inline.
When testUndoRedo() is called at the end, it:

• Undoes once, and runs the first assert again
• Redoes once, and runs the second assert again

This works with any number of asserts. For instance, if you have three asserts, the sequence will be:

• Initial run: run 1, run 2, run 3
• Undo/redo test: undo, run 2, undo, run 1; then redo, run 2, redo, run 3

If your asserts need to be different during the undo/redo test, you can register asserts without immediately running them, by passing true for onlyRunOnReplay when calling endUndoGroupWithAsserts().

If you do a lot of undo/redo testing, I recommend setting up your model and UndoAssertManager in your test class's setUp method, and to run checkTestDidRun() in your tearDown() method, to ensure you don't forget to call testUndoRedo().