import copy import importlib.util import json import sys import tempfile import types import unittest from pathlib import Path from unittest.mock import patch class _FakeResponse: def __init__(self, payload, status=200): self.payload = copy.deepcopy(payload) self.status = status class _FakeRequest: def __init__(self, payload=None, error=None): self.payload = payload self.error = error async def json(self): if self.error is not None: raise self.error return self.payload class _Routes: def get(self, _path): return lambda handler: handler def post(self, _path): return lambda handler: handler def _load_modules(): repo = Path(__file__).resolve().parents[1] package_name = "dist_serverless_routes_testpkg" for name in list(sys.modules): if name == package_name or name.startswith(f"{package_name}."): del sys.modules[name] root = types.ModuleType(package_name) root.__path__ = [str(repo)] sys.modules[package_name] = root api_pkg = types.ModuleType(f"{package_name}.api") api_pkg.__path__ = [str(repo / "api")] sys.modules[f"{package_name}.api"] = api_pkg utils_pkg = types.ModuleType(f"{package_name}.utils") utils_pkg.__path__ = [str(repo / "utils")] sys.modules[f"{package_name}.utils"] = utils_pkg orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration") orchestration_pkg.__path__ = [str(repo / "api" / "orchestration")] sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg logging_module = types.ModuleType(f"{package_name}.utils.logging") logging_module.log = lambda *_args, **_kwargs: None logging_module.debug_log = lambda *_args, **_kwargs: None sys.modules[logging_module.__name__] = logging_module constants_module = types.ModuleType(f"{package_name}.utils.constants") constants_module.HEARTBEAT_TIMEOUT = 30 sys.modules[constants_module.__name__] = constants_module def load_file(module_name, path): spec = importlib.util.spec_from_file_location(module_name, 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) return module config_module = load_file(f"{package_name}.utils.config", repo / "utils" / "config.py") serverless_config = load_file( f"{package_name}.utils.serverless_config", repo / "utils" / "serverless_config.py" ) runtime = types.ModuleType(f"{package_name}.api.orchestration.serverless_runtime") runtime.secret_store = None runtime.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False) runtime.get_secret_store = lambda: runtime.secret_store runtime.get_manager = lambda: runtime.manager sys.modules[runtime.__name__] = runtime import aiohttp # Ensure the real dependency is loaded; never replace global aiohttp. previous_server = sys.modules.get("server") server = types.ModuleType("server") server.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes())) sys.modules["server"] = server network = types.ModuleType(f"{package_name}.utils.network") async def handle_api_error(_request, error, status=500): return _FakeResponse({"status": "error", "message": str(error)}, status) network.handle_api_error = handle_api_error network.normalize_host = lambda value: value sys.modules[network.__name__] = network routes = load_file(f"{package_name}.api.config_routes", repo / "api" / "config_routes.py") routes.web = types.SimpleNamespace( json_response=lambda payload, status=200: _FakeResponse(payload, status) ) if previous_server is None: sys.modules.pop("server", None) else: sys.modules["server"] = previous_server return config_module, serverless_config, runtime, routes config_module, serverless_config, runtime_module, config_routes = _load_modules() class _RunningTask: def done(self): return False class _ConfigRouteFixture(unittest.IsolatedAsyncioTestCase): def setUp(self): self.tempdir = tempfile.TemporaryDirectory() self.root = Path(self.tempdir.name) config_module.CONFIG_FILE = str(self.root / "gpu_config.json") config_module.invalidate_config_cache() self.config = config_module.get_default_config() self.config["workers"] = [] self.assertTrue(config_module.save_config(self.config)) self.secrets = serverless_config.SecretStore(self.root / "private" / "credentials.json") runtime_module.secret_store = self.secrets runtime_module.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False) def tearDown(self): config_module.invalidate_config_cache() self.tempdir.cleanup() def read_config_file(self): return json.loads(Path(config_module.CONFIG_FILE).read_text(encoding="utf-8")) class ServerlessConfigRouteTests(_ConfigRouteFixture): async def test_get_config_redacts_keys_and_reports_only_configured_boolean(self): self.secrets.set("pool-1", "private-api-key") self.config["workers"] = [ {"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-1", "api_key": "inline-api-key", "name": "GPU"} ] Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8") config_module.invalidate_config_cache() response = await config_routes.get_config_endpoint(_FakeRequest()) self.assertEqual(response.status, 200) worker = response.payload["workers"][0] self.assertTrue(worker["api_key_configured"]) self.assertNotIn("api_key", worker) self.assertNotIn("private-api-key", json.dumps(response.payload)) self.assertNotIn("inline-api-key", json.dumps(response.payload)) async def test_get_config_redacts_normalized_secret_field_aliases(self): aliases = { "apiKey": "camel-api-key", "APIKEY": "upper-api-key", "clearApiKey": "camel-clear-key", "CLEAR_API_KEY": "upper-clear-key", "runpod_CREDENTIAL": "provider-credential", } self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", **aliases, "nested": {"credential_note": "nested-credential"}, }] Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8") config_module.invalidate_config_cache() response = await config_routes.get_config_endpoint(_FakeRequest()) self.assertEqual(response.status, 200) worker = response.payload["workers"][0] for alias, value in aliases.items(): self.assertNotIn(alias, worker) self.assertNotIn(value, json.dumps(response.payload)) self.assertNotIn("credential_note", worker["nested"]) self.assertNotIn("nested-credential", json.dumps(response.payload)) async def test_bulk_config_rejects_normalized_secret_field_aliases(self): for secret_field in ("apiKey", "APIKEY", "clearApiKey", "RUNPOD_CREDENTIAL"): with self.subTest(secret_field=secret_field): secret = f"private-{secret_field}" response = await config_routes.update_config_endpoint(_FakeRequest({ "workers": [{"id": "pool-1", "type": "local", secret_field: secret}] })) self.assertEqual(response.status, 400) self.assertNotIn(secret, json.dumps(response.payload)) self.assertEqual(self.read_config_file()["workers"], []) async def test_create_serverless_worker_without_local_fields_and_store_key_privately(self): secret = "runpod-secret-never-returned" payload = { "worker_id": "pool-1", "name": "GPU pool", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", "jobs_per_run": 3, "concurrency": 2, "timeout_seconds": 600, "enabled": True, "api_key": secret, } response = await config_routes.update_worker_endpoint(_FakeRequest(payload)) self.assertEqual(response.status, 200) self.assertTrue(response.payload["api_key_configured"]) self.assertNotIn(secret, json.dumps(response.payload)) saved = self.read_config_file()["workers"][0] self.assertEqual(saved["endpoint_id"], "ep-123") self.assertEqual(saved["jobs_per_run"], 3) # Legacy field from older browser tabs is accepted but not persisted. self.assertNotIn("concurrency", saved) self.assertEqual(saved["timeout_seconds"], 600) for obsolete in ("host", "port", "cuda_device", "extra_args", "api_key", "clear_api_key"): self.assertNotIn(obsolete, saved) self.assertEqual(self.secrets.get("pool-1"), secret) async def test_media_profile_persists_and_is_public_without_changing_credentials(self): for profile in ("legacy-image-v1", "comfy-media-v1"): with self.subTest(profile=profile): response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "media_profile": profile, "api_key": "saved-key", "enabled": False, })) self.assertEqual(response.status, 200) self.assertEqual(self.read_config_file()["workers"][0]["media_profile"], profile) config_module.invalidate_config_cache() public = await config_routes.get_config_endpoint(_FakeRequest()) self.assertEqual(public.status, 200) self.assertEqual(public.payload["workers"][0]["media_profile"], profile) self.assertTrue(public.payload["workers"][0]["api_key_configured"]) self.assertNotIn("saved-key", json.dumps(public.payload)) self.assertEqual(self.secrets.get("pool-1"), "saved-key") async def test_old_clients_omit_media_profile_without_downgrading_saved_profile(self): for stored in ({}, {"media_profile": "comfy-media-v1"}): with self.subTest(stored=stored): worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", **stored} self.config["workers"] = [worker] self.assertTrue(config_module.save_config(self.config)) response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "name": "Old client save", "jobs_per_run": 2, })) self.assertEqual(response.status, 200) saved = self.read_config_file()["workers"][0] self.assertEqual({k: v for k, v in saved.items() if k == "media_profile"}, stored) # The bulk route is also used by older clients with no profile field. response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{ "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "jobs_per_run": 3, }]})) self.assertEqual(response.status, 200) saved = self.read_config_file()["workers"][0] self.assertEqual({k: v for k, v in saved.items() if k == "media_profile"}, stored) async def test_invalid_media_profile_is_rejected_by_worker_and_bulk_routes(self): for profile in ("auto", "comfy-media-v2", None, False, 1, [], {}): with self.subTest(profile=profile): response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "media_profile": profile, "api_key": "never-save", })) self.assertEqual(response.status, 400) response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{ "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "media_profile": profile, }]})) self.assertEqual(response.status, 400) self.assertEqual(self.read_config_file()["workers"], []) self.assertIsNone(self.secrets.get("pool-1")) async def test_active_pool_blocks_media_profile_changes_in_both_directions(self): for old_profile, new_profile in ((None, "comfy-media-v1"), ("comfy-media-v1", "legacy-image-v1")): with self.subTest(old_profile=old_profile): worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"} if old_profile is not None: worker["media_profile"] = old_profile self.config["workers"] = [worker] self.assertTrue(config_module.save_config(self.config)) original = self.read_config_file() with patch.object(config_routes, "_serverless_pool_is_active", return_value=True): response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "media_profile": new_profile, })) self.assertEqual(response.status, 409) response = await config_routes.update_config_endpoint(_FakeRequest({ "workers": [dict(worker, media_profile=new_profile)], })) self.assertEqual(response.status, 409) self.assertEqual(self.read_config_file(), original) async def test_enable_only_update_preserves_serverless_settings_and_saved_key(self): self.secrets.set("pool-1", "saved-secret") self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "provider": "runpod", "name": "Pool", "endpoint_id": "ep-123", "jobs_per_run": 4, "concurrency": 3, "timeout_seconds": 900, "enabled": False, }] self.assertTrue(config_module.save_config(self.config)) response = await config_routes.update_worker_endpoint( _FakeRequest({"worker_id": "pool-1", "enabled": True}) ) worker = self.read_config_file()["workers"][0] self.assertTrue(worker["enabled"]) self.assertEqual(worker["endpoint_id"], "ep-123") self.assertEqual(worker["jobs_per_run"], 4) self.assertNotIn("concurrency", worker) self.assertEqual(worker["timeout_seconds"], 900) self.assertEqual(self.secrets.get("pool-1"), "saved-secret") self.assertTrue(response.payload["api_key_configured"]) async def test_invalid_serverless_record_is_not_persisted_or_secret_written(self): secret = "must-not-leak" response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", "jobs_per_run": True, "api_key": secret, })) self.assertEqual(response.status, 400) self.assertNotIn(secret, json.dumps(response.payload)) self.assertEqual(self.read_config_file()["workers"], []) self.assertIsNone(self.secrets.get("pool-1")) async def test_json_decode_error_is_generic_and_does_not_echo_body(self): secret = "body-secret-value" response = await config_routes.update_worker_endpoint( _FakeRequest(error=ValueError(f"bad json near {secret}")) ) self.assertEqual(response.status, 400) self.assertNotIn(secret, json.dumps(response.payload)) async def test_bulk_config_rejects_api_keys_at_any_depth_without_persisting(self): secret = "bulk-secret-never-write" response = await config_routes.update_config_endpoint(_FakeRequest({ "workers": [{ "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", "nested": {"api_key": secret}, }] })) self.assertEqual(response.status, 400) self.assertNotIn(secret, json.dumps(response.payload)) self.assertEqual(self.read_config_file()["workers"], []) self.assertFalse(self.secrets.path.exists()) async def test_bulk_update_validates_and_returns_redacted_serverless_config(self): self.secrets.set("pool-1", "saved-secret") response = await config_routes.update_config_endpoint(_FakeRequest({ "workers": [{ "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", "name": "Pool", "enabled": True, "host": "obsolete-host", "port": 8188, "cuda_device": 0, }] })) self.assertEqual(response.status, 200) self.assertTrue(response.payload["config"]["workers"][0]["api_key_configured"]) saved = self.read_config_file()["workers"][0] for obsolete in ("host", "port", "cuda_device"): self.assertNotIn(obsolete, saved) self.assertNotIn("api_key", response.payload["config"]["workers"][0]) self.assertEqual(self.secrets.get("pool-1"), "saved-secret") async def test_bulk_type_conversion_clears_secret_and_serverless_fields(self): self.secrets.set("pool-1", "bulk-clear-key") self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", }] self.assertTrue(config_module.save_config(self.config)) response = await config_routes.update_config_endpoint(_FakeRequest({ "workers": [{ "id": "pool-1", "type": "local", "port": 8188, "provider": "runpod", "endpoint_id": "legacy-ep", "api_key_configured": True, "media_profile": "comfy-media-v1", }] })) self.assertEqual(response.status, 200) self.assertIsNone(self.secrets.get("pool-1")) saved = self.read_config_file()["workers"][0] for field in ("provider", "endpoint_id", "api_key_configured", "media_profile"): self.assertNotIn(field, saved) async def test_bulk_invalid_serverless_record_does_not_replace_configuration(self): existing = {"id": "local-1", "type": "local", "port": 8188} self.config["workers"] = [existing] self.assertTrue(config_module.save_config(self.config)) response = await config_routes.update_config_endpoint(_FakeRequest({ "workers": [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}] })) self.assertEqual(response.status, 400) self.assertEqual(self.read_config_file()["workers"], [existing]) async def test_delete_serverless_worker_clears_saved_key(self): self.secrets.set("pool-1", "delete-me") self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123" }] self.assertTrue(config_module.save_config(self.config)) response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) self.assertEqual(response.status, 200) self.assertEqual(self.read_config_file()["workers"], []) self.assertIsNone(self.secrets.get("pool-1")) async def test_type_conversion_clears_secret_and_removes_serverless_fields(self): self.secrets.set("pool-1", "old-key") self.config["workers"] = [{ "id": "pool-1", "name": "Pool", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, "timeout_seconds": 600, "media_profile": "comfy-media-v1", }] self.assertTrue(config_module.save_config(self.config)) response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "type": "local", "name": "Local", "port": 8188, "cuda_device": 0, "enabled": False, })) self.assertEqual(response.status, 200) self.assertIsNone(self.secrets.get("pool-1")) worker = self.read_config_file()["workers"][0] self.assertEqual(worker["type"], "local") for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "media_profile"): self.assertNotIn(serverless_field, worker) async def test_enabled_toggle_allowed_during_active_run_but_endpoint_change_is_rejected(self): self.secrets.set("pool-1", "saved-key") self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", "enabled": False, }] self.assertTrue(config_module.save_config(self.config)) runtime_module.manager = types.SimpleNamespace( tasks={"run-1": _RunningTask()}, runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}}, pool_is_active=lambda worker_id: worker_id == "pool-1", ) enabled = await config_routes.update_worker_endpoint( _FakeRequest({"worker_id": "pool-1", "enabled": True}) ) blocked = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "endpoint_id": "ep-456", })) self.assertEqual(enabled.status, 200) self.assertEqual(blocked.status, 409) self.assertEqual(self.read_config_file()["workers"][0]["endpoint_id"], "ep-123") async def test_blank_api_key_keeps_saved_key_and_clear_api_key_deletes_it(self): self.secrets.set("pool-1", "saved-key") self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", }] self.assertTrue(config_module.save_config(self.config)) keep = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "api_key": "", })) clear = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "clear_api_key": True, })) self.assertTrue(keep.payload["api_key_configured"]) self.assertFalse(clear.payload["api_key_configured"]) self.assertIsNone(self.secrets.get("pool-1")) async def test_config_save_refuses_inline_api_keys_and_invalid_serverless_records(self): with self.subTest(case="inline key"): config = config_module.get_default_config() config["workers"] = [{"id": "pool-1", "type": "local", "api_key": "private"}] self.assertFalse(config_module.save_config(config)) with self.subTest(case="invalid serverless record"): config = config_module.get_default_config() config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}] self.assertFalse(config_module.save_config(config)) async def test_active_pool_blocks_credential_changes_type_conversion_and_deletion(self): self.secrets.set("pool-1", "saved-key") self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", }] self.assertTrue(config_module.save_config(self.config)) runtime_module.manager = types.SimpleNamespace( tasks={"run-1": _RunningTask()}, runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}}, pool_is_active=lambda worker_id: worker_id == "pool-1", ) credential_change = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "api_key": "new-key", })) type_change = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "type": "local", "port": 8188, "cuda_device": 0, })) deletion = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) self.assertEqual(credential_change.status, 409) self.assertEqual(type_change.status, 409) self.assertEqual(deletion.status, 409) self.assertEqual(self.secrets.get("pool-1"), "saved-key") self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless") async def test_secret_write_is_rolled_back_when_config_persistence_fails(self): secret = "write-then-rollback" with patch.object(config_module, "save_config", return_value=False): response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "api_key": secret, })) self.assertEqual(response.status, 500) self.assertIsNone(self.secrets.get("pool-1")) self.assertEqual(self.read_config_file()["workers"], []) self.assertEqual(config_module.load_config()["workers"], []) async def test_bulk_secret_delete_is_rolled_back_when_config_persistence_fails(self): self.secrets.set("pool-1", "bulk-restore-key") self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", }] self.assertTrue(config_module.save_config(self.config)) with patch.object(config_module, "save_config", return_value=False): response = await config_routes.update_config_endpoint(_FakeRequest({ "workers": [{"id": "pool-1", "type": "local", "port": 8188}] })) self.assertEqual(response.status, 500) self.assertEqual(self.secrets.get("pool-1"), "bulk-restore-key") self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless") async def test_secret_replacement_is_rolled_back_when_config_persistence_fails(self): self.secrets.set("pool-1", "prior-key") self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", }] self.assertTrue(config_module.save_config(self.config)) with patch.object(config_module, "save_config", return_value=False): response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "api_key": "replacement-key", })) self.assertEqual(response.status, 500) self.assertEqual(self.secrets.get("pool-1"), "prior-key") self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless") async def test_delete_secret_is_rolled_back_when_config_persistence_fails(self): self.secrets.set("pool-1", "delete-restore-key") self.config["workers"] = [{ "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", }] self.assertTrue(config_module.save_config(self.config)) with patch.object(config_module, "save_config", return_value=False): response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) self.assertEqual(response.status, 500) self.assertEqual(self.secrets.get("pool-1"), "delete-restore-key") self.assertEqual(self.read_config_file()["workers"][0]["id"], "pool-1") async def test_regular_worker_delete_does_not_require_posix_secret_storage(self): self.config["workers"] = [{"id": "local-1", "type": "local", "port": 8188}] self.assertTrue(config_module.save_config(self.config)) with patch.object(config_routes, "_serverless_secret_store", side_effect=RuntimeError("POSIX only")): response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "local-1"})) self.assertEqual(response.status, 200) self.assertEqual(self.read_config_file()["workers"], []) async def test_enabling_serverless_without_key_is_rejected(self): response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "enabled": True, })) self.assertEqual(response.status, 400) self.assertEqual(self.read_config_file()["workers"], []) async def test_bulk_endpoint_change_is_blocked_during_active_run(self): self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"}] self.assertTrue(config_module.save_config(self.config)) runtime_module.manager = types.SimpleNamespace( tasks={"run-1": _RunningTask()}, runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}}, pool_is_active=lambda worker_id: worker_id == "pool-1", ) updated = copy.deepcopy(self.config["workers"]) updated[0]["endpoint_id"] = "ep-other" response = await config_routes.update_config_endpoint(_FakeRequest({"workers": updated})) self.assertEqual(response.status, 409) self.assertEqual(self.read_config_file()["workers"][0]["endpoint_id"], "ep-123") async def test_non_serverless_update_behavior_is_preserved(self): self.config["workers"] = [{"id": "local-1", "name": "Local", "type": "local", "port": 8188}] self.assertTrue(config_module.save_config(self.config)) response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "local-1", "name": "Renamed", "enabled": True, })) self.assertEqual(response.status, 200) worker = self.read_config_file()["workers"][0] self.assertEqual(worker["name"], "Renamed") self.assertTrue(worker["enabled"]) async def test_config_normalization_removes_legacy_local_fields_from_serverless_records(self): raw = { "workers": [{ "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "host": "localhost", "port": 8188, "cuda_device": 2, "extra_args": "--old", }] } normalized = config_module._merge_with_defaults(raw, config_module.get_default_config()) worker = normalized["workers"][0] for obsolete in ("host", "port", "cuda_device", "extra_args"): self.assertNotIn(obsolete, worker) class PersistedPoolGuardTests(_ConfigRouteFixture): """Use the actual manager/ledger across restart, not a mocked idle result.""" def setUp(self): super().setUp() from test_serverless_runtime import runtime from unittest.mock import AsyncMock self.runtime = runtime self.session_factory = AsyncMock() self.ledger_path = self.root / "jobs.json" self.manager = runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory) runtime_module.manager = self.manager self.secrets.set("pool-1", "saved-key") self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}] self.assertTrue(config_module.save_config(self.config)) async def test_retained_jobs_block_destructive_changes_after_restart(self): original = self.read_config_file() for status in ("UNKNOWN", "CANCEL_REQUESTED", "IN_QUEUE", "IN_PROGRESS", "DOWNLOADING", "SUBMITTING"): for restart in (False, True): with self.subTest(status=status, restart=restart): self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status=status, job_id=None if status in {"UNKNOWN", "SUBMITTING"} else "owned-id") if restart: self.manager = self.runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory) runtime_module.manager = self.manager else: self.manager.tasks["retained"] = types.SimpleNamespace(done=lambda: True) for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"type": "local", "port": 8188}, {"media_profile": "comfy-media-v1"}): response = await config_routes.update_worker_endpoint(_FakeRequest(dict(fields, worker_id="pool-1"))) self.assertEqual(response.status, 409, fields) response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) self.assertEqual(response.status, 409) worker = original["workers"][0] for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")], [dict(worker, media_profile="comfy-media-v1")]): response = await config_routes.update_config_endpoint(_FakeRequest({"workers": workers})) self.assertEqual(response.status, 409, workers) self.assertEqual(self.read_config_file(), original) self.assertEqual(self.secrets.get("pool-1"), "saved-key") self.session_factory.assert_not_awaited() async def test_retained_jobs_allow_same_or_omitted_media_profile(self): self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status="UNKNOWN") for profile in ("legacy-image-v1", "comfy-media-v1"): with self.subTest(profile=profile): worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"} # An absent legacy field and an explicit legacy selection are equivalent. if profile != "legacy-image-v1": worker["media_profile"] = profile self.config["workers"] = [worker] self.assertTrue(config_module.save_config(self.config)) for fields in ({"media_profile": profile}, {}): response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "name": "Same profile", **fields, })) self.assertEqual(response.status, 200) response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{ "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", **fields, }]})) self.assertEqual(response.status, 200) self.assertEqual(self.read_config_file()["workers"][0]["media_profile"], profile) self.session_factory.assert_not_awaited() async def test_retained_jobs_allow_images_per_run_and_legacy_concurrency_saves(self): # Upgrading from a config with concurrency=1 must not trip the busy guard, # and the image count stays editable while old jobs are unresolved. self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status="UNKNOWN") response = await config_routes.update_worker_endpoint(_FakeRequest({ "worker_id": "pool-1", "endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, })) self.assertEqual(response.status, 200) saved = self.read_config_file()["workers"][0] self.assertEqual(saved["jobs_per_run"], 2) self.assertNotIn("concurrency", saved) response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [dict(saved, jobs_per_run=3, concurrency=5)]})) self.assertEqual(response.status, 200) self.assertEqual(self.read_config_file()["workers"][0]["jobs_per_run"], 3) async def test_retained_job_allows_disable_and_rename_but_not_key_clear(self): self.manager.ledger.update("retained", worker_id="pool-1", status="UNKNOWN") response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "enabled": False, "name": "Paused"})) self.assertEqual(response.status, 200) self.assertEqual(self.read_config_file()["workers"][0]["name"], "Paused") response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "clear_api_key": True})) self.assertEqual(response.status, 409) self.assertEqual(self.secrets.get("pool-1"), "saved-key") async def test_terminal_or_other_pool_records_do_not_block_deletion(self): self.manager.ledger.update("done", worker_id="pool-1", status="RESOLVED") self.manager.ledger.update("other", worker_id="other", status="UNKNOWN") response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) self.assertEqual(response.status, 200) self.assertIsNone(self.secrets.get("pool-1")) if __name__ == "__main__": unittest.main()