1# Test utilities for fetching & caching assets 2# 3# Copyright 2024 Red Hat, Inc. 4# 5# This work is licensed under the terms of the GNU GPL, version 2 or 6# later. See the COPYING file in the top-level directory. 7 8import hashlib 9import logging 10import os 11import subprocess 12import sys 13import unittest 14import urllib.request 15from pathlib import Path 16from shutil import copyfileobj 17 18 19# Instances of this class must be declared as class level variables 20# starting with a name "ASSET_". This enables the pre-caching logic 21# to easily find all referenced assets and download them prior to 22# execution of the tests. 23class Asset: 24 25 def __init__(self, url, hashsum): 26 self.url = url 27 self.hash = hashsum 28 cache_dir_env = os.getenv('QEMU_TEST_CACHE_DIR') 29 if cache_dir_env: 30 self.cache_dir = Path(cache_dir_env, "download") 31 else: 32 self.cache_dir = Path(Path("~").expanduser(), 33 ".cache", "qemu", "download") 34 self.cache_file = Path(self.cache_dir, hashsum) 35 self.log = logging.getLogger('qemu-test') 36 37 def __repr__(self): 38 return "Asset: url=%s hash=%s cache=%s" % ( 39 self.url, self.hash, self.cache_file) 40 41 def _check(self, cache_file): 42 if self.hash is None: 43 return True 44 if len(self.hash) == 64: 45 sum_prog = 'sha256sum' 46 elif len(self.hash) == 128: 47 sum_prog = 'sha512sum' 48 else: 49 raise Exception("unknown hash type") 50 51 checksum = subprocess.check_output( 52 [sum_prog, str(cache_file)]).split()[0] 53 return self.hash == checksum.decode("utf-8") 54 55 def valid(self): 56 return self.cache_file.exists() and self._check(self.cache_file) 57 58 def fetch(self): 59 if not self.cache_dir.exists(): 60 self.cache_dir.mkdir(parents=True, exist_ok=True) 61 62 if self.valid(): 63 self.log.debug("Using cached asset %s for %s", 64 self.cache_file, self.url) 65 return str(self.cache_file) 66 67 if os.environ.get("QEMU_TEST_NO_DOWNLOAD", False): 68 raise Exception("Asset cache is invalid and downloads disabled") 69 70 self.log.info("Downloading %s to %s...", self.url, self.cache_file) 71 tmp_cache_file = self.cache_file.with_suffix(".download") 72 73 try: 74 resp = urllib.request.urlopen(self.url) 75 except Exception as e: 76 self.log.error("Unable to download %s: %s", self.url, e) 77 raise 78 79 try: 80 with tmp_cache_file.open("wb+") as dst: 81 copyfileobj(resp, dst) 82 except: 83 tmp_cache_file.unlink() 84 raise 85 try: 86 # Set these just for informational purposes 87 os.setxattr(str(tmp_cache_file), "user.qemu-asset-url", 88 self.url.encode('utf8')) 89 os.setxattr(str(tmp_cache_file), "user.qemu-asset-hash", 90 self.hash.encode('utf8')) 91 except Exception as e: 92 self.log.debug("Unable to set xattr on %s: %s", tmp_cache_file, e) 93 pass 94 95 if not self._check(tmp_cache_file): 96 tmp_cache_file.unlink() 97 raise Exception("Hash of %s does not match %s" % 98 (self.url, self.hash)) 99 tmp_cache_file.replace(self.cache_file) 100 101 self.log.info("Cached %s at %s" % (self.url, self.cache_file)) 102 return str(self.cache_file) 103 104 def precache_test(test): 105 log = logging.getLogger('qemu-test') 106 log.setLevel(logging.DEBUG) 107 handler = logging.StreamHandler(sys.stdout) 108 handler.setLevel(logging.DEBUG) 109 formatter = logging.Formatter( 110 '%(asctime)s - %(name)s - %(levelname)s - %(message)s') 111 handler.setFormatter(formatter) 112 log.addHandler(handler) 113 for name, asset in vars(test.__class__).items(): 114 if name.startswith("ASSET_") and type(asset) == Asset: 115 log.info("Attempting to cache '%s'" % asset) 116 asset.fetch() 117 log.removeHandler(handler) 118 119 def precache_suite(suite): 120 for test in suite: 121 if isinstance(test, unittest.TestSuite): 122 Asset.precache_suite(test) 123 elif isinstance(test, unittest.TestCase): 124 Asset.precache_test(test) 125 126 def precache_suites(path, cacheTstamp): 127 loader = unittest.loader.defaultTestLoader 128 tests = loader.loadTestsFromNames([path], None) 129 130 with open(cacheTstamp, "w") as fh: 131 Asset.precache_suite(tests) 132