Add support for global artifacts in ArtifactedPlugin (#107)

setup.cfg

@@ -45,5 +45,6 @@ markers = only
 asyncio_mode = auto
 filterwarnings = 
 	ignore:Deprecated
+	ignore::DeprecationWarning:dessert
 addopts = 
 	--basetemp=./__tmpdir__

tests/core/scenario_result/test_scenario_result.py

@@ -314,7 +314,7 @@ def test_scenario_result_attach_incorrect_artifact(*, virtual_scenario: VirtualS
 
     with then:
         assert exc.type is TypeError
-        assert str(exc.value) == "artifact must be an instance of Artifact"
+        assert str(exc.value) == "Expected an instance of Artifact, got dict"
 
 
 def test_scenario_result_get_artifacts(*, virtual_scenario: VirtualScenario):

tests/core/test_report.py

@@ -2,7 +2,9 @@
 from unittest.mock import Mock
 
 from baby_steps import given, then, when
+from pytest import raises
 
+from vedro import MemoryArtifact
 from vedro.core import ExcInfo, Report, ScenarioResult, VirtualScenario
 
 
@@ -255,3 +257,45 @@ def test_repr_with_interrupted():
     with then:
         assert res == ("<Report total=0 passed=0 failed=0 skipped=0"
                        " interrupted=<class 'KeyboardInterrupt'>>")
+
+
+def test_report_attach_artifact():
+    with given:
+        artifact = MemoryArtifact("name", "text/plain", b"")
+        report = Report()
+
+    with when:
+        res = report.attach(artifact)
+
+    with then:
+        assert res is None
+
+
+def test_report_attach_incorrect_artifact():
+    with given:
+        artifact = {}
+        report = Report()
+
+    with when, raises(BaseException) as exc:
+        report.attach(artifact)
+
+    with then:
+        assert exc.type is TypeError
+        assert str(exc.value) == "Expected an instance of Artifact, got dict"
+
+
+def test_report_get_artifacts():
+    with given:
+        report = Report()
+
+        artifact1 = MemoryArtifact("name1", "text/plain", b"")
+        report.attach(artifact1)
+
+        artifact2 = MemoryArtifact("name2", "text/plain", b"")
+        report.attach(artifact2)
+
+    with when:
+        artifacts = report.artifacts
+
+    with then:
+        assert artifacts == [artifact1, artifact2]

tests/core/test_step_result.py

@@ -207,7 +207,7 @@ def test_step_result_attach_incorrect_artifact(*, virtual_step: VirtualStep):
 
     with then:
         assert exc.type is TypeError
-        assert str(exc.value) == "artifact must be an instance of Artifact"
+        assert str(exc.value) == "Expected an instance of Artifact, got dict"
 
 
 def test_step_result_get_artifacts(*, virtual_step: VirtualStep):

tests/plugins/artifacted/_utils.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 from time import monotonic_ns
 from typing import Optional
+from unittest.mock import patch
 
 import pytest
 
@@ -11,14 +12,19 @@
 from vedro import Scenario as VedroScenario
 from vedro.core import Dispatcher, VirtualScenario, VirtualStep
 from vedro.events import ArgParsedEvent, ArgParseEvent, ConfigLoadedEvent
-from vedro.plugins.artifacted import Artifacted, ArtifactedPlugin, MemoryArtifact
+from vedro.plugins.artifacted import Artifacted, ArtifactedPlugin, ArtifactManager, MemoryArtifact
 
 
 @pytest.fixture()
 def dispatcher() -> Dispatcher:
     return Dispatcher()
 
 
+@pytest.fixture()
+def global_artifacts() -> deque:
+    return deque()
+
+
 @pytest.fixture()
 def scenario_artifacts() -> deque:
     return deque()
@@ -35,9 +41,24 @@ def project_dir(tmp_path: Path) -> Path:
 
 
 @pytest.fixture()
-def artifacted(dispatcher: Dispatcher, scenario_artifacts: deque,
+def artifacts_dir(project_dir: Path):
+    artifacts_dir = project_dir / "artifacts/"
+    artifacts_dir.mkdir(exist_ok=True)
+    return artifacts_dir
+
+
+@pytest.fixture()
+def artifact_manager(artifacts_dir: Path, project_dir: Path):
+    return ArtifactManager(artifacts_dir, project_dir)
+
+
+@pytest.fixture()
+def artifacted(dispatcher: Dispatcher,
+               global_artifacts: deque,
+               scenario_artifacts: deque,
                step_artifacts: deque) -> ArtifactedPlugin:
     artifacted = ArtifactedPlugin(Artifacted,
+                                  global_artifacts=global_artifacts,
                                   scenario_artifacts=scenario_artifacts,
                                   step_artifacts=step_artifacts)
     artifacted.subscribe(dispatcher)
@@ -72,9 +93,35 @@ class Config(_Config):
 
 
 async def fire_arg_parsed_event(dispatcher: Dispatcher, *,
-                                save_artifacts: bool = False,
+                                save_artifacts: bool = Artifacted.save_artifacts,
+                                add_artifact_details: bool = Artifacted.add_artifact_details,
                                 artifacts_dir: Optional[Path] = None) -> None:
     await dispatcher.fire(ArgParseEvent(ArgumentParser()))
 
-    namespace = Namespace(save_artifacts=save_artifacts, artifacts_dir=artifacts_dir)
+    namespace = Namespace(
+        save_artifacts=save_artifacts,
+        add_artifact_details=add_artifact_details,
+        artifacts_dir=artifacts_dir,
+    )
     await dispatcher.fire(ArgParsedEvent(namespace))
+
+
+def patch_rmtree(exception: Optional[Exception] = None):
+    return patch("shutil.rmtree", side_effect=exception)
+
+
+def patch_copy2(exception: Optional[Exception] = None):
+    return patch("shutil.copy2", side_effect=exception)
+
+
+def patch_write_bytes(exception: Optional[Exception] = None):
+    return patch("pathlib.Path.write_bytes", side_effect=exception)
+
+
+def patch_mkdir(exception: Optional[Exception] = None):
+    return patch("pathlib.Path.mkdir", side_effect=exception)
+
+
+def make_artifacted_plugin(artifact_manager: ArtifactManager) -> ArtifactedPlugin:
+    return ArtifactedPlugin(Artifacted,
+                            artifact_manager_factory=lambda *args: artifact_manager)

tests/plugins/artifacted/test_artifact_manager.py

@@ -0,0 +1,187 @@
+from pathlib import Path
+from typing import Type
+from unittest.mock import call
+
+import pytest
+from baby_steps import given, then, when
+from pytest import raises
+
+from vedro.core import Artifact
+from vedro.plugins.artifacted import ArtifactManager
+
+from ._utils import (
+    artifact_manager,
+    artifacts_dir,
+    create_file_artifact,
+    create_memory_artifact,
+    patch_copy2,
+    patch_mkdir,
+    patch_rmtree,
+    patch_write_bytes,
+    project_dir,
+)
+
+__all__ = ("project_dir", "artifacts_dir", "artifact_manager")  # fixtures
+
+
+def test_cleanup_artifacts(*, artifact_manager: ArtifactManager, artifacts_dir: Path):
+    with given:
+        artifacts_dir.mkdir(parents=True, exist_ok=True)
+
+    with when:
+        artifact_manager.cleanup_artifacts()
+
+    with then:
+        assert not artifacts_dir.exists()
+
+
+def test_cleanup_artifacts_file_not_found(*, artifact_manager: ArtifactManager,
+                                          artifacts_dir: Path):
+    with given:
+        artifacts_dir.mkdir(parents=True, exist_ok=True)
+
+    with when, patch_rmtree(FileNotFoundError()) as mock:
+        artifact_manager.cleanup_artifacts()
+
+    with then:
+        # no exception raised
+        assert mock.mock_calls == [call(artifacts_dir)]
+
+
+@pytest.mark.parametrize("exc_type", [PermissionError, OSError])
+def test_cleanup_artifacts_os_error(exc_type: Type[Exception], *,
+                                    artifact_manager: ArtifactManager, artifacts_dir: Path):
+    with given:
+        artifacts_dir.mkdir(parents=True, exist_ok=True)
+
+    with when, \
+         patch_rmtree(exc_type()) as mock, \
+         raises(Exception) as exc:
+        artifact_manager.cleanup_artifacts()
+
+    with then:
+        assert exc.type is exc_type
+        assert "Failed to clean up artifacts directory" in str(exc.value)
+
+        assert mock.mock_calls == [call(artifacts_dir)]
+
+
+def test_save_memory_artifact(*, artifact_manager: ArtifactManager, artifacts_dir: Path):
+    with given:
+        file_content = "Hello, World!"
+        artifact = create_memory_artifact(file_content)
+
+    with when:
+        artifact_path = artifact_manager.save_artifact(artifact, artifacts_dir)
+
+    with then:
+        assert artifact_path.exists()
+        assert artifact_path.read_bytes() == file_content.encode()
+
+
+@pytest.mark.parametrize(("exc_type", "exc_msg"), [
+    (PermissionError, "Permission denied when writing to"),
+    (OSError, "Failed to write MemoryArtifact to"),
+])
+def test_save_memory_artifact_os_error(exc_type: Type[Exception], exc_msg: str, *,
+                                       artifact_manager: ArtifactManager, artifacts_dir: Path):
+    with given:
+        file_content = "Hello, World!"
+        artifact = create_memory_artifact(file_content)
+
+    with when, \
+         patch_write_bytes(exc_type()) as mock, \
+         raises(Exception) as exc:
+        artifact_manager.save_artifact(artifact, artifacts_dir)
+
+    with then:
+        assert exc.type is exc_type
+        assert exc_msg in str(exc.value)
+
+        assert mock.mock_calls == [call(file_content.encode())]
+
+
+def test_save_file_artifact(*, artifact_manager: ArtifactManager,
+                            artifacts_dir: Path, tmp_path: Path):
+    with given:
+        file_path = tmp_path / "file_artifact.txt"
+        file_content = "Hello, World!"
+        artifact = create_file_artifact(file_path, file_content)
+
+    with when:
+        artifact_path = artifact_manager.save_artifact(artifact, artifacts_dir)
+
+    with then:
+        assert artifact_path.exists()
+        assert artifact_path.read_bytes() == file_content.encode()
+
+        assert artifact.path.exists()
+        assert artifact.path.read_bytes() == file_content.encode()
+
+
+@pytest.mark.parametrize(("exc_type", "exc_msg"), [
+    (FileNotFoundError, "Source file"),
+    (PermissionError, "Permission denied when copying from "),
+    (OSError, "Failed to copy FileArtifact from"),
+])
+def test_save_file_artifact_os_error(exc_type: Type[Exception], exc_msg: str, *,
+                                     artifact_manager: ArtifactManager,
+                                     artifacts_dir: Path,
+                                     tmp_path: Path):
+    with given:
+        file_path = tmp_path / "file_artifact.txt"
+        file_content = "Hello, World!"
+        artifact = create_file_artifact(file_path, file_content)
+
+    with when, \
+         patch_copy2(exc_type()) as mock, \
+         raises(Exception) as exc:
+        artifact_manager.save_artifact(artifact, artifacts_dir)
+
+    with then:
+        assert exc.type is exc_type
+        assert exc_msg in str(exc.value)
+
+        assert artifact.path.exists()
+        assert artifact.path.read_bytes() == file_content.encode()
+
+        assert mock.mock_calls == [call(artifact.path, artifacts_dir / artifact.name)]
+
+
+@pytest.mark.parametrize("exc_type", [PermissionError, OSError])
+def test_save_artifact_directory_mkdir_failure(exc_type: Type[Exception], *,
+                                               artifact_manager: ArtifactManager,
+                                               artifacts_dir: Path):
+    with given:
+        artifacts_dir.rmdir()  # artifacts_dir is automatically created by the fixture
+
+        file_content = "Hello, World!"
+        artifact = create_memory_artifact(file_content)
+
+    with when, \
+         patch_mkdir(exc_type()) as mock, \
+         raises(Exception) as exc:
+        artifact_manager.save_artifact(artifact, artifacts_dir)
+
+    with then:
+        assert exc.type is exc_type
+        assert "Failed to create directory" in str(exc.value)
+
+        assert mock.mock_calls == [call(parents=True, exist_ok=True)]
+
+
+def test_save_artifact_unknown_type(*, artifact_manager: ArtifactManager, artifacts_dir: Path):
+    with given:
+        class UnknownArtifact(Artifact):
+            pass
+
+        artifact = UnknownArtifact()
+
+    with when, raises(BaseException) as exc:
+        artifact_manager.save_artifact(artifact, artifacts_dir)
+
+    with then:
+        assert exc.type is TypeError
+        assert str(exc.value) == (
+            f"Can't save artifact to '{artifacts_dir}': unknown type 'UnknownArtifact'"
+        )