3 files changed, 83 insertions, 5 deletions
diff --git a/test/py/tests/test_android/test_ab.py b/test/py/tests/test_android/test_ab.py
new file mode 100644
index 00000000000..c79cb07fda3
--- /dev/null
+++ b/test/py/tests/test_android/test_ab.py
@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: GPL-2.0
+# (C) Copyright 2018 Texas Instruments, <www.ti.com>
+
+# Test A/B update commands.
+
+import os
+import pytest
+import u_boot_utils
+
+class ABTestDiskImage(object):
+    """Disk Image used by the A/B tests."""
+
+    def __init__(self, u_boot_console):
+        """Initialize a new ABTestDiskImage object.
+
+        Args:
+            u_boot_console: A U-Boot console.
+
+        Returns:
+            Nothing.
+        """
+
+        filename = 'test_ab_disk_image.bin'
+
+        persistent = u_boot_console.config.persistent_data_dir + '/' + filename
+        self.path = u_boot_console.config.result_dir  + '/' + filename
+
+        with u_boot_utils.persistent_file_helper(u_boot_console.log, persistent):
+            if os.path.exists(persistent):
+                u_boot_console.log.action('Disk image file ' + persistent +
+                    ' already exists')
+            else:
+                u_boot_console.log.action('Generating ' + persistent)
+                fd = os.open(persistent, os.O_RDWR | os.O_CREAT)
+                os.ftruncate(fd, 524288)
+                os.close(fd)
+                cmd = ('sgdisk', persistent)
+                u_boot_utils.run_and_log(u_boot_console, cmd)
+
+                cmd = ('sgdisk', '--new=1:64:512', '--change-name=1:misc',
+                    persistent)
+                u_boot_utils.run_and_log(u_boot_console, cmd)
+                cmd = ('sgdisk', '--load-backup=' + persistent)
+                u_boot_utils.run_and_log(u_boot_console, cmd)
+
+        cmd = ('cp', persistent, self.path)
+        u_boot_utils.run_and_log(u_boot_console, cmd)
+
+di = None
+@pytest.fixture(scope='function')
+def ab_disk_image(u_boot_console):
+    global di
+    if not di:
+        di = ABTestDiskImage(u_boot_console)
+    return di
+
+@pytest.mark.boardspec('sandbox')
+@pytest.mark.buildconfigspec('android_ab')
+@pytest.mark.buildconfigspec('cmd_ab_select')
+@pytest.mark.requiredtool('sgdisk')
+def test_ab(ab_disk_image, u_boot_console):
+    """Test the 'ab_select' command."""
+
+    u_boot_console.run_command('host bind 0 ' + ab_disk_image.path)
+
+    output = u_boot_console.run_command('ab_select slot_name host 0#misc')
+    assert 're-initializing A/B metadata' in output
+    assert 'Attempting slot a, tries remaining 7' in output
+    output = u_boot_console.run_command('printenv slot_name')
+    assert 'slot_name=a' in output
+
+    output = u_boot_console.run_command('ab_select slot_name host 0:1')
+    assert 'Attempting slot b, tries remaining 7' in output
+    output = u_boot_console.run_command('printenv slot_name')
+    assert 'slot_name=b' in output
diff --git a/test/py/tests/test_avb.py b/test/py/tests/test_avb.py
index 2bb75ed6e2a..81324234355 100644
--- a/test/py/tests/test_avb.py
+++ b/test/py/tests/test_avb.py
@@ -8,7 +8,7 @@
 This tests Android Verified Boot 2.0 support in U-boot:
 
 For additional details about how to build proper vbmeta partition
-check doc/README.avb2
+check doc/android/avb2.txt
 
 For configuration verification:
 - Corrupt boot partition and check for failure
diff --git a/test/py/tests/test_gpt.py b/test/py/tests/test_gpt.py
index c8edb336422..229d7eb2c21 100644
--- a/test/py/tests/test_gpt.py
+++ b/test/py/tests/test_gpt.py
@@ -40,16 +40,19 @@ class GptTestDiskImage(object):
                 fd = os.open(persistent, os.O_RDWR | os.O_CREAT)
                 os.ftruncate(fd, 4194304)
                 os.close(fd)
-                cmd = ('sgdisk', '-U', '375a56f7-d6c9-4e81-b5f0-09d41ca89efe',
+                cmd = ('sgdisk',
+                    '--disk-guid=375a56f7-d6c9-4e81-b5f0-09d41ca89efe',
                     persistent)
                 u_boot_utils.run_and_log(u_boot_console, cmd)
                 # part1 offset 1MB size 1MB
-                cmd = ('sgdisk', '--new=1:2048:4095', '-c 1:part1', persistent)
+                cmd = ('sgdisk', '--new=1:2048:4095', '--change-name=1:part1',
+                    persistent)
                 # part2 offset 2MB size 1.5MB
                 u_boot_utils.run_and_log(u_boot_console, cmd)
-                cmd = ('sgdisk', '--new=2:4096:7167', '-c 2:part2', persistent)
+                cmd = ('sgdisk', '--new=2:4096:7167', '--change-name=2:part2',
+                    persistent)
                 u_boot_utils.run_and_log(u_boot_console, cmd)
-                cmd = ('sgdisk', '-l', persistent)
+                cmd = ('sgdisk', '--load-backup=' + persistent)
                 u_boot_utils.run_and_log(u_boot_console, cmd)
 
         cmd = ('cp', persistent, self.path)