chongming-utils-task/tests/unit/test_buckets.py

"""分桶策略测试。

最重要的一条是**跨进程稳定性**：内置 ``hash()`` 带进程级随机盐，若图省事用它，
同一个键在不同进程里会落到不同的桶，分布式下"同桶任务在同一台机器"这条约定会
**静默**失效。这里用真实子进程与不同的 ``PYTHONHASHSEED`` 把它钉死。
"""

from __future__ import annotations

import os
import subprocess
import sys

import pytest

from chongming_task.policies import BucketPolicy


def _run_with_hash_seed(script: str, seed: str) -> str:
    """在指定 ``PYTHONHASHSEED`` 的子进程里跑一段脚本，返回其标准输出。"""
    environment = {**os.environ, "PYTHONHASHSEED": seed}
    completed = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        env=environment,
        check=True,
        timeout=60,
    )
    return completed.stdout.strip()


# ---------------------------------------------------------------------------
# 稳定性（分桶能用于分布式的全部前提）
# ---------------------------------------------------------------------------


def test_same_key_always_maps_to_the_same_bucket() -> None:
    policy = BucketPolicy(64)
    for _ in range(50):
        assert policy.bucket_for("tenant-42") == policy.bucket_for("tenant-42")


def test_assignment_is_stable_across_processes_with_different_hash_seeds() -> None:
    """内置 ``hash()`` 带随机盐，本实现必须不受它影响。"""
    script = (
        "from chongming_task.policies import BucketPolicy;"
        "print(' '.join(str(BucketPolicy(64).bucket_for(f'tenant-{i}')) for i in range(20)))"
    )
    first = _run_with_hash_seed(script, "0")
    second = _run_with_hash_seed(script, "1")
    assert first == second, "同一个键在不同进程里落到了不同的桶"

    local = " ".join(str(BucketPolicy(64).bucket_for(f"tenant-{i}")) for i in range(20))
    assert first == local, "子进程与本进程的结果不一致"


def test_builtin_hash_would_have_been_unusable() -> None:
    """留一条反例，说明为什么不能图省事用内置 ``hash()``。

    不是"可能不一样"，而是**确实不一样**——这就是那个坑本身。
    """
    script = "print(' '.join(str(hash(f'k{i}') % 64) for i in range(20)))"
    assert _run_with_hash_seed(script, "0") != _run_with_hash_seed(script, "1")


# ---------------------------------------------------------------------------
# 取值范围与分布
# ---------------------------------------------------------------------------


def test_bucket_is_always_within_range() -> None:
    for num_buckets in (1, 2, 7, 64, 1000):
        policy = BucketPolicy(num_buckets)
        for index in range(200):
            bucket = policy.bucket_for(f"key-{index}")
            assert 0 <= bucket < num_buckets


def test_single_bucket_policy_puts_everything_in_bucket_zero() -> None:
    policy = BucketPolicy(1)
    assert all(policy.bucket_for(key) == 0 for key in ("a", "b", "c"))


def test_distribution_is_reasonably_uniform() -> None:
    """分桶明显倾斜会让某些机器空转、另一些排队。

    这里只做一个粗检：64 个桶、4096 个键，每个桶的落点不应当偏离均值太多。
    """
    policy = BucketPolicy(64)
    counts = [0] * 64
    for index in range(4096):
        counts[policy.bucket_for(f"tenant-{index}")] += 1

    expected = 4096 / 64
    assert all(count > 0 for count in counts), "有桶完全没分到键，分布明显异常"
    assert max(counts) < expected * 3, f"分布过于集中：max={max(counts)}"


def test_similar_keys_do_not_cluster() -> None:
    """``tenant-1`` 与 ``tenant-2`` 这类只差一个字符的键必须被打散。"""
    policy = BucketPolicy(64)
    buckets = {policy.bucket_for(f"tenant-{index}") for index in range(50)}
    assert len(buckets) > 20, "相邻键挤在少数几个桶里，哈希的雪崩性不足"


def test_policy_validates_bucket_range() -> None:
    policy = BucketPolicy(8)
    assert policy.is_valid(0)
    assert policy.is_valid(7)
    assert not policy.is_valid(8)
    assert not policy.is_valid(-1)


def test_policy_exposes_all_buckets() -> None:
    assert list(BucketPolicy(4).buckets) == [0, 1, 2, 3]


# ---------------------------------------------------------------------------
# 参数校验
# ---------------------------------------------------------------------------


def test_num_buckets_must_be_positive() -> None:
    with pytest.raises(ValueError, match="num_buckets"):
        BucketPolicy(0)


def test_empty_key_is_rejected() -> None:
    """空键会让所有"没填租户"的任务挤进同一个桶。"""
    with pytest.raises(ValueError, match="不能为空"):
        BucketPolicy().bucket_for("")


def test_policy_is_hashable_and_comparable() -> None:
    """策略会作为配置传递，值语义让它能被比较与当作字典键。"""
    assert BucketPolicy(64) == BucketPolicy(64)
    assert BucketPolicy(64) != BucketPolicy(128)
    assert len({BucketPolicy(64), BucketPolicy(64)}) == 1