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 fetchable(self): 69 return not os.environ.get("QEMU_TEST_NO_DOWNLOAD", False) 70 71 def available(self): 72 return self.valid() or self.fetchable() 73 74 def _wait_for_other_download(self, tmp_cache_file): 75 # Another thread already seems to download the asset, so wait until 76 # it is done, while also checking the size to see whether it is stuck 77 try: 78 current_size = tmp_cache_file.stat().st_size 79 new_size = current_size 80 except: 81 if os.path.exists(self.cache_file): 82 return True 83 raise 84 waittime = lastchange = 600 85 while waittime > 0: 86 sleep(1) 87 waittime -= 1 88 try: 89 new_size = tmp_cache_file.stat().st_size 90 except: 91 if os.path.exists(self.cache_file): 92 return True 93 raise 94 if new_size != current_size: 95 lastchange = waittime 96 current_size = new_size 97 elif lastchange - waittime > 90: 98 return False 99 100 self.log.debug("Time out while waiting for %s!", tmp_cache_file) 101 raise 102 103 def fetch(self): 104 if not self.cache_dir.exists(): 105 self.cache_dir.mkdir(parents=True, exist_ok=True) 106 107 if self.valid(): 108 self.log.debug("Using cached asset %s for %s", 109 self.cache_file, self.url) 110 return str(self.cache_file) 111 112 if not self.fetchable(): 113 raise Exception("Asset cache is invalid and downloads disabled") 114 115 self.log.info("Downloading %s to %s...", self.url, self.cache_file) 116 tmp_cache_file = self.cache_file.with_suffix(".download") 117 118 for retries in range(3): 119 try: 120 with tmp_cache_file.open("xb") as dst: 121 with urllib.request.urlopen(self.url) as resp: 122 copyfileobj(resp, dst) 123 break 124 except FileExistsError: 125 self.log.debug("%s already exists, " 126 "waiting for other thread to finish...", 127 tmp_cache_file) 128 if self._wait_for_other_download(tmp_cache_file): 129 return str(self.cache_file) 130 self.log.debug("%s seems to be stale, " 131 "deleting and retrying download...", 132 tmp_cache_file) 133 tmp_cache_file.unlink() 134 continue 135 except Exception as e: 136 self.log.error("Unable to download %s: %s", self.url, e) 137 tmp_cache_file.unlink() 138 raise 139 140 try: 141 # Set these just for informational purposes 142 os.setxattr(str(tmp_cache_file), "user.qemu-asset-url", 143 self.url.encode('utf8')) 144 os.setxattr(str(tmp_cache_file), "user.qemu-asset-hash", 145 self.hash.encode('utf8')) 146 except Exception as e: 147 self.log.debug("Unable to set xattr on %s: %s", tmp_cache_file, e) 148 pass 149 150 if not self._check(tmp_cache_file): 151 tmp_cache_file.unlink() 152 raise Exception("Hash of %s does not match %s" % 153 (self.url, self.hash)) 154 tmp_cache_file.replace(self.cache_file) 155 # Remove write perms to stop tests accidentally modifying them 156 os.chmod(self.cache_file, stat.S_IRUSR | stat.S_IRGRP) 157 158 self.log.info("Cached %s at %s" % (self.url, self.cache_file)) 159 return str(self.cache_file) 160 161 def precache_test(test): 162 log = logging.getLogger('qemu-test') 163 log.setLevel(logging.DEBUG) 164 handler = logging.StreamHandler(sys.stdout) 165 handler.setLevel(logging.DEBUG) 166 formatter = logging.Formatter( 167 '%(asctime)s - %(name)s - %(levelname)s - %(message)s') 168 handler.setFormatter(formatter) 169 log.addHandler(handler) 170 for name, asset in vars(test.__class__).items(): 171 if name.startswith("ASSET_") and type(asset) == Asset: 172 log.info("Attempting to cache '%s'" % asset) 173 asset.fetch() 174 log.removeHandler(handler) 175 176 def precache_suite(suite): 177 for test in suite: 178 if isinstance(test, unittest.TestSuite): 179 Asset.precache_suite(test) 180 elif isinstance(test, unittest.TestCase): 181 Asset.precache_test(test) 182 183 def precache_suites(path, cacheTstamp): 184 loader = unittest.loader.defaultTestLoader 185 tests = loader.loadTestsFromNames([path], None) 186 187 with open(cacheTstamp, "w") as fh: 188 Asset.precache_suite(tests) 189