zhifu-platform/tests/test_sse_limits.py

import asyncio
from uuid import uuid4

import pytest
from pydantic import ValidationError
from test_after_sales import scenario  # noqa: F401

from zhifu.core.config import Settings
from zhifu.modules.conversations.stream import ADMIT


@pytest.mark.parametrize("global_limit,actor_limit", [(0, 3), (100001, 3), (200, 0), (3, 4)])
def test_invalid_sse_limits(global_limit, actor_limit):
    with pytest.raises(ValidationError):
        Settings(sse_global_limit=global_limit, sse_actor_limit=actor_limit)


async def test_configurable_sse_admission_is_atomic_and_shared(scenario):  # noqa: F811
    state = scenario["app"].state.state
    prefix = "verify:sse:" + uuid4().hex
    keys = [prefix, prefix + ":actor1", prefix + ":actor2"]
    try:
        results = await asyncio.gather(*(
            state.eval(ADMIT, 2, keys[0], keys[1 + index % 2], uuid4().hex, 5, 3)
            for index in range(40)
        ))
        assert sum(results) == 5
        assert await state.zcard(keys[0]) == 5
        assert all(count <= 3 for count in await asyncio.gather(
            state.zcard(keys[1]), state.zcard(keys[2])
        ))
        # 使用真实 Redis 到期分数验证异常中断的额度回收。
        for key in keys:
            members = await state.zrange(key, 0, -1)
            if members:
                await state.zadd(key, dict.fromkeys(members, 0))
        assert await state.eval(ADMIT, 2, keys[0], keys[1], uuid4().hex, 3000, 3) == 1
        assert await state.zcard(keys[0]) == 1
    finally:
        await state.delete(*keys)