agent: complete plan05 closeout

agent/nix_builder_autoscaler/tests/test_reservations_api.py

@@ -3,24 +3,29 @@
 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() -> tuple[TestClient, StateDB, FakeClock, MetricsRegistry]:
+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)
+    app = create_app(db, config, clock, metrics, reconcile_now=reconcile_now)
     return TestClient(app), db, clock, metrics
 
 
@@ -120,6 +125,20 @@ def test_health_ready_returns_ok_when_no_checks() -> None:
     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)
@@ -150,3 +169,67 @@ def test_release_nonexistent_returns_404() -> None:
     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