1 files changed, 12 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.