1 files changed, 22 insertions, 6 deletions
diff --git a/src/utils/test-utils.ts b/src/utils/test-utils.ts
index 4b2efff..e620748 100644
--- a/src/utils/test-utils.ts
+++ b/src/utils/test-utils.ts
@@ -5,38 +5,54 @@ import { avrInstruction } from '../cpu/instruction';
 const BREAK_OPCODE = 0x9598;
 
 export function asmProgram(source: string) {
-  const { bytes, errors, lines } = assemble(source);
+  const { bytes, errors, lines, labels } = assemble(source);
   if (errors.length) {
     throw new Error('Assembly failed: ' + errors);
   }
-  return { program: new Uint16Array(bytes.buffer), lines, instructionCount: lines.length };
+  return { program: new Uint16Array(bytes.buffer), lines, instructionCount: lines.length, labels };
 }
 
+const defaultOnBreak = () => {
+  throw new Error('BREAK instruction encountered');
+};
+
 export class TestProgramRunner {
-  constructor(private readonly cpu: CPU, private readonly onBreak?: (cpu: CPU) => void) {}
+  constructor(
+    private readonly cpu: CPU,
+    private readonly onBreak: (cpu: CPU) => void = defaultOnBreak
+  ) {}
 
   runInstructions(count: number) {
     const { cpu, onBreak } = this;
     for (let i = 0; i < count; i++) {
       if (cpu.progMem[cpu.pc] === BREAK_OPCODE) {
         onBreak?.(cpu);
-        throw new Error('BREAK instruction encountered');
       }
       avrInstruction(cpu);
       cpu.tick();
     }
   }
 
-  runToBreak(maxIterations = 5000) {
+  runUntil(predicate: (cpu: CPU) => boolean, maxIterations = 5000) {
     const { cpu, onBreak } = this;
     for (let i = 0; i < maxIterations; i++) {
       if (cpu.progMem[cpu.pc] === BREAK_OPCODE) {
         onBreak?.(cpu);
+      }
+      if (predicate(cpu)) {
         return;
       }
       avrInstruction(cpu);
       cpu.tick();
     }
-    throw new Error('Program ran for too long without a BREAK instruction');
+    throw new Error('Test program ran for too long, check your predicate');
+  }
+
+  runToBreak() {
+    this.runUntil((cpu) => cpu.progMem[cpu.pc] === BREAK_OPCODE);
+  }
+
+  runToAddress(byteAddr: number) {
+    this.runUntil((cpu) => cpu.pc * 2 === byteAddr);
   }
 }