1# SPDX-License-Identifier: GPL-2.0
2# (C) Copyright 2018 Texas Instruments, <www.ti.com>
3
4# Test A/B update commands.
5
6import os
7import pytest
8import u_boot_utils
9
10class ABTestDiskImage(object):
11    """Disk Image used by the A/B tests."""
12
13    def __init__(self, u_boot_console):
14        """Initialize a new ABTestDiskImage object.
15
16        Args:
17            u_boot_console: A U-Boot console.
18
19        Returns:
20            Nothing.
21        """
22
23        filename = 'test_ab_disk_image.bin'
24
25        persistent = u_boot_console.config.persistent_data_dir + '/' + filename
26        self.path = u_boot_console.config.result_dir  + '/' + filename
27
28        with u_boot_utils.persistent_file_helper(u_boot_console.log, persistent):
29            if os.path.exists(persistent):
30                u_boot_console.log.action('Disk image file ' + persistent +
31                    ' already exists')
32            else:
33                u_boot_console.log.action('Generating ' + persistent)
34                fd = os.open(persistent, os.O_RDWR | os.O_CREAT)
35                os.ftruncate(fd, 524288)
36                os.close(fd)
37                cmd = ('sgdisk', persistent)
38                u_boot_utils.run_and_log(u_boot_console, cmd)
39
40                cmd = ('sgdisk', '--new=1:64:512', '--change-name=1:misc',
41                    persistent)
42                u_boot_utils.run_and_log(u_boot_console, cmd)
43                cmd = ('sgdisk', '--load-backup=' + persistent)
44                u_boot_utils.run_and_log(u_boot_console, cmd)
45
46        cmd = ('cp', persistent, self.path)
47        u_boot_utils.run_and_log(u_boot_console, cmd)
48
49di = None
50@pytest.fixture(scope='function')
51def ab_disk_image(u_boot_console):
52    global di
53    if not di:
54        di = ABTestDiskImage(u_boot_console)
55    return di
56
57@pytest.mark.boardspec('sandbox')
58@pytest.mark.buildconfigspec('android_ab')
59@pytest.mark.buildconfigspec('cmd_ab_select')
60@pytest.mark.requiredtool('sgdisk')
61def test_ab(ab_disk_image, u_boot_console):
62    """Test the 'ab_select' command."""
63
64    u_boot_console.run_command('host bind 0 ' + ab_disk_image.path)
65
66    output = u_boot_console.run_command('ab_select slot_name host 0#misc')
67    assert 're-initializing A/B metadata' in output
68    assert 'Attempting slot a, tries remaining 7' in output
69    output = u_boot_console.run_command('printenv slot_name')
70    assert 'slot_name=a' in output
71
72    output = u_boot_console.run_command('ab_select slot_name host 0:1')
73    assert 'Attempting slot b, tries remaining 7' in output
74    output = u_boot_console.run_command('printenv slot_name')
75    assert 'slot_name=b' in output
76