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 stat 12import sys 13import unittest 14import urllib.request 15from time import sleep 16from pathlib import Path 17from shutil import copyfileobj 18 19 20# Instances of this class must be declared as class level variables 21# starting with a name "ASSET_". This enables the pre-caching logic 22# to easily find all referenced assets and download them prior to 23# execution of the tests. 24class Asset: 25 26 def __init__(self, url, hashsum): 27 self.url = url 28 self.hash = hashsum 29 cache_dir_env = os.getenv('QEMU_TEST_CACHE_DIR') 30 if cache_dir_env: 31 self.cache_dir = Path(cache_dir_env, "download") 32 else: 33 self.cache_dir = Path(Path("~").expanduser(), 34 ".cache", "qemu", "download") 35 self.cache_file = Path(self.cache_dir, hashsum) 36 self.log = logging.getLogger('qemu-test') 37 38 def __repr__(self): 39 return "Asset: url=%s hash=%s cache=%s" % ( 40 self.url, self.hash, self.cache_file) 41 42 def __str__(self): 43 return str(self.cache_file) 44 45 def _check(self, cache_file): 46 if self.hash is None: 47 return True 48 if len(self.hash) == 64: 49 hl = hashlib.sha256() 50 elif len(self.hash) == 128: 51 hl = hashlib.sha512() 52 else: 53 raise Exception("unknown hash type") 54 55 # Calculate the hash of the file: 56 with open(cache_file, 'rb') as file: 57 while True: 58 chunk = file.read(1 << 20) 59 if not chunk: 60 break 61 hl.update(chunk) 62 63 return self.hash == hl.hexdigest() 64 65 def valid(self): 66 return self.cache_file.exists() and self._check(self.cache_file) 67 68 def _wait_for_other_download(self, tmp_cache_file): 69 # Another thread already seems to download the asset, so wait until 70 # it is done, while also checking the size to see whether it is stuck 71 try: 72 current_size = tmp_cache_file.stat().st_size 73 new_size = current_size 74 except: 75 if os.path.exists(self.cache_file): 76 return True 77 raise 78 waittime = lastchange = 600 79 while waittime > 0: 80 sleep(1) 81 waittime -= 1 82 try: 83 new_size = tmp_cache_file.stat().st_size 84 except: 85 if os.path.exists(self.cache_file): 86 return True 87 raise 88 if new_size != current_size: 89 lastchange = waittime 90 current_size = new_size 91 elif lastchange - waittime > 90: 92 return False 93 94 self.log.debug("Time out while waiting for %s!", tmp_cache_file) 95 raise 96 97 def fetch(self): 98 if not self.cache_dir.exists(): 99 self.cache_dir.mkdir(parents=True, exist_ok=True) 100 101 if self.valid(): 102 self.log.debug("Using cached asset %s for %s", 103 self.cache_file, self.url) 104 return str(self.cache_file) 105 106 if os.environ.get("QEMU_TEST_NO_DOWNLOAD", False): 107 raise Exception("Asset cache is invalid and downloads disabled") 108 109 self.log.info("Downloading %s to %s...", self.url, self.cache_file) 110 tmp_cache_file = self.cache_file.with_suffix(".download") 111 112 for retries in range(3): 113 try: 114 with tmp_cache_file.open("xb") as dst: 115 with urllib.request.urlopen(self.url) as resp: 116 copyfileobj(resp, dst) 117 break 118 except FileExistsError: 119 self.log.debug("%s already exists, " 120 "waiting for other thread to finish...", 121 tmp_cache_file) 122 if self._wait_for_other_download(tmp_cache_file): 123 return str(self.cache_file) 124 self.log.debug("%s seems to be stale, " 125 "deleting and retrying download...", 126 tmp_cache_file) 127 tmp_cache_file.unlink() 128 continue 129 except Exception as e: 130 self.log.error("Unable to download %s: %s", self.url, e) 131 tmp_cache_file.unlink() 132 raise 133 134 try: 135 # Set these just for informational purposes 136 os.setxattr(str(tmp_cache_file), "user.qemu-asset-url", 137 self.url.encode('utf8')) 138 os.setxattr(str(tmp_cache_file), "user.qemu-asset-hash", 139 self.hash.encode('utf8')) 140 except Exception as e: 141 self.log.debug("Unable to set xattr on %s: %s", tmp_cache_file, e) 142 pass 143 144 if not self._check(tmp_cache_file): 145 tmp_cache_file.unlink() 146 raise Exception("Hash of %s does not match %s" % 147 (self.url, self.hash)) 148 tmp_cache_file.replace(self.cache_file) 149 # Remove write perms to stop tests accidentally modifying them 150 os.chmod(self.cache_file, stat.S_IRUSR | stat.S_IRGRP) 151 152 self.log.info("Cached %s at %s" % (self.url, self.cache_file)) 153 return str(self.cache_file) 154 155 def precache_test(test): 156 log = logging.getLogger('qemu-test') 157 log.setLevel(logging.DEBUG) 158 handler = logging.StreamHandler(sys.stdout) 159 handler.setLevel(logging.DEBUG) 160 formatter = logging.Formatter( 161 '%(asctime)s - %(name)s - %(levelname)s - %(message)s') 162 handler.setFormatter(formatter) 163 log.addHandler(handler) 164 for name, asset in vars(test.__class__).items(): 165 if name.startswith("ASSET_") and type(asset) == Asset: 166 log.info("Attempting to cache '%s'" % asset) 167 asset.fetch() 168 log.removeHandler(handler) 169 170 def precache_suite(suite): 171 for test in suite: 172 if isinstance(test, unittest.TestSuite): 173 Asset.precache_suite(test) 174 elif isinstance(test, unittest.TestCase): 175 Asset.precache_test(test) 176 177 def precache_suites(path, cacheTstamp): 178 loader = unittest.loader.defaultTestLoader 179 tests = loader.loadTestsFromNames([path], None) 180 181 with open(cacheTstamp, "w") as fh: 182 Asset.precache_suite(tests) 183