77 lines
2.2 KiB
Python
77 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
import time
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
|
|
|
|
def uuidv7() -> str:
|
|
ts_ms = int(time.time() * 1000)
|
|
rand_a = random.getrandbits(12)
|
|
rand_b = random.getrandbits(62)
|
|
high = (ts_ms & 0xFFFFFFFFFFFF) << 16 | 0x7000 | rand_a
|
|
low = (0b10 << 62) | rand_b
|
|
return str(uuid.UUID(int=(high << 64) | low))
|
|
|
|
|
|
def now_iso(offset_sec: int = 0) -> str:
|
|
return datetime.now(UTC).replace(microsecond=0).isoformat()
|
|
|
|
|
|
def make_event(
|
|
check_id: str = "disk_root",
|
|
status: str = "ok",
|
|
exit_code: int = 0,
|
|
observed_at: str | None = None,
|
|
output: str | None = None,
|
|
incident_key: str | None = None,
|
|
notifications_enabled: bool = True,
|
|
) -> dict:
|
|
ev = {
|
|
"event_id": uuidv7(),
|
|
"check_id": check_id,
|
|
"observed_at": observed_at or now_iso(),
|
|
"status": status,
|
|
"exit_code": exit_code,
|
|
"duration_ms": 12,
|
|
"output_truncated": False,
|
|
"notifications_enabled": notifications_enabled,
|
|
}
|
|
if output is not None:
|
|
ev["output"] = output
|
|
if incident_key is not None:
|
|
ev["incident_key"] = incident_key
|
|
return ev
|
|
|
|
|
|
def make_heartbeat(agent_id: str = "agent-1") -> dict:
|
|
return {
|
|
"agent_id": agent_id,
|
|
"observed_at": now_iso(),
|
|
"hostname": agent_id,
|
|
"features": {"push": True, "metrics": True},
|
|
"labels": {},
|
|
}
|
|
|
|
|
|
def agent_token(agent_id: str = "agent-1") -> str:
|
|
safe = "".join(ch if ch.isalnum() else "0" for ch in agent_id)
|
|
return f"test-agent-key-{safe}-00000000000000000000000000000000"
|
|
|
|
|
|
def agent_headers(agent_id: str = "agent-1") -> dict[str, str]:
|
|
return {"Authorization": f"Bearer {agent_token(agent_id)}"}
|
|
|
|
|
|
async def register_agent(
|
|
client, ui_headers: dict[str, str], agent_id: str = "agent-1"
|
|
) -> dict[str, str]:
|
|
headers = agent_headers(agent_id)
|
|
r = await client.post("/api/v1/heartbeat", json=make_heartbeat(agent_id), headers=headers)
|
|
assert r.status_code == 202
|
|
r = await client.post(f"/api/v1/agents/{agent_id}/accept", headers=ui_headers)
|
|
assert r.status_code == 200
|
|
assert r.json()["count"] == 1
|
|
return headers
|