142 lines
4.5 KiB
Python
142 lines
4.5 KiB
Python
from pathlib import Path
|
|
import stat
|
|
import uuid
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from flashops_agent.client import AgentUnauthorized, ControlPlaneClient
|
|
from flashops_agent.state import AgentState, StateStore
|
|
|
|
|
|
def test_control_plane_client_registers_and_sends_heartbeat():
|
|
requests = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request)
|
|
if request.url.path == "/api/v1/agent/register":
|
|
return httpx.Response(
|
|
201,
|
|
json={
|
|
"agent_id": "agent_test",
|
|
"host_id": "host_test",
|
|
"token": "raw-token",
|
|
"heartbeat_interval_s": 5,
|
|
},
|
|
)
|
|
if request.url.path == "/api/v1/agent/agent_test/heartbeat":
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"ok": True,
|
|
"server_time": "2026-07-27T00:00:00Z",
|
|
"commands": [],
|
|
},
|
|
)
|
|
return httpx.Response(404)
|
|
|
|
with ControlPlaneClient(
|
|
"http://control.test",
|
|
enrollment_token="enroll-test",
|
|
transport=httpx.MockTransport(handler),
|
|
) as client:
|
|
registered = client.register(
|
|
{
|
|
"host_name": "pytest-host",
|
|
"agent_version": "0.1.0",
|
|
"os_family": "linux",
|
|
}
|
|
)
|
|
heartbeat = client.heartbeat(
|
|
registered["agent_id"],
|
|
registered["token"],
|
|
{"healthy": True},
|
|
)
|
|
|
|
assert heartbeat["ok"] is True
|
|
assert requests[0].headers["x-flashops-enrollment-token"] == "enroll-test"
|
|
assert requests[1].headers["authorization"] == "Bearer raw-token"
|
|
|
|
|
|
def test_control_plane_client_rejects_unauthorized_heartbeat():
|
|
def handler(_: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(401, json={"detail": "invalid"})
|
|
|
|
with ControlPlaneClient(
|
|
"http://control.test",
|
|
transport=httpx.MockTransport(handler),
|
|
) as client:
|
|
with pytest.raises(AgentUnauthorized):
|
|
client.heartbeat("agent_test", "bad-token", {"healthy": True})
|
|
|
|
|
|
def test_state_store_round_trip_uses_private_permissions(tmp_path: Path):
|
|
path = tmp_path / "agent-state.json"
|
|
store = StateStore(path)
|
|
expected = AgentState(
|
|
server="http://control.test",
|
|
agent_id="agent_test",
|
|
host_id="host_test",
|
|
token="raw-token",
|
|
heartbeat_interval_s=5.0,
|
|
)
|
|
|
|
store.save(expected)
|
|
|
|
assert store.load() == expected
|
|
assert stat.S_IMODE(path.stat().st_mode) == 0o600
|
|
|
|
|
|
def test_control_plane_client_runs_step_lease_protocol():
|
|
requests = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request)
|
|
path = request.url.path
|
|
if path.endswith("/lease"):
|
|
return httpx.Response(
|
|
200,
|
|
json={"step_id": "step_test", "attempt": 1, "dry_run": True},
|
|
)
|
|
if path.endswith("/ack"):
|
|
return httpx.Response(200, json={"step_id": "step_test", "state": "RUNNING"})
|
|
if path.endswith("/renew"):
|
|
return httpx.Response(200, json={"step_id": "step_test", "state": "RUNNING"})
|
|
if path.endswith("/events"):
|
|
return httpx.Response(200, json={"step_id": "step_test", "accepted": 1})
|
|
if path.endswith("/complete"):
|
|
return httpx.Response(200, json={"step_id": "step_test", "state": "SUCCEEDED"})
|
|
return httpx.Response(404)
|
|
|
|
token = "raw-token"
|
|
keys = [str(uuid.uuid4()) for _ in range(4)]
|
|
with ControlPlaneClient(
|
|
"http://control.test", transport=httpx.MockTransport(handler)
|
|
) as client:
|
|
lease = client.lease("agent_test", token)
|
|
assert lease is not None
|
|
client.ack("agent_test", token, "step_test", 1, keys[0])
|
|
client.renew("agent_test", token, "step_test", 1, keys[1])
|
|
client.report_events(
|
|
"agent_test",
|
|
token,
|
|
"step_test",
|
|
1,
|
|
[{"message": "hello"}],
|
|
keys[2],
|
|
)
|
|
client.complete(
|
|
"agent_test",
|
|
token,
|
|
"step_test",
|
|
1,
|
|
status="SUCCEEDED",
|
|
exit_code=0,
|
|
result={"simulated": True},
|
|
idempotency_key=keys[3],
|
|
)
|
|
|
|
assert len(requests) == 5
|
|
assert requests[0].headers["authorization"] == "Bearer raw-token"
|
|
assert requests[1].headers["idempotency-key"] == keys[0]
|