Scope gallery image endpoints by owner (#3001)

tests/test_gallery_image_endpoint_owner_scope.py

@@ -0,0 +1,126 @@
+"""Owner-scope regression for gallery image endpoint selection.
+
+The image editor/upscale proxies select ``ModelEndpoint`` rows and may copy the
+row's stored ``api_key`` for OpenAI-compatible image endpoints. That lookup must
+only consider endpoints visible to the caller, otherwise users sharing the same
+base URL can borrow another account's private image API key.
+"""
+
+from types import SimpleNamespace
+
+import routes.gallery_routes as gallery_routes
+
+
+class _Predicate:
+    def __init__(self, check):
+        self._check = check
+
+    def __call__(self, row):
+        return self._check(row)
+
+    def __or__(self, other):
+        return _Predicate(lambda row: self(row) or other(row))
+
+
+class _Column:
+    def __init__(self, name):
+        self.name = name
+
+    def __eq__(self, value):
+        return _Predicate(lambda row: getattr(row, self.name) == value)
+
+
+class _ModelEndpoint:
+    base_url = _Column("base_url")
+    model_type = _Column("model_type")
+    is_enabled = _Column("is_enabled")
+    owner = _Column("owner")
+
+
+class _Query:
+    def __init__(self, rows):
+        self._rows = list(rows)
+
+    def filter(self, *predicates):
+        self._rows = [row for row in self._rows if all(pred(row) for pred in predicates)]
+        return self
+
+    def all(self):
+        return list(self._rows)
+
+
+class _DB:
+    def __init__(self, rows):
+        self._rows = rows
+
+    def query(self, model):
+        assert model is _ModelEndpoint
+        return _Query(self._rows)
+
+
+def _ep(base_url, owner, *, enabled=True, model_type="image", api_key="sk-secret"):
+    return SimpleNamespace(
+        base_url=base_url,
+        owner=owner,
+        is_enabled=enabled,
+        model_type=model_type,
+        api_key=api_key,
+    )
+
+
+def _patch_model(monkeypatch):
+    monkeypatch.setattr(gallery_routes, "ModelEndpoint", _ModelEndpoint)
+
+
+URL = "https://api.example.com/v1"
+
+
+def test_first_visible_image_endpoint_rejects_another_owner(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [_ep(URL, "bob")]
+
+    assert gallery_routes._first_visible_image_endpoint(_DB(rows), "alice") is None
+
+
+def test_first_visible_image_endpoint_prefers_callers_own_row(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [_ep(URL, None, api_key="shared"), _ep(URL, "alice", api_key="own")]
+
+    ep = gallery_routes._first_visible_image_endpoint(_DB(rows), "alice")
+
+    assert ep is not None
+    assert ep.owner == "alice"
+    assert ep.api_key == "own"
+
+
+def test_visible_image_endpoint_for_base_rejects_same_url_other_owner(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [_ep(URL, "bob")]
+
+    assert gallery_routes._visible_image_endpoint_for_base(_DB(rows), URL, "alice") is None
+
+
+def test_visible_image_endpoint_for_base_allows_shared_or_own(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [
+        _ep("https://other.example/v1", "alice"),
+        _ep(URL, None, api_key="shared"),
+        _ep(URL, "alice", api_key="own"),
+    ]
+
+    ep = gallery_routes._visible_image_endpoint_for_base(_DB(rows), "https://api.example.com", "alice")
+
+    assert ep is not None
+    assert ep.owner == "alice"
+    assert ep.api_key == "own"
+    assert ep.base_url == URL
+
+
+def test_image_endpoint_owner_filter_is_noop_in_single_user_mode(monkeypatch):
+    _patch_model(monkeypatch)
+    rows = [_ep(URL, "bob")]
+
+    ep = gallery_routes._visible_image_endpoint_for_base(_DB(rows), URL, None)
+
+    assert ep is not None
+    assert ep.owner == "bob"