Repository: NucciTheBoss/cleantest (language: python)

Issue:
Add support for file object uploads and downloads from test environments
Self explanatory - to test Apptainer, I need to be able to upload definition files into the test provider. Here is what needs to be done:

- [x] Create `File` class somewhere under `cleantest.hook`.
- [x] Enable support for uploading files via the `StartEnv` hook.
- [x] Enable support for downloading files via the `StopEnv` hook.


Discussion / hints:
Added `Dir` class for managing directories too.

Fix patch:
diff --git a/tox.ini b/tox.ini
index bdf51f2..242e746 100644
--- a/tox.ini
+++ b/tox.ini
@@ -106,3 +106,11 @@ deps =
   pytest
 commands =
   pytest -v --tb native --log-cli-level=ERROR {[vars]tst_path}/packages/snap
+
+[testenv:files]
+description = Run functional tests to test file and directory injections.
+deps =
+  -r {toxinidir}/requirements.txt
+  pytest
+commands =
+  pytest -v --tb native --log-cli-level=ERROR {[vars]tst_path}/files


Test patch:
diff --git a/src/cleantest/hooks/data/directory.py b/src/cleantest/hooks/data/directory.py
index 4631351..491e29c 100644
--- a/src/cleantest/hooks/data/directory.py
+++ b/src/cleantest/hooks/data/directory.py
@@ -8,7 +8,8 @@
 import pathlib
 import tarfile
 import tempfile
-import zipfile
+from io import BytesIO
+from typing import Iterable
 
 from .file import File
 
@@ -25,24 +26,16 @@ class DirectoryNotFoundError(Exception):
     ...
 
 
-class Dir(File):
-    def __init__(
-        self, src: str, dest: str, compression: str = "gz", overwrite: bool = False
-    ) -> None:
-        self.src = pathlib.Path(src)
-        self.dest = pathlib.Path(dest)
-        self.compression = compression
-        self.overwrite = overwrite
-        self._data = None
+def _strip_tar(tar: tarfile.TarFile, n_components: int = 1) -> Iterable[tarfile.TarInfo]:
+    for member in tar.getmembers():
+        member_path = pathlib.Path(member.path)
+        member.path = member_path.relative_to(*member_path.parts[:n_components])
+        yield member
 
-        if self.src.is_file():
-            raise NotADirectoryError(f"{self.src} is a file. Use File class instead.")
 
-        if self.compression not in ["gz", "zip", "bz2", "xz", None]:
-            raise DirectoryError(
-                f"{self.compression} not valid compression format. "
-                f"Format can be one of the following: {['gz', 'zip', 'bz2', 'xz', None]}"
-            )
+class Dir(File):
+    def __init__(self, src: str, dest: str, overwrite: bool = False) -> None:
+        super().__init__(src, dest, overwrite)
 
     def dump(self) -> None:
         if self.dest.exists() and self.overwrite is False:
@@ -53,32 +46,20 @@ def dump(self) -> None:
         if self._data is None:
             DirectoryError("Nothing to write.")
 
-        if self.compression in ["gz", "bz2", "xz", None]:
-            protocol = "r" if self.compression is None else f"r:{self.compression}"
-            with tarfile.open(self._data, protocol) as tar:
-                tar.extractall(self.dest)
-        else:
-            with zipfile.ZipFile(self._data, "r") as zip_out:
-                zip_out.extractall(self.dest)
+        with tarfile.open(fileobj=BytesIO(self._data), mode="r:gz") as tar:
+            tar.extractall(self.dest, members=_strip_tar(tar))
 
     def load(self) -> None:
         if self.src.exists() is False:
             raise FileNotFoundError(f"Could not find {self.src}.")
 
+        if self.src.is_file():
+            raise NotADirectoryError(f"{self.src} is a file. Use File class instead.")
+
         old_dir = os.getcwd()
         os.chdir(os.sep.join(str(self.src).split(os.sep)[:-1]))
         archive_path = pathlib.Path(tempfile.gettempdir()).joinpath(self.src.name)
-        if self.compression in ["gz", "bz2", "xz", None]:
-            protocol = "w" if self.compression is None else f"w:{self.compression}"
-            with tarfile.open(archive_path, protocol) as tar:
-                tar.add(self.src.name)
-        else:
-            data = []
-            for root, directory, file in os.walk(os.getcwd()):
-                for filename in file:
-                    data.append(os.path.join(root, filename))
-            with zipfile.ZipFile(archive_path, "w") as zip_in:
-                for file in data:
-                    zip_in.write(file)
+        with tarfile.open(archive_path, "w:gz") as tar:
+            tar.add(self.src.name)
         os.chdir(old_dir)
         self._data = archive_path.read_bytes()
diff --git a/src/cleantest/hooks/data/file.py b/src/cleantest/hooks/data/file.py
index a0b7160..34712d4 100644
--- a/src/cleantest/hooks/data/file.py
+++ b/src/cleantest/hooks/data/file.py
@@ -9,7 +9,7 @@
 import tarfile
 import tempfile
 import textwrap
-import zipfile
+from io import BytesIO
 
 from cleantest.meta import Injectable
 
@@ -23,24 +23,12 @@ class InjectableModeError(Exception):
 
 
 class File(Injectable):
-    def __init__(
-        self, src: str, dest: str, compression: str = "gz", overwrite: bool = False
-    ) -> None:
+    def __init__(self, src: str, dest: str, overwrite: bool = False) -> None:
         self.src = pathlib.Path(src)
         self.dest = pathlib.Path(dest)
-        self.compression = compression
         self.overwrite = overwrite
         self._data = None
 
-        if self.src.is_dir():
-            raise FileError(f"{self.src} is a directory. Use Dir class instead.")
-
-        if self.compression not in ["gz", "zip", "bz2", "xz", None]:
-            raise FileError(
-                f"{self.compression} not valid compression format. "
-                f"Format can be one of the following: {['gz', 'zip', 'bz2', 'xz', None]}"
-            )
-
     def dump(self) -> None:
         if self.dest.exists() and self.overwrite is False:
             raise FileExistsError(
@@ -50,28 +38,24 @@ def dump(self) -> None:
         if self._data is None:
             raise FileError(f"Nothing to write.")
 
-        if self.compression in ["gz", "bz2", "xz", None]:
-            protocol = "r" if self.compression is None else f"r:{self.compression}"
-            with tarfile.open(self._data, protocol) as tar:
-                tar.extract(self.src.name, self.dest)
-        else:
-            with zipfile.ZipFile(self._data, "r") as zip_out:
-                zip_out.extract(self.src.name, self.dest)
+        with tarfile.open(
+            fileobj=BytesIO(self._data), mode="r:gz"
+        ) as tar, tempfile.TemporaryDirectory() as tmp_dir:
+            tar.extractall(tmp_dir)
+            self.dest.write_bytes(pathlib.Path(tmp_dir).joinpath(self.src.name).read_bytes())
 
     def load(self) -> None:
         if self.src.exists() is False:
             raise FileNotFoundError(f"Could not find {self.src}.")
 
+        if self.src.is_dir():
+            raise FileError(f"{self.src} is a directory. Use Dir class instead.")
+
         old_dir = os.getcwd()
         os.chdir(os.sep.join(str(self.src).split(os.sep)[:-1]))
         archive_path = pathlib.Path(tempfile.gettempdir()).joinpath(self.src.name)
-        if self.compression in ["gz", "bz2", "xz", None]:
-            protocol = "w" if self.compression is None else f"w:{self.compression}"
-            with tarfile.open(archive_path, protocol) as tar:
-                tar.add(self.src.name)
-        else:
-            with zipfile.ZipFile(archive_path, "w") as zip_out:
-                zip_out.write(self.src.name)
+        with tarfile.open(archive_path, "w:gz") as tar:
+            tar.add(self.src.name)
         os.chdir(old_dir)
         self._data = archive_path.read_bytes()
 
@@ -83,7 +67,7 @@ def __injectable__(self, path: str, verification_hash: str, mode: str) -> str:
                 
                 from {self.__module__} import {self.__class__.__name__}
                 
-                holder = {self.__class__.__name__}._load("{path}")
+                holder = {self.__class__.__name__}._load("{path}", "{verification_hash}")
                 holder.dump()
                 """
             ).strip("\n")
@@ -91,12 +75,14 @@ def __injectable__(self, path: str, verification_hash: str, mode: str) -> str:
             return textwrap.dedent(
                 f"""
                 #!/usr/bin/env python3
+                import json
+                import sys
                 
                 from {self.__module__} import {self.__class__.__name__}
                 
                 holder = {self.__class__.__name__}._load("{path}", "{verification_hash}")
                 holder.load()
-                holder._dump()
+                print(json.dumps(holder._dump()._asdict()), file=sys.stdout)
                 """
             ).strip("\n")
         else:
diff --git a/src/cleantest/hooks/startenv.py b/src/cleantest/hooks/startenv.py
index 777e7f5..a2ac6a0 100644
--- a/src/cleantest/hooks/startenv.py
+++ b/src/cleantest/hooks/startenv.py
@@ -13,7 +13,7 @@ class StartEnvHook:
     def __init__(
         self,
         name: str = "default",
-        packages: List[object] = [],
+        packages: List[Injectable] = [],
         upload: List[Injectable] = [],
     ) -> None:
         """Hook run at the start of the test environment.
diff --git a/src/cleantest/hooks/stopenv.py b/src/cleantest/hooks/stopenv.py
index 64f0d62..829067c 100644
--- a/src/cleantest/hooks/stopenv.py
+++ b/src/cleantest/hooks/stopenv.py
@@ -10,10 +10,12 @@
 
 
 class StopEnvHook:
-    def __init__(self, download: List[Injectable]) -> None:
+    def __init__(self, name: str = "default", download: List[Injectable] = []) -> None:
         """Hook run before stopping test environment.
 
         Args:
+            name (str): Unique name of hook.
             download (List[Injectable]): Artifacts to download from test environment.
         """
+        self.name = name
         self.download = download
diff --git a/src/cleantest/provider/_handler/lxd_handler.py b/src/cleantest/provider/_handler/lxd_handler.py
index a5274c6..1a1db58 100644
--- a/src/cleantest/provider/_handler/lxd_handler.py
+++ b/src/cleantest/provider/_handler/lxd_handler.py
@@ -8,12 +8,20 @@
 import json
 import pathlib
 import re
+import tempfile
 from concurrent.futures import ProcessPoolExecutor
 from typing import Any, Callable, Dict, List
 
 from pylxd import Client
 
-from cleantest.meta import BaseEntrypoint, BaseHandler, BasePackage, CleantestInfo, Result
+from cleantest.meta import (
+    BaseEntrypoint,
+    BaseHandler,
+    BasePackage,
+    CleantestInfo,
+    Injectable,
+    Result,
+)
 
 
 class InstanceMetadata:
@@ -87,15 +95,24 @@ def _exists(self, instance: InstanceMetadata) -> InstanceMetadata:
 
     def _handle_start_env_hooks(self, instance: InstanceMetadata) -> None:
         start_env_hooks = self._clean_config.get_start_env_hooks()
-        while len(start_env_hooks) > 0:
+        while start_env_hooks:
             hook = start_env_hooks.pop()
             instance = self._client.instances.get(instance.name)
             if hook.packages is not None:
                 for pkg in hook.packages:
                     self._handle_package_install(instance, pkg)
-
-    def _handle_stop_env_hooks(self) -> None:
-        ...
+            if hook.upload is not None:
+                for artifact in hook.upload:
+                    self._handle_artifact_upload(instance, artifact)
+
+    def _handle_stop_env_hooks(self, instance: InstanceMetadata) -> None:
+        stop_env_hooks = self._clean_config.get_stop_env_hooks()
+        while stop_env_hooks:
+            hook = stop_env_hooks.pop()
+            instance = self._client.instances.get(instance.name)
+            if hook.download is not None:
+                for artifact in hook.download:
+                    self._handle_artifact_download(instance, artifact)
 
     def _handle_package_install(self, instance: Any, pkg: BasePackage) -> None:
         dispatch = {"charmlib": lambda x: self._env.add(json.loads(x))}
@@ -113,6 +130,39 @@ def _handle_package_install(self, instance: Any, pkg: BasePackage) -> None:
         if pkg.__class__.__name__.lower() in dispatch:
             dispatch[pkg.__class__.__name__.lower()](result.stdout)
 
+    def _handle_artifact_upload(self, instance: Any, artifact: Injectable) -> None:
+        artifact.load()
+        dump_data = artifact._dump()
+        data_path = pathlib.Path(dump_data.path)
+        instance.execute(["mkdir", "-p", "/root/init/data"])
+        instance.files.put(f"/root/init/data/{data_path.name}", data_path.read_bytes())
+        instance.files.put(
+            "/root/init/data/dump",
+            artifact.__injectable__(
+                f"/root/init/data/{data_path.name}", dump_data.hash, mode="upload"
+            ),
+        )
+        instance.execute(["python3", "/root/init/data/dump"])
+
+    def _handle_artifact_download(self, instance: Any, artifact: Injectable) -> None:
+        dump_data = artifact._dump()
+        data_path = pathlib.Path(dump_data.path)
+        instance.execute(["mkdir", "-p", "/root/post/data"])
+        instance.files.put(f"/root/post/data/{data_path.name}", data_path.read_bytes())
+        instance.files.put(
+            "/root/post/data/load",
+            artifact.__injectable__(
+                f"/root/post/data/{data_path.name}", dump_data.hash, mode="download"
+            ),
+        )
+        result = json.loads(instance.execute(["python3", "/root/post/data/load"]).stdout)
+        data = instance.files.get(result["path"])
+        with tempfile.NamedTemporaryFile() as fout:
+            handler = pathlib.Path(fout.name)
+            handler.write_bytes(data)
+            holder = artifact.__class__._load(str(handler), result["hash"])
+            holder.dump()
+
 
 class Serial(BaseEntrypoint, LXDHandler):
     def __init__(self, attr: Dict[str, Any], func: Callable) -> None:
@@ -129,6 +179,7 @@ def run(self) -> Dict[str, Result]:
                 self._make_testlet(self._func, self._func_name, [re.compile(r"^@lxd\(([^)]+)\)")]),
                 i,
             )
+            self._handle_stop_env_hooks(i)
             if self._preserve is False:
                 self._teardown(i)
             results.update({i.name: self._process(result)})
@@ -160,6 +211,7 @@ def _target(self, i: InstanceMetadata) -> Dict[str, Result]:
             self._make_testlet(self._func, self._func_name, [re.compile(r"^@lxd\(([^)]+)\)")]),
             i,
         )
+        self._handle_stop_env_hooks(i)
         if self._preserve is False:
             self._teardown(i)
 
diff --git a/tests/files/greeting.txt b/tests/files/greeting.txt
new file mode 100644
index 0000000..8da7bb2
--- /dev/null
+++ b/tests/files/greeting.txt
@@ -0,0 +1,1 @@
+Hello there General Kenobi
\ No newline at end of file
diff --git a/tests/files/greetings/greeting_1.txt b/tests/files/greetings/greeting_1.txt
new file mode 100644
index 0000000..b08588d
--- /dev/null
+++ b/tests/files/greetings/greeting_1.txt
@@ -0,0 +1,1 @@
+Good to see you.
\ No newline at end of file
diff --git a/tests/files/greetings/greeting_2.txt b/tests/files/greetings/greeting_2.txt
new file mode 100644
index 0000000..8ae9a6e
--- /dev/null
+++ b/tests/files/greetings/greeting_2.txt
@@ -0,0 +1,1 @@
+What's cooking good-looking?
\ No newline at end of file
diff --git a/tests/files/greetings/greeting_3.txt b/tests/files/greetings/greeting_3.txt
new file mode 100644
index 0000000..4295a58
--- /dev/null
+++ b/tests/files/greetings/greeting_3.txt
@@ -0,0 +1,1 @@
+Stop! You have violated the law!
\ No newline at end of file
diff --git a/tests/files/test_upload_download.py b/tests/files/test_upload_download.py
new file mode 100644
index 0000000..b3f32c8
--- /dev/null
+++ b/tests/files/test_upload_download.py
@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+# Copyright 2022 Jason C. Nucciarone, Canonical Ltd.
+# See LICENSE file for licensing details.
+
+"""Test that cleantest can successfully upload and download files."""
+
+import os
+import pathlib
+import tempfile
+
+from cleantest import Configure
+from cleantest.hooks import StartEnvHook, StopEnvHook
+from cleantest.hooks.data import Dir, File
+from cleantest.provider import lxd
+
+root = os.path.dirname(os.path.realpath(__file__))
+config = Configure()
+start_hook = StartEnvHook(
+    name="upload_artifact",
+    upload=[
+        File(os.path.join(root, "greeting.txt"), "/root/greeting.txt"),
+        Dir(os.path.join(root, "greetings"), "/root/greetings"),
+    ],
+)
+stop_hook = StopEnvHook(
+    name="download_artifact",
+    download=[
+        File("/root/dump.txt", os.path.join(tempfile.gettempdir(), "dump.txt")),
+        Dir("/root/dump", os.path.join(tempfile.gettempdir(), "dump")),
+    ],
+)
+config.register_hook(start_hook)
+config.register_hook(stop_hook)
+
+
+@lxd(image="jammy-amd64", preserve=True)
+def work_on_artifacts():
+    import os
+    import pathlib
+    import sys
+
+    print(pathlib.Path("/root/greeting.txt").is_file(), file=sys.stdout)
+    print(pathlib.Path("/root/dump").is_dir(), file=sys.stdout)
+
+    pathlib.Path("/root/dump.txt").write_text("Dumped like a sack of rocks")
+    os.mkdir("/root/dump")
+    pathlib.Path("/root/dump/dump_1.txt").write_text("Oh I have been dumped again!")
+
+
+class TestUploadDownload:
+    def test_upload_download(self) -> None:
+        work_on_artifacts()
+        assert pathlib.Path(tempfile.gettempdir()).joinpath("dump.txt").is_file() is True
+        assert pathlib.Path(tempfile.gettempdir()).joinpath("dump").is_dir() is True