xref: /openbmc/u-boot/test/py/u_boot_utils.py (revision fbe502e9)
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 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[], 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    return ram_base
240
241class PersistentFileHelperCtxMgr(object):
242    """A context manager for Python's "with" statement, which ensures that any
243    generated file is deleted (and hence regenerated) if its mtime is older
244    than the mtime of the Python module which generated it, and gets an mtime
245    newer than the mtime of the Python module which generated after it is
246    generated. Objects of this type should be created by factory function
247    persistent_file_helper rather than directly."""
248
249    def __init__(self, log, filename):
250        """Initialize a new object.
251
252        Args:
253            log: The Logfile object to log to.
254            filename: The filename of the generated file.
255
256        Returns:
257            Nothing.
258        """
259
260        self.log = log
261        self.filename = filename
262
263    def __enter__(self):
264        frame = inspect.stack()[1]
265        module = inspect.getmodule(frame[0])
266        self.module_filename = module.__file__
267        self.module_timestamp = os.path.getmtime(self.module_filename)
268
269        if os.path.exists(self.filename):
270            filename_timestamp = os.path.getmtime(self.filename)
271            if filename_timestamp < self.module_timestamp:
272                self.log.action('Removing stale generated file ' +
273                    self.filename)
274                os.unlink(self.filename)
275
276    def __exit__(self, extype, value, traceback):
277        if extype:
278            try:
279                os.path.unlink(self.filename)
280            except:
281                pass
282            return
283        logged = False
284        for i in range(20):
285            filename_timestamp = os.path.getmtime(self.filename)
286            if filename_timestamp > self.module_timestamp:
287                break
288            if not logged:
289                self.log.action(
290                    'Waiting for generated file timestamp to increase')
291                logged = True
292            os.utime(self.filename)
293            time.sleep(0.1)
294
295def persistent_file_helper(u_boot_log, filename):
296    """Manage the timestamps and regeneration of a persistent generated
297    file. This function creates a context manager for Python's "with"
298    statement
299
300    Usage:
301        with persistent_file_helper(u_boot_console.log, filename):
302            code to generate the file, if it's missing.
303
304    Args:
305        u_boot_log: u_boot_console.log.
306        filename: The filename of the generated file.
307
308    Returns:
309        A context manager object.
310    """
311
312    return PersistentFileHelperCtxMgr(u_boot_log, filename)
313
314def crc32(u_boot_console, address, count):
315    """Helper function used to compute the CRC32 value of a section of RAM.
316
317    Args:
318        u_boot_console: A U-Boot console connection.
319        address: Address where data starts.
320        count: Amount of data to use for calculation.
321
322    Returns:
323        CRC32 value
324    """
325
326    bcfg = u_boot_console.config.buildconfig
327    has_cmd_crc32 = bcfg.get('config_cmd_crc32', 'n') == 'y'
328    assert has_cmd_crc32, 'Cannot compute crc32 without CONFIG_CMD_CRC32.'
329    output = u_boot_console.run_command('crc32 %08x %x' % (address, count))
330
331    m = re.search('==> ([0-9a-fA-F]{8})$', output)
332    assert m, 'CRC32 operation failed.'
333
334    return m.group(1)
335