chore(repo): initialize team collaboration repository

2026-07-27 20:40:12 +08:00
commit c91a64fddb
109 changed files with 21121 additions and 0 deletions
@@ -0,0 +1,141 @@
+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]