zhifu-platform/tests/test_tenant_dispatch.py

import asyncio
import secrets
from uuid import UUID, uuid4

import pytest
from sqlalchemy import select
from test_after_sales import scenario  # noqa: F401

from zhifu.modules.agent_runtime.models import AgentRun
from zhifu.modules.agent_runtime.worker import claim
from zhifu.modules.identity.access_models import Principal, StoreGrant
from zhifu.modules.identity.auth import password_hasher

pytestmark = pytest.mark.integration


async def test_large_tenant_queue_does_not_block_another_tenant(scenario):  # noqa: F811
    sessions = scenario["app"].state.sessions
    client = scenario["client"]
    tenant = scenario["tenant"]
    password = secrets.token_urlsafe(24)
    encoded = await asyncio.to_thread(password_hasher.hash, password)
    users = [uuid4() for _ in range(12)]
    async with sessions.begin() as session:
        for user in users:
            session.add(
                Principal(
                    id=user.bytes,
                    tenant_id=tenant.bytes,
                    username=user.hex,
                    password_hash=encoded,
                    kind="CUSTOMER",
                )
            )
        await session.flush()
        for user in users:
            session.add(
                StoreGrant(
                    tenant_id=tenant.bytes,
                    principal_id=user.bytes,
                    store_id=scenario["store"].bytes,
                    permission="orders:read",
                )
            )
    headers_by_run = {}

    async def enqueue(headers, store_id):
        response = await client.post(
            "/openapi/v1/agent/runs",
            headers={**headers, "Idempotency-Key": uuid4().hex},
            json={"store_id": store_id, "message": "查询订单"},
        )
        assert response.status_code == 202, response.text
        identifier = UUID(response.json()["run_id"]).bytes
        headers_by_run[identifier] = headers
        return identifier

    async def cancel(identifier):
        response = await client.post(
            f"/openapi/v1/agent/runs/{UUID(bytes=identifier)}/cancel",
            headers=headers_by_run[identifier],
        )
        assert response.status_code == 200, response.text

    # 使用真实登录和入队接口建立 36 个任务，每账户遵守 3 个未结束任务上限。
    for user in users:
        login = await client.post(
            "/openapi/v1/auth/sessions",
            json={
                "tenant_id": str(tenant),
                "username": user.hex,
                "password": password,
            },
        )
        assert login.status_code == 200, login.text
        headers = {"Authorization": "Bearer " + login.json()["access_token"]}
        for _ in range(3):
            await enqueue(headers, str(scenario["store"]))
        rejected = await client.post(
            "/openapi/v1/agent/runs",
            headers={**headers, "Idempotency-Key": uuid4().hex},
            json={"store_id": str(scenario["store"]), "message": "超额请求"},
        )
        assert rejected.status_code == 429
    first = await claim(sessions, 2)
    second = await claim(sessions, 2)
    assert first and second
    assert await claim(sessions, 2) is None
    identity = (await scenario["call"]("GET", "/openapi/v1/auth/me", "foreign")).json()
    other_tenant = UUID(identity["tenant_id"]).bytes
    other_runs = [
        await enqueue(scenario["headers"]["foreign"], identity["grants"][0]["store_id"])
        for _ in range(3)
    ]
    # A 已达到并发上限且还有 34 个更早的等待任务，B 仍应立即获得下一次空闲机会。
    third = await claim(sessions, 2)
    assert third and third[0] == other_runs[0]
    for identifier in (first[0], second[0], third[0]):
        await cancel(identifier)
    sequence = []
    for _ in range(4):
        selected = await claim(sessions, 2)
        assert selected
        async with sessions() as session:
            row = await session.get(AgentRun, selected[0])
            sequence.append(row.tenant_id)
        await cancel(selected[0])
    assert sequence == [tenant.bytes, other_tenant, tenant.bytes, other_tenant]
    async with sessions() as session:
        remaining = (
            await session.scalars(
                select(AgentRun).where(
                    AgentRun.tenant_id == tenant.bytes,
                    AgentRun.status == "QUEUED",
                )
            )
        ).all()
        assert len(remaining) == 32