xref: /openbmc/u-boot/test/py/u_boot_utils.py (revision e3963c09)
1# SPDX-License-Identifier: GPL-2.0
2# Copyright (c) 2016, NVIDIA CORPORATION. All rights reserved.
3
4# Utility code shared across multiple tests.
5
6import hashlib
7import inspect
8import os
9import os.path
10import pytest
11import sys
12import time
13import re
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 range(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 range(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[], or a string.
162            If a string, note that it is split up so that quoted spaces
163            will not be preserved. E.g. "fred and" becomes ['"fred', 'and"']
164        ignore_errors: Indicate whether to ignore errors. If True, the function
165            will simply return if the command cannot be executed or exits with
166            an error code, otherwise an exception will be raised if such
167            problems occur.
168
169    Returns:
170        The output as a string.
171    """
172    if isinstance(cmd, str):
173        cmd = cmd.split()
174    runner = u_boot_console.log.get_runner(cmd[0], sys.stdout)
175    output = runner.run(cmd, ignore_errors=ignore_errors)
176    runner.close()
177    return output
178
179def run_and_log_expect_exception(u_boot_console, cmd, retcode, msg):
180    """Run a command that is expected to fail.
181
182    This runs a command and checks that it fails with the expected return code
183    and exception method. If not, an exception is raised.
184
185    Args:
186        u_boot_console: A console connection to U-Boot.
187        cmd: The command to run, as an array of argv[].
188        retcode: Expected non-zero return code from the command.
189        msg: String that should be contained within the command's output.
190    """
191    try:
192        runner = u_boot_console.log.get_runner(cmd[0], sys.stdout)
193        runner.run(cmd)
194    except Exception as e:
195        assert(retcode == runner.exit_status)
196        assert(msg in runner.output)
197    else:
198        raise Exception("Expected an exception with retcode %d message '%s',"
199                        "but it was not raised" % (retcode, msg))
200    finally:
201        runner.close()
202
203ram_base = None
204def find_ram_base(u_boot_console):
205    """Find the running U-Boot's RAM location.
206
207    Probe the running U-Boot to determine the address of the first bank
208    of RAM. This is useful for tests that test reading/writing RAM, or
209    load/save files that aren't associated with some standard address
210    typically represented in an environment variable such as
211    ${kernel_addr_r}. The value is cached so that it only needs to be
212    actively read once.
213
214    Args:
215        u_boot_console: A console connection to U-Boot.
216
217    Returns:
218        The address of U-Boot's first RAM bank, as an integer.
219    """
220
221    global ram_base
222    if u_boot_console.config.buildconfig.get('config_cmd_bdi', 'n') != 'y':
223        pytest.skip('bdinfo command not supported')
224    if ram_base == -1:
225        pytest.skip('Previously failed to find RAM bank start')
226    if ram_base is not None:
227        return ram_base
228
229    with u_boot_console.log.section('find_ram_base'):
230        response = u_boot_console.run_command('bdinfo')
231        for l in response.split('\n'):
232            if '-> start' in l or 'memstart    =' in l:
233                ram_base = int(l.split('=')[1].strip(), 16)
234                break
235        if ram_base is None:
236            ram_base = -1
237            raise Exception('Failed to find RAM bank start in `bdinfo`')
238
239    # We don't want ram_base to be zero as some functions test if the given
240    # address is NULL (0). Let's add 2MiB then (size of an ARM LPAE/v8 section).
241
242    if ram_base == 0:
243        ram_base += 1024 * 1024 * 2
244
245    return ram_base
246
247class PersistentFileHelperCtxMgr(object):
248    """A context manager for Python's "with" statement, which ensures that any
249    generated file is deleted (and hence regenerated) if its mtime is older
250    than the mtime of the Python module which generated it, and gets an mtime
251    newer than the mtime of the Python module which generated after it is
252    generated. Objects of this type should be created by factory function
253    persistent_file_helper rather than directly."""
254
255    def __init__(self, log, filename):
256        """Initialize a new object.
257
258        Args:
259            log: The Logfile object to log to.
260            filename: The filename of the generated file.
261
262        Returns:
263            Nothing.
264        """
265
266        self.log = log
267        self.filename = filename
268
269    def __enter__(self):
270        frame = inspect.stack()[1]
271        module = inspect.getmodule(frame[0])
272        self.module_filename = module.__file__
273        self.module_timestamp = os.path.getmtime(self.module_filename)
274
275        if os.path.exists(self.filename):
276            filename_timestamp = os.path.getmtime(self.filename)
277            if filename_timestamp < self.module_timestamp:
278                self.log.action('Removing stale generated file ' +
279                    self.filename)
280                os.unlink(self.filename)
281
282    def __exit__(self, extype, value, traceback):
283        if extype:
284            try:
285                os.path.unlink(self.filename)
286            except:
287                pass
288            return
289        logged = False
290        for i in range(20):
291            filename_timestamp = os.path.getmtime(self.filename)
292            if filename_timestamp > self.module_timestamp:
293                break
294            if not logged:
295                self.log.action(
296                    'Waiting for generated file timestamp to increase')
297                logged = True
298            os.utime(self.filename)
299            time.sleep(0.1)
300
301def persistent_file_helper(u_boot_log, filename):
302    """Manage the timestamps and regeneration of a persistent generated
303    file. This function creates a context manager for Python's "with"
304    statement
305
306    Usage:
307        with persistent_file_helper(u_boot_console.log, filename):
308            code to generate the file, if it's missing.
309
310    Args:
311        u_boot_log: u_boot_console.log.
312        filename: The filename of the generated file.
313
314    Returns:
315        A context manager object.
316    """
317
318    return PersistentFileHelperCtxMgr(u_boot_log, filename)
319
320def crc32(u_boot_console, address, count):
321    """Helper function used to compute the CRC32 value of a section of RAM.
322
323    Args:
324        u_boot_console: A U-Boot console connection.
325        address: Address where data starts.
326        count: Amount of data to use for calculation.
327
328    Returns:
329        CRC32 value
330    """
331
332    bcfg = u_boot_console.config.buildconfig
333    has_cmd_crc32 = bcfg.get('config_cmd_crc32', 'n') == 'y'
334    assert has_cmd_crc32, 'Cannot compute crc32 without CONFIG_CMD_CRC32.'
335    output = u_boot_console.run_command('crc32 %08x %x' % (address, count))
336
337    m = re.search('==> ([0-9a-fA-F]{8})$', output)
338    assert m, 'CRC32 operation failed.'
339
340    return m.group(1)
341