1 files changed, 6 insertions, 2 deletions
diff --git a/tests/test_vba.py b/tests/test_vba.py
index 13f1c2c..9c12cc1 100644
--- a/tests/test_vba.py
+++ b/tests/test_vba.py
@@ -6,15 +6,19 @@ import unittest
 
 import pokemontools.vba.vba as vba
 
+from pokemontools.vba.battle import (
+    Battle,
+    BattleException,
+)
+
 try:
     import pokemontools.vba.vba_autoplayer as autoplayer
 except ImportError:
     import pokemontools.vba.autoplayer as autoplayer
+autoplayer.vba = vba
 
 import pokemontools.vba.keyboard as keyboard
 
-autoplayer.vba = vba
-
 def setup_wram():
     """
     Loads up some default addresses. Should eventually be replaced with the