Initial public release

tests/anyimage/server/support.py

@@ -0,0 +1,56 @@
+from tests.support.paths import PROJECT_ROOT
+import sys
+from pathlib import Path
+
+
+def load_server_module(name):
+    addon_directory = PROJECT_ROOT / "src" / "anyimage"
+    sys.path.insert(0, str(addon_directory))
+    try:
+        return __import__(f"server.{name}", fromlist=[name])
+    finally:
+        sys.path.remove(str(addon_directory))
+
+
+class FakeModels:
+    def __init__(self, directory="BEN2-ONNX"):
+        self.model_directory = Path(directory)
+
+    def directory(self, _model):
+        return self.model_directory
+
+    def get_background(self, _model_directory, _device):
+        from server.models.onnx_ben2 import create_session
+
+        return create_session(self.model_directory, "cpu"), 0.0
+
+
+    def get_upscale(self, _model_directory, _device):
+        return object()
+
+
+class FakeJobContext:
+    def __init__(self, directory, job_type, models=None):
+        self.directory = Path(directory)
+        self.job_type = job_type
+        self.models = models or FakeModels()
+        self.status = {}
+        self.logs = []
+
+    def resource(self, name):
+        if name != "model_manager":
+            raise KeyError(name)
+        return self.models
+
+    def progress(self, progress, message, **details):
+        self.status = {
+            "progress": progress,
+            "message": message,
+            **details,
+        }
+
+    def log(self, message, *, level="INFO"):
+        self.logs.append((level, message))
+
+    def check_cancelled(self):
+        return None