nix-builder-autoscaler/agent/nix_builder_autoscaler/tests/test_reservations_api.py
Abel Luck 44bc99ab85
All checks were successful
buildbot/nix-eval Build done.
buildbot/nix-build Build done.
buildbot/nix-effects Build done.
add termination cooldown for slot scale-down
2026-02-27 18:37:58 +01:00

247 lines
8.9 KiB
Python

"""Reservations API unit tests."""
from __future__ import annotations
from datetime import UTC, datetime
from typing import Any
from fastapi.testclient import TestClient
from nix_builder_autoscaler.api import create_app
from nix_builder_autoscaler.config import AppConfig, CapacityConfig
from nix_builder_autoscaler.metrics import MetricsRegistry
from nix_builder_autoscaler.models import SlotState
from nix_builder_autoscaler.providers.clock import FakeClock
from nix_builder_autoscaler.state_db import StateDB
def _make_client(
*,
reconcile_now: Any = None, # noqa: ANN401
) -> tuple[TestClient, StateDB, FakeClock, MetricsRegistry]:
clock = FakeClock()
db = StateDB(":memory:", clock=clock)
db.init_schema()
db.init_slots("slot", 3, "x86_64-linux", "all")
config = AppConfig(capacity=CapacityConfig(reservation_ttl_seconds=1200))
metrics = MetricsRegistry()
app = create_app(db, config, clock, metrics, reconcile_now=reconcile_now)
return TestClient(app), db, clock, metrics
def test_create_reservation_returns_200() -> None:
client, _, _, _ = _make_client()
response = client.post("/v1/reservations", json={"system": "x86_64-linux", "reason": "test"})
assert response.status_code == 200
body = response.json()
assert body["reservation_id"].startswith("resv_")
assert body["phase"] == "pending"
assert body["system"] == "x86_64-linux"
assert "created_at" in body
assert "expires_at" in body
def test_get_reservation_returns_current_phase() -> None:
client, _, _, _ = _make_client()
created = client.post("/v1/reservations", json={"system": "x86_64-linux", "reason": "test"})
reservation_id = created.json()["reservation_id"]
response = client.get(f"/v1/reservations/{reservation_id}")
assert response.status_code == 200
assert response.json()["phase"] == "pending"
def test_release_reservation_moves_to_released() -> None:
client, _, _, _ = _make_client()
created = client.post("/v1/reservations", json={"system": "x86_64-linux", "reason": "test"})
reservation_id = created.json()["reservation_id"]
response = client.post(f"/v1/reservations/{reservation_id}/release")
assert response.status_code == 200
assert response.json()["phase"] == "released"
def test_unknown_reservation_returns_404() -> None:
client, _, _, _ = _make_client()
response = client.get("/v1/reservations/resv_nonexistent")
assert response.status_code == 404
body = response.json()
assert body["error"]["code"] == "not_found"
assert "request_id" in body
def test_malformed_body_returns_422() -> None:
client, _, _, _ = _make_client()
response = client.post("/v1/reservations", json={"invalid": "data"})
assert response.status_code == 422
def test_list_reservations_returns_all() -> None:
client, _, _, _ = _make_client()
client.post("/v1/reservations", json={"system": "x86_64-linux", "reason": "a"})
client.post("/v1/reservations", json={"system": "x86_64-linux", "reason": "b"})
response = client.get("/v1/reservations")
assert response.status_code == 200
assert len(response.json()) == 2
def test_list_reservations_filters_by_phase() -> None:
client, _, _, _ = _make_client()
created = client.post("/v1/reservations", json={"system": "x86_64-linux", "reason": "test"})
reservation_id = created.json()["reservation_id"]
client.post(f"/v1/reservations/{reservation_id}/release")
response = client.get("/v1/reservations?phase=released")
assert response.status_code == 200
body = response.json()
assert len(body) == 1
assert body[0]["phase"] == "released"
def test_list_slots_returns_all_slots() -> None:
client, _, _, _ = _make_client()
response = client.get("/v1/slots")
assert response.status_code == 200
assert len(response.json()) == 3
def test_state_summary_returns_counts() -> None:
client, _, _, _ = _make_client()
response = client.get("/v1/state/summary")
assert response.status_code == 200
body = response.json()
assert body["slots"]["total"] == 3
assert body["slots"]["empty"] == 3
def test_effective_config_returns_capacity_and_scheduler() -> None:
client, _, _, _ = _make_client()
response = client.get("/v1/config/effective")
assert response.status_code == 200
body = response.json()
assert body["capacity"]["max_slots"] == 8
assert body["capacity"]["idle_scale_down_seconds"] == 900
assert body["capacity"]["termination_cooldown_seconds"] == 180
assert body["scheduler"]["tick_seconds"] == 3.0
assert body["scheduler"]["reconcile_seconds"] == 15.0
def test_health_live_returns_ok() -> None:
client, _, _, _ = _make_client()
response = client.get("/health/live")
assert response.status_code == 200
assert response.json()["status"] == "ok"
def test_health_ready_returns_ok_when_no_checks() -> None:
client, _, _, _ = _make_client()
response = client.get("/health/ready")
assert response.status_code == 200
assert response.json()["status"] == "ok"
def test_health_ready_degraded_when_ready_check_fails() -> None:
clock = FakeClock()
db = StateDB(":memory:", clock=clock)
db.init_schema()
db.init_slots("slot", 3, "x86_64-linux", "all")
config = AppConfig(capacity=CapacityConfig(reservation_ttl_seconds=1200))
metrics = MetricsRegistry()
app = create_app(db, config, clock, metrics, ready_check=lambda: False)
client = TestClient(app)
response = client.get("/health/ready")
assert response.status_code == 503
assert response.json()["status"] == "degraded"
def test_metrics_returns_prometheus_text() -> None:
client, _, _, metrics = _make_client()
metrics.counter("autoscaler_test_counter", {}, 1.0)
response = client.get("/metrics")
assert response.status_code == 200
assert "text/plain" in response.headers["content-type"]
assert "autoscaler_test_counter" in response.text
def test_capacity_hint_accepted() -> None:
client, _, _, _ = _make_client()
response = client.post(
"/v1/hints/capacity",
json={
"builder": "buildbot",
"queued": 2,
"running": 4,
"system": "x86_64-linux",
"timestamp": datetime(2026, 1, 1, tzinfo=UTC).isoformat(),
},
)
assert response.status_code == 200
assert response.json()["status"] == "accepted"
def test_release_nonexistent_returns_404() -> None:
client, _, _, _ = _make_client()
response = client.post("/v1/reservations/resv_nonexistent/release")
assert response.status_code == 404
assert response.json()["error"]["code"] == "not_found"
def test_admin_drain_success() -> None:
client, db, _, _ = _make_client()
db.update_slot_state("slot001", SlotState.LAUNCHING, instance_id="i-test")
db.update_slot_state("slot001", SlotState.BOOTING)
db.update_slot_state("slot001", SlotState.BINDING, instance_ip="100.64.0.1")
db.update_slot_state("slot001", SlotState.READY)
response = client.post("/v1/admin/drain", json={"slot_id": "slot001"})
assert response.status_code == 200
assert response.json()["state"] == "draining"
slot = db.get_slot("slot001")
assert slot is not None
assert slot["state"] == SlotState.DRAINING.value
def test_admin_drain_invalid_state_returns_409() -> None:
client, _, _, _ = _make_client()
response = client.post("/v1/admin/drain", json={"slot_id": "slot001"})
assert response.status_code == 409
assert response.json()["error"]["code"] == "invalid_state"
def test_admin_unquarantine_success() -> None:
client, db, _, _ = _make_client()
db.update_slot_state("slot001", SlotState.ERROR, instance_id="i-bad")
response = client.post("/v1/admin/unquarantine", json={"slot_id": "slot001"})
assert response.status_code == 200
assert response.json()["state"] == "empty"
slot = db.get_slot("slot001")
assert slot is not None
assert slot["state"] == SlotState.EMPTY.value
assert slot["instance_id"] is None
def test_admin_unquarantine_invalid_state_returns_409() -> None:
client, _, _, _ = _make_client()
response = client.post("/v1/admin/unquarantine", json={"slot_id": "slot001"})
assert response.status_code == 409
assert response.json()["error"]["code"] == "invalid_state"
def test_admin_reconcile_now_not_configured_returns_503() -> None:
client, _, _, _ = _make_client()
response = client.post("/v1/admin/reconcile-now")
assert response.status_code == 503
assert response.json()["error"]["code"] == "not_configured"
def test_admin_reconcile_now_success() -> None:
called = {"value": False}
def _reconcile_now() -> dict[str, object]:
called["value"] = True
return {"triggered": True}
client, _, _, _ = _make_client(reconcile_now=_reconcile_now)
response = client.post("/v1/admin/reconcile-now")
assert response.status_code == 200
assert response.json()["status"] == "accepted"
assert response.json()["triggered"] is True
assert called["value"] is True