feat: add in-app LM Remote configuration
This commit is contained in:
@@ -0,0 +1,149 @@
|
||||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
|
||||
import {
|
||||
RemoteConfigFormError,
|
||||
buildConfigDraft,
|
||||
buildConnectionDraft,
|
||||
didEffectiveRemoteUrlChange,
|
||||
isConfigWritable,
|
||||
mappingsToRows,
|
||||
} from "../../web/comfyui/remote_config_utils.js";
|
||||
|
||||
test("buildConfigDraft normalizes form values and skips blank mappings", () => {
|
||||
assert.deepEqual(
|
||||
buildConfigDraft({
|
||||
remoteUrl: " http://manager.local:8188/ ",
|
||||
timeout: "45",
|
||||
mappingRows: [
|
||||
{ remote: "/data/loras/", local: "/mnt/loras" },
|
||||
{ remote: "", local: "" },
|
||||
],
|
||||
}),
|
||||
{
|
||||
remote_url: "http://manager.local:8188",
|
||||
timeout: 45,
|
||||
path_mappings: { "/data/loras": "/mnt/loras" },
|
||||
}
|
||||
);
|
||||
});
|
||||
|
||||
test("buildConfigDraft allows an empty URL to disable remote mode", () => {
|
||||
const draft = buildConfigDraft({
|
||||
remoteUrl: "",
|
||||
timeout: 30,
|
||||
mappingRows: [],
|
||||
});
|
||||
assert.equal(draft.remote_url, "");
|
||||
});
|
||||
|
||||
test("buildConfigDraft rejects unsafe URLs and invalid mappings", () => {
|
||||
assert.throws(
|
||||
() =>
|
||||
buildConfigDraft({
|
||||
remoteUrl: "ftp://manager.local",
|
||||
timeout: 30,
|
||||
mappingRows: [],
|
||||
}),
|
||||
(error) =>
|
||||
error instanceof RemoteConfigFormError && error.field === "remote_url"
|
||||
);
|
||||
assert.throws(
|
||||
() =>
|
||||
buildConfigDraft({
|
||||
remoteUrl: "http://manager.local",
|
||||
timeout: 30,
|
||||
mappingRows: [{ remote: "/remote", local: "" }],
|
||||
}),
|
||||
(error) =>
|
||||
error instanceof RemoteConfigFormError && error.field === "path_mappings"
|
||||
);
|
||||
assert.throws(
|
||||
() =>
|
||||
buildConfigDraft({
|
||||
remoteUrl: "http://manager.local",
|
||||
timeout: 301,
|
||||
mappingRows: [],
|
||||
}),
|
||||
(error) =>
|
||||
error instanceof RemoteConfigFormError && error.field === "timeout"
|
||||
);
|
||||
});
|
||||
|
||||
test("buildConfigDraft rejects normalized duplicate path prefixes", () => {
|
||||
assert.throws(
|
||||
() =>
|
||||
buildConfigDraft({
|
||||
remoteUrl: "http://manager.local",
|
||||
timeout: 30,
|
||||
mappingRows: [
|
||||
{ remote: "/data/loras", local: "/mnt/a" },
|
||||
{ remote: "\\data\\loras\\", local: "/mnt/b" },
|
||||
],
|
||||
}),
|
||||
/Duplicate remote path prefix/
|
||||
);
|
||||
});
|
||||
|
||||
test("buildConfigDraft preserves __proto__ as a path mapping key", () => {
|
||||
const draft = buildConfigDraft({
|
||||
remoteUrl: "http://manager.local",
|
||||
timeout: 30,
|
||||
mappingRows: [{ remote: "__proto__", local: "/mnt/prototype" }],
|
||||
});
|
||||
|
||||
assert.equal(Object.hasOwn(draft.path_mappings, "__proto__"), true);
|
||||
assert.equal(draft.path_mappings.__proto__, "/mnt/prototype");
|
||||
assert.equal(
|
||||
JSON.parse(JSON.stringify(draft.path_mappings)).__proto__,
|
||||
"/mnt/prototype"
|
||||
);
|
||||
});
|
||||
|
||||
test("connection tests use effective environment-managed values", () => {
|
||||
assert.deepEqual(
|
||||
buildConnectionDraft({
|
||||
remoteUrl: "http://stored.local:8188",
|
||||
timeout: "30",
|
||||
effective: {
|
||||
remote_url: "https://managed.local:443",
|
||||
timeout: 12,
|
||||
},
|
||||
overrides: {
|
||||
remote_url: "LM_REMOTE_URL",
|
||||
timeout: "LM_REMOTE_TIMEOUT",
|
||||
},
|
||||
}),
|
||||
{ remote_url: "https://managed.local:443", timeout: 12 }
|
||||
);
|
||||
});
|
||||
|
||||
test("mappingsToRows creates editable row values", () => {
|
||||
assert.deepEqual(mappingsToRows({ "/remote": "/local" }), [
|
||||
{ remote: "/remote", local: "/local" },
|
||||
]);
|
||||
});
|
||||
|
||||
test("didEffectiveRemoteUrlChange only tracks the effective URL", () => {
|
||||
assert.equal(
|
||||
didEffectiveRemoteUrlChange(
|
||||
{ effective: { remote_url: "http://old:8188", timeout: 30 } },
|
||||
{ effective: { remote_url: "http://new:8188", timeout: 30 } }
|
||||
),
|
||||
true
|
||||
);
|
||||
assert.equal(
|
||||
didEffectiveRemoteUrlChange(
|
||||
{ effective: { remote_url: "http://same:8188", timeout: 30 } },
|
||||
{ effective: { remote_url: "http://same:8188", timeout: 60 } }
|
||||
),
|
||||
false
|
||||
);
|
||||
});
|
||||
|
||||
test("isConfigWritable only disables saving for an explicit false value", () => {
|
||||
assert.equal(isConfigWritable({ storage: { writable: false } }), false);
|
||||
assert.equal(isConfigWritable({ storage: { writable: true } }), true);
|
||||
assert.equal(isConfigWritable({ storage: {} }), true);
|
||||
assert.equal(isConfigWritable(null), true);
|
||||
});
|
||||
@@ -0,0 +1,318 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def config_module():
|
||||
module_name = "lm_remote_config_test_module"
|
||||
module_path = Path(__file__).resolve().parents[1] / "config.py"
|
||||
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
yield module
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
|
||||
def write_json(path: Path, data: object) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(data), encoding="utf-8")
|
||||
|
||||
|
||||
def test_legacy_config_loads_and_first_save_migrates(config_module, tmp_path):
|
||||
legacy = tmp_path / "package" / "config.json"
|
||||
user = tmp_path / "user" / "config.json"
|
||||
original = {
|
||||
"remote_url": "http://legacy.local:8188",
|
||||
"timeout": 30,
|
||||
"path_mappings": {},
|
||||
"future_setting": {"preserve": True},
|
||||
}
|
||||
write_json(legacy, original)
|
||||
|
||||
config = config_module.RemoteConfig(user, legacy, environ={})
|
||||
state = config.as_dict()
|
||||
assert state["storage"]["source"] == "legacy"
|
||||
assert state["effective"]["remote_url"] == "http://legacy.local:8188"
|
||||
|
||||
config.save(
|
||||
{
|
||||
"remote_url": "http://new.local:8188/",
|
||||
"timeout": 45,
|
||||
"path_mappings": {"/remote": "/local"},
|
||||
},
|
||||
expected_revision=state["revision"],
|
||||
)
|
||||
|
||||
assert json.loads(legacy.read_text(encoding="utf-8")) == original
|
||||
saved = json.loads(user.read_text(encoding="utf-8"))
|
||||
assert saved["remote_url"] == "http://new.local:8188"
|
||||
assert saved["future_setting"] == {"preserve": True}
|
||||
assert config.as_dict()["storage"]["source"] == "user"
|
||||
|
||||
|
||||
def test_environment_values_override_without_rewriting_stored_values(
|
||||
config_module, tmp_path
|
||||
):
|
||||
user = tmp_path / "config.json"
|
||||
write_json(
|
||||
user,
|
||||
{
|
||||
"remote_url": "http://stored.local:8188",
|
||||
"timeout": 30,
|
||||
"path_mappings": {},
|
||||
},
|
||||
)
|
||||
config = config_module.RemoteConfig(
|
||||
user,
|
||||
tmp_path / "missing.json",
|
||||
environ={
|
||||
"LM_REMOTE_URL": "https://managed.local:443/",
|
||||
"LM_REMOTE_TIMEOUT": "12",
|
||||
},
|
||||
)
|
||||
|
||||
state = config.as_dict()
|
||||
assert state["configured"]["remote_url"] == "http://stored.local:8188"
|
||||
assert state["effective"]["remote_url"] == "https://managed.local:443"
|
||||
assert state["effective"]["timeout"] == 12
|
||||
assert state["overrides"] == {
|
||||
"remote_url": "LM_REMOTE_URL",
|
||||
"timeout": "LM_REMOTE_TIMEOUT",
|
||||
}
|
||||
|
||||
|
||||
def test_explicit_config_environment_reports_explicit_source(config_module, tmp_path):
|
||||
explicit = tmp_path / "managed" / "remote.json"
|
||||
config = config_module.RemoteConfig(
|
||||
legacy_config_file=tmp_path / "missing-legacy.json",
|
||||
environ={"LM_REMOTE_CONFIG": str(explicit)},
|
||||
)
|
||||
initial = config.as_dict()
|
||||
config.save(
|
||||
{"remote_url": "http://manager.local", "timeout": 30, "path_mappings": {}},
|
||||
expected_revision=initial["revision"],
|
||||
)
|
||||
|
||||
assert config.as_dict()["storage"]["source"] == "explicit"
|
||||
|
||||
|
||||
def test_explicit_missing_target_does_not_fall_back_to_legacy(config_module, tmp_path):
|
||||
explicit = tmp_path / "managed" / "remote.json"
|
||||
legacy = tmp_path / "package" / "config.json"
|
||||
legacy_data = {
|
||||
"remote_url": "http://legacy.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {},
|
||||
}
|
||||
write_json(legacy, legacy_data)
|
||||
|
||||
config = config_module.RemoteConfig(
|
||||
legacy_config_file=legacy,
|
||||
environ={"LM_REMOTE_CONFIG": str(explicit)},
|
||||
)
|
||||
state = config.as_dict()
|
||||
|
||||
assert state["storage"]["source"] == "explicit"
|
||||
assert state["configured"]["remote_url"] == ""
|
||||
config.save(
|
||||
{"remote_url": "http://explicit.local", "timeout": 30, "path_mappings": {}},
|
||||
expected_revision=state["revision"],
|
||||
)
|
||||
assert json.loads(legacy.read_text(encoding="utf-8")) == legacy_data
|
||||
assert (
|
||||
json.loads(explicit.read_text(encoding="utf-8"))["remote_url"]
|
||||
== "http://explicit.local"
|
||||
)
|
||||
|
||||
|
||||
def test_invalid_environment_timeout_is_ignored(config_module, tmp_path):
|
||||
config = config_module.RemoteConfig(
|
||||
tmp_path / "missing-user.json",
|
||||
tmp_path / "missing-legacy.json",
|
||||
environ={"LM_REMOTE_TIMEOUT": "not-a-number"},
|
||||
)
|
||||
state = config.as_dict()
|
||||
assert state["effective"]["timeout"] == 30
|
||||
assert state["overrides"]["timeout"] is None
|
||||
assert any("LM_REMOTE_TIMEOUT" in warning for warning in state["warnings"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
("remote_url", "ftp://manager.local"),
|
||||
("remote_url", "http://user:secret@manager.local"),
|
||||
("remote_url", "http://manager.local?query=yes"),
|
||||
("timeout", True),
|
||||
("timeout", 0),
|
||||
("timeout", 301),
|
||||
("path_mappings", []),
|
||||
],
|
||||
)
|
||||
def test_validation_rejects_invalid_values(config_module, field, value):
|
||||
candidate = {
|
||||
"remote_url": "http://manager.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {},
|
||||
}
|
||||
candidate[field] = value
|
||||
with pytest.raises(config_module.ConfigValidationError) as caught:
|
||||
config_module.validate_config(candidate)
|
||||
assert caught.value.field == field
|
||||
|
||||
|
||||
def test_mapping_is_longest_first_and_path_boundary_aware(config_module, tmp_path):
|
||||
user = tmp_path / "config.json"
|
||||
write_json(
|
||||
user,
|
||||
{
|
||||
"remote_url": "http://manager.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {
|
||||
"/models": "/mnt/general",
|
||||
"/models/special": "/mnt/special",
|
||||
},
|
||||
},
|
||||
)
|
||||
config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={})
|
||||
assert (
|
||||
config.map_path("/models/special/a.safetensors") == "/mnt/special/a.safetensors"
|
||||
)
|
||||
assert (
|
||||
config.map_path("/models/base.safetensors") == "/mnt/general/base.safetensors"
|
||||
)
|
||||
assert config.map_path("/models-old/a.safetensors") == "/models-old/a.safetensors"
|
||||
|
||||
|
||||
def test_mapping_rejects_separator_only_prefix_but_allows_root(config_module):
|
||||
candidate = {
|
||||
"remote_url": "http://manager.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {"////": "/mnt/invalid"},
|
||||
}
|
||||
with pytest.raises(config_module.ConfigValidationError) as caught:
|
||||
config_module.validate_config(candidate)
|
||||
assert caught.value.field == "path_mappings"
|
||||
|
||||
root_mapping = config_module.validate_config(
|
||||
{**candidate, "path_mappings": {"/": "/mnt/root"}}
|
||||
)
|
||||
snapshot = config_module.ConfigSnapshot(
|
||||
1,
|
||||
root_mapping["remote_url"],
|
||||
root_mapping["timeout"],
|
||||
config_module._normalize_mappings(root_mapping["path_mappings"]),
|
||||
)
|
||||
assert (
|
||||
snapshot.map_path("/models/a.safetensors") == "/mnt/root/models/a.safetensors"
|
||||
)
|
||||
|
||||
|
||||
def test_snapshot_mapping_does_not_change_after_reload(config_module, tmp_path):
|
||||
user = tmp_path / "config.json"
|
||||
write_json(
|
||||
user,
|
||||
{
|
||||
"remote_url": "http://one.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {"/remote": "/local-one"},
|
||||
},
|
||||
)
|
||||
config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={})
|
||||
original = config.snapshot
|
||||
revision = config.as_dict()["revision"]
|
||||
|
||||
config.save(
|
||||
{
|
||||
"remote_url": "http://two.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {"/remote": "/local-two"},
|
||||
},
|
||||
expected_revision=revision,
|
||||
)
|
||||
|
||||
assert original.map_path("/remote/a.safetensors") == "/local-one/a.safetensors"
|
||||
assert config.map_path("/remote/a.safetensors") == "/local-two/a.safetensors"
|
||||
|
||||
|
||||
def test_storage_reports_unwritable_when_parent_is_not_a_directory(
|
||||
config_module, tmp_path
|
||||
):
|
||||
blocking_file = tmp_path / "not-a-directory"
|
||||
blocking_file.write_text("blocked", encoding="utf-8")
|
||||
config = config_module.RemoteConfig(
|
||||
blocking_file / "config.json",
|
||||
tmp_path / "missing-legacy.json",
|
||||
environ={},
|
||||
)
|
||||
|
||||
assert config.as_dict()["storage"]["writable"] is False
|
||||
|
||||
|
||||
def test_stale_revision_does_not_overwrite_external_change(config_module, tmp_path):
|
||||
user = tmp_path / "config.json"
|
||||
initial = {"remote_url": "http://one.local", "timeout": 30, "path_mappings": {}}
|
||||
write_json(user, initial)
|
||||
config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={})
|
||||
revision = config.as_dict()["revision"]
|
||||
write_json(
|
||||
user,
|
||||
{"remote_url": "http://two.local", "timeout": 30, "path_mappings": {}},
|
||||
)
|
||||
|
||||
with pytest.raises(config_module.ConfigConflictError):
|
||||
config.save(initial, expected_revision=revision)
|
||||
assert (
|
||||
json.loads(user.read_text(encoding="utf-8"))["remote_url"] == "http://two.local"
|
||||
)
|
||||
|
||||
|
||||
def test_atomic_write_failure_preserves_file_and_live_snapshot(
|
||||
config_module, tmp_path, monkeypatch
|
||||
):
|
||||
user = tmp_path / "config.json"
|
||||
initial = {"remote_url": "http://one.local", "timeout": 30, "path_mappings": {}}
|
||||
write_json(user, initial)
|
||||
config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={})
|
||||
before = config.snapshot
|
||||
revision = config.as_dict()["revision"]
|
||||
|
||||
def fail_replace(source, target):
|
||||
raise OSError("simulated replace failure")
|
||||
|
||||
monkeypatch.setattr(config_module.os, "replace", fail_replace)
|
||||
with pytest.raises(OSError, match="simulated"):
|
||||
config.save(
|
||||
{"remote_url": "http://two.local", "timeout": 50, "path_mappings": {}},
|
||||
expected_revision=revision,
|
||||
)
|
||||
|
||||
assert json.loads(user.read_text(encoding="utf-8")) == initial
|
||||
assert config.snapshot == before
|
||||
assert not list(tmp_path.glob(".config.json.*.tmp"))
|
||||
|
||||
|
||||
def test_invalid_json_can_be_repaired_from_loaded_revision(config_module, tmp_path):
|
||||
user = tmp_path / "config.json"
|
||||
user.write_text("{not valid", encoding="utf-8")
|
||||
config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={})
|
||||
state = config.as_dict()
|
||||
assert state["effective"]["remote_url"] == ""
|
||||
assert state["warnings"]
|
||||
|
||||
config.save(
|
||||
{"remote_url": "http://fixed.local", "timeout": 30, "path_mappings": {}},
|
||||
expected_revision=state["revision"],
|
||||
)
|
||||
assert (
|
||||
json.loads(user.read_text(encoding="utf-8"))["remote_url"]
|
||||
== "http://fixed.local"
|
||||
)
|
||||
@@ -0,0 +1,617 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def modules():
|
||||
root = Path(__file__).resolve().parents[1]
|
||||
package_name = "lm_remote_proxy_test_package"
|
||||
package = types.ModuleType(package_name)
|
||||
package.__path__ = [str(root)]
|
||||
sys.modules[package_name] = package
|
||||
|
||||
loaded = {}
|
||||
for name in ("config", "remote_client", "proxy"):
|
||||
full_name = f"{package_name}.{name}"
|
||||
spec = importlib.util.spec_from_file_location(full_name, root / f"{name}.py")
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[full_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
loaded[name] = module
|
||||
|
||||
yield types.SimpleNamespace(**loaded)
|
||||
for name in ("proxy", "remote_client", "config"):
|
||||
sys.modules.pop(f"{package_name}.{name}", None)
|
||||
sys.modules.pop(package_name, None)
|
||||
|
||||
|
||||
class DummyContent:
|
||||
def __init__(self, body: bytes):
|
||||
self.body = body
|
||||
self.offset = 0
|
||||
|
||||
async def read(self, size: int) -> bytes:
|
||||
chunk = self.body[self.offset : self.offset + size]
|
||||
self.offset += len(chunk)
|
||||
return chunk
|
||||
|
||||
|
||||
class FragmentedContent:
|
||||
def __init__(self, chunks: list[bytes]):
|
||||
self.chunks = list(chunks)
|
||||
|
||||
async def read(self, size: int) -> bytes:
|
||||
if not self.chunks:
|
||||
return b""
|
||||
chunk = self.chunks.pop(0)
|
||||
if len(chunk) <= size:
|
||||
return chunk
|
||||
self.chunks.insert(0, chunk[size:])
|
||||
return chunk[:size]
|
||||
|
||||
|
||||
class DummyRequest:
|
||||
def __init__(
|
||||
self,
|
||||
method="GET",
|
||||
path="/",
|
||||
payload=None,
|
||||
headers=None,
|
||||
*,
|
||||
raw_body: bytes | None = None,
|
||||
content_length: int | None | object = ...,
|
||||
):
|
||||
self.method = method
|
||||
self.path = path
|
||||
self._payload = payload
|
||||
self.headers = headers or {}
|
||||
self.query_string = ""
|
||||
self.can_read_body = payload is not None
|
||||
self.content_type = (
|
||||
"application/json" if payload is not None else "application/octet-stream"
|
||||
)
|
||||
self._body = (
|
||||
raw_body
|
||||
if raw_body is not None
|
||||
else json.dumps(payload).encode("utf-8")
|
||||
if payload is not None
|
||||
else b""
|
||||
)
|
||||
self.content = DummyContent(self._body)
|
||||
if content_length is ...:
|
||||
self.content_length = len(self._body) if self._body else None
|
||||
else:
|
||||
self.content_length = content_length
|
||||
if raw_body is not None:
|
||||
self.content_type = "application/json"
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
async def read(self):
|
||||
return self._body
|
||||
|
||||
|
||||
def response_json(response: web.Response) -> dict:
|
||||
return json.loads(response.body.decode("utf-8"))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_proxy(modules, tmp_path, monkeypatch):
|
||||
config = modules.config.RemoteConfig(
|
||||
tmp_path / "user" / "config.json",
|
||||
tmp_path / "missing-legacy.json",
|
||||
environ={},
|
||||
)
|
||||
monkeypatch.setattr(modules.proxy, "remote_config", config)
|
||||
monkeypatch.setattr(modules.remote_client, "remote_config", config)
|
||||
modules.proxy.RemoteLoraClient._instance = None
|
||||
modules.proxy._proxy_sessions = {}
|
||||
modules.proxy._proxy_session_lock = None
|
||||
modules.proxy._active_proxy_websockets = set()
|
||||
modules.proxy._active_proxy_websockets_lock = None
|
||||
return modules.proxy, config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_endpoint_saves_and_hot_enables(isolated_proxy):
|
||||
proxy, config = isolated_proxy
|
||||
get_response = await proxy._handle_config(DummyRequest())
|
||||
initial = response_json(get_response)
|
||||
assert initial["configured"]["remote_url"] == ""
|
||||
|
||||
put_response = await proxy._handle_config(
|
||||
DummyRequest(
|
||||
"PUT",
|
||||
proxy._CONFIG_ROUTE,
|
||||
{
|
||||
"revision": initial["revision"],
|
||||
"config": {
|
||||
"remote_url": "http://manager.local:8188/",
|
||||
"timeout": 40,
|
||||
"path_mappings": {"/remote": "/local"},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
payload = response_json(put_response)
|
||||
assert put_response.status == 200
|
||||
assert payload["effective"]["remote_url"] == "http://manager.local:8188"
|
||||
assert payload["restart_required"] is False
|
||||
assert config.is_configured
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_falls_through_when_disabled_then_uses_new_url(
|
||||
isolated_proxy, monkeypatch
|
||||
):
|
||||
proxy, config = isolated_proxy
|
||||
|
||||
async def local_handler(request):
|
||||
return web.Response(text="local")
|
||||
|
||||
disabled_response = await proxy.lm_remote_proxy_middleware(
|
||||
DummyRequest(path="/loras"), local_handler
|
||||
)
|
||||
assert disabled_response.text == "local"
|
||||
|
||||
state = config.as_dict()
|
||||
config.save(
|
||||
{"remote_url": "http://manager.local", "timeout": 30, "path_mappings": {}},
|
||||
expected_revision=state["revision"],
|
||||
)
|
||||
captured = {}
|
||||
|
||||
async def fake_proxy_http(request, snapshot):
|
||||
captured["url"] = snapshot.remote_url
|
||||
return web.Response(text="remote")
|
||||
|
||||
monkeypatch.setattr(proxy, "_proxy_http", fake_proxy_http)
|
||||
enabled_response = await proxy.lm_remote_proxy_middleware(
|
||||
DummyRequest(path="/loras"), local_handler
|
||||
)
|
||||
assert enabled_response.text == "remote"
|
||||
assert captured["url"] == "http://manager.local"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_endpoint_rejects_changes_to_environment_managed_field(
|
||||
modules, tmp_path, monkeypatch
|
||||
):
|
||||
config = modules.config.RemoteConfig(
|
||||
tmp_path / "user.json",
|
||||
tmp_path / "legacy.json",
|
||||
environ={"LM_REMOTE_URL": "http://managed.local"},
|
||||
)
|
||||
monkeypatch.setattr(modules.proxy, "remote_config", config)
|
||||
initial = config.as_dict()
|
||||
response = await modules.proxy._handle_config(
|
||||
DummyRequest(
|
||||
"PUT",
|
||||
modules.proxy._CONFIG_ROUTE,
|
||||
{
|
||||
"revision": initial["revision"],
|
||||
"config": {
|
||||
"remote_url": "http://changed.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
payload = response_json(response)
|
||||
assert response.status == 409
|
||||
assert payload["field"] == "remote_url"
|
||||
assert "LM_REMOTE_URL" in payload["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conflict_reload_rotates_runtime_generation(isolated_proxy, monkeypatch):
|
||||
proxy, config = isolated_proxy
|
||||
initial = config.as_dict()
|
||||
config._config_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
config._config_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"remote_url": "http://external.local",
|
||||
"timeout": 20,
|
||||
"path_mappings": {},
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
retired = []
|
||||
rotated = []
|
||||
|
||||
async def fake_retire(generation):
|
||||
retired.append(generation)
|
||||
|
||||
async def fake_rotate(generation):
|
||||
rotated.append(generation)
|
||||
|
||||
monkeypatch.setattr(proxy, "_retire_proxy_sessions", fake_retire)
|
||||
monkeypatch.setattr(proxy, "_rotate_active_websockets", fake_rotate)
|
||||
response = await proxy._handle_config(
|
||||
DummyRequest(
|
||||
"PUT",
|
||||
proxy._CONFIG_ROUTE,
|
||||
{
|
||||
"revision": initial["revision"],
|
||||
"config": {
|
||||
"remote_url": "http://browser.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
payload = response_json(response)
|
||||
|
||||
assert response.status == 409
|
||||
assert payload["latest"]["effective"]["remote_url"] == "http://external.local"
|
||||
assert retired == [config.generation]
|
||||
assert rotated == [config.generation]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_test_uses_unsaved_draft(isolated_proxy, monkeypatch):
|
||||
proxy, _ = isolated_proxy
|
||||
captured = {}
|
||||
|
||||
async def fake_test(remote_url, timeout):
|
||||
captured.update(remote_url=remote_url, timeout=timeout)
|
||||
return 17
|
||||
|
||||
monkeypatch.setattr(proxy, "_perform_connection_test", fake_test)
|
||||
response = await proxy._handle_test_connection(
|
||||
DummyRequest(
|
||||
"POST",
|
||||
proxy._TEST_CONNECTION_ROUTE,
|
||||
{"remote_url": "http://draft.local:8188/", "timeout": 8},
|
||||
)
|
||||
)
|
||||
payload = response_json(response)
|
||||
assert response.status == 200
|
||||
assert payload["latency_ms"] == 17
|
||||
assert captured == {"remote_url": "http://draft.local:8188", "timeout": 8}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_test_rejects_empty_url(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
response = await proxy._handle_test_connection(
|
||||
DummyRequest(
|
||||
"POST",
|
||||
proxy._TEST_CONNECTION_ROUTE,
|
||||
{"remote_url": "", "timeout": 30},
|
||||
)
|
||||
)
|
||||
assert response.status == 400
|
||||
assert response_json(response)["field"] == "remote_url"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_response_reader_accepts_fragmented_json(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
response = types.SimpleNamespace(
|
||||
content=FragmentedContent([b'{"sta', b'tus":"', b'ok"}'])
|
||||
)
|
||||
|
||||
assert await proxy._read_small_json(response) == {"status": "ok"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_response_reader_enforces_hard_limit(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
response = types.SimpleNamespace(
|
||||
content=FragmentedContent([b" " * proxy._MAX_TEST_RESPONSE, b" "])
|
||||
)
|
||||
|
||||
with pytest.raises(proxy._ConnectionTestError, match="unexpectedly large"):
|
||||
await proxy._read_small_json(response)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_loop_header_is_rejected(isolated_proxy):
|
||||
proxy, config = isolated_proxy
|
||||
state = config.as_dict()
|
||||
config.save(
|
||||
{"remote_url": "http://manager.local", "timeout": 30, "path_mappings": {}},
|
||||
expected_revision=state["revision"],
|
||||
)
|
||||
|
||||
async def local_handler(request):
|
||||
return web.Response(text="local")
|
||||
|
||||
response = await proxy.lm_remote_proxy_middleware(
|
||||
DummyRequest(
|
||||
path="/api/lm/health-check",
|
||||
headers={proxy._PROXY_HOP_HEADER: "1"},
|
||||
),
|
||||
local_handler,
|
||||
)
|
||||
assert response.status == 508
|
||||
|
||||
|
||||
def test_register_proxy_is_available_while_unconfigured(isolated_proxy):
|
||||
proxy, config = isolated_proxy
|
||||
assert not config.is_configured
|
||||
app = web.Application()
|
||||
proxy.register_proxy(app)
|
||||
assert proxy.lm_remote_proxy_middleware in app.middlewares
|
||||
assert proxy._cleanup_proxy_session in app.on_shutdown
|
||||
|
||||
|
||||
def test_register_proxy_keeps_existing_security_middleware_first(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
|
||||
@web.middleware
|
||||
async def security_guard(request, handler):
|
||||
return await handler(request)
|
||||
|
||||
app = web.Application(middlewares=[security_guard])
|
||||
proxy.register_proxy(app)
|
||||
|
||||
assert list(app.middlewares) == [security_guard, proxy.lm_remote_proxy_middleware]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunked_config_body_is_hard_limited(isolated_proxy):
|
||||
proxy, _ = isolated_proxy
|
||||
response = await proxy._handle_config(
|
||||
DummyRequest(
|
||||
"PUT",
|
||||
proxy._CONFIG_ROUTE,
|
||||
raw_body=b" " * (proxy._MAX_CONFIG_BODY + 1),
|
||||
content_length=None,
|
||||
)
|
||||
)
|
||||
|
||||
assert response.status == 413
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_client_fetches_all_pages_beyond_server_cap(
|
||||
isolated_proxy, modules, monkeypatch
|
||||
):
|
||||
_, config = isolated_proxy
|
||||
client = modules.remote_client.RemoteLoraClient()
|
||||
calls = []
|
||||
|
||||
async def fake_get_json(path, params=None, *, snapshot=None):
|
||||
page = int(params["page"])
|
||||
calls.append((page, int(params["page_size"]), snapshot.generation))
|
||||
start = (page - 1) * 100
|
||||
count = 100 if page == 1 else 55
|
||||
return {
|
||||
"items": [
|
||||
{"file_name": f"model-{index}"} for index in range(start, start + count)
|
||||
],
|
||||
"total_pages": 2,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(client, "_get_json", fake_get_json)
|
||||
items = await client._get_lora_list_cached(snapshot=config.snapshot)
|
||||
|
||||
assert len(items) == 155
|
||||
assert items[-1]["file_name"] == "model-154"
|
||||
assert [call[:2] for call in calls] == [(1, 100), (2, 100)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("method_name", "expected_path"),
|
||||
[
|
||||
("_get_lora_list_cached", "/api/lm/loras/list"),
|
||||
("_get_checkpoint_list_cached", "/api/lm/checkpoints/list"),
|
||||
],
|
||||
)
|
||||
async def test_successful_empty_listing_is_cached(
|
||||
isolated_proxy, modules, monkeypatch, method_name, expected_path
|
||||
):
|
||||
_, config = isolated_proxy
|
||||
client = modules.remote_client.RemoteLoraClient()
|
||||
calls = []
|
||||
|
||||
async def fake_get_all_pages(path, *, snapshot):
|
||||
calls.append((path, snapshot.generation))
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(client, "_get_all_pages", fake_get_all_pages)
|
||||
cached_method = getattr(client, method_name)
|
||||
|
||||
assert await cached_method(snapshot=config.snapshot) == []
|
||||
assert await cached_method(snapshot=config.snapshot) == []
|
||||
assert calls == [(expected_path, config.generation)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_lookup_never_returns_another_generations_cache(
|
||||
isolated_proxy, modules, monkeypatch
|
||||
):
|
||||
_, config = isolated_proxy
|
||||
client = modules.remote_client.RemoteLoraClient()
|
||||
client._lora_cache = [{"file_name": "new-generation"}]
|
||||
client._lora_cache_generation = config.generation
|
||||
old_snapshot = modules.config.ConfigSnapshot(
|
||||
config.generation - 1,
|
||||
"http://old.local",
|
||||
30,
|
||||
(("/remote", "/old-local"),),
|
||||
)
|
||||
|
||||
async def fail_fetch(path, *, snapshot):
|
||||
raise OSError("old server unavailable")
|
||||
|
||||
monkeypatch.setattr(client, "_get_all_pages", fail_fetch)
|
||||
|
||||
assert await client._get_lora_list_cached(snapshot=old_snapshot) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lora_info_maps_with_the_fetch_generation(
|
||||
isolated_proxy, modules, monkeypatch
|
||||
):
|
||||
_, config = isolated_proxy
|
||||
initial = config.as_dict()
|
||||
config.save(
|
||||
{
|
||||
"remote_url": "http://one.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {"/remote": "/local-one"},
|
||||
},
|
||||
expected_revision=initial["revision"],
|
||||
)
|
||||
client = modules.remote_client.RemoteLoraClient()
|
||||
|
||||
async def fake_list(*, snapshot=None):
|
||||
latest = config.as_dict()
|
||||
config.save(
|
||||
{
|
||||
"remote_url": "http://two.local",
|
||||
"timeout": 30,
|
||||
"path_mappings": {"/remote": "/local-two"},
|
||||
},
|
||||
expected_revision=latest["revision"],
|
||||
)
|
||||
return [
|
||||
{
|
||||
"file_name": "portrait",
|
||||
"file_path": "/remote/portrait.safetensors",
|
||||
"folder": "",
|
||||
"civitai": {},
|
||||
}
|
||||
]
|
||||
|
||||
monkeypatch.setattr(client, "_get_lora_list_cached", fake_list)
|
||||
monkeypatch.setattr(client, "_relative_lora_path", lambda path, folder: path)
|
||||
|
||||
relative, _ = await client.get_lora_info("portrait")
|
||||
|
||||
assert relative == "/local-one/portrait.safetensors"
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_session_does_not_retain_remote_cookies(
|
||||
isolated_proxy, monkeypatch
|
||||
):
|
||||
proxy, config = isolated_proxy
|
||||
captured = {}
|
||||
|
||||
def create_session(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return FakeSession()
|
||||
|
||||
monkeypatch.setattr(proxy.aiohttp, "ClientSession", create_session)
|
||||
lease = proxy._proxy_session_lease(config.snapshot)
|
||||
await lease.__aenter__()
|
||||
await lease.__aexit__(None, None, None)
|
||||
|
||||
assert isinstance(captured["cookie_jar"], proxy.aiohttp.DummyCookieJar)
|
||||
await proxy._close_all_proxy_sessions()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_session_rotation_waits_for_inflight_request(
|
||||
isolated_proxy, monkeypatch
|
||||
):
|
||||
proxy, config = isolated_proxy
|
||||
created = []
|
||||
|
||||
def create_session(*args, **kwargs):
|
||||
session = FakeSession()
|
||||
created.append(session)
|
||||
return session
|
||||
|
||||
monkeypatch.setattr(proxy.aiohttp, "ClientSession", create_session)
|
||||
old_snapshot = config.snapshot
|
||||
old_lease = proxy._proxy_session_lease(old_snapshot)
|
||||
old_session = await old_lease.__aenter__()
|
||||
|
||||
state = config.as_dict()
|
||||
new_snapshot = config.save(
|
||||
{"remote_url": "http://new.local", "timeout": 30, "path_mappings": {}},
|
||||
expected_revision=state["revision"],
|
||||
)
|
||||
await proxy._retire_proxy_sessions(new_snapshot.generation)
|
||||
assert old_session.closed is False
|
||||
|
||||
new_lease = proxy._proxy_session_lease(new_snapshot)
|
||||
new_session = await new_lease.__aenter__()
|
||||
assert new_session is not old_session
|
||||
await new_lease.__aexit__(None, None, None)
|
||||
assert new_session.closed is False
|
||||
|
||||
await old_lease.__aexit__(None, None, None)
|
||||
assert old_session.closed is True
|
||||
assert new_session.closed is False
|
||||
|
||||
await proxy._close_all_proxy_sessions()
|
||||
assert new_session.closed is True
|
||||
|
||||
|
||||
class FakeWebSocket:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
self.close_code = None
|
||||
|
||||
async def close(self, *, code=None, message=None):
|
||||
self.closed = True
|
||||
self.close_code = code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_rotation_closes_only_retired_generation(isolated_proxy):
|
||||
proxy, config = isolated_proxy
|
||||
current_generation = config.generation
|
||||
old_bridge = proxy._ActiveWebSocket(
|
||||
current_generation - 1, FakeWebSocket(), FakeSession()
|
||||
)
|
||||
current_bridge = proxy._ActiveWebSocket(
|
||||
current_generation, FakeWebSocket(), FakeSession()
|
||||
)
|
||||
proxy._active_proxy_websockets.update({old_bridge, current_bridge})
|
||||
|
||||
await proxy._rotate_active_websockets(current_generation)
|
||||
|
||||
assert old_bridge.local_ws.closed is True
|
||||
assert old_bridge.local_ws.close_code == 1012
|
||||
assert old_bridge.session.closed is True
|
||||
assert current_bridge.local_ws.closed is False
|
||||
assert proxy._active_proxy_websockets == {current_bridge}
|
||||
|
||||
await proxy._rotate_active_websockets(None)
|
||||
assert current_bridge.local_ws.closed is True
|
||||
assert current_bridge.session.closed is True
|
||||
|
||||
|
||||
def test_mapped_local_path_is_resolved_against_comfy_lora_roots(
|
||||
modules, tmp_path, monkeypatch
|
||||
):
|
||||
root = tmp_path / "models" / "loras"
|
||||
fake_folder_paths = types.SimpleNamespace(
|
||||
get_folder_paths=lambda model_type: [str(root)]
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", fake_folder_paths)
|
||||
relative = modules.remote_client.RemoteLoraClient._relative_lora_path(
|
||||
str(root / "styles" / "portrait.safetensors"),
|
||||
"wrong-remote-folder",
|
||||
)
|
||||
assert relative == "styles/portrait.safetensors"
|
||||
Reference in New Issue
Block a user