464 lines
21 KiB
Python
464 lines
21 KiB
Python
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={})
|
|
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 ServerlessConfigRouteTests(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={})
|
|
|
|
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"))
|
|
|
|
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_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)
|
|
self.assertEqual(saved["concurrency"], 2)
|
|
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_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.assertEqual(worker["concurrency"], 3)
|
|
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", "concurrency": 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,
|
|
}]
|
|
}))
|
|
|
|
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"):
|
|
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,
|
|
}]
|
|
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"):
|
|
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"}}},
|
|
)
|
|
|
|
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"}}},
|
|
)
|
|
|
|
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_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_concurrency_change_is_blocked_during_active_run(self):
|
|
self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}]
|
|
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"}}},
|
|
)
|
|
updated = copy.deepcopy(self.config["workers"])
|
|
updated[0]["concurrency"] = 2
|
|
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": updated}))
|
|
self.assertEqual(response.status, 409)
|
|
self.assertEqual(self.read_config_file()["workers"][0]["concurrency"], 1)
|
|
|
|
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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|