diff --git a/Sources/CPU.swift b/Sources/CPU.swift
index 1bfc643..bd282f6 100644
--- a/Sources/CPU.swift
+++ b/Sources/CPU.swift
@@ -103,18 +103,20 @@ class CPU {
     }
 
     func run() {
-        run {}
+        run(onCycle: {}, onComplete: {})
     }
 
-    func run(callback: @escaping () -> ()) {
+    func run(onCycle: @escaping () -> (), onComplete: @escaping () -> ())  {
         let opcodes = OPCODES_MAP
         Timer.scheduledTimer(withTimeInterval: 0.00007, repeats: true) { [self] timer in
-            processOpcodes(callback: callback, opcodes: opcodes, timer: timer)
+            processOpcodes(onCycle: onCycle, opcodes: opcodes, timer: timer) {
+                onComplete()
+            }
         }
     }
 
-    func processOpcodes(callback: () -> (), opcodes: [UInt8: OpCode], timer: Timer) {
-        callback()
+    func processOpcodes(onCycle: () -> (), opcodes: [UInt8: OpCode], timer: Timer, onComplete: () -> ()) {
+        onCycle()
         let code = memRead(programCounter)
         programCounter += 1
 
@@ -278,9 +280,7 @@ class CPU {
         /// BRK
         case 0x00:
             timer.invalidate()
-            SDL_DestroyWindow(window)
-            SDL_Quit()
-            exit(0)
+            onComplete()
         default: fatalError("TODO!")
         }
 
diff --git a/Sources/main.swift b/Sources/main.swift
index 63be3b7..a5b29f7 100644
--- a/Sources/main.swift
+++ b/Sources/main.swift
@@ -124,7 +124,7 @@ cpu.reset()
 var screenState = [UInt8](repeating: 0, count: 32 * 3 * 32)
 var rng = SystemRandomNumberGenerator()
 
-cpu.run() {
+cpu.run(onCycle: {
     handleUserInput(cpu, event: &event)
     cpu.memWrite(0xfe, data: UInt8.random(in: 1...16, using: &rng))
 
@@ -133,7 +133,11 @@ cpu.run() {
         SDL_RenderCopy(canvas, texture, nil, nil)
         SDL_RenderPresent(canvas)
     }
-}
+}, onComplete: {
+    SDL_DestroyWindow(window)
+    SDL_Quit()
+    exit(0)
+})
 
 // Infinite loop otherwise the program will exit prematurely
 RunLoop.main.run()
diff --git a/Tests/CPU.swift b/Tests/CPU.swift
index 27cb284..b0eb7c1 100644
--- a/Tests/CPU.swift
+++ b/Tests/CPU.swift
@@ -3,8 +3,13 @@ import XCTest
 
 class CPUTests: XCTestCase {
     func test_0xa9_lda_immediate_load_data() {
+        let expect = XCTestExpectation()
         let cpu = CPU()
-        cpu.loadAndRun([0xa9, 0x05, 0x00])
+        cpu.load([0xa9, 0x05, 0x00])
+        cpu.reset()
+        cpu.run(onCycle: {}, onComplete: { expect.fulfill() })
+
+        wait(for: [expect], timeout: 5.0)
 
         XCTAssertEqual(cpu.register_a, 0x05)
         XCTAssert(cpu.status.rawValue & 0b0000_0010 == 0b00)
@@ -12,61 +17,88 @@ class CPUTests: XCTestCase {
     }
 
     func test_lda_from_memory() {
+        let expect = XCTestExpectation()
         let cpu = CPU()
         cpu.memWrite(0x10, data: 0x55)
 
-        cpu.loadAndRun([0xa5, 0x10, 0x00])
+        cpu.load([0xa5, 0x10, 0x00])
+        cpu.reset()
+        cpu.run(onCycle: {}, onComplete: { expect.fulfill() })
+
+        wait(for: [expect], timeout: 5.0)
 
         XCTAssertEqual(cpu.register_a, 0x55)
     }
 
     func test_0xa9_lda_zero_flag() {
+        let expect = XCTestExpectation()
         let cpu = CPU()
-        cpu.loadAndRun([0xa9, 0x00, 0x00])
+        cpu.load([0xa9, 0x00, 0x00])
+        cpu.reset()
+        cpu.run(onCycle: {}, onComplete: { expect.fulfill() })
+
+        wait(for: [expect], timeout: 5.0)
+
         XCTAssert(cpu.status.rawValue & 0b0000_0010 == 0b10)
     }
 
     func test_0xa9_lda_neg_flag() {
+        let expect = XCTestExpectation()
         let cpu = CPU()
-        cpu.loadAndRun([0xa9, 0xFF, 0x00])
+        cpu.load([0xa9, 0xFF, 0x00])
+        cpu.reset()
+        cpu.run(onCycle: {}, onComplete: { expect.fulfill() })
+
+        wait(for: [expect], timeout: 5.0)
         XCTAssert(cpu.status.rawValue & 0b1000_0000 == 0b1000_0000)
     }
 
     func test_0xaa_tax_move_a_to_x() {
+        let expect = XCTestExpectation()
         let cpu = CPU()
 
         cpu.load([0xaa, 0x00])
         cpu.reset()
         cpu.register_a = 10
-        cpu.run()
+        cpu.run(onCycle: {}, onComplete: { expect.fulfill() })
+        wait(for: [expect], timeout: 5.0)
+
         XCTAssertEqual(cpu.register_x, 10)
     }
 
     func test_0xe8_inx_incriment_x() {
+        let expect = XCTestExpectation()
         let cpu = CPU()
 
         cpu.load([0xe8, 0x00])
         cpu.reset()
         cpu.register_x = 10
-        cpu.run()
+        cpu.run(onCycle: {}, onComplete: { expect.fulfill() })
+        wait(for: [expect], timeout: 5.0)
 
         XCTAssertEqual(cpu.register_x, 11)
     }
 
     func test_5_ops_working_together() {
+        let expect = XCTestExpectation()
         let cpu = CPU()
-        cpu.loadAndRun([0xa9, 0xc0, 0xaa, 0xe8, 0x00])
+        cpu.load([0xa9, 0xc0, 0xaa, 0xe8, 0x00])
+        cpu.reset()
+        cpu.run(onCycle: {}, onComplete: { expect.fulfill() })
+        wait(for: [expect], timeout: 5.0)
 
         XCTAssertEqual(cpu.register_x, 0xc1)
    }
 
     func test_inx_overflow() {
+        let expect = XCTestExpectation()
         let cpu = CPU()
 
         cpu.load([0xe8, 0xe8, 0x00])
         cpu.reset()
         cpu.register_x = 0xff
-        cpu.run()
+        cpu.run(onCycle: {}, onComplete: { expect.fulfill() })
+        wait(for: [expect], timeout: 5.0)
         XCTAssertEqual(cpu.register_x, 1)
     }
 }