From 19b844c47289aab64774cf7b6f579cc8f7def5bc Mon Sep 17 00:00:00 2001
From: Leonardo Schwarz <leonardo.schwarz@fgcz.ethz.ch>
Date: Fri, 5 Jul 2024 17:10:06 +0200
Subject: [PATCH] add GenericWriteFile and missing unit test

---
 .../persistence/imzml/imzml_write_file.py     |  3 +-
 src/depiction/persistence/types.py            | 13 +++-
 .../imzml/test_imzml_write_file.py            | 74 +++++++++++++++++++
 3 files changed, 87 insertions(+), 3 deletions(-)
 create mode 100644 tests/unit/persistence/imzml/test_imzml_write_file.py

diff --git a/src/depiction/persistence/imzml/imzml_write_file.py b/src/depiction/persistence/imzml/imzml_write_file.py
index eb3fdcf..403bc02 100644
--- a/src/depiction/persistence/imzml/imzml_write_file.py
+++ b/src/depiction/persistence/imzml/imzml_write_file.py
@@ -4,12 +4,13 @@
 from typing import TYPE_CHECKING
 
 from depiction.persistence.imzml.imzml_writer import ImzmlWriter
+from depiction.persistence.types import GenericWriteFile
 
 if TYPE_CHECKING:
     from depiction.persistence.imzml.imzml_mode_enum import ImzmlModeEnum
 
 
-class ImzmlWriteFile:
+class ImzmlWriteFile(GenericWriteFile):
     """A handle for a .imzML file that is to be written.
 
     Args:
diff --git a/src/depiction/persistence/types.py b/src/depiction/persistence/types.py
index f620d0a..1cdbfe9 100644
--- a/src/depiction/persistence/types.py
+++ b/src/depiction/persistence/types.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Self, Protocol
+
 import numpy as np
 
 if TYPE_CHECKING:
@@ -11,7 +12,7 @@
 from tqdm import tqdm
 
 from collections.abc import Generator
-from contextlib import contextmanager
+from contextlib import contextmanager, AbstractContextManager
 from functools import cached_property
 from pathlib import Path
 from typing import Optional, TextIO
@@ -225,4 +226,12 @@ def progress_fn(x: Sequence[int]) -> Sequence[int]:
 
 
 class GenericWriteFile(Protocol):
-    pass
+    @property
+    def imzml_mode(self) -> ImzmlModeEnum:
+        """The imzml mode of the .imzML file."""
+        ...
+
+    @contextmanager
+    def writer(self) -> AbstractContextManager[GenericWriter]:
+        """Opens the .imzML file for writing and yields an `ImzmlWriter` instance."""
+        ...
diff --git a/tests/unit/persistence/imzml/test_imzml_write_file.py b/tests/unit/persistence/imzml/test_imzml_write_file.py
new file mode 100644
index 0000000..40a989c
--- /dev/null
+++ b/tests/unit/persistence/imzml/test_imzml_write_file.py
@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+from pytest_mock import MockerFixture
+
+from depiction.persistence import ImzmlModeEnum, ImzmlWriteFile, ImzmlWriter
+
+mock_path = Path("/dev/null/test.imzML")
+mock_imzml_mode = ImzmlModeEnum.CONTINUOUS
+
+
+@pytest.fixture
+def mock_write_file() -> ImzmlWriteFile:
+    return ImzmlWriteFile(path=mock_path, imzml_mode=mock_imzml_mode)
+
+
+def test_imzml_file(mock_write_file: ImzmlWriteFile) -> None:
+    assert mock_write_file.imzml_file == mock_path
+
+
+def test_imzml_file_when_invalid() -> None:
+    file = ImzmlWriteFile(path="/dev/null/test.txt", imzml_mode=mock_imzml_mode)
+    with pytest.raises(ValueError):
+        _ = file.imzml_file
+
+
+def test_ibd_file(mock_write_file: ImzmlWriteFile) -> None:
+    assert mock_write_file.ibd_file == Path("/dev/null/test.ibd")
+
+
+def test_imzml_mode(mock_write_file: ImzmlWriteFile) -> None:
+    assert mock_write_file.imzml_mode == mock_imzml_mode
+
+
+def test_writer_when_success(mocker: MockerFixture, mock_write_file: ImzmlWriteFile) -> None:
+    mock_imzml_file = mocker.MagicMock(name="mock_imzml_file", spec=Path)
+    mock_imzml_file.exists.return_value = False
+    mocker.patch.object(Path, "exists", return_value=False)
+    mock_open = mocker.patch.object(ImzmlWriter, "open")
+    with mock_write_file.writer() as writer:
+        assert writer == mock_open.return_value
+    mock_open.assert_called_once_with(path=mock_write_file.imzml_file, imzml_mode=mock_imzml_mode)
+    mock_open.return_value.close.assert_called_once_with()
+
+
+def test_writer_when_mode_x_file_exists(mocker: MockerFixture, mock_write_file: ImzmlWriteFile) -> None:
+    mocker.patch.object(Path, "exists", return_value=True)
+    with pytest.raises(ValueError):
+        with mock_write_file.writer():
+            pass
+
+
+def test_writer_when_mode_w_file_exists(mocker: MockerFixture) -> None:
+    mock_write_file = ImzmlWriteFile(path=mock_path, imzml_mode=mock_imzml_mode, write_mode="w")
+    mocker.patch.object(Path, "exists", return_value=True)
+    mock_unlink = mocker.patch.object(Path, "unlink")
+    mock_open = mocker.patch.object(ImzmlWriter, "open")
+    with mock_write_file.writer() as writer:
+        assert writer == mock_open.return_value
+    assert mock_unlink.mock_calls == [mocker.call(), mocker.call()]
+    mock_open.assert_called_once_with(path=mock_write_file.imzml_file, imzml_mode=mock_imzml_mode)
+    mock_open.return_value.close.assert_called_once_with()
+
+
+def test_repr(mock_write_file: ImzmlWriteFile) -> None:
+    assert (
+        repr(mock_write_file) == f"ImzmlWriteFile(path={mock_path!r}, imzml_mode={mock_imzml_mode!r}, write_mode='x')"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main()