Initial public release

tests/ai/test_background_alpha.py

@@ -0,0 +1,36 @@
+from unittest.mock import patch
+
+import numpy as np
+from PIL import Image
+import pytest
+
+from server.jobs import remove_background
+from server.models import onnx_ben2
+from tests.anyimage.server.support import FakeJobContext
+
+
+@pytest.mark.parametrize("failure", ["cancel", "prediction", "sequence", "output_kind"])
+def test_alpha_job_failure_does_not_write_result(tmp_path, failure):
+    source = tmp_path / "input.png"
+    Image.new("RGB", (2, 2)).save(source)
+    context = FakeJobContext(tmp_path / "job", "remove-background")
+    parameters = {"input": str(source), "model": "BEN2_BASE", "device": "cpu", "output_kind": "alpha"}
+    if failure == "sequence":
+        frames = tmp_path / "frames"
+        frames.mkdir()
+        for i in (1, 2):
+            Image.new("RGB", (2, 2)).save(frames / f"{i}.png")
+        parameters["input"] = str(frames)
+    if failure == "output_kind":
+        parameters["output_kind"] = "invalid"
+
+    def infer(*args, **kwargs):
+        if failure == "prediction":
+            raise RuntimeError("Prediction failed")
+        context.check_cancelled = lambda: (_ for _ in ()).throw(RuntimeError("Cancelled"))
+        return np.ones((2, 2), np.float32)
+
+    with patch.object(onnx_ben2, "create_session", return_value=object()), patch.object(onnx_ben2, "infer_alpha", side_effect=infer):
+        with pytest.raises((RuntimeError, ValueError)):
+            remove_background.run(context, parameters)
+    assert not (context.directory / "alpha.npy").exists()