|
1 | 1 | import json |
2 | 2 | import os |
| 3 | +import shutil |
| 4 | +from collections import namedtuple |
3 | 5 | from contextlib import contextmanager |
4 | 6 | from pathlib import Path, PurePath |
| 7 | +from zipfile import ZipFile |
5 | 8 |
|
6 | 9 | import pytest |
| 10 | +from jinja2 import Environment, FileSystemLoader |
7 | 11 | from packaging.requirements import Requirement |
8 | 12 |
|
9 | 13 | from hab import Resolver, Site |
@@ -111,6 +115,141 @@ def resolver(request): |
111 | 115 | return request.getfixturevalue(test_map[request.param]) |
112 | 116 |
|
113 | 117 |
|
| 118 | +Distro = namedtuple("Distro", ["name", "version", "inc_version", "distros"]) |
| 119 | + |
| 120 | + |
| 121 | +class DistroInfo(namedtuple("DistroInfo", ["root", "versions", "zip_root"])): |
| 122 | + default_versions = ( |
| 123 | + ("dist_a", "0.1", True, None), |
| 124 | + ("dist_a", "0.2", False, ["dist_b"]), |
| 125 | + ("dist_a", "1.0", False, None), |
| 126 | + ("dist_b", "0.5", False, None), |
| 127 | + ("dist_b", "0.6", False, None), |
| 128 | + ) |
| 129 | + |
| 130 | + @classmethod |
| 131 | + def dist_version(cls, distro, version): |
| 132 | + return f"{distro}_v{version}" |
| 133 | + |
| 134 | + @classmethod |
| 135 | + def hab_json(cls, distro, version=None, distros=None): |
| 136 | + data = {"name": distro} |
| 137 | + if version: |
| 138 | + data["version"] = version |
| 139 | + if distros: |
| 140 | + data["distros"] = distros |
| 141 | + return json.dumps(data, indent=4) |
| 142 | + |
| 143 | + @classmethod |
| 144 | + def generate(cls, root, versions=None, zip_created=None, zip_root=None): |
| 145 | + if versions is None: |
| 146 | + versions = cls.default_versions |
| 147 | + if zip_root is None: |
| 148 | + zip_root = root |
| 149 | + |
| 150 | + versions = {(x[0], x[1]): Distro(*x) for x in versions} |
| 151 | + |
| 152 | + for version in versions.values(): |
| 153 | + name = cls.dist_version(version.name, version.version) |
| 154 | + filename = root / f"{name}.zip" |
| 155 | + ver = version.version if version.inc_version else None |
| 156 | + with ZipFile(filename, "w") as zf: |
| 157 | + # Make the .zip file larger than the remotezip initial_buffer_size |
| 158 | + # so testing of partial archive reading is forced use multiple requests |
| 159 | + zf.writestr("data.txt", "-" * 64 * 1024) |
| 160 | + zf.writestr( |
| 161 | + ".hab.json", |
| 162 | + cls.hab_json(version.name, version=ver, distros=version.distros), |
| 163 | + ) |
| 164 | + zf.writestr("file_a.txt", "File A inside the distro.") |
| 165 | + zf.writestr("folder/file_b.txt", "File B inside the distro.") |
| 166 | + if zip_created: |
| 167 | + zip_created(zf) |
| 168 | + |
| 169 | + # Create a correctly named .zip file that doesn't have a .hab.json file |
| 170 | + # to test for .zip files that are not distros. |
| 171 | + with ZipFile(root / "not_valid_v0.1.zip", "w") as zf: |
| 172 | + zf.writestr("README.txt", "This file is not a hab distro zip.") |
| 173 | + |
| 174 | + return cls(root, versions, zip_root) |
| 175 | + |
| 176 | + |
| 177 | +@pytest.fixture(scope="session") |
| 178 | +def distro_finder_info(tmp_path_factory): |
| 179 | + """Returns a DistroInfo instance with extracted distros ready for hab. |
| 180 | +
|
| 181 | + This is useful for using an existing hab distro structure as your download server. |
| 182 | + """ |
| 183 | + root = tmp_path_factory.mktemp("_distro_finder") |
| 184 | + |
| 185 | + def zip_created(zf): |
| 186 | + """Extract all contents zip into a distro folder structure.""" |
| 187 | + filename = Path(zf.filename).stem |
| 188 | + distro, version = filename.split("_v") |
| 189 | + zf.extractall(root / distro / version) |
| 190 | + |
| 191 | + return DistroInfo.generate(root, zip_created=zip_created) |
| 192 | + |
| 193 | + |
| 194 | +@pytest.fixture(scope="session") |
| 195 | +def zip_distro(tmp_path_factory): |
| 196 | + """Returns a DistroInfo instance for a zip folder structure. |
| 197 | +
|
| 198 | + This is useful if the zip files are locally accessible or if your hab download |
| 199 | + server supports `HTTP range requests`_. For example if you are using Amazon S3. |
| 200 | +
|
| 201 | + .. _HTTP range requests: |
| 202 | + https://developer.mozilla.org/en-US/docs/Web/HTTP/Range_requests |
| 203 | + """ |
| 204 | + root = tmp_path_factory.mktemp("_zip_distro_files") |
| 205 | + return DistroInfo.generate(root) |
| 206 | + |
| 207 | + |
| 208 | +@pytest.fixture(scope="session") |
| 209 | +def zip_distro_sidecar(tmp_path_factory): |
| 210 | + """Returns a DistroInfo instance for a zip folder structure with sidecar |
| 211 | + `.hab.json` files. |
| 212 | +
|
| 213 | + This is useful when your hab download server does not support HTTP range requests. |
| 214 | + """ |
| 215 | + root = tmp_path_factory.mktemp("_zip_distro_sidecar_files") |
| 216 | + |
| 217 | + def zip_created(zf): |
| 218 | + """Extract the .hab.json from the zip to a sidecar file.""" |
| 219 | + filename = Path(zf.filename).stem |
| 220 | + sidecar = root / f"{filename}.hab.json" |
| 221 | + path = zf.extract(".hab.json", root) |
| 222 | + shutil.move(path, sidecar) |
| 223 | + |
| 224 | + return DistroInfo.generate(root, zip_created=zip_created) |
| 225 | + |
| 226 | + |
| 227 | +@pytest.fixture(scope="session") |
| 228 | +def _zip_distro_s3(tmp_path_factory): |
| 229 | + """The files used by `zip_distro_s3` only generated once per test.""" |
| 230 | + root = tmp_path_factory.mktemp("_zip_distro_s3_files") |
| 231 | + bucket_root = root / "hab-test-bucket" |
| 232 | + bucket_root.mkdir() |
| 233 | + return DistroInfo.generate(bucket_root, zip_root=root) |
| 234 | + |
| 235 | + |
| 236 | +@pytest.fixture() |
| 237 | +def zip_distro_s3(_zip_distro_s3, monkeypatch): |
| 238 | + """Returns a DistroInfo instance for a s3 zip cloud based folder structure. |
| 239 | +
|
| 240 | + This is used to simulate using an aws s3 cloud storage bucket to host hab |
| 241 | + distro zip files. |
| 242 | + """ |
| 243 | + from cloudpathlib import implementation_registry |
| 244 | + from cloudpathlib.local import LocalS3Client, local_s3_implementation |
| 245 | + |
| 246 | + from hab.distro_finders import s3_zip |
| 247 | + |
| 248 | + monkeypatch.setitem(implementation_registry, "s3", local_s3_implementation) |
| 249 | + monkeypatch.setattr(s3_zip, "S3Client", LocalS3Client) |
| 250 | + return _zip_distro_s3 |
| 251 | + |
| 252 | + |
114 | 253 | class Helpers(object): |
115 | 254 | """A collection of reusable functions that tests can use.""" |
116 | 255 |
|
@@ -204,6 +343,36 @@ def compare_files(generated, check): |
204 | 343 | cache[i] == check[i] |
205 | 344 | ), f"Difference on line: {i} between the generated cache and {generated}." |
206 | 345 |
|
| 346 | + @staticmethod |
| 347 | + def render_template(template, dest, **kwargs): |
| 348 | + """Render a jinja template in from the test templates directory. |
| 349 | +
|
| 350 | + Args: |
| 351 | + template (str): The name of the template file in the templates dir. |
| 352 | + dest (os.PathLike): The destination filename to write the output. |
| 353 | + **kwargs: All kwargs are used to render the template. |
| 354 | + """ |
| 355 | + environment = Environment( |
| 356 | + loader=FileSystemLoader(str(Path(__file__).parent / "templates")), |
| 357 | + trim_blocks=True, |
| 358 | + lstrip_blocks=True, |
| 359 | + ) |
| 360 | + template = environment.get_template(template) |
| 361 | + |
| 362 | + text = template.render(**kwargs).rstrip() + "\n" |
| 363 | + with dest.open("w") as fle: |
| 364 | + fle.write(text) |
| 365 | + |
| 366 | + @classmethod |
| 367 | + def render_resolver(cls, site_template, dest, **kwargs): |
| 368 | + """Calls `render_template` and constructs a Resolver instance for it.""" |
| 369 | + # Build the hab site |
| 370 | + site_file = dest / "site.json" |
| 371 | + cls.render_template(site_template, site_file, **kwargs) |
| 372 | + |
| 373 | + site = Site([site_file]) |
| 374 | + return Resolver(site) |
| 375 | + |
207 | 376 |
|
208 | 377 | @pytest.fixture |
209 | 378 | def helpers(): |
|
0 commit comments