4/4: Add testing for the new hab install and DistroFinder features

tests/conftest.py

@@ -1,9 +1,13 @@
 import json
 import os
+import shutil
+from collections import namedtuple
 from contextlib import contextmanager
 from pathlib import Path, PurePath
+from zipfile import ZipFile
 
 import pytest
+from jinja2 import Environment, FileSystemLoader
 from packaging.requirements import Requirement
 
 from hab import Resolver, Site
@@ -111,6 +115,110 @@ def resolver(request):
     return request.getfixturevalue(test_map[request.param])
 
 
+Distro = namedtuple("Distro", ["name", "version", "inc_version", "distros"])
+
+
+class DistroInfo(namedtuple("DistroInfo", ["root", "versions"])):
+    default_versions = (
+        ("dist_a", "0.1", True, None),
+        ("dist_a", "0.2", False, ["dist_b"]),
+        ("dist_a", "1.0", False, None),
+        ("dist_b", "0.5", False, None),
+        ("dist_b", "0.6", False, None),
+    )
+
+    @classmethod
+    def dist_version(cls, distro, version):
+        return f"{distro}_v{version}"
+
+    @classmethod
+    def hab_json(cls, distro, version=None, distros=None):
+        data = {"name": distro}
+        if version:
+            data["version"] = version
+        if distros:
+            data["distros"] = distros
+        return json.dumps(data, indent=4)
+
+    @classmethod
+    def generate(cls, root, versions=None, zip_created=None):
+        if versions is None:
+            versions = cls.default_versions
+
+        versions = {(x[0], x[1]): Distro(*x) for x in versions}
+
+        for version in versions.values():
+            name = cls.dist_version(version.name, version.version)
+            filename = root / f"{name}.zip"
+            ver = version.version if version.inc_version else None
+            with ZipFile(filename, "w") as zf:
+                zf.writestr(
+                    ".hab.json",
+                    cls.hab_json(version.name, version=ver, distros=version.distros),
+                )
+                zf.writestr("file_a.txt", "File A inside the distro.")
+                zf.writestr("folder/file_b.txt", "File B inside the distro.")
+                if zip_created:
+                    zip_created(zf)
+
+        # Create a correctly named .zip file that doesn't have a .hab.json file
+        # to test for .zip files that are not distros.
+        with ZipFile(root / "not_valid_v0.1.zip", "w") as zf:
+            zf.writestr("README.txt", "This file is not a hab distro zip.")
+
+        return cls(root, versions)
+
+
+@pytest.fixture(scope="session")
+def distro_finder_info(tmp_path_factory):
+    """Returns a DistroInfo instance with extracted distros ready for hab.
+
+    This is useful for using an existing hab distro structure as your download server.
+    """
+    root = tmp_path_factory.mktemp("_distro_finder")
+
+    def zip_created(zf):
+        """Extract all contents zip into a distro folder structure."""
+        filename = Path(zf.filename).stem
+        distro, version = filename.split("_v")
+        zf.extractall(root / distro / version)
+
+    return DistroInfo.generate(root, zip_created=zip_created)
+
+
+@pytest.fixture(scope="session")
+def zip_distro(tmp_path_factory):
+    """Returns a DistroInfo instance for a zip folder structure.
+
+    This is useful if the zip files are locally accessible or if your hab download
+    server supports `HTTP range requests`_. For example if you are using Amazon S3.
+
+    .. _HTTP range requests:
+       https://developer.mozilla.org/en-US/docs/Web/HTTP/Range_requests
+    """
+    root = tmp_path_factory.mktemp("_zip_distro")
+    return DistroInfo.generate(root)
+
+
+@pytest.fixture(scope="session")
+def zip_distro_sidecar(tmp_path_factory):
+    """Returns a DistroInfo instance for a zip folder structure with sidecar
+    `.hab.json` files.
+
+    This is useful when your hab download server does not support HTTP range requests.
+    """
+    root = tmp_path_factory.mktemp("_zip_distro_sidecar")
+
+    def zip_created(zf):
+        """Extract the .hab.json from the zip to a sidecar file."""
+        filename = Path(zf.filename).stem
+        sidecar = root / f"{filename}.hab.json"
+        path = zf.extract(".hab.json", root)
+        shutil.move(path, sidecar)
+
+    return DistroInfo.generate(root, zip_created=zip_created)
+
+
 class Helpers(object):
     """A collection of reusable functions that tests can use."""
 
@@ -204,6 +312,36 @@ def compare_files(generated, check):
                 cache[i] == check[i]
             ), f"Difference on line: {i} between the generated cache and {generated}."
 
+    @staticmethod
+    def render_template(template, dest, **kwargs):
+        """Render a jinja template in from the test templates directory.
+
+        Args:
+            template (str): The name of the template file in the templates dir.
+            dest (os.PathLike): The destination filename to write the output.
+            **kwargs: All kwargs are used to render the template.
+        """
+        environment = Environment(
+            loader=FileSystemLoader(str(Path(__file__).parent / "templates")),
+            trim_blocks=True,
+            lstrip_blocks=True,
+        )
+        template = environment.get_template(template)
+
+        text = template.render(**kwargs).rstrip() + "\n"
+        with dest.open("w") as fle:
+            fle.write(text)
+
+    @classmethod
+    def render_resolver(cls, site_template, dest, **kwargs):
+        """Calls `render_template` and constructs a Resolver instance for it."""
+        # Build the hab site
+        site_file = dest / "site.json"
+        cls.render_template(site_template, site_file, **kwargs)
+
+        site = Site([site_file])
+        return Resolver(site)
+
 
 @pytest.fixture
 def helpers():

tests/site/site_distro_finder.json

@@ -0,0 +1,32 @@
+{
+    "set":
+    {
+        "distro_paths":
+        [
+            [
+                "hab.distro_finders.distro_finder:DistroFinder",
+                "hab testable/download/path"
+            ],
+            [
+                "hab.distro_finders.distro_finder:DistroFinder",
+                "hab testing/downloads",
+                {
+                    "site": "for testing only, do not specify site"
+                }
+            ]
+        ],
+        "downloads":
+        {
+            "cache_root": "hab testable/download/path",
+            "distros":
+            [
+                [
+                    "hab.distro_finders.df_zip:DistroFinderZip",
+                    "network_server/distro/source"
+                ]
+            ],
+            "install_root": "{relative_root}/distros",
+            "relative_path": "{{distro_name}}_v{{version}}"
+        }
+    }
+}

tests/site/site_distro_finder_empty.json

@@ -0,0 +1,7 @@
+{
+    "set": {
+        "downloads": {
+            "cache_root": ""
+        }
+    }
+}

tests/templates/site_distro_finder.json

@@ -0,0 +1,20 @@
+{
+    "set": {
+        "config_paths": [
+            "{relative_root}/configs"
+        ],
+        "distro_paths": [
+            "{relative_root}/distros/*"
+        ],
+        "downloads": {
+            "cache_root": "{relative_root}/downloads",
+            "distros": [
+                [
+                    "hab.distro_finders.distro_finder:DistroFinder",
+                    "{{ zip_root }}/*"
+                ]
+            ],
+            "install_root": "{relative_root}/distros"
+        }
+    }
+}

tests/templates/site_distro_zip.json

@@ -0,0 +1,20 @@
+{
+    "set": {
+        "config_paths": [
+            "{relative_root}/configs"
+        ],
+        "distro_paths": [
+            "{relative_root}/distros/*"
+        ],
+        "downloads": {
+            "cache_root": "{relative_root}/downloads",
+            "distros": [
+                [
+                    "hab.distro_finders.df_zip:DistroFinderZip",
+                    "{{ zip_root }}"
+                ]
+            ],
+            "install_root": "{relative_root}/distros"
+        }
+    }
+}

tests/templates/site_distro_zip_sidecar.json

@@ -0,0 +1,20 @@
+{
+    "set": {
+        "config_paths": [
+            "{relative_root}/configs"
+        ],
+        "distro_paths": [
+            "{relative_root}/distros/*"
+        ],
+        "downloads": {
+            "cache_root": "{relative_root}/downloads",
+            "distros": [
+                [
+                    "hab.distro_finders.zip_sidecar:DistroFinderZipSidecar",
+                    "{{ zip_root }}"
+                ]
+            ],
+            "install_root": "{relative_root}/distros"
+        }
+    }
+}

tests/templates/site_download.json

@@ -0,0 +1,20 @@
+{
+    "set": {
+        "config_paths": [
+            "{relative_root}/configs"
+        ],
+        "distro_paths": [
+            "{relative_root}/distros/*"
+        ],
+        "downloads": {
+            "cache_root": "{relative_root}/downloads",
+            "distros": [
+                [
+                    "hab.distro_finders.df_zip:DistroFinderZip",
+                    "{{ zip_root }}"
+                ]
+            ],
+            "install_root": "{relative_root}/distros"
+        }
+    }
+}