diff --git a/test/py/tests/test_extension.py b/test/py/tests/test_extension.py
index 2a3c511..d0840f7 100644
--- a/test/py/tests/test_extension.py
+++ b/test/py/tests/test_extension.py
@@ -13,43 +13,43 @@
 SANDBOX_DTB='arch/sandbox/dts/sandbox.dtb'
 OVERLAY_DIR='arch/sandbox/dts/'
 
-def load_dtb(u_boot_console):
-    u_boot_console.log.action('Loading devicetree to RAM...')
-    u_boot_console.run_command('host load hostfs - $fdt_addr_r %s' % (os.path.join(u_boot_console.config.build_dir, SANDBOX_DTB)))
-    u_boot_console.run_command('fdt addr $fdt_addr_r')
+def load_dtb(ubman):
+    ubman.log.action('Loading devicetree to RAM...')
+    ubman.run_command('host load hostfs - $fdt_addr_r %s' % (os.path.join(ubman.config.build_dir, SANDBOX_DTB)))
+    ubman.run_command('fdt addr $fdt_addr_r')
 
 @pytest.mark.buildconfigspec('cmd_fdt')
 @pytest.mark.boardspec('sandbox')
-def test_extension(u_boot_console):
+def test_extension(ubman):
     """Test the 'extension' command."""
 
-    load_dtb(u_boot_console)
+    load_dtb(ubman)
 
-    output = u_boot_console.run_command('extension list')
+    output = ubman.run_command('extension list')
     # extension_bootdev_hunt may have already run.
     # Without reboot we cannot make any assumption here.
     # assert('No extension' in output)
 
-    output = u_boot_console.run_command('extension scan')
+    output = ubman.run_command('extension scan')
     assert output == 'Found 2 extension board(s).'
 
-    output = u_boot_console.run_command('extension list')
+    output = ubman.run_command('extension list')
     assert('overlay0.dtbo' in output)
     assert('overlay1.dtbo' in output)
 
-    u_boot_console.run_command_list([
+    ubman.run_command_list([
         'setenv extension_overlay_addr %s' % (overlay_addr),
-        'setenv extension_overlay_cmd \'host load hostfs - ${extension_overlay_addr} %s${extension_overlay_name}\'' % (os.path.join(u_boot_console.config.build_dir, OVERLAY_DIR))])
+        'setenv extension_overlay_cmd \'host load hostfs - ${extension_overlay_addr} %s${extension_overlay_name}\'' % (os.path.join(ubman.config.build_dir, OVERLAY_DIR))])
 
-    output = u_boot_console.run_command('extension apply 0')
+    output = ubman.run_command('extension apply 0')
     assert('bytes read' in output)
 
-    output = u_boot_console.run_command('fdt print')
+    output = ubman.run_command('fdt print')
     assert('button3' in output)
 
-    output = u_boot_console.run_command('extension apply all')
+    output = ubman.run_command('extension apply all')
     assert('bytes read' in output)
 
-    output = u_boot_console.run_command('fdt print')
+    output = ubman.run_command('fdt print')
     assert('button4' in output)