diff --git a/hal_test.go b/hal_test.go
index edcf6f5..642320c 100644
--- a/hal_test.go
+++ b/hal_test.go
@@ -25,12 +25,14 @@ func TestAddition(t *testing.T) {
 	module, err := hal.NewHALModule(program, 256, 2, false)
 	assert.NoError(t, err)
 
-	module.IO[0] = 10
+	module.IO[0] = hal.NewMockIO(10, 10)
+	outputMock := hal.NewMockIO()
+	module.IO[1] = outputMock
 
 	err = module.Run()
 	assert.NoError(t, err)
 
-	assert.Equal(t, float64(20), module.IO[1])
+	assert.Equal(t, float64(20), outputMock.Outputs[0])
 }
 
 func TestMax(t *testing.T) {
@@ -104,7 +106,7 @@ func TestNewton1(t *testing.T) {
 		"64 SUB 40",
 		"65 JUMP 03",
 		"66 LOAD 10",
-		"67 OUT 1",
+		"67 OUT 0",
 		"68 STOP",
 	}
 
@@ -138,7 +140,7 @@ func TestNewton1(t *testing.T) {
 		return x
 	}
 
-	assert.Equal(t, calculateNewton(1), module.IO[1])
+	assert.Equal(t, calculateNewton(1), outputMock.Outputs[0])
 }
 
 func TestNewton2(t *testing.T) {
@@ -191,16 +193,19 @@ func TestNewton2(t *testing.T) {
 		"54 JUMP 03",
 
 		"60 LOAD 10", // Load the result
-		"61 OUT 1",   // Print it
+		"61 OUT 0",   // Print it
 		"62 STOP",
 	}
 
 	program, err := parser.ParseProgram(input)
 	assert.NoError(t, err)
 
-	module, err := hal.NewHALModule(program, 256, 2, false)
+	module, err := hal.NewHALModule(program, 256, 1, false)
 	assert.NoError(t, err)
 
+	outputMock := hal.NewMockIO()
+	module.IO[0] = outputMock
+
 	err = module.Run()
 	assert.NoError(t, err)
 
@@ -222,5 +227,5 @@ func TestNewton2(t *testing.T) {
 		return x
 	}
 
-	assert.Equal(t, calculateNewton(1), module.IO[1])
+	assert.Equal(t, calculateNewton(1), outputMock.Outputs[0])
 }