"""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