110 lines
4.7 KiB
Python
110 lines
4.7 KiB
Python
import pytest
|
|
|
|
from app.gateway.worker_pool.discovery import discover, recommend
|
|
from app.gateway.worker_pool.settings import PoolSettings
|
|
|
|
|
|
def test_recommend_leaves_headroom_and_caps_starter():
|
|
result = recommend(32, 128 * 1024)
|
|
assert result["desired"] == 4
|
|
assert result["total_cpu"] == 8
|
|
assert result["total_memory_mb"] == 8192
|
|
|
|
|
|
@pytest.mark.parametrize("cpu,memory", [(1, 1024), (2, 2048), (0, 0)])
|
|
def test_small_host_has_no_viable_worker(cpu, memory):
|
|
assert recommend(cpu, memory)["desired"] == 0
|
|
|
|
|
|
def test_other_services_and_connections_limit_recommendation():
|
|
assert recommend(16, 32768, reserved_cpu=14)["desired"] == 1
|
|
assert recommend(32, 32768, reserved_memory_mb=30000)["desired"] == 0
|
|
assert recommend(32, 32768, connection_budget=12)["desired"] == 1
|
|
|
|
|
|
def test_discovery_excludes_own_workers_and_returns_no_secrets(monkeypatch):
|
|
def request(self, method, path):
|
|
if path == "/info":
|
|
return {"NCPU": 16, "MemTotal": 32768 * 1024**2}
|
|
if path == "/containers/json":
|
|
return [{"Id": "main", "Labels": {}}, {"Id": "worker", "Labels": {"deerflow.worker-pool": "mine"}}]
|
|
assert path == "/containers/main/json"
|
|
return {"Name": "/main", "Config": {"Env": ["PASSWORD=secret"]}, "HostConfig": {"NanoCpus": 2_000_000_000, "Memory": 4096 * 1024**2}, "NetworkSettings": {"Networks": {"bridge": {}, "app": {}}}}
|
|
monkeypatch.setattr("app.gateway.worker_pool.discovery.DockerEngine.request", request)
|
|
result = discover(PoolSettings(), "mine")
|
|
assert result["containers"] == [{"name": "main", "networks": ["app"]}]
|
|
assert result["recommendation"]["desired"] == 4
|
|
assert "secret" not in str(result)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_discovery_does_not_forward_saved_password_to_new_host(tmp_path, monkeypatch):
|
|
from types import SimpleNamespace
|
|
|
|
from app.gateway.worker_pool import router
|
|
from app.gateway.worker_pool.settings import SettingsStore
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(host="old-host", username="root", host_key_sha256="SHA256:old", password="original-secret"))
|
|
async def controller(request): return SimpleNamespace(store=store)
|
|
monkeypatch.setattr(router, "controller", controller)
|
|
seen = []
|
|
def inspect(settings, instance):
|
|
seen.append(settings.password)
|
|
return {"ok": True}
|
|
monkeypatch.setattr("app.gateway.worker_pool.discovery.discover", inspect)
|
|
await router.discover_server(PoolSettings(host="new-host", username="root", host_key_sha256="SHA256:new"), None)
|
|
await router.discover_server(PoolSettings(host="old-host", username="root", host_key_sha256="SHA256:old"), None)
|
|
assert seen == [None, "original-secret"]
|
|
|
|
|
|
def test_ssh_host_key_mismatch_stops_before_authentication(monkeypatch):
|
|
import paramiko
|
|
|
|
from app.gateway.worker_pool.docker import DockerEngine, DockerError
|
|
authenticated = []
|
|
class Client:
|
|
def set_missing_host_key_policy(self, policy): self.policy = policy
|
|
def connect(self, *args, **kwargs):
|
|
from types import SimpleNamespace
|
|
self.policy.missing_host_key(self, "host", SimpleNamespace(asbytes=lambda: b"wrong-key"))
|
|
authenticated.append(kwargs["password"])
|
|
def close(self): pass
|
|
monkeypatch.setattr(paramiko, "SSHClient", Client)
|
|
with pytest.raises(DockerError, match="指纹不匹配"):
|
|
DockerEngine(PoolSettings(host="host", password="secret", host_key_sha256="SHA256:expected")).request("GET", "/info")
|
|
assert not authenticated
|
|
|
|
|
|
def test_snapshot_worker_runtime_commits_main_container_without_pausing(monkeypatch):
|
|
from app.gateway.worker_pool.docker import DockerEngine
|
|
|
|
calls = []
|
|
|
|
def request(self, method, path, body=None, **kwargs):
|
|
calls.append((method, path, body))
|
|
return {"Id": "sha256:worker-runtime-image"}
|
|
|
|
monkeypatch.setattr(DockerEngine, "request", request)
|
|
image = DockerEngine(PoolSettings()).snapshot_worker_runtime(
|
|
{"Id": "main-container-id"}, "instance123",
|
|
)
|
|
assert image == "sha256:worker-runtime-image"
|
|
assert calls == [(
|
|
"POST",
|
|
"/commit?container=main-container-id&repo=deerflow-worker-runtime-instance123&tag=main-contain&pause=false",
|
|
None,
|
|
)]
|
|
|
|
|
|
def test_snapshot_worker_runtime_reuses_image_for_tiny_writable_layer(monkeypatch):
|
|
from app.gateway.worker_pool.docker import DockerEngine
|
|
|
|
def unexpected_request(*args, **kwargs):
|
|
pytest.fail("a tiny writable layer should reuse the immutable image")
|
|
|
|
monkeypatch.setattr(DockerEngine, "request", unexpected_request)
|
|
image = DockerEngine(PoolSettings()).snapshot_worker_runtime(
|
|
{"Id": "main-container-id", "Image": "sha256:base-image", "SizeRw": 65536}, "instance123",
|
|
)
|
|
assert image == "sha256:base-image"
|