tests/test_allocator.py

"""Unit tests for the multi-tenant space allocator.

Updated for the 2026-07-05 tier / approval rewrite.

Removed (deprecated)
--------------------
- ``test_assign_user_*`` \u2014 ``assign_user`` is now a compatibility
  shim that does NOT allocate.  Users no longer bind to a single
  Space; bots are placed per-tier via ``assign_bot``.

Added
-----
- ``test_packable_capacity_subtracts_baseline_cpu``
- ``test_assign_bot_uses_free_tier_space_for_free_user``
- ``test_assign_bot_rejects_tier_mismatch_via_pending_approval``
- ``test_request_space_creates_pending_approval_row``
- ``test_approve_space_flips_status_to_provisioning``
- ``test_reject_space_flips_status_to_error``
- ``test_create_space_manual_skips_approval``
- ``test_user_hard_cpu_limit_internal_only``
"""
from __future__ import annotations

import os

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 asyncio
from datetime import datetime

import pytest
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from models import (
    PLAN_TIER_FREE,
    PLAN_TIER_PAID,
    Base,
    BotInstance,
    BotStatus,
    DeploymentMode,
    HostSpace,
    SpaceStatus,
    User,
    UserRole,
)
from allocator import (
    ALLOCATOR,
    BOT_PORT_MAX,
    BOT_PORT_MIN,
    AllocationError,
    Capacity,
    SpacePendingApprovalError,
    space_allocable_capacity,
    user_hard_cpu_limit,
)


@pytest.fixture
async def db():
    """Per-test in-memory SQLite session."""
    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()


# ---- Factories -------------------------------------------------------------

async def _make_space(
    db, *, cpu=2.0, ram=16 * 1024, used_cpu=0.0, used_ram=0,
    status=SpaceStatus.READY, is_local=True, name="host-local",
    tier=PLAN_TIER_FREE, hardware="cpu-basic",
    baseline_cpu=0.1, baseline_ram=128,
):
    sp = HostSpace(
        name=name, tier=tier, hardware_tier=hardware,
        cpu_capacity_cores=cpu, ram_capacity_mb=ram,
        disk_capacity_mb=50 * 1024,
        baseline_cpu_reserved=baseline_cpu,
        baseline_ram_reserved_mb=baseline_ram,
        cpu_used_cores=used_cpu, ram_used_mb=used_ram,
        status=status, is_local=is_local,
    )
    db.add(sp)
    await db.commit()
    await db.refresh(sp)
    return sp


async def _make_user(
    db, *, cpu=0.5, ram=1024, storage=3072, space=None,
    tier=PLAN_TIER_FREE,
):
    u = User(
        username=f"u{abs(hash((cpu, ram, storage, space, tier))) % 9999}",
        password_hash="x",
        role=UserRole.VIEWER,
        cpu_quota_cores=cpu, ram_quota_mb=ram, storage_quota_mb=storage,
        plan_tier=tier,
        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_bot(db, owner, *, cpu=0.2, ram=256):
    import uuid
    b = BotInstance(
        owner_id=owner.id, name="Test", slug=f"test-{uuid.uuid4().hex[:8]}",
        status=BotStatus.PENDING,
        cpu_cores=cpu, ram_mb=ram,
        deployment_mode=DeploymentMode.MULTITENANT,
        env_vars_json="{}",
        space_state="on_dataset",
        created_at=datetime.utcnow(),
        updated_at=datetime.utcnow(),
    )
    db.add(b)
    await db.commit()
    await db.refresh(b)
    return b


# ---- Capacity math --------------------------------------------------------

def test_capacity_fits_when_under_limit():
    cap = Capacity(cpu=2.0, ram_mb=1024)
    used = Capacity(cpu=0.5, ram_mb=256)
    need = Capacity(cpu=0.5, ram_mb=256)
    assert need.fits_in(used, cap) is True


def test_capacity_rejects_overflow():
    cap = Capacity(cpu=2.0, ram_mb=1024)
    used = Capacity(cpu=1.5, ram_mb=900)
    need = Capacity(cpu=1.0, ram_mb=200)
    assert need.fits_in(used, cap) is False


def test_user_hard_cpu_limit_is_internal():
    """The factor used at runtime is always < the advertised quota so a
    single bot can never spend the orchestrator's headroom. Test only
    asserts the factor relationship \u2014 the actual number must NOT be
    leaked to UI strings."""
    assert user_hard_cpu_limit(0.4) == pytest.approx(0.38, abs=1e-4)
    assert user_hard_cpu_limit(0.4) < 0.4


def test_space_allocable_subtracts_baseline_cpu():
    """Packable capacity = raw - baseline.  With a 0.1 baseline on a
    2.0-CPU Space we get 1.9 CPU available for bots."""
    cap = space_allocable_capacity(
        cap_cores=2.0, cap_ram_mb=16384, cap_disk_mb=51200,
        baseline_cpu=0.1, baseline_ram_mb=128,
    )
    assert cap.cpu == pytest.approx(1.9, abs=1e-4)
    assert cap.ram_mb == 16256


# ---- Bot placement (tier-aware) ------------------------------------------

async def test_assign_bot_to_existing_tier_space(db):
    sp = await _make_space(
        db, cpu=4.0, ram=8 * 1024, baseline_cpu=0.0, baseline_ram=0,
    )
    u = await _make_user(db, cpu=0.5, ram=1024, tier=PLAN_TIER_FREE)
    b = await _make_bot(db, u, cpu=0.2, ram=256)
    space = await ALLOCATOR.assign_bot(db, b, u)
    assert space.id == sp.id
    assert b.space_id == sp.id
    assert b.port is not None
    assert BOT_PORT_MIN <= b.port <= BOT_PORT_MAX


async def test_assign_bot_refuses_tier_mismatch_via_pending_approval(db):
    """Paid user has no paid Space; allocator should request one in
    PENDING_APPROVAL and raise ``SpacePendingApprovalError``."""
    await _make_space(db, cpu=2.0, ram=8 * 1024, tier=PLAN_TIER_FREE,
                      baseline_cpu=0.0)
    u = await _make_user(db, cpu=0.4, ram=512, tier=PLAN_TIER_PAID)
    b = await _make_bot(db, u, cpu=0.2, ram=256)
    # Force require_space_approval=True so the new row sits as pending.
    from config import PANEL_CONFIG
    original = PANEL_CONFIG.require_space_approval
    PANEL_CONFIG.require_space_approval = True
    try:
        from sqlalchemy import select
        with pytest.raises(SpacePendingApprovalError) as exc_info:
            await ALLOCATOR.assign_bot(db, b, u)
        pending = (
            await db.execute(select(HostSpace).where(
                HostSpace.status == SpaceStatus.PENDING_APPROVAL
            ))
        ).scalar_one_or_none()
        assert pending is not None
        assert pending.tier == PLAN_TIER_PAID
        assert pending.id == exc_info.value.space.id
    finally:
        PANEL_CONFIG.require_space_approval = original


async def test_assign_bot_rejects_over_user_quota(db):
    """Two bots claiming more CPU than the user's quota: combined
    reservation overflows the soft limit."""
    sp = await _make_space(db, cpu=4.0, ram=8 * 1024,
                           baseline_cpu=0.0, baseline_ram=0)
    u = await _make_user(db, cpu=0.5, ram=1024, tier=PLAN_TIER_FREE, space=sp)
    b1 = await _make_bot(db, u, cpu=0.3, ram=256)
    b2 = await _make_bot(db, u, cpu=0.3, ram=256)  # 0.3 + 0.3 = 0.6 > 0.5 quota
    await ALLOCATOR.assign_bot(db, b1, u)
    # Refresh the owner so the just-allocated b1 is visible in
    # ``owner.bots`` (the relationship is selectin-loaded at user
    # creation time and won't pick up new rows until the next refresh).
    await db.refresh(u, attribute_names=["bots"])
    with pytest.raises(AllocationError, match="User CPU quota"):
        await ALLOCATOR.assign_bot(db, b2, u)


async def test_release_bot_returns_capacity(db):
    sp = await _make_space(db, cpu=4.0, ram=8 * 1024,
                           baseline_cpu=0.0, baseline_ram=0)
    u = await _make_user(db, cpu=0.5, ram=1024, tier=PLAN_TIER_FREE, space=sp)
    b = await _make_bot(db, u, cpu=0.3, ram=512)
    await ALLOCATOR.assign_bot(db, b, u)
    await db.refresh(sp)
    cpu_before = sp.cpu_used_cores
    ram_before = sp.ram_used_mb
    await ALLOCATOR.release_bot(db, b)
    await db.refresh(sp)
    assert sp.cpu_used_cores == max(0.0, cpu_before - 0.3)
    assert sp.ram_used_mb == max(0, ram_before - 512)
    assert b.port is None


# ---- Approval flow --------------------------------------------------------

async def test_request_space_creates_pending_approval_row(db):
    from config import PANEL_CONFIG
    original = PANEL_CONFIG.require_space_approval
    PANEL_CONFIG.require_space_approval = True
    try:
        space = await ALLOCATOR.request_space(
            db, tier=PLAN_TIER_PAID, requested_by_user_id=None,
        )
        assert space.status == SpaceStatus.PENDING_APPROVAL
        assert space.tier == PLAN_TIER_PAID
        assert space.requested_at is not None
        assert space.approved_at is None
        assert space.baseline_cpu_reserved == PANEL_CONFIG.space_baseline_cpu_reserved
    finally:
        PANEL_CONFIG.require_space_approval = original


async def test_approve_space_flips_status_to_provisioning(db):
    space = await ALLOCATOR.request_space(
        db, tier=PLAN_TIER_FREE, requested_by_user_id=None,
    )
    approved = await ALLOCATOR.approve_space(
        db, space.id, approved_by_user_id=42,
    )
    assert approved.status == SpaceStatus.PROVISIONING
    assert approved.approved_at is not None
    assert approved.approved_by_user_id == 42


async def test_reject_space_flips_status_to_error(db):
    space = await ALLOCATOR.request_space(
        db, tier=PLAN_TIER_FREE, requested_by_user_id=None,
    )
    rejected = await ALLOCATOR.reject_space(
        db, space.id, reason="redundant tier",
    )
    assert rejected.status == SpaceStatus.ERROR
    assert "redundant tier" in (rejected.operator_note or "")


async def test_approve_non_pending_space_is_rejected(db):
    space = await ALLOCATOR.create_space_manual(
        db, tier=PLAN_TIER_FREE, requested_by_user_id=1,
    )
    assert space.status == SpaceStatus.PROVISIONING
    # Approving a non-pending Space is wrong -> AllocationError.
    with pytest.raises(AllocationError):
        await ALLOCATOR.approve_space(db, space.id, approved_by_user_id=1)


async def test_create_space_manual_starts_in_provisioning(db):
    space = await ALLOCATOR.create_space_manual(
        db, tier=PLAN_TIER_FREE, requested_by_user_id=1,
        operator_note="test fixture",
    )
    assert space.status == SpaceStatus.PROVISIONING
    assert space.approved_at is not None
    assert space.approved_by_user_id == 1
    assert space.operator_note == "test fixture"


# ---- Port allocation ------------------------------------------------------

async def test_distinct_bots_get_distinct_ports(db):
    sp = await _make_space(db, cpu=4.0, ram=8 * 1024,
                           baseline_cpu=0.0, baseline_ram=0)
    u = await _make_user(db, cpu=2.0, ram=4096, tier=PLAN_TIER_FREE, space=sp)
    b1 = await _make_bot(db, u, cpu=0.1, ram=128)
    b2 = await _make_bot(db, u, cpu=0.1, ram=128)
    await ALLOCATOR.assign_bot(db, b1, u)
    await ALLOCATOR.assign_bot(db, b2, u)
    assert b1.port != b2.port
    assert BOT_PORT_MIN <= b1.port <= BOT_PORT_MAX
    assert BOT_PORT_MIN <= b2.port <= BOT_PORT_MAX