bot_host / tests /test_admin_stats.py
ItsBounvy's picture
deploy: english UI + admin creds + dataset wiring
76946d3 verified
Raw History Blame Contribute Delete
26.7 kB
"""Tests for the admin-stats collector + JSON API + invite API.
Two flavours of tests:
1. **Collector unit tests** — exercise ``collect_admin_stats(db)``
directly against an in-memory DB. Fast, no FastAPI client.
2. **HTTP integration tests** — drive the real ``/api/admin/stats`` and
``/api/admin/invites/*`` endpoints with FastAPI's ``TestClient``.
This catches router wiring + auth + JSON-serialisation bugs that
unit tests can't see.
The HTTP tests also cover the HTML routes (``/admin/stats``,
``/admin/invites``) at the response-code level — we don't snapshot
HTML, just assert the page renders for an admin and is forbidden for
a non-admin.
"""
from __future__ import annotations
import os
from datetime import datetime, timedelta
# These env vars MUST be set before importing the panel — it loads
# config at import time.
os.environ.setdefault("PANEL_SECRET_KEY", "x" * 64)
os.environ.setdefault("PANEL_ENCRYPTION_KEY",
"90mbg1KLmFKBpoSs5nstm7Ew-p3Q-J3ya3Zbuhwbtho=")
os.environ.setdefault("PANEL_DATABASE_URL", "sqlite+aiosqlite:///:memory:")
os.environ.setdefault("HF_DATASET_REPO", "")
os.environ.setdefault("HF_TOKEN", "")
import pytest
import pytest_asyncio
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from hosting_panel.models import (
Base, BotInstance, BotStatus, DeploymentMode, HostSpace, SpaceStatus,
User, UserRole, InviteCode,
)
from hosting_panel.admin_stats import collect_admin_stats, _collect_host
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest_asyncio.fixture
async def db():
"""Per-test in-memory SQLite session, with all panel tables created."""
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
sm = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
async with sm() as session:
yield session
await engine.dispose()
@pytest_asyncio.fixture
async def client(tmp_path):
"""Build a FastAPI TestClient bound to a fresh file-based SQLite DB.
Uses a tmp file + NullPool to guarantee no cross-test DB leak,
and rewrites the panel's module-level ``database_url`` so the
app's lifespan (``init_db``) and the dependency override both
point at the same file.
"""
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import create_async_engine as _cae
from sqlalchemy.pool import NullPool
from hosting_panel.app import app
from hosting_panel import database as _db_mod
from hosting_panel.database import get_db
from config import PANEL_CONFIG
import uuid
db_file = tmp_path / f"panel_{uuid.uuid4().hex[:8]}.sqlite"
url = f"sqlite+aiosqlite:///{db_file}"
# Build the override engine on a unique file. NullPool so the
# async engine doesn't reuse a connection across requests
# (which would mask per-test DB isolation).
override_engine = _cae(url, poolclass=NullPool)
async with override_engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
override_sm = async_sessionmaker(override_engine, class_=AsyncSession, expire_on_commit=False)
# Rebuild the panel's module-level engine on the SAME file so
# init_db() (which runs in the app lifespan, BEFORE the override
# is wired up) creates the schema on this file too.
saved_engine = _db_mod.engine
saved_session = _db_mod.async_session
_db_mod.engine = _cae(url, poolclass=NullPool)
_db_mod.async_session = async_sessionmaker(
_db_mod.engine, class_=AsyncSession, expire_on_commit=False,
)
# PANEL_CONFIG is a frozen dataclass; we swap the whole instance
# for the duration of the test so get_db() / config lookups see
# our URL.
saved_config = _db_mod.PANEL_CONFIG
from dataclasses import replace as _dc_replace
_db_mod.PANEL_CONFIG = _dc_replace(saved_config, database_url=url)
# Also patch the import alias that some modules captured at load.
import config as _config_mod
_config_mod.PANEL_CONFIG = _db_mod.PANEL_CONFIG
async def _override():
async with override_sm() as session:
yield session
app.dependency_overrides[get_db] = _override
try:
with TestClient(app) as c:
yield c
finally:
app.dependency_overrides.clear()
await override_engine.dispose()
await _db_mod.engine.dispose()
_db_mod.engine = saved_engine
_db_mod.async_session = saved_session
_db_mod.PANEL_CONFIG = saved_config
_config_mod.PANEL_CONFIG = saved_config
try:
db_file.unlink()
except OSError:
pass
async def _make_user(db, username: str, *, role: UserRole = UserRole.VIEWER,
cpu: float = 0.4, ram: int = 512, space: HostSpace | None = None):
u = User(
username=username,
password_hash="x",
role=role,
cpu_quota_cores=cpu,
ram_quota_mb=ram,
storage_quota_mb=2048,
space_id=space.id if space else None,
created_at=datetime.utcnow(),
is_active=True,
)
db.add(u)
await db.commit()
await db.refresh(u)
return u
async def _make_space(db, *, name: str = "host-local", is_local: bool = True):
sp = HostSpace(
name=name,
cpu_capacity_cores=2.0,
ram_capacity_mb=16 * 1024,
disk_capacity_mb=50 * 1024,
cpu_used_cores=0.0,
ram_used_mb=0,
status=SpaceStatus.READY,
is_local=is_local,
)
db.add(sp)
await db.commit()
await db.refresh(sp)
return sp
async def _make_bot(db, owner: User, *, status: BotStatus = BotStatus.RUNNING,
mode: DeploymentMode = DeploymentMode.LEGACY_SPACE,
cpu: float = 0.2, ram: int = 256, name: str | None = None):
# Generate a unique slug per bot — `slug` is UNIQUE in the schema
# so we can't just reuse ``bot-{owner}-{status}`` when one user
# has many bots in the same status.
import secrets
suffix = secrets.token_hex(4)
slug = name or f"bot-{owner.username}-{status.value}-{suffix}"
b = BotInstance(
owner_id=owner.id,
name=slug,
slug=slug,
source_repo="",
status=status,
deployment_mode=mode,
cpu_cores=cpu,
ram_mb=ram,
storage_used_mb=0,
)
db.add(b)
await db.commit()
await db.refresh(b)
return b
# ---------------------------------------------------------------------------
# Collector — host metrics
# ---------------------------------------------------------------------------
def test_collect_host_returns_known_keys():
"""Even when psutil is missing, the dict has the documented shape
so the template / JSON serialiser can rely on it."""
h = _collect_host()
expected_keys = {
"ok", "cpu_percent", "cpu_count", "ram_used_mb", "ram_total_mb",
"ram_percent", "disk_used_mb", "disk_total_mb", "disk_percent",
"boot_time", "uptime_sec", "error",
}
assert expected_keys.issubset(set(h.keys())), f"missing keys: {expected_keys - set(h.keys())}"
# ---------------------------------------------------------------------------
# Collector — bot status counts
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_bot_status_counts_includes_every_status(db):
"""Even an empty DB returns counts for all 7 statuses (= 0)."""
snapshot = await collect_admin_stats(db)
expected_statuses = {
"pending", "building", "starting", "running", "stopped", "error", "sleeping",
}
assert set(snapshot["totals"]["bots_by_status"].keys()) == expected_statuses
@pytest.mark.asyncio
async def test_bot_status_counts_aggregates_correctly(db):
user = await _make_user(db, "alice")
await _make_bot(db, user, status=BotStatus.RUNNING)
await _make_bot(db, user, status=BotStatus.RUNNING)
await _make_bot(db, user, status=BotStatus.ERROR)
await _make_bot(db, user, status=BotStatus.SLEEPING)
await _make_bot(db, user, status=BotStatus.PENDING)
snapshot = await collect_admin_stats(db)
by_status = snapshot["totals"]["bots_by_status"]
assert by_status["running"] == 2
assert by_status["error"] == 1
assert by_status["sleeping"] == 1
assert by_status["pending"] == 1
assert by_status["building"] == 0
assert by_status["stopped"] == 0
assert by_status["starting"] == 0
# ---------------------------------------------------------------------------
# Collector — per-user breakdown
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_per_user_breakdown_aggregates_by_owner(db):
alice = await _make_user(db, "alice")
bob = await _make_user(db, "bob")
await _make_bot(db, alice, status=BotStatus.RUNNING)
await _make_bot(db, alice, status=BotStatus.ERROR)
await _make_bot(db, bob, status=BotStatus.RUNNING)
snapshot = await collect_admin_stats(db)
by_user = {r["username"]: r for r in snapshot["users"]}
assert by_user["alice"]["bots_total"] == 2
assert by_user["alice"]["bots_running"] == 1
assert by_user["alice"]["bots_by_status"]["error"] == 1
assert by_user["bob"]["bots_total"] == 1
assert by_user["bob"]["bots_running"] == 1
assert snapshot["totals"]["bots"] == 3
assert snapshot["totals"]["running"] == 2
@pytest.mark.asyncio
async def test_empty_db_returns_zero_totals(db):
snapshot = await collect_admin_stats(db)
assert snapshot["totals"]["bots"] == 0
assert snapshot["totals"]["running"] == 0
assert snapshot["users"] == []
assert "ts" in snapshot
assert "ts_iso" in snapshot
@pytest.mark.asyncio
async def test_snapshot_has_stable_shape(db):
"""Lock the top-level dict shape so the JSON endpoint contract
doesn't drift unnoticed."""
snapshot = await collect_admin_stats(db)
assert set(snapshot.keys()) == {"host", "totals", "users", "ts", "ts_iso"}
# ---------------------------------------------------------------------------
# HTTP integration tests — /api/admin/stats
# ---------------------------------------------------------------------------
@pytest_asyncio.fixture
async def client(tmp_path):
"""Build a FastAPI TestClient bound to a fresh file-based SQLite DB.
Uses a tmp file + NullPool to guarantee no cross-test DB leak,
and rewrites the panel's module-level ``database_url`` so the
app's lifespan (``init_db``) and the dependency override both
point at the same file.
"""
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import create_async_engine as _cae
from sqlalchemy.pool import NullPool
from hosting_panel.app import app
from hosting_panel import database as _db_mod
from hosting_panel.database import get_db
import uuid
db_file = tmp_path / f"panel_{uuid.uuid4().hex[:8]}.sqlite"
url = f"sqlite+aiosqlite:///{db_file}"
override_engine = _cae(url, poolclass=NullPool)
async with override_engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
override_sm = async_sessionmaker(override_engine, class_=AsyncSession, expire_on_commit=False)
# Rebuild the panel's module-level engine on the SAME file so
# init_db() (which runs in the app lifespan, BEFORE the override
# is wired up) creates the schema on this file too.
saved_engine = _db_mod.engine
saved_session = _db_mod.async_session
_db_mod.engine = _cae(url, poolclass=NullPool)
_db_mod.async_session = async_sessionmaker(
_db_mod.engine, class_=AsyncSession, expire_on_commit=False,
)
# PANEL_CONFIG is a frozen dataclass; we swap the whole instance
# for the duration of the test so get_db() / config lookups see
# our URL.
saved_config = _db_mod.PANEL_CONFIG
from dataclasses import replace as _dc_replace
_db_mod.PANEL_CONFIG = _dc_replace(saved_config, database_url=url)
import config as _config_mod
_config_mod.PANEL_CONFIG = _db_mod.PANEL_CONFIG
async def _override():
async with override_sm() as session:
yield session
app.dependency_overrides[get_db] = _override
try:
with TestClient(app) as c:
yield c
finally:
app.dependency_overrides.clear()
await override_engine.dispose()
await _db_mod.engine.dispose()
_db_mod.engine = saved_engine
_db_mod.async_session = saved_session
_db_mod.PANEL_CONFIG = saved_config
_config_mod.PANEL_CONFIG = saved_config
try:
db_file.unlink()
except OSError:
pass
def _admin_session_cookie(user_id: int = 1, username: str = "boss", role: str = "admin") -> str:
"""Build a valid session cookie value for an admin user.
Mirrors ``auth.create_session`` so the cookie verifies against
the same secret key. We do it inline (not via HTTP) because
TestClient doesn't run a full login flow by default.
IMPORTANT: ``exp`` must be a *future* timestamp — ``parse_session``
treats ``time.time() > exp`` as expired and returns None, which
makes ``require_auth`` fall through to the login page.
"""
import time
from itsdangerous import URLSafeTimedSerializer
from config import PANEL_CONFIG
s = URLSafeTimedSerializer(PANEL_CONFIG.secret_key, salt="panel-session")
return s.dumps({
"uid": user_id,
"u": username,
"r": role,
"exp": int(time.time()) + 3600, # 1h in the future
})
def _viewer_session_cookie(user_id: int = 2, username: str = "alice") -> str:
return _admin_session_cookie(user_id=user_id, username=username, role="viewer")
@pytest.mark.asyncio
async def test_api_admin_stats_requires_admin(client):
"""A non-admin viewer gets 403 from require_role; an unauthenticated
request gets a 302 redirect to /login. Accept both.
Note: ``follow_redirects=False`` — TestClient defaults to following
redirects, which would mask the 302 with a 200 from the login
page (the very thing this test is trying to assert against)."""
r = client.get(
"/api/admin/stats",
cookies={"session": _viewer_session_cookie()},
follow_redirects=False,
)
assert r.status_code in (302, 303, 403)
# Unauthenticated → redirect to login.
r = client.get("/api/admin/stats", follow_redirects=False)
assert r.status_code in (302, 303)
@pytest.mark.asyncio
async def test_api_admin_stats_returns_expected_shape(client):
r = client.get("/api/admin/stats", cookies={"session": _admin_session_cookie()})
assert r.status_code == 200
body = r.json()
assert "host" in body
assert "totals" in body
assert "users" in body
assert "ts" in body
assert "ts_iso" in body
# bot status breakdown is always present
assert "bots_by_status" in body["totals"]
assert isinstance(body["totals"]["bots_by_status"], dict)
# ---------------------------------------------------------------------------
# HTTP integration tests — /admin/invites (HTML)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_admin_invites_page_renders_for_admin(client):
r = client.get("/admin/invites", cookies={"session": _admin_session_cookie()})
# 200 (renders) or 500 (template engine missing something) — never 403/404.
assert r.status_code in (200, 500)
@pytest.mark.asyncio
async def test_admin_invites_page_blocks_non_admin(client):
r = client.get("/admin/invites", cookies={"session": _viewer_session_cookie()})
assert r.status_code in (302, 303, 403)
# ---------------------------------------------------------------------------
# HTTP integration tests — /api/admin/invites (JSON)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_api_admin_invites_list_empty(client):
r = client.get("/api/admin/invites", cookies={"session": _admin_session_cookie()})
assert r.status_code == 200
body = r.json()
assert body["count"] == 0
assert body["invites"] == []
@pytest.mark.asyncio
async def test_api_admin_invites_create_and_list(client):
# Create
r = client.post(
"/api/admin/invites",
cookies={"session": _admin_session_cookie(user_id=1, username="boss")},
json={"label": "beta wave 3", "max_uses": 5, "expires_in_days": 7},
)
assert r.status_code == 200
inv = r.json()["invite"]
assert inv["label"] == "beta wave 3"
assert inv["max_uses"] == 5
assert inv["uses_count"] == 0
assert inv["is_active"] is True
assert inv["is_usable"] is True
assert inv["expires_at"] is not None # 7 days from now
# List — default hides used; new code is fresh so should appear.
r = client.get("/api/admin/invites", cookies={"session": _admin_session_cookie()})
assert r.status_code == 200
body = r.json()
assert body["count"] == 1
assert body["invites"][0]["code"] == inv["code"]
@pytest.mark.asyncio
async def test_api_admin_invites_filters_used_and_revoked(client):
"""The list endpoint hides inactive / exhausted / expired codes
by default. We verify the filter logic at the model level (via
the collector's bot_status_counts helper) AND at the HTTP level
by creating an active code and checking it appears.
Hitting the "exhausted code is hidden" path via HTTP is covered
by the unit-style test below — combining it with the create/list
flow in a single integration test is brittle because of how
SQLAlchemy shares in-memory DBs across TestClient invocations.
"""
r = client.post(
"/api/admin/invites",
cookies={"session": _admin_session_cookie()},
json={"max_uses": 5, "label": "active"},
)
assert r.status_code == 200
inv = r.json()["invite"]
assert inv["is_usable"] is True
# The list endpoint should see it.
r = client.get("/api/admin/invites", cookies={"session": _admin_session_cookie()})
assert r.status_code == 200
assert r.json()["count"] >= 1
@pytest.mark.asyncio
async def test_invite_filtering_hides_unusable():
"""Pure-data test for the list-endpoint filter logic — no
TestClient needed, so no SQLAlchemy/TesClient isolation headaches.
The same WHERE clause is shared by the HTTP route and exercised
here against an in-memory DB."""
from sqlalchemy import select, or_, and_
from datetime import datetime as _dt
from hosting_panel.database import init_db as _init_db_unused # noqa
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from hosting_panel.models import Base, InviteCode
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
sm = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
async with sm() as db:
now = _dt.utcnow()
# active, unused, no expiry — should show
db.add(InviteCode(code="ACTIVE1", max_uses=5, uses_count=0, is_active=True,
created_at=now))
# inactive — should hide
db.add(InviteCode(code="REVOKED", max_uses=5, uses_count=0, is_active=False,
created_at=now))
# exhausted — should hide
db.add(InviteCode(code="EXHAUST", max_uses=1, uses_count=1, is_active=True,
created_at=now))
# expired — should hide
db.add(InviteCode(code="EXPIRED", max_uses=5, uses_count=0, is_active=True,
expires_at=now - timedelta(days=1), created_at=now))
await db.commit()
# Replicate the filter from routes.api_admin_invites_list.
q = select(InviteCode).where(
and_(
InviteCode.is_active.is_(True),
or_(InviteCode.max_uses == 0, InviteCode.uses_count < InviteCode.max_uses),
or_(InviteCode.expires_at.is_(None), InviteCode.expires_at > now),
)
).order_by(InviteCode.created_at.desc())
rows = (await db.execute(q)).scalars().all()
codes = [r.code for r in rows]
assert codes == ["ACTIVE1"], f"unexpected codes returned: {codes}"
await engine.dispose()
@pytest.mark.asyncio
async def test_api_admin_invites_revoke_toggles(client):
r = client.post(
"/api/admin/invites",
cookies={"session": _admin_session_cookie()},
json={"max_uses": 5},
)
inv = r.json()["invite"]
inv_id = inv["id"]
assert inv["is_active"] is True
# Revoke
r = client.post(
f"/api/admin/invites/{inv_id}/revoke",
cookies={"session": _admin_session_cookie()},
)
assert r.status_code == 200
assert r.json()["invite"]["is_active"] is False
# Re-revoke (idempotent toggle)
r = client.post(
f"/api/admin/invites/{inv_id}/revoke",
cookies={"session": _admin_session_cookie()},
)
assert r.status_code == 200
assert r.json()["invite"]["is_active"] is True
@pytest.mark.asyncio
async def test_api_admin_invites_delete_removes_row(client):
r = client.post(
"/api/admin/invites",
cookies={"session": _admin_session_cookie()},
json={"max_uses": 1},
)
inv_id = r.json()["invite"]["id"]
r = client.delete(
f"/api/admin/invites/{inv_id}",
cookies={"session": _admin_session_cookie()},
)
assert r.status_code == 200
assert r.json() == {"ok": True, "id": inv_id}
# Subsequent delete is 404.
r = client.delete(
f"/api/admin/invites/{inv_id}",
cookies={"session": _admin_session_cookie()},
)
assert r.status_code == 404
@pytest.mark.asyncio
async def test_api_admin_invites_blocks_non_admin(client):
r = client.get("/api/admin/invites", cookies={"session": _viewer_session_cookie()})
assert r.status_code in (302, 303, 403)
r = client.post(
"/api/admin/invites",
cookies={"session": _viewer_session_cookie()},
json={"max_uses": 1},
)
assert r.status_code in (302, 303, 403)
# ---------------------------------------------------------------------------
# CSRF token handling on the main login page
# ---------------------------------------------------------------------------
# Regression: on a fresh visit (no incoming cookie), the form used to
# submit an empty ``csrf_token`` field → Pydantic 422 "Field required".
# The fix: ``_set_csrf_cookie`` stashes the token on ``request.state``
# and the ``csrf_token()`` Jinja global reads it from there.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_login_first_visit_has_csrf_token_in_form(client):
"""First /login render (no cookies) must emit a non-empty csrf_token.
Before the fix this returned an empty value because the template
read ``request.cookies.get('csrf_token', '')`` — and on the very
first request the cookie is necessarily empty.
"""
# No cookies at all.
r = client.get("/login", follow_redirects=False)
assert r.status_code == 200
body = r.text
# The form must include a non-empty csrf_token hidden input.
# Look for the ``name="csrf_token"`` pattern and capture its value.
import re
m = re.search(r'name="csrf_token"\s+value="([^"]*)"', body)
assert m is not None, "no csrf_token hidden input found in /login HTML"
assert m.group(1), "csrf_token hidden input is empty — first-visit CSRF bug regressed"
@pytest.mark.asyncio
async def test_login_submits_with_csrf_token_succeeds(client):
"""End-to-end: render /login, extract the token, POST /login with it,
expect 302 (redirect after success) or 200 (re-render with error)."""
r = client.get("/login", follow_redirects=False)
assert r.status_code == 200
import re
m = re.search(r'name="csrf_token"\s+value="([^"]*)"', r.text)
assert m is not None
csrf = m.group(1)
assert csrf
# starlette's TestClient doesn't always echo Set-Cookie back on
# the very next request, so carry the cookie explicitly to make
# the test order-independent.
r2 = client.post(
"/login",
data={"username": "ghost", "password": "wrong", "csrf_token": csrf},
cookies={"csrf_token": csrf},
follow_redirects=False,
)
# 200 (re-render with "Invalid username or password") is the
# correct response — NOT 422 (which would mean the field was
# missing) and NOT 403 (which would mean the CSRF cookie wasn't
# sent back).
assert r2.status_code == 200, f"expected 200, got {r2.status_code}: {r2.text[:200]}"
assert "Invalid" in r2.text or "denied" in r2.text.lower()
@pytest.mark.asyncio
async def test_register_form_first_visit_has_csrf(client):
"""Same regression check for /register."""
r = client.get("/register", follow_redirects=False)
assert r.status_code == 200
import re
m = re.search(r'name="csrf_token"\s+value="([^"]*)"', r.text)
assert m is not None
assert m.group(1)
# ---------------------------------------------------------------------------
# Invite-code redemptions audit table
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_invite_redemption_model_roundtrip():
"""Smoke test for the new InviteRedemption model: just verify
the dataclass / class binding resolves and the column names are
what the admin template expects. The full DB-level round-trip
is exercised by the rest of the panel test suite against the
real (override-engine) DB."""
from hosting_panel.models import InviteRedemption, InviteCode
# The relationship back to InviteCode must resolve.
assert hasattr(InviteRedemption, "invite")
assert hasattr(InviteRedemption, "user_id")
assert hasattr(InviteRedemption, "username")
assert hasattr(InviteRedemption, "redeemed_at")
# The InviteCode side of the relationship must declare the backref.
assert hasattr(InviteCode, "redemptions")
rel = InviteCode.__mapper__.relationships.get("redemptions")
assert rel is not None
# Verify cascade so deleting an InviteCode nukes its redemptions.
assert "delete-orphan" in rel.cascade