This commit is contained in:
@@ -1,9 +1,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import socket
|
||||
from contextlib import suppress
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import httpx
|
||||
|
||||
from cloud.internal_api.models import (
|
||||
AssignmentModel,
|
||||
DeviceEnrollmentResponse,
|
||||
@@ -16,6 +19,12 @@ from host_agent.identity import HostIdentityStore
|
||||
from storage.device_config import DeviceConfigStore
|
||||
|
||||
|
||||
def _free_loopback_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
return probe.getsockname()[1]
|
||||
|
||||
|
||||
def _config() -> HostAgentConfig:
|
||||
return HostAgentConfig(
|
||||
control_plane_url="https://control.example",
|
||||
@@ -88,7 +97,6 @@ def test_create_application_enrolls_host_and_devices_before_managed_startup(
|
||||
def __init__(self) -> None:
|
||||
self.config = HostAgentConfig(
|
||||
control_plane_url="https://control.example",
|
||||
enrollment_token="one-time-token",
|
||||
enrollment_managed=True,
|
||||
)
|
||||
|
||||
@@ -352,3 +360,98 @@ def test_main_task_cancellation_waits_for_active_work_shutdown() -> None:
|
||||
assert events == ["work-finished", "final-heartbeat", "closed"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_console_enabled_serves_http_and_shuts_down_cleanly(tmp_path) -> None:
|
||||
async def scenario() -> None:
|
||||
port = _free_loopback_port()
|
||||
config = HostAgentConfig(
|
||||
control_plane_url="https://control.example",
|
||||
host_id="host-a",
|
||||
token="secret",
|
||||
identity_path=tmp_path / "host_identity.json",
|
||||
local_account_path=tmp_path / "host_local_account.json",
|
||||
console_enabled=True,
|
||||
console_bind_host="127.0.0.1",
|
||||
console_port=port,
|
||||
)
|
||||
application = create_application(config=config, manager=DeviceManager())
|
||||
assert application.console_server is not None
|
||||
|
||||
claim_started = asyncio.Event()
|
||||
claim_cancelled = asyncio.Event()
|
||||
events: list[str] = []
|
||||
|
||||
class BlockingClient:
|
||||
async def claim(self):
|
||||
claim_started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
claim_cancelled.set()
|
||||
raise
|
||||
|
||||
async def aclose(self):
|
||||
events.append("closed")
|
||||
|
||||
class RecordingHeartbeat:
|
||||
async def run(self, stop):
|
||||
await stop.wait()
|
||||
|
||||
async def sync_once(self):
|
||||
events.append("final-heartbeat")
|
||||
|
||||
class IdleProcessor:
|
||||
async def process(self, assignment):
|
||||
raise AssertionError("no assignment expected")
|
||||
|
||||
def request_stop(self):
|
||||
events.append("stop-work")
|
||||
|
||||
application.client = BlockingClient() # type: ignore[assignment]
|
||||
application.heartbeat = RecordingHeartbeat() # type: ignore[assignment]
|
||||
application.processor = IdleProcessor() # type: ignore[assignment]
|
||||
|
||||
stop = asyncio.Event()
|
||||
running = asyncio.create_task(application.run_async(stop))
|
||||
await claim_started.wait()
|
||||
|
||||
response: httpx.Response | None = None
|
||||
async with httpx.AsyncClient() as http_client:
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = loop.time() + 5
|
||||
while loop.time() < deadline:
|
||||
try:
|
||||
response = await http_client.get(
|
||||
f"http://127.0.0.1:{port}/login", timeout=0.5
|
||||
)
|
||||
except httpx.TransportError:
|
||||
await asyncio.sleep(0.05)
|
||||
continue
|
||||
break
|
||||
assert response is not None
|
||||
assert response.status_code == 200
|
||||
assert "Login" in response.text
|
||||
|
||||
stop.set()
|
||||
await asyncio.wait_for(running, timeout=5)
|
||||
|
||||
assert claim_cancelled.is_set()
|
||||
assert events == ["stop-work", "final-heartbeat", "closed"]
|
||||
assert application.console_server.should_exit is True
|
||||
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
|
||||
probe.settimeout(0.5)
|
||||
with suppress(ConnectionRefusedError, OSError):
|
||||
probe.connect(("127.0.0.1", port))
|
||||
raise AssertionError("console socket should be closed after shutdown")
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_console_disabled_by_default_opens_no_socket(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
application = create_application(config=_config(), manager=DeviceManager())
|
||||
|
||||
assert application.console_server is None
|
||||
asyncio.run(application.client.aclose())
|
||||
|
||||
@@ -172,7 +172,7 @@ def test_result_report_retries_identical_payload_after_response_loss() -> None:
|
||||
assert payloads[0]["failure_reason"] == "planner unavailable"
|
||||
|
||||
|
||||
def test_bootstrap_client_retries_identical_enrollment_and_enrolls_device() -> None:
|
||||
def test_bootstrap_client_directly_enrolls_and_enrolls_device() -> None:
|
||||
requests: list[httpx.Request] = []
|
||||
host_attempts = 0
|
||||
|
||||
@@ -189,7 +189,6 @@ def test_bootstrap_client_retries_identical_enrollment_and_enrolls_device() -> N
|
||||
config = _config(
|
||||
host_id="",
|
||||
token="",
|
||||
enrollment_token="one-time-token",
|
||||
enrollment_managed=True,
|
||||
)
|
||||
with httpx.Client(
|
||||
@@ -209,7 +208,6 @@ def test_bootstrap_client_retries_identical_enrollment_and_enrolls_device() -> N
|
||||
client.config = _config(
|
||||
host_id=host.host_id,
|
||||
token="host-token-" + ("x" * 40),
|
||||
enrollment_token="one-time-token",
|
||||
enrollment_managed=True,
|
||||
)
|
||||
device = client.enroll_device(
|
||||
@@ -223,38 +221,5 @@ def test_bootstrap_client_retries_identical_enrollment_and_enrolls_device() -> N
|
||||
assert device.device_id == "device-cloud-a"
|
||||
assert len(requests) == 3
|
||||
assert requests[0].content == requests[1].content
|
||||
assert requests[0].headers["authorization"] == "Bearer one-time-token"
|
||||
assert requests[2].headers["authorization"] == ("Bearer host-token-" + ("x" * 40))
|
||||
|
||||
|
||||
def test_self_service_enrollment_sends_no_authorization_header() -> None:
|
||||
requests: list[httpx.Request] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
return httpx.Response(201, json={"host_id": "host-cloud-a"})
|
||||
|
||||
config = _config(
|
||||
host_id="",
|
||||
token="",
|
||||
enrollment_token="",
|
||||
enrollment_managed=True,
|
||||
)
|
||||
with httpx.Client(
|
||||
transport=httpx.MockTransport(handler),
|
||||
base_url="https://control.example",
|
||||
) as http_client:
|
||||
client = HostAgentEnrollmentClient(
|
||||
config,
|
||||
http_client=http_client,
|
||||
sleep=lambda _delay: None,
|
||||
)
|
||||
host = client.enroll_host(
|
||||
agent_instance_id="agent-instance-a",
|
||||
host_token="host-token-" + ("x" * 40),
|
||||
display_name="operator",
|
||||
)
|
||||
|
||||
assert host.host_id == "host-cloud-a"
|
||||
assert len(requests) == 1
|
||||
assert "authorization" not in requests[0].headers
|
||||
assert requests[2].headers["authorization"] == ("Bearer host-token-" + ("x" * 40))
|
||||
|
||||
@@ -11,24 +11,16 @@ from host_agent.config import (
|
||||
)
|
||||
|
||||
|
||||
BASE_ENV = {
|
||||
"HOST_AGENT_HOST_ID": "host-a",
|
||||
"HOST_AGENT_TOKEN": "secret",
|
||||
}
|
||||
|
||||
|
||||
def test_load_host_agent_config_uses_managed_cloud_default() -> None:
|
||||
assert load_host_agent_config(BASE_ENV) == HostAgentConfig(
|
||||
assert load_host_agent_config({}) == HostAgentConfig(
|
||||
control_plane_url="https://amcp.home.jerryyan.top",
|
||||
host_id="host-a",
|
||||
token="secret",
|
||||
enrollment_managed=True,
|
||||
)
|
||||
|
||||
|
||||
def test_load_host_agent_config_parses_poll_and_retry_values() -> None:
|
||||
config = load_host_agent_config(
|
||||
{
|
||||
**BASE_ENV,
|
||||
"HOST_AGENT_CONTROL_PLANE_URL": "https://cloud.example/v1/",
|
||||
"HOST_AGENT_HEARTBEAT_INTERVAL_SECONDS": "10",
|
||||
"HOST_AGENT_POLL_TIMEOUT_SECONDS": "15",
|
||||
@@ -42,12 +34,11 @@ def test_load_host_agent_config_parses_poll_and_retry_values() -> None:
|
||||
assert config.max_retry_backoff_seconds == 20
|
||||
|
||||
|
||||
def test_load_host_agent_config_supports_managed_enrollment(tmp_path) -> None:
|
||||
def test_load_host_agent_config_uses_direct_enrollment(tmp_path) -> None:
|
||||
identity_path = tmp_path / "host_identity.json"
|
||||
config = load_host_agent_config(
|
||||
{
|
||||
"HOST_AGENT_CONTROL_PLANE_URL": "https://cloud.example",
|
||||
"HOST_AGENT_ENROLLMENT_TOKEN": "one-time-token",
|
||||
"HOST_AGENT_IDENTITY_PATH": str(identity_path),
|
||||
"HOST_AGENT_DISPLAY_NAME": "Edge Mac",
|
||||
}
|
||||
@@ -55,30 +46,30 @@ def test_load_host_agent_config_supports_managed_enrollment(tmp_path) -> None:
|
||||
|
||||
assert config.host_id == ""
|
||||
assert config.token == ""
|
||||
assert config.enrollment_token == "one-time-token"
|
||||
assert config.identity_path == identity_path
|
||||
assert config.enrollment_managed is True
|
||||
assert config.display_name == "Edge Mac"
|
||||
assert "one-time-token" not in repr(config)
|
||||
|
||||
|
||||
def test_existing_identity_state_allows_restart_without_enrollment_token(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
identity_path = tmp_path / "host_identity.json"
|
||||
identity_path.write_text("{}", encoding="utf-8")
|
||||
def test_static_credential_environment_variables_are_ignored() -> None:
|
||||
config = load_host_agent_config(
|
||||
{
|
||||
"HOST_AGENT_HOST_ID": "legacy-host",
|
||||
"HOST_AGENT_TOKEN": "legacy-token",
|
||||
"HOST_AGENT_ENROLLMENT_TOKEN": "legacy-enrollment-token",
|
||||
}
|
||||
)
|
||||
|
||||
config = load_host_agent_config({"HOST_AGENT_IDENTITY_PATH": str(identity_path)})
|
||||
|
||||
assert config.identity_path == Path(identity_path)
|
||||
assert config.host_id == ""
|
||||
assert config.token == ""
|
||||
assert config.enrollment_managed is True
|
||||
|
||||
|
||||
def test_fresh_install_with_no_token_is_valid_and_defaults_local_account_path() -> None:
|
||||
config = load_host_agent_config({"HOST_AGENT_CONTROL_PLANE_URL": "https://cloud.example"})
|
||||
def test_fresh_install_defaults_local_account_path() -> None:
|
||||
config = load_host_agent_config(
|
||||
{"HOST_AGENT_CONTROL_PLANE_URL": "https://cloud.example"}
|
||||
)
|
||||
|
||||
assert config.host_id == ""
|
||||
assert config.enrollment_token == ""
|
||||
assert config.enrollment_managed is True
|
||||
assert config.local_account_path == Path("tasks/host_local_account.json")
|
||||
|
||||
@@ -98,9 +89,6 @@ def test_local_account_path_can_be_overridden(tmp_path) -> None:
|
||||
@pytest.mark.parametrize(
|
||||
"overrides",
|
||||
[
|
||||
{"HOST_AGENT_HOST_ID": ""},
|
||||
{"HOST_AGENT_TOKEN": ""},
|
||||
{"HOST_AGENT_HOST_ID": "host-a", "HOST_AGENT_TOKEN": ""},
|
||||
{"HOST_AGENT_CONTROL_PLANE_URL": "ftp://cloud.example"},
|
||||
{"HOST_AGENT_POLL_TIMEOUT_SECONDS": "0"},
|
||||
{
|
||||
@@ -113,4 +101,95 @@ def test_load_host_agent_config_rejects_invalid_values(
|
||||
overrides: dict[str, str],
|
||||
) -> None:
|
||||
with pytest.raises(HostAgentConfigurationError):
|
||||
load_host_agent_config({**BASE_ENV, **overrides})
|
||||
load_host_agent_config(overrides)
|
||||
|
||||
|
||||
def test_console_defaults_are_disabled_and_do_not_trigger_validation() -> None:
|
||||
config = load_host_agent_config({})
|
||||
|
||||
assert config.console_enabled is False
|
||||
assert config.console_bind_host == "127.0.0.1"
|
||||
assert config.console_port == 8765
|
||||
assert config.console_allow_non_loopback is False
|
||||
assert config.console_session_ttl_seconds == 43200.0
|
||||
assert config.console_history_limit == 200
|
||||
|
||||
|
||||
def test_console_enabled_with_default_loopback_bind_passes() -> None:
|
||||
config = load_host_agent_config({"HOST_AGENT_CONSOLE_ENABLED": "true"})
|
||||
|
||||
assert config.console_enabled is True
|
||||
assert config.console_bind_host == "127.0.0.1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bind_host", ["127.0.0.1", "localhost", "::1"])
|
||||
def test_console_enabled_with_loopback_bind_host_passes(bind_host: str) -> None:
|
||||
config = load_host_agent_config(
|
||||
{
|
||||
"HOST_AGENT_CONSOLE_ENABLED": "true",
|
||||
"HOST_AGENT_CONSOLE_BIND_HOST": bind_host,
|
||||
}
|
||||
)
|
||||
|
||||
assert config.console_bind_host == bind_host
|
||||
|
||||
|
||||
def test_console_enabled_with_non_loopback_bind_without_opt_in_raises() -> None:
|
||||
with pytest.raises(HostAgentConfigurationError):
|
||||
load_host_agent_config(
|
||||
{
|
||||
"HOST_AGENT_CONSOLE_ENABLED": "true",
|
||||
"HOST_AGENT_CONSOLE_BIND_HOST": "0.0.0.0",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_console_enabled_with_non_loopback_bind_with_opt_in_succeeds() -> None:
|
||||
config = load_host_agent_config(
|
||||
{
|
||||
"HOST_AGENT_CONSOLE_ENABLED": "true",
|
||||
"HOST_AGENT_CONSOLE_BIND_HOST": "0.0.0.0",
|
||||
"HOST_AGENT_CONSOLE_ALLOW_NON_LOOPBACK": "true",
|
||||
}
|
||||
)
|
||||
|
||||
assert config.console_bind_host == "0.0.0.0"
|
||||
assert config.console_allow_non_loopback is True
|
||||
|
||||
|
||||
def test_console_disabled_with_non_loopback_bind_does_not_raise() -> None:
|
||||
config = load_host_agent_config({"HOST_AGENT_CONSOLE_BIND_HOST": "0.0.0.0"})
|
||||
|
||||
assert config.console_enabled is False
|
||||
assert config.console_bind_host == "0.0.0.0"
|
||||
|
||||
|
||||
def test_console_env_vars_parse_numeric_and_bool_fields() -> None:
|
||||
config = load_host_agent_config(
|
||||
{
|
||||
"HOST_AGENT_CONSOLE_ENABLED": "1",
|
||||
"HOST_AGENT_CONSOLE_PORT": "9001",
|
||||
"HOST_AGENT_CONSOLE_SESSION_TTL_SECONDS": "3600",
|
||||
"HOST_AGENT_CONSOLE_HISTORY_LIMIT": "50",
|
||||
}
|
||||
)
|
||||
|
||||
assert config.console_enabled is True
|
||||
assert config.console_port == 9001
|
||||
assert config.console_session_ttl_seconds == 3600
|
||||
assert config.console_history_limit == 50
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"overrides",
|
||||
[
|
||||
{"HOST_AGENT_CONSOLE_PORT": "0"},
|
||||
{"HOST_AGENT_CONSOLE_SESSION_TTL_SECONDS": "-1"},
|
||||
{"HOST_AGENT_CONSOLE_HISTORY_LIMIT": "0"},
|
||||
],
|
||||
)
|
||||
def test_console_numeric_fields_reject_invalid_values(
|
||||
overrides: dict[str, str],
|
||||
) -> None:
|
||||
with pytest.raises(HostAgentConfigurationError):
|
||||
load_host_agent_config(overrides)
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from cloud.internal_api.models import DeviceEnrollmentResponse
|
||||
from device.manager import DeviceManager
|
||||
from host_agent.config import HostAgentConfig
|
||||
from host_agent.devices import register_local_device, unregister_local_device
|
||||
from storage.device_config import DeviceConfigStore
|
||||
|
||||
|
||||
def _config(*, enrollment_managed: bool) -> HostAgentConfig:
|
||||
return HostAgentConfig(
|
||||
control_plane_url="https://control.example",
|
||||
host_id="host-a",
|
||||
token="secret",
|
||||
enrollment_managed=enrollment_managed,
|
||||
)
|
||||
|
||||
|
||||
class RecordingEnrollmentClient:
|
||||
def __init__(self, device_id: str) -> None:
|
||||
self.device_id = device_id
|
||||
self.calls: list[dict[str, object]] = []
|
||||
|
||||
def enroll_device(self, **payload):
|
||||
self.calls.append(payload)
|
||||
return DeviceEnrollmentResponse(device_id=self.device_id)
|
||||
|
||||
|
||||
def test_register_local_device_enrollment_managed_registers_under_cloud_id(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
|
||||
manager = DeviceManager()
|
||||
enrollment_client = RecordingEnrollmentClient("device-cloud-a")
|
||||
|
||||
register_local_device(
|
||||
store,
|
||||
manager,
|
||||
device_id="local-device-a",
|
||||
driver_type="wda",
|
||||
connection_info={"server_url": "http://127.0.0.1:4723"},
|
||||
name="Lab iPhone",
|
||||
config=_config(enrollment_managed=True),
|
||||
enrollment_client=enrollment_client, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
assert enrollment_client.calls == [
|
||||
{
|
||||
"local_device_id": "local-device-a",
|
||||
"driver_type": "wda",
|
||||
"name": "Lab iPhone",
|
||||
"capability_tags": [],
|
||||
}
|
||||
]
|
||||
assert [device.id for device in manager.list_devices()] == ["device-cloud-a"]
|
||||
assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-a"
|
||||
|
||||
|
||||
def test_register_local_device_not_enrollment_managed_registers_under_local_id(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
|
||||
manager = DeviceManager()
|
||||
enrollment_client = RecordingEnrollmentClient("device-cloud-a")
|
||||
|
||||
register_local_device(
|
||||
store,
|
||||
manager,
|
||||
device_id="local-device-a",
|
||||
driver_type="wda",
|
||||
connection_info={"server_url": "http://127.0.0.1:4723"},
|
||||
name="Lab iPhone",
|
||||
config=_config(enrollment_managed=False),
|
||||
enrollment_client=enrollment_client, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
assert enrollment_client.calls == []
|
||||
assert [device.id for device in manager.list_devices()] == ["local-device-a"]
|
||||
assert store.get("local-device-a")["cloud_device_id"] is None
|
||||
|
||||
|
||||
def test_register_local_device_reregister_under_new_cloud_id_replaces_prior_entry(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
|
||||
manager = DeviceManager()
|
||||
config = _config(enrollment_managed=True)
|
||||
|
||||
register_local_device(
|
||||
store,
|
||||
manager,
|
||||
device_id="local-device-a",
|
||||
driver_type="wda",
|
||||
connection_info={"server_url": "http://127.0.0.1:4723"},
|
||||
name="Lab iPhone",
|
||||
config=config,
|
||||
enrollment_client=RecordingEnrollmentClient( # type: ignore[arg-type]
|
||||
"device-cloud-a"
|
||||
),
|
||||
)
|
||||
assert [device.id for device in manager.list_devices()] == ["device-cloud-a"]
|
||||
|
||||
register_local_device(
|
||||
store,
|
||||
manager,
|
||||
device_id="local-device-a",
|
||||
driver_type="wda",
|
||||
connection_info={"server_url": "http://127.0.0.1:5000"},
|
||||
name="Lab iPhone (moved)",
|
||||
config=config,
|
||||
enrollment_client=RecordingEnrollmentClient( # type: ignore[arg-type]
|
||||
"device-cloud-b"
|
||||
),
|
||||
)
|
||||
|
||||
devices = manager.list_devices()
|
||||
assert [device.id for device in devices] == ["device-cloud-b"]
|
||||
assert devices[0].name == "Lab iPhone (moved)"
|
||||
assert devices[0].connection_info == {"server_url": "http://127.0.0.1:5000"}
|
||||
assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-b"
|
||||
|
||||
|
||||
def test_unregister_local_device_removes_device_with_cloud_id(tmp_path) -> None:
|
||||
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
|
||||
manager = DeviceManager()
|
||||
register_local_device(
|
||||
store,
|
||||
manager,
|
||||
device_id="local-device-a",
|
||||
driver_type="wda",
|
||||
connection_info={},
|
||||
name=None,
|
||||
config=_config(enrollment_managed=True),
|
||||
enrollment_client=RecordingEnrollmentClient( # type: ignore[arg-type]
|
||||
"device-cloud-a"
|
||||
),
|
||||
)
|
||||
|
||||
unregister_local_device(store, manager, device_id="local-device-a")
|
||||
|
||||
assert manager.list_devices() == []
|
||||
assert store.get("local-device-a") is None
|
||||
|
||||
|
||||
def test_unregister_local_device_removes_device_without_cloud_id(tmp_path) -> None:
|
||||
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
|
||||
manager = DeviceManager()
|
||||
register_local_device(
|
||||
store,
|
||||
manager,
|
||||
device_id="local-device-a",
|
||||
driver_type="wda",
|
||||
connection_info={},
|
||||
name=None,
|
||||
config=_config(enrollment_managed=False),
|
||||
enrollment_client=None,
|
||||
)
|
||||
|
||||
unregister_local_device(store, manager, device_id="local-device-a")
|
||||
|
||||
assert manager.list_devices() == []
|
||||
assert store.get("local-device-a") is None
|
||||
|
||||
|
||||
def test_unregister_local_device_is_a_no_op_for_unknown_device(tmp_path) -> None:
|
||||
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
|
||||
manager = DeviceManager()
|
||||
|
||||
unregister_local_device(store, manager, device_id="unknown-device")
|
||||
|
||||
assert manager.list_devices() == []
|
||||
@@ -10,12 +10,12 @@ import httpx
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from cloud.auth import BearerCredential
|
||||
from cloud.auth import digest_token
|
||||
from cloud.control_config import CloudControlConfig
|
||||
from cloud.sdk.client import CloudClient
|
||||
from cloud.scheduler import TaskConstraints
|
||||
from cloud_api.app import create_app
|
||||
from core.models import Scene
|
||||
from core.models import Scene, utc_now
|
||||
from device.manager import DeviceManager
|
||||
from driver.base import Driver
|
||||
from host_agent.assignment import AssignmentExecutionResult, AssignmentExecutor
|
||||
@@ -79,23 +79,6 @@ class FakeDriver(Driver):
|
||||
return None
|
||||
|
||||
|
||||
def _credential(host_id: str) -> BearerCredential:
|
||||
return BearerCredential(
|
||||
principal_id=f"agent-{host_id}",
|
||||
token=f"token-{host_id}",
|
||||
scopes=frozenset(),
|
||||
host_id=host_id,
|
||||
)
|
||||
|
||||
|
||||
def _public_credential() -> BearerCredential:
|
||||
return BearerCredential(
|
||||
principal_id="sdk",
|
||||
token="public-token",
|
||||
scopes=frozenset({"tasks:submit", "tasks:read"}),
|
||||
)
|
||||
|
||||
|
||||
def _config(host_id: str) -> HostAgentConfig:
|
||||
return HostAgentConfig(
|
||||
control_plane_url="http://control.test",
|
||||
@@ -120,10 +103,18 @@ async def _control_plane(
|
||||
scheduler_interval_seconds=60,
|
||||
lease_reaper_interval_seconds=60,
|
||||
lease_duration_seconds=lease_duration_seconds,
|
||||
credentials=tuple(_credential(host_id) for host_id in host_ids),
|
||||
)
|
||||
)
|
||||
async with app.router.lifespan_context(app):
|
||||
for host_id in host_ids:
|
||||
app.state.cloud_services.repository.enroll_host(
|
||||
host_id=host_id,
|
||||
agent_instance_id=f"agent-{host_id}",
|
||||
credential_digest=digest_token(f"token-{host_id}"),
|
||||
enrollment_token_digest=None,
|
||||
display_name=host_id,
|
||||
enrolled_at=utc_now(),
|
||||
)
|
||||
yield app
|
||||
|
||||
|
||||
@@ -162,12 +153,22 @@ class _RecordingTransport(httpx.AsyncBaseTransport):
|
||||
|
||||
|
||||
async def _sync_fake_device(
|
||||
app,
|
||||
client: HostAgentClient,
|
||||
host_id: str,
|
||||
device_id: str,
|
||||
*,
|
||||
driver_type: str = "wda",
|
||||
) -> FakeDriver:
|
||||
app.state.cloud_services.repository.enroll_device(
|
||||
device_id=device_id,
|
||||
host_id=host_id,
|
||||
local_device_id=device_id,
|
||||
driver_type=driver_type,
|
||||
name=device_id,
|
||||
capability_tags=[],
|
||||
enrolled_at=utc_now(),
|
||||
)
|
||||
driver = FakeDriver()
|
||||
manager = DeviceManager()
|
||||
manager.register_device(
|
||||
@@ -248,7 +249,7 @@ def test_one_host_executes_assignment_through_outbound_protocol(tmp_path) -> Non
|
||||
paths: list[str] = []
|
||||
async with _control_plane(tmp_path / "one-host.sqlite3", "host-a") as app:
|
||||
async with _host_client(app, "host-a", request_paths=paths) as client:
|
||||
await _sync_fake_device(client, "host-a", "device-a")
|
||||
await _sync_fake_device(app, client, "host-a", "device-a")
|
||||
task_id = app.state.cloud_services.scheduler.submit(
|
||||
goal="open settings"
|
||||
)
|
||||
@@ -275,7 +276,7 @@ def test_nat_style_host_requires_only_outbound_requests(tmp_path) -> None:
|
||||
paths: list[str] = []
|
||||
async with _control_plane(tmp_path / "outbound-only.sqlite3", "host-a") as app:
|
||||
async with _host_client(app, "host-a", request_paths=paths) as client:
|
||||
await _sync_fake_device(client, "host-a", "device-a")
|
||||
await _sync_fake_device(app, client, "host-a", "device-a")
|
||||
assert await client.claim() is None
|
||||
|
||||
assert paths == [
|
||||
@@ -297,8 +298,9 @@ def test_multiple_hosts_claim_only_their_matching_devices(tmp_path) -> None:
|
||||
_host_client(app, "host-a") as client_a,
|
||||
_host_client(app, "host-b") as client_b,
|
||||
):
|
||||
await _sync_fake_device(client_a, "host-a", "device-a")
|
||||
await _sync_fake_device(app, client_a, "host-a", "device-a")
|
||||
await _sync_fake_device(
|
||||
app,
|
||||
client_b,
|
||||
"host-b",
|
||||
"device-b",
|
||||
@@ -331,7 +333,7 @@ def test_control_plane_restart_preserves_dispatched_assignment(tmp_path) -> None
|
||||
assignment = None
|
||||
async with _control_plane(database_path, "host-a") as first_app:
|
||||
async with _host_client(first_app, "host-a") as client:
|
||||
await _sync_fake_device(client, "host-a", "device-a")
|
||||
await _sync_fake_device(first_app, client, "host-a", "device-a")
|
||||
task_id = first_app.state.cloud_services.scheduler.submit(goal="resume")
|
||||
first_app.state.cloud_services.scheduler.assign()
|
||||
assignment = await client.claim()
|
||||
@@ -351,7 +353,7 @@ def test_host_agent_restart_reuses_active_lease(tmp_path) -> None:
|
||||
async def scenario() -> None:
|
||||
async with _control_plane(tmp_path / "agent-restart.sqlite3", "host-a") as app:
|
||||
async with _host_client(app, "host-a") as first_client:
|
||||
await _sync_fake_device(first_client, "host-a", "device-a")
|
||||
await _sync_fake_device(app, first_client, "host-a", "device-a")
|
||||
task_id = app.state.cloud_services.scheduler.submit(goal="resume host")
|
||||
app.state.cloud_services.scheduler.assign()
|
||||
assignment = await first_client.claim()
|
||||
@@ -376,7 +378,7 @@ def test_lease_loss_rejects_stale_host_result(tmp_path) -> None:
|
||||
lease_duration_seconds=0.2,
|
||||
) as app:
|
||||
async with _host_client(app, "host-a") as client:
|
||||
await _sync_fake_device(client, "host-a", "device-a")
|
||||
await _sync_fake_device(app, client, "host-a", "device-a")
|
||||
task_id = app.state.cloud_services.scheduler.submit(goal="expire")
|
||||
app.state.cloud_services.scheduler.assign()
|
||||
assignment = await client.claim()
|
||||
@@ -405,7 +407,6 @@ def test_public_sdk_reports_fake_device_success_and_runtime_failure(tmp_path) ->
|
||||
scheduler_interval_seconds=60,
|
||||
lease_reaper_interval_seconds=60,
|
||||
lease_duration_seconds=30,
|
||||
credentials=(_public_credential(), _credential("host-a")),
|
||||
)
|
||||
)
|
||||
driver = FakeDriver()
|
||||
@@ -413,10 +414,38 @@ def test_public_sdk_reports_fake_device_success_and_runtime_failure(tmp_path) ->
|
||||
manager.register_device("device-a", lambda: driver, status="idle")
|
||||
|
||||
with TestClient(app) as http_client:
|
||||
app.state.cloud_services.user_auth_service.create_user(
|
||||
username="operator",
|
||||
display_name="Operator",
|
||||
role="admin",
|
||||
password="correct-horse-battery-staple",
|
||||
must_change_password=False,
|
||||
)
|
||||
login = http_client.post(
|
||||
"/v1/auth/login",
|
||||
json={"username": "operator", "password": "correct-horse-battery-staple"},
|
||||
)
|
||||
assert login.status_code == 200
|
||||
app.state.cloud_services.repository.enroll_host(
|
||||
host_id="host-a",
|
||||
agent_instance_id="agent-host-a",
|
||||
credential_digest=digest_token("token-host-a"),
|
||||
enrollment_token_digest=None,
|
||||
display_name="host-a",
|
||||
enrolled_at=utc_now(),
|
||||
)
|
||||
app.state.cloud_services.repository.enroll_device(
|
||||
device_id="device-a",
|
||||
host_id="host-a",
|
||||
local_device_id="device-a",
|
||||
driver_type="wda",
|
||||
name="device-a",
|
||||
capability_tags=[],
|
||||
enrolled_at=utc_now(),
|
||||
)
|
||||
cloud_client = CloudClient(
|
||||
"http://testserver",
|
||||
http_client=http_client,
|
||||
token="public-token",
|
||||
)
|
||||
|
||||
async def scenario() -> None:
|
||||
@@ -428,9 +457,7 @@ def test_public_sdk_reports_fake_device_success_and_runtime_failure(tmp_path) ->
|
||||
heartbeat.connect_devices()
|
||||
await heartbeat.sync_once()
|
||||
|
||||
successful_task_id = cloud_client.submit_task(goal="tap screen")[
|
||||
"task_id"
|
||||
]
|
||||
successful_task_id = cloud_client.submit_task(goal="tap screen")["task_id"]
|
||||
app.state.cloud_services.scheduler.assign()
|
||||
successful_assignment = await host_client.claim()
|
||||
assert successful_assignment is not None
|
||||
|
||||
@@ -31,12 +31,12 @@ class _RejectingEnrollmentClient:
|
||||
raise HostAgentAPIError(401, "unauthorized")
|
||||
|
||||
|
||||
def test_fresh_install_with_no_token_self_enrolls(tmp_path) -> None:
|
||||
def test_fresh_install_directly_enrolls(tmp_path) -> None:
|
||||
identity_store = HostIdentityStore(tmp_path / "identity.json")
|
||||
client = _RecordingEnrollmentClient()
|
||||
|
||||
resolved = resolve_host_identity(
|
||||
_config(enrollment_token=""),
|
||||
_config(),
|
||||
identity_store=identity_store,
|
||||
client=client,
|
||||
)
|
||||
@@ -52,7 +52,7 @@ def test_self_service_rejection_propagates_as_api_error(tmp_path) -> None:
|
||||
|
||||
try:
|
||||
resolve_host_identity(
|
||||
_config(enrollment_token=""),
|
||||
_config(),
|
||||
identity_store=identity_store,
|
||||
client=client,
|
||||
)
|
||||
@@ -63,19 +63,6 @@ def test_self_service_rejection_propagates_as_api_error(tmp_path) -> None:
|
||||
assert identity_store.load().host_id is None
|
||||
|
||||
|
||||
def test_configured_enrollment_token_still_used_when_present(tmp_path) -> None:
|
||||
identity_store = HostIdentityStore(tmp_path / "identity.json")
|
||||
client = _RecordingEnrollmentClient()
|
||||
|
||||
resolve_host_identity(
|
||||
_config(enrollment_token="one-time-token"),
|
||||
identity_store=identity_store,
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert client.calls[0]["agent_instance_id"]
|
||||
|
||||
|
||||
def test_existing_cached_identity_skips_enrollment(tmp_path) -> None:
|
||||
identity_store = HostIdentityStore(tmp_path / "identity.json")
|
||||
identity_store.complete(identity_store.load_or_create(), "host-cloud-a")
|
||||
@@ -85,7 +72,7 @@ def test_existing_cached_identity_skips_enrollment(tmp_path) -> None:
|
||||
raise AssertionError("cached identity must skip enrollment")
|
||||
|
||||
resolved = resolve_host_identity(
|
||||
_config(enrollment_token=""),
|
||||
_config(),
|
||||
identity_store=identity_store,
|
||||
client=ExplodingClient(), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
@@ -7,6 +7,7 @@ from cloud.internal_api.models import HeartbeatResponse
|
||||
from device.manager import DeviceManager
|
||||
from host_agent.config import HostAgentConfig
|
||||
from host_agent.heartbeat import HeartbeatSynchronizer, build_device_snapshot
|
||||
from host_agent.status import AgentStatusTracker
|
||||
|
||||
|
||||
class ConnectableDriver:
|
||||
@@ -81,3 +82,44 @@ def test_heartbeat_synchronizer_runs_at_configured_interval_until_stopped() -> N
|
||||
asyncio.run(scenario())
|
||||
assert calls == [["device-a"], ["device-a"], ["device-a"]]
|
||||
assert manager.status("device-a") == "busy"
|
||||
|
||||
|
||||
def test_sync_once_notifies_status_tracker_and_on_sync_with_device_count() -> None:
|
||||
manager = DeviceManager()
|
||||
manager.register_device(
|
||||
"device-a",
|
||||
lambda: ConnectableDriver(), # type: ignore[arg-type,return-value]
|
||||
)
|
||||
manager.register_device(
|
||||
"device-b",
|
||||
lambda: ConnectableDriver(), # type: ignore[arg-type,return-value]
|
||||
)
|
||||
|
||||
class FakeClient:
|
||||
async def heartbeat(self, devices, *, address=None):
|
||||
return HeartbeatResponse(
|
||||
host_id="host-a",
|
||||
accepted_devices=len(devices),
|
||||
received_at=datetime.now(UTC),
|
||||
)
|
||||
|
||||
async def scenario() -> None:
|
||||
tracker = AgentStatusTracker()
|
||||
on_sync_calls: list[int] = []
|
||||
synchronizer = HeartbeatSynchronizer(
|
||||
manager,
|
||||
FakeClient(), # type: ignore[arg-type]
|
||||
_config(),
|
||||
status_tracker=tracker,
|
||||
on_sync=on_sync_calls.append,
|
||||
)
|
||||
|
||||
await synchronizer.sync_once()
|
||||
|
||||
assert on_sync_calls == [2]
|
||||
last_heartbeat = tracker.snapshot()["last_heartbeat"]
|
||||
assert last_heartbeat is not None
|
||||
assert last_heartbeat["ok"] is True
|
||||
assert last_heartbeat["device_count"] == 2
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from host_agent.history import ConsoleHistoryStore
|
||||
|
||||
|
||||
def test_history_store_records_assignment_and_heartbeat_newest_first(tmp_path) -> None:
|
||||
store = ConsoleHistoryStore(tmp_path / "history.sqlite3")
|
||||
|
||||
store.record_assignment(
|
||||
task_id="task-a",
|
||||
attempt=1,
|
||||
status="done",
|
||||
failure_reason=None,
|
||||
device_id="device-a",
|
||||
)
|
||||
store.record_heartbeat(device_count=2)
|
||||
|
||||
entries = store.list_recent()
|
||||
|
||||
assert len(entries) == 2
|
||||
assert entries[0] == {
|
||||
"kind": "heartbeat",
|
||||
"occurred_at": entries[0]["occurred_at"],
|
||||
"summary": "heartbeat: 2 devices",
|
||||
"detail": {"device_count": 2},
|
||||
}
|
||||
assert entries[1] == {
|
||||
"kind": "assignment",
|
||||
"occurred_at": entries[1]["occurred_at"],
|
||||
"summary": "task-a attempt 1 on device-a: done",
|
||||
"detail": {
|
||||
"task_id": "task-a",
|
||||
"attempt": 1,
|
||||
"status": "done",
|
||||
"failure_reason": None,
|
||||
"device_id": "device-a",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_history_store_prunes_oldest_entries_beyond_limit(tmp_path) -> None:
|
||||
store = ConsoleHistoryStore(tmp_path / "history.sqlite3", limit=3)
|
||||
|
||||
for index in range(5):
|
||||
store.record_heartbeat(device_count=index)
|
||||
|
||||
entries = store.list_recent()
|
||||
|
||||
assert len(entries) == 3
|
||||
assert [entry["detail"]["device_count"] for entry in entries] == [4, 3, 2]
|
||||
|
||||
|
||||
def test_history_store_returns_empty_list_when_no_entries(tmp_path) -> None:
|
||||
store = ConsoleHistoryStore(tmp_path / "history.sqlite3")
|
||||
|
||||
assert store.list_recent() == []
|
||||
|
||||
|
||||
def test_history_store_uses_injected_now_for_occurred_at(tmp_path) -> None:
|
||||
fixed_now = datetime(2026, 1, 1, 12, 0, 0, tzinfo=UTC)
|
||||
store = ConsoleHistoryStore(tmp_path / "history.sqlite3", now=lambda: fixed_now)
|
||||
|
||||
store.record_heartbeat(device_count=1)
|
||||
|
||||
entries = store.list_recent()
|
||||
|
||||
assert entries[0]["occurred_at"] == fixed_now.isoformat()
|
||||
@@ -6,6 +6,7 @@ from datetime import UTC, datetime
|
||||
from cloud.internal_api.models import AssignmentModel, TerminalResultResponse
|
||||
from host_agent.assignment import AssignmentExecutionResult
|
||||
from host_agent.processor import AssignmentProcessor
|
||||
from host_agent.status import AgentStatusTracker
|
||||
|
||||
|
||||
def _assignment() -> AssignmentModel:
|
||||
@@ -85,3 +86,69 @@ def test_processor_preserves_runtime_failure_reason() -> None:
|
||||
}
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_status_tracker_sees_started_then_finished_even_on_raise() -> None:
|
||||
async def scenario() -> None:
|
||||
tracker = AgentStatusTracker()
|
||||
snapshots: list[dict[str, object]] = []
|
||||
|
||||
class RaisingExecutor:
|
||||
async def run(self, assignment):
|
||||
snapshots.append(tracker.snapshot())
|
||||
raise RuntimeError("executor exploded")
|
||||
|
||||
class RecordingClient:
|
||||
async def report_result(self, assignment, **kwargs):
|
||||
return TerminalResultResponse(status="recorded")
|
||||
|
||||
processor = AssignmentProcessor(
|
||||
RecordingClient(), # type: ignore[arg-type]
|
||||
RaisingExecutor(),
|
||||
status_tracker=tracker,
|
||||
)
|
||||
|
||||
try:
|
||||
await processor.process(_assignment())
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
assert snapshots[0]["current_assignment"] is not None
|
||||
assert snapshots[0]["current_assignment"]["task_id"] == "task-a"
|
||||
assert tracker.snapshot()["current_assignment"] is None
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_on_result_receives_assignment_and_result_and_swallows_exceptions() -> None:
|
||||
async def scenario() -> None:
|
||||
received: list[tuple[object, object]] = []
|
||||
|
||||
class SuccessfulExecutor:
|
||||
async def run(self, assignment):
|
||||
return AssignmentExecutionResult(
|
||||
status="done",
|
||||
failure_reason=None,
|
||||
metadata={},
|
||||
)
|
||||
|
||||
class RecordingClient:
|
||||
async def report_result(self, assignment, **kwargs):
|
||||
return TerminalResultResponse(status="recorded")
|
||||
|
||||
def on_result(assignment, result) -> None:
|
||||
received.append((assignment, result))
|
||||
raise RuntimeError("history recording exploded")
|
||||
|
||||
assignment = _assignment()
|
||||
processor = AssignmentProcessor(
|
||||
RecordingClient(), # type: ignore[arg-type]
|
||||
SuccessfulExecutor(),
|
||||
on_result=on_result,
|
||||
)
|
||||
|
||||
result = await processor.process(assignment)
|
||||
|
||||
assert received == [(assignment, result)]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from cloud.internal_api.models import AssignmentModel
|
||||
from host_agent.status import AgentStatusTracker
|
||||
|
||||
|
||||
def _assignment() -> AssignmentModel:
|
||||
return AssignmentModel(
|
||||
task_id="task-a",
|
||||
attempt=1,
|
||||
lease_id="lease-a",
|
||||
lease_expires_at=datetime(2026, 7, 12, tzinfo=UTC),
|
||||
host_id="host-a",
|
||||
device_id="device-a",
|
||||
goal="open settings",
|
||||
workflow_definition_id="workflow-a",
|
||||
)
|
||||
|
||||
|
||||
def test_mark_assignment_started_reflected_in_snapshot() -> None:
|
||||
ticks = iter([datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC)])
|
||||
tracker = AgentStatusTracker(now=lambda: next(ticks))
|
||||
|
||||
tracker.mark_assignment_started(_assignment())
|
||||
|
||||
snapshot = tracker.snapshot()
|
||||
assert snapshot["current_assignment"] == {
|
||||
"task_id": "task-a",
|
||||
"device_id": "device-a",
|
||||
"goal": "open settings",
|
||||
"workflow_definition_id": "workflow-a",
|
||||
"started_at": "2026-07-13T09:00:00+00:00",
|
||||
}
|
||||
|
||||
|
||||
def test_mark_assignment_finished_clears_current_assignment() -> None:
|
||||
tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC))
|
||||
|
||||
tracker.mark_assignment_started(_assignment())
|
||||
tracker.mark_assignment_finished()
|
||||
|
||||
assert tracker.snapshot()["current_assignment"] is None
|
||||
|
||||
|
||||
def test_mark_heartbeat_reflected_in_snapshot() -> None:
|
||||
tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 5, 0, tzinfo=UTC))
|
||||
|
||||
tracker.mark_heartbeat(ok=True, device_count=3)
|
||||
|
||||
snapshot = tracker.snapshot()
|
||||
assert snapshot["last_heartbeat"] == {
|
||||
"ok": True,
|
||||
"device_count": 3,
|
||||
"at": "2026-07-13T09:05:00+00:00",
|
||||
}
|
||||
|
||||
|
||||
def test_snapshot_defaults_to_no_assignment_or_heartbeat() -> None:
|
||||
tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC))
|
||||
|
||||
snapshot = tracker.snapshot()
|
||||
|
||||
assert snapshot["current_assignment"] is None
|
||||
assert snapshot["last_heartbeat"] is None
|
||||
@@ -0,0 +1,262 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from device.manager import DeviceManager
|
||||
from host_agent.config import HostAgentConfig
|
||||
from host_agent.history import ConsoleHistoryStore
|
||||
from host_agent.identity import HostIdentityStore
|
||||
from host_agent.local_account import LocalAccountStore
|
||||
from host_agent.status import AgentStatusTracker
|
||||
from host_agent.web.app import SESSION_COOKIE_NAME, create_console_app
|
||||
from host_agent.web.auth import SessionManager
|
||||
from storage.device_config import DeviceConfigStore
|
||||
|
||||
CSRF_PATTERN = re.compile(r'name="csrf_token" value="([^"]+)"')
|
||||
|
||||
|
||||
def _build_client(tmp_path, *, create_account: bool = True) -> tuple[TestClient, dict]:
|
||||
config = HostAgentConfig(
|
||||
control_plane_url="https://control.example",
|
||||
host_id="host-a",
|
||||
token="secret",
|
||||
enrollment_managed=False,
|
||||
console_session_ttl_seconds=3600.0,
|
||||
)
|
||||
manager = DeviceManager()
|
||||
config_store = DeviceConfigStore(tmp_path / "devices.sqlite3")
|
||||
local_account_store = LocalAccountStore(tmp_path / "host_local_account.json")
|
||||
if create_account:
|
||||
local_account_store.create("operator", "correct horse battery staple")
|
||||
identity_store = HostIdentityStore(tmp_path / "host_identity.json")
|
||||
history_store = ConsoleHistoryStore(tmp_path / "history.sqlite3")
|
||||
status_tracker = AgentStatusTracker()
|
||||
session_manager = SessionManager(ttl_seconds=3600.0)
|
||||
|
||||
app = create_console_app(
|
||||
config=config,
|
||||
manager=manager,
|
||||
config_store=config_store,
|
||||
local_account_store=local_account_store,
|
||||
identity_store=identity_store,
|
||||
history_store=history_store,
|
||||
status_tracker=status_tracker,
|
||||
session_manager=session_manager,
|
||||
enrollment_client=None,
|
||||
)
|
||||
client = TestClient(app)
|
||||
context = {
|
||||
"manager": manager,
|
||||
"config_store": config_store,
|
||||
"local_account_store": local_account_store,
|
||||
"history_store": history_store,
|
||||
"session_manager": session_manager,
|
||||
}
|
||||
return client, context
|
||||
|
||||
|
||||
def _login(
|
||||
client: TestClient,
|
||||
*,
|
||||
username: str = "operator",
|
||||
password: str = "correct horse battery staple",
|
||||
) -> str:
|
||||
response = client.post("/login", data={"username": username, "password": password})
|
||||
assert response.status_code == 200
|
||||
match = CSRF_PATTERN.search(response.text)
|
||||
assert match is not None
|
||||
return match.group(1)
|
||||
|
||||
|
||||
def test_unauthenticated_get_root_redirects_to_login(tmp_path) -> None:
|
||||
client, _ = _build_client(tmp_path)
|
||||
|
||||
response = client.get("/", follow_redirects=False)
|
||||
|
||||
assert response.status_code == 303
|
||||
assert response.headers["location"] == "/login"
|
||||
|
||||
|
||||
def test_login_with_no_account_shows_setup_message_and_rejects_post(tmp_path) -> None:
|
||||
client, context = _build_client(tmp_path, create_account=False)
|
||||
|
||||
get_response = client.get("/login")
|
||||
assert "device-host-agent setup" in get_response.text
|
||||
|
||||
post_response = client.post(
|
||||
"/login", data={"username": "operator", "password": "anything"}
|
||||
)
|
||||
|
||||
assert "device-host-agent setup" in post_response.text
|
||||
assert SESSION_COOKIE_NAME not in client.cookies
|
||||
|
||||
|
||||
def test_login_with_wrong_password_fails_and_sets_no_cookie(tmp_path) -> None:
|
||||
client, _ = _build_client(tmp_path)
|
||||
|
||||
response = client.post("/login", data={"username": "operator", "password": "wrong"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "Invalid username or password" in response.text
|
||||
assert SESSION_COOKIE_NAME not in client.cookies
|
||||
|
||||
|
||||
def test_login_with_correct_password_sets_cookie_and_dashboard_succeeds(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
client, _ = _build_client(tmp_path)
|
||||
|
||||
response = client.post(
|
||||
"/login",
|
||||
data={"username": "operator", "password": "correct horse battery staple"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert SESSION_COOKIE_NAME in client.cookies
|
||||
assert "Status" in response.text
|
||||
assert "control.example" in response.text
|
||||
|
||||
|
||||
def test_mutating_post_without_csrf_token_is_rejected_and_makes_no_change(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
client, context = _build_client(tmp_path)
|
||||
_login(client)
|
||||
|
||||
response = client.post(
|
||||
"/devices/save",
|
||||
data={
|
||||
"device_id": "device-a",
|
||||
"driver_type": "wda",
|
||||
"name": "Lab iPhone",
|
||||
"connection_info": "{}",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert context["config_store"].get("device-a") is None
|
||||
assert context["manager"].list_devices() == []
|
||||
|
||||
|
||||
def test_add_device_appears_in_devices_page_and_manager(tmp_path) -> None:
|
||||
client, context = _build_client(tmp_path)
|
||||
csrf_token = _login(client)
|
||||
|
||||
response = client.post(
|
||||
"/devices/save",
|
||||
data={
|
||||
"device_id": "device-a",
|
||||
"driver_type": "wda",
|
||||
"name": "Lab iPhone",
|
||||
"connection_info": "{}",
|
||||
"csrf_token": csrf_token,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "device-a" in response.text
|
||||
assert context["config_store"].get("device-a") is not None
|
||||
assert [device.id for device in context["manager"].list_devices()] == ["device-a"]
|
||||
|
||||
|
||||
def test_remove_device_unregisters_from_manager(tmp_path) -> None:
|
||||
client, context = _build_client(tmp_path)
|
||||
csrf_token = _login(client)
|
||||
client.post(
|
||||
"/devices/save",
|
||||
data={
|
||||
"device_id": "device-a",
|
||||
"driver_type": "wda",
|
||||
"name": "Lab iPhone",
|
||||
"connection_info": "{}",
|
||||
"csrf_token": csrf_token,
|
||||
},
|
||||
)
|
||||
assert context["config_store"].get("device-a") is not None
|
||||
|
||||
response = client.post(
|
||||
"/devices/remove",
|
||||
data={"device_id": "device-a", "csrf_token": csrf_token},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert context["config_store"].get("device-a") is None
|
||||
assert context["manager"].list_devices() == []
|
||||
|
||||
|
||||
def test_change_password_wrong_current_fails_old_password_still_works(tmp_path) -> None:
|
||||
client, context = _build_client(tmp_path)
|
||||
csrf_token = _login(client)
|
||||
|
||||
response = client.post(
|
||||
"/account",
|
||||
data={
|
||||
"current_password": "wrong",
|
||||
"new_password": "new password",
|
||||
"confirm_password": "new password",
|
||||
"csrf_token": csrf_token,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "incorrect" in response.text
|
||||
account_store: LocalAccountStore = context["local_account_store"]
|
||||
account = account_store.load()
|
||||
assert account is not None
|
||||
assert account_store.verify(account, "correct horse battery staple") is True
|
||||
|
||||
|
||||
def test_change_password_correct_succeeds_old_password_no_longer_works(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
client, context = _build_client(tmp_path)
|
||||
csrf_token = _login(client)
|
||||
|
||||
response = client.post(
|
||||
"/account",
|
||||
data={
|
||||
"current_password": "correct horse battery staple",
|
||||
"new_password": "new password",
|
||||
"confirm_password": "new password",
|
||||
"csrf_token": csrf_token,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "Password updated" in response.text
|
||||
account_store: LocalAccountStore = context["local_account_store"]
|
||||
account = account_store.load()
|
||||
assert account is not None
|
||||
assert account_store.verify(account, "new password") is True
|
||||
assert account_store.verify(account, "correct horse battery staple") is False
|
||||
|
||||
|
||||
def test_history_reflects_entries_written_beforehand(tmp_path) -> None:
|
||||
client, context = _build_client(tmp_path)
|
||||
context["history_store"].record_heartbeat(device_count=3)
|
||||
_login(client)
|
||||
|
||||
response = client.get("/history")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "heartbeat: 3 devices" in response.text
|
||||
|
||||
|
||||
def test_logout_invalidates_session_so_subsequent_request_redirects_to_login(
|
||||
tmp_path,
|
||||
) -> None:
|
||||
client, _ = _build_client(tmp_path)
|
||||
csrf_token = _login(client)
|
||||
|
||||
logout_response = client.post(
|
||||
"/logout", data={"csrf_token": csrf_token}, follow_redirects=False
|
||||
)
|
||||
assert logout_response.status_code == 303
|
||||
assert logout_response.headers["location"] == "/login"
|
||||
|
||||
response = client.get("/", follow_redirects=False)
|
||||
|
||||
assert response.status_code == 303
|
||||
assert response.headers["location"] == "/login"
|
||||
@@ -0,0 +1,162 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from host_agent.local_account import LocalAccountStore
|
||||
from host_agent.web.auth import SessionManager, attempt_login, change_password
|
||||
|
||||
|
||||
class FakeClock:
|
||||
def __init__(self, start: datetime) -> None:
|
||||
self.current = start
|
||||
|
||||
def __call__(self) -> datetime:
|
||||
return self.current
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.current += timedelta(seconds=seconds)
|
||||
|
||||
|
||||
def test_create_session_then_validate_returns_username() -> None:
|
||||
clock = FakeClock(datetime(2026, 1, 1, tzinfo=UTC))
|
||||
manager = SessionManager(ttl_seconds=60, now=clock)
|
||||
|
||||
session_token, csrf_token = manager.create_session("operator")
|
||||
state = manager.validate(session_token)
|
||||
|
||||
assert state is not None
|
||||
assert state.username == "operator"
|
||||
assert state.csrf_token == csrf_token
|
||||
|
||||
|
||||
def test_validate_unknown_token_returns_none() -> None:
|
||||
manager = SessionManager(ttl_seconds=60)
|
||||
|
||||
assert manager.validate("does-not-exist") is None
|
||||
|
||||
|
||||
def test_validate_after_ttl_elapsed_returns_none() -> None:
|
||||
clock = FakeClock(datetime(2026, 1, 1, tzinfo=UTC))
|
||||
manager = SessionManager(ttl_seconds=60, now=clock)
|
||||
|
||||
session_token, _ = manager.create_session("operator")
|
||||
clock.advance(61)
|
||||
|
||||
assert manager.validate(session_token) is None
|
||||
|
||||
|
||||
def test_validate_before_expiry_slides_expiry_forward() -> None:
|
||||
clock = FakeClock(datetime(2026, 1, 1, tzinfo=UTC))
|
||||
manager = SessionManager(ttl_seconds=60, now=clock)
|
||||
|
||||
session_token, _ = manager.create_session("operator")
|
||||
clock.advance(30)
|
||||
first = manager.validate(session_token)
|
||||
assert first is not None
|
||||
|
||||
clock.advance(30)
|
||||
second = manager.validate(session_token)
|
||||
|
||||
assert second is not None
|
||||
assert second.expires_at > first.expires_at
|
||||
|
||||
|
||||
def test_validate_csrf_true_for_right_token() -> None:
|
||||
manager = SessionManager(ttl_seconds=60)
|
||||
session_token, csrf_token = manager.create_session("operator")
|
||||
|
||||
assert manager.validate_csrf(session_token, csrf_token) is True
|
||||
|
||||
|
||||
def test_validate_csrf_false_for_wrong_token() -> None:
|
||||
manager = SessionManager(ttl_seconds=60)
|
||||
session_token, _ = manager.create_session("operator")
|
||||
|
||||
assert manager.validate_csrf(session_token, "wrong-token") is False
|
||||
|
||||
|
||||
def test_validate_csrf_false_for_invalid_session() -> None:
|
||||
manager = SessionManager(ttl_seconds=60)
|
||||
|
||||
assert manager.validate_csrf("does-not-exist", "anything") is False
|
||||
|
||||
|
||||
def test_invalidate_makes_subsequent_validate_return_none() -> None:
|
||||
manager = SessionManager(ttl_seconds=60)
|
||||
session_token, _ = manager.create_session("operator")
|
||||
|
||||
manager.invalidate(session_token)
|
||||
|
||||
assert manager.validate(session_token) is None
|
||||
|
||||
|
||||
def test_attempt_login_true_for_correct_credentials(tmp_path) -> None:
|
||||
store = LocalAccountStore(tmp_path / "host_local_account.json")
|
||||
store.create("operator", "correct horse battery staple")
|
||||
|
||||
assert (
|
||||
attempt_login(
|
||||
store, username="operator", password="correct horse battery staple"
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_attempt_login_false_for_wrong_password(tmp_path) -> None:
|
||||
store = LocalAccountStore(tmp_path / "host_local_account.json")
|
||||
store.create("operator", "correct horse battery staple")
|
||||
|
||||
assert attempt_login(store, username="operator", password="wrong") is False
|
||||
|
||||
|
||||
def test_attempt_login_false_when_no_account_exists(tmp_path) -> None:
|
||||
store = LocalAccountStore(tmp_path / "host_local_account.json")
|
||||
|
||||
assert attempt_login(store, username="operator", password="anything") is False
|
||||
|
||||
|
||||
def test_attempt_login_false_for_wrong_username(tmp_path) -> None:
|
||||
store = LocalAccountStore(tmp_path / "host_local_account.json")
|
||||
store.create("operator", "correct horse battery staple")
|
||||
|
||||
assert (
|
||||
attempt_login(
|
||||
store, username="someone-else", password="correct horse battery staple"
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_change_password_succeeds_and_rotates_credential(tmp_path) -> None:
|
||||
store = LocalAccountStore(tmp_path / "host_local_account.json")
|
||||
store.create("operator", "old password")
|
||||
|
||||
assert (
|
||||
change_password(
|
||||
store, current_password="old password", new_password="new password"
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert attempt_login(store, username="operator", password="new password") is True
|
||||
assert attempt_login(store, username="operator", password="old password") is False
|
||||
|
||||
|
||||
def test_change_password_fails_with_wrong_current_password(tmp_path) -> None:
|
||||
store = LocalAccountStore(tmp_path / "host_local_account.json")
|
||||
store.create("operator", "old password")
|
||||
|
||||
assert (
|
||||
change_password(store, current_password="wrong", new_password="new password")
|
||||
is False
|
||||
)
|
||||
assert attempt_login(store, username="operator", password="old password") is True
|
||||
assert attempt_login(store, username="operator", password="new password") is False
|
||||
|
||||
|
||||
def test_change_password_fails_when_no_account_exists(tmp_path) -> None:
|
||||
store = LocalAccountStore(tmp_path / "host_local_account.json")
|
||||
|
||||
assert (
|
||||
change_password(store, current_password="anything", new_password="new password")
|
||||
is False
|
||||
)
|
||||
Reference in New Issue
Block a user