xref: /openbmc/u-boot/test/py/u_boot_utils.py (revision 729c2db7) 
1# Copyright (c) 2016, NVIDIA CORPORATION. All rights reserved.
2#
3# SPDX-License-Identifier: GPL-2.0
4
5# Utility code shared across multiple tests.
6
7import hashlib
8import os
9import os.path
10import pytest
11import sys
12import time
13import pytest
14
15def md5sum_data(data):
16    """Calculate the MD5 hash of some data.
17
18    Args:
19        data: The data to hash.
20
21    Returns:
22        The hash of the data, as a binary string.
23    """
24
25    h = hashlib.md5()
26    h.update(data)
27    return h.digest()
28
29def md5sum_file(fn, max_length=None):
30    """Calculate the MD5 hash of the contents of a file.
31
32    Args:
33        fn: The filename of the file to hash.
34        max_length: The number of bytes to hash. If the file has more
35            bytes than this, they will be ignored. If None or omitted, the
36            entire file will be hashed.
37
38    Returns:
39        The hash of the file content, as a binary string.
40    """
41
42    with open(fn, 'rb') as fh:
43        if max_length:
44            params = [max_length]
45        else:
46            params = []
47        data = fh.read(*params)
48    return md5sum_data(data)
49
50class PersistentRandomFile(object):
51    """Generate and store information about a persistent file containing
52    random data."""
53
54    def __init__(self, u_boot_console, fn, size):
55        """Create or process the persistent file.
56
57        If the file does not exist, it is generated.
58
59        If the file does exist, its content is hashed for later comparison.
60
61        These files are always located in the "persistent data directory" of
62        the current test run.
63
64        Args:
65            u_boot_console: A console connection to U-Boot.
66            fn: The filename (without path) to create.
67            size: The desired size of the file in bytes.
68
69        Returns:
70            Nothing.
71        """
72
73        self.fn = fn
74
75        self.abs_fn = u_boot_console.config.persistent_data_dir + '/' + fn
76
77        if os.path.exists(self.abs_fn):
78            u_boot_console.log.action('Persistent data file ' + self.abs_fn +
79                ' already exists')
80            self.content_hash = md5sum_file(self.abs_fn)
81        else:
82            u_boot_console.log.action('Generating ' + self.abs_fn +
83                ' (random, persistent, %d bytes)' % size)
84            data = os.urandom(size)
85            with open(self.abs_fn, 'wb') as fh:
86                fh.write(data)
87            self.content_hash = md5sum_data(data)
88
89def attempt_to_open_file(fn):
90    """Attempt to open a file, without throwing exceptions.
91
92    Any errors (exceptions) that occur during the attempt to open the file
93    are ignored. This is useful in order to test whether a file (in
94    particular, a device node) exists and can be successfully opened, in order
95    to poll for e.g. USB enumeration completion.
96
97    Args:
98        fn: The filename to attempt to open.
99
100    Returns:
101        An open file handle to the file, or None if the file could not be
102            opened.
103    """
104
105    try:
106        return open(fn, 'rb')
107    except:
108        return None
109
110def wait_until_open_succeeds(fn):
111    """Poll until a file can be opened, or a timeout occurs.
112
113    Continually attempt to open a file, and return when this succeeds, or
114    raise an exception after a timeout.
115
116    Args:
117        fn: The filename to attempt to open.
118
119    Returns:
120        An open file handle to the file.
121    """
122
123    for i in xrange(100):
124        fh = attempt_to_open_file(fn)
125        if fh:
126            return fh
127        time.sleep(0.1)
128    raise Exception('File could not be opened')
129
130def wait_until_file_open_fails(fn, ignore_errors):
131    """Poll until a file cannot be opened, or a timeout occurs.
132
133    Continually attempt to open a file, and return when this fails, or
134    raise an exception after a timeout.
135
136    Args:
137        fn: The filename to attempt to open.
138        ignore_errors: Indicate whether to ignore timeout errors. If True, the
139            function will simply return if a timeout occurs, otherwise an
140            exception will be raised.
141
142    Returns:
143        Nothing.
144    """
145
146    for i in xrange(100):
147        fh = attempt_to_open_file(fn)
148        if not fh:
149            return
150        fh.close()
151        time.sleep(0.1)
152    if ignore_errors:
153        return
154    raise Exception('File can still be opened')
155
156def run_and_log(u_boot_console, cmd, ignore_errors=False):
157    """Run a command and log its output.
158
159    Args:
160        u_boot_console: A console connection to U-Boot.
161        cmd: The command to run, as an array of argv[].
162        ignore_errors: Indicate whether to ignore errors. If True, the function
163            will simply return if the command cannot be executed or exits with
164            an error code, otherwise an exception will be raised if such
165            problems occur.
166
167    Returns:
168        The output as a string.
169    """
170
171    runner = u_boot_console.log.get_runner(cmd[0], sys.stdout)
172    output = runner.run(cmd, ignore_errors=ignore_errors)
173    runner.close()
174    return output
175
176def cmd(u_boot_console, cmd_str):
177    """Run a single command string and log its output.
178
179    Args:
180        u_boot_console: A console connection to U-Boot.
181        cmd: The command to run, as a string.
182
183    Returns:
184        The output as a string.
185    """
186    return run_and_log(u_boot_console, cmd_str.split())
187
188def run_and_log_expect_exception(u_boot_console, cmd, retcode, msg):
189    """Run a command which is expected to fail.
190
191    This runs a command and checks that it fails with the expected return code
192    and exception method. If not, an exception is raised.
193
194    Args:
195        u_boot_console: A console connection to U-Boot.
196        cmd: The command to run, as an array of argv[].
197        retcode: Expected non-zero return code from the command.
198        msg: String which should be contained within the command's output.
199    """
200    try:
201        runner = u_boot_console.log.get_runner(cmd[0], sys.stdout)
202        runner.run(cmd)
203    except Exception as e:
204        assert(msg in runner.output)
205    else:
206        raise Exception('Expected exception, but not raised')
207    finally:
208        runner.close()
209
210ram_base = None
211def find_ram_base(u_boot_console):
212    """Find the running U-Boot's RAM location.
213
214    Probe the running U-Boot to determine the address of the first bank
215    of RAM. This is useful for tests that test reading/writing RAM, or
216    load/save files that aren't associated with some standard address
217    typically represented in an environment variable such as
218    ${kernel_addr_r}. The value is cached so that it only needs to be
219    actively read once.
220
221    Args:
222        u_boot_console: A console connection to U-Boot.
223
224    Returns:
225        The address of U-Boot's first RAM bank, as an integer.
226    """
227
228    global ram_base
229    if u_boot_console.config.buildconfig.get('config_cmd_bdi', 'n') != 'y':
230        pytest.skip('bdinfo command not supported')
231    if ram_base == -1:
232        pytest.skip('Previously failed to find RAM bank start')
233    if ram_base is not None:
234        return ram_base
235
236    with u_boot_console.log.section('find_ram_base'):
237        response = u_boot_console.run_command('bdinfo')
238        for l in response.split('\n'):
239            if '-> start' in l or 'memstart    =' in l:
240                ram_base = int(l.split('=')[1].strip(), 16)
241                break
242        if ram_base is None:
243            ram_base = -1
244            raise Exception('Failed to find RAM bank start in `bdinfo`')
245
246    return ram_base
247