2 files changed, 24 insertions, 4 deletions
diff --git a/src/1.s b/src/1.s
index d4cffa2..4f4ae38 100644
--- a/src/1.s
+++ b/src/1.s
@@ -33,6 +33,8 @@
 .include "../include/globals.s"
 .include "../include/print.s"
 
+.include "../test/defines.s"
+
 ;; "Variables" used by this program.
 .scope Vars
     ;; After each iteration, this will hold the tenth of the number that has
@@ -99,8 +101,18 @@
     lda #%00011110
     sta PPU::MASK
 @loop:
+    ;; Skip everything if the `done` flag is set.
+    bit Globals::m_flags
+    bvs @end
+
     jsr compute_next
     jmp @loop
+
+@end:
+    ;; Run tests now that everything has been done.
+    CALL_TESTS_ON_DONE
+@halt:
+    jmp @halt
 .endproc
 
 ;; Iterate over the row pointed by `m_address` and compute its value. The tenth
@@ -153,10 +165,6 @@
 ;; address to the next row. This function will also set the `done` flag whenever
 ;; we reach the $ED control byte.
 .proc compute_next
-    ;; Skip everything if the `done` flag is set.
-    bit Globals::m_flags
-    bvs @end
-
     ;; Call `compute` and `accumulate_number`, which work in tandem in order to
     ;; fetch the value for the current row and add it into the `m_sum`
     ;; accumulator.
diff --git a/src/2.s b/src/2.s
index e4f0361..353cb16 100644
--- a/src/2.s
+++ b/src/2.s
@@ -38,6 +38,8 @@
 .include "../include/globals.s"
 .include "../include/print.s"
 
+.include "../test/defines.s"
+
 ;; "Variables" used by this program.
 .scope Vars
     ;; Maximum number for each color according to Part 1.
@@ -127,8 +129,18 @@
     sta PPU::MASK
 
 @loop:
+    ;; Skip everything if the `done` flag is set.
+    bit Globals::m_flags
+    bvs @end
+
     jsr compute_next
     jmp @loop
+
+@end:
+    ;; Run tests now that everything has been done.
+    CALL_TESTS_ON_DONE
+@halt:
+    jmp @halt
 .endproc
 
 ;; Compute the next game if needed.