xref: /qemu/tests/functional/qemu_test/utils.py (revision cfcb4484fc78cbbd835e2880add561e1fbef6796)
1850a1951SThomas Huth# Utilities for python-based QEMU tests
2850a1951SThomas Huth#
3850a1951SThomas Huth# Copyright 2024 Red Hat, Inc.
4850a1951SThomas Huth#
5850a1951SThomas Huth# Authors:
6850a1951SThomas Huth#  Thomas Huth <thuth@redhat.com>
7850a1951SThomas Huth#
8850a1951SThomas Huth# This work is licensed under the terms of the GNU GPL, version 2 or
9850a1951SThomas Huth# later.  See the COPYING file in the top-level directory.
10850a1951SThomas Huth
11d5674412SThomas Huthimport gzip
12e2e9fd25SThomas Huthimport lzma
13e2e9fd25SThomas Huthimport os
14e2e9fd25SThomas Huthimport shutil
15*cfcb4484SDaniel P. Berrangé
16*cfcb4484SDaniel P. Berrangéfrom .archive import tar_extract as archive_extract
17*cfcb4484SDaniel P. Berrangéfrom .archive import cpio_extract
18850a1951SThomas Huth
19f7d6b772SThomas Huth"""
20f7d6b772SThomas HuthRound up to next power of 2
21f7d6b772SThomas Huth"""
22f7d6b772SThomas Huthdef pow2ceil(x):
23f7d6b772SThomas Huth    return 1 if x == 0 else 2**(x - 1).bit_length()
24f7d6b772SThomas Huth
25f7d6b772SThomas Huthdef file_truncate(path, size):
26f7d6b772SThomas Huth    if size != os.path.getsize(path):
27f7d6b772SThomas Huth        with open(path, 'ab+') as fd:
28f7d6b772SThomas Huth            fd.truncate(size)
29f7d6b772SThomas Huth
30f7d6b772SThomas Huth"""
31f7d6b772SThomas HuthExpand file size to next power of 2
32f7d6b772SThomas Huth"""
33f7d6b772SThomas Huthdef image_pow2ceil_expand(path):
34f7d6b772SThomas Huth        size = os.path.getsize(path)
35f7d6b772SThomas Huth        size_aligned = pow2ceil(size)
36f7d6b772SThomas Huth        if size != size_aligned:
37f7d6b772SThomas Huth            with open(path, 'ab+') as fd:
38f7d6b772SThomas Huth                fd.truncate(size_aligned)
39f7d6b772SThomas Huth
40d5674412SThomas Huthdef gzip_uncompress(gz_path, output_path):
41d5674412SThomas Huth    if os.path.exists(output_path):
42d5674412SThomas Huth        return
43d5674412SThomas Huth    with gzip.open(gz_path, 'rb') as gz_in:
44d5674412SThomas Huth        try:
45d5674412SThomas Huth            with open(output_path, 'wb') as raw_out:
46d5674412SThomas Huth                shutil.copyfileobj(gz_in, raw_out)
47d5674412SThomas Huth        except:
48d5674412SThomas Huth            os.remove(output_path)
49d5674412SThomas Huth            raise
50d5674412SThomas Huth
51e2e9fd25SThomas Huthdef lzma_uncompress(xz_path, output_path):
52e2e9fd25SThomas Huth    if os.path.exists(output_path):
53e2e9fd25SThomas Huth        return
54e2e9fd25SThomas Huth    with lzma.open(xz_path, 'rb') as lzma_in:
55e2e9fd25SThomas Huth        try:
56e2e9fd25SThomas Huth            with open(output_path, 'wb') as raw_out:
57e2e9fd25SThomas Huth                shutil.copyfileobj(lzma_in, raw_out)
58e2e9fd25SThomas Huth        except:
59e2e9fd25SThomas Huth            os.remove(output_path)
60e2e9fd25SThomas Huth            raise
61