mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
fix(scheduling): 负载均衡模式与无亲和键场景统一走随机排序
重构 CandidateSorter.shuffle_keys_by_internal_priority 中同优先级 Key 的排序逻辑,将三分支简化为两分支: - 随机排序:TTL=0 / 负载均衡模式 / 无 affinity_key - 哈希确定性排序:缓存亲和模式且有 affinity_key 移除了原先"无 affinity_key 时按 ID 排序"的冗余分支,新增 对应单元测试覆盖三种场景。
This commit is contained in:
93
tests/unit/test_candidate_sorter_key_shuffle.py
Normal file
93
tests/unit/test_candidate_sorter_key_shuffle.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import cast
|
||||
from unittest.mock import patch
|
||||
|
||||
from src.models.database import ProviderAPIKey
|
||||
from src.services.scheduling.candidate_sorter import CandidateSorter
|
||||
from src.services.scheduling.scheduling_config import SchedulingConfig
|
||||
from src.services.scheduling.utils import affinity_hash
|
||||
|
||||
|
||||
def _make_key(key_id: str, internal_priority: int = 1) -> ProviderAPIKey:
|
||||
return cast(
|
||||
ProviderAPIKey,
|
||||
SimpleNamespace(
|
||||
id=key_id,
|
||||
internal_priority=internal_priority,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _reverse_in_place(items: list[object]) -> None:
|
||||
items.reverse()
|
||||
|
||||
|
||||
def test_shuffle_keys_random_in_load_balance_mode_even_with_affinity_key() -> None:
|
||||
config = SchedulingConfig(
|
||||
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
|
||||
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_LOAD_BALANCE,
|
||||
)
|
||||
sorter = CandidateSorter(config)
|
||||
|
||||
keys = [_make_key("k1"), _make_key("k2"), _make_key("k3")]
|
||||
|
||||
with patch(
|
||||
"src.services.scheduling.candidate_sorter.random.shuffle",
|
||||
side_effect=_reverse_in_place,
|
||||
) as shuffle_mock:
|
||||
result = sorter.shuffle_keys_by_internal_priority(
|
||||
keys,
|
||||
affinity_key="affinity-1",
|
||||
use_random=False,
|
||||
)
|
||||
|
||||
assert [k.id for k in result] == ["k3", "k2", "k1"]
|
||||
assert shuffle_mock.call_count == 1
|
||||
|
||||
|
||||
def test_shuffle_keys_random_when_affinity_key_absent() -> None:
|
||||
config = SchedulingConfig(
|
||||
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
|
||||
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_CACHE_AFFINITY,
|
||||
)
|
||||
sorter = CandidateSorter(config)
|
||||
|
||||
keys = [_make_key("k1"), _make_key("k2"), _make_key("k3")]
|
||||
|
||||
with patch(
|
||||
"src.services.scheduling.candidate_sorter.random.shuffle",
|
||||
side_effect=_reverse_in_place,
|
||||
) as shuffle_mock:
|
||||
result = sorter.shuffle_keys_by_internal_priority(
|
||||
keys,
|
||||
affinity_key=None,
|
||||
use_random=False,
|
||||
)
|
||||
|
||||
assert [k.id for k in result] == ["k3", "k2", "k1"]
|
||||
assert shuffle_mock.call_count == 1
|
||||
|
||||
|
||||
def test_shuffle_keys_still_hashes_with_affinity_in_non_load_balance_mode() -> None:
|
||||
config = SchedulingConfig(
|
||||
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
|
||||
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_CACHE_AFFINITY,
|
||||
)
|
||||
sorter = CandidateSorter(config)
|
||||
|
||||
keys = [_make_key("k1"), _make_key("k2"), _make_key("k3")]
|
||||
affinity_key = "affinity-1"
|
||||
|
||||
expected = sorted(keys, key=lambda k: affinity_hash(affinity_key, k.id))
|
||||
|
||||
with patch("src.services.scheduling.candidate_sorter.random.shuffle") as shuffle_mock:
|
||||
result = sorter.shuffle_keys_by_internal_priority(
|
||||
keys,
|
||||
affinity_key=affinity_key,
|
||||
use_random=False,
|
||||
)
|
||||
|
||||
assert [k.id for k in result] == [k.id for k in expected]
|
||||
shuffle_mock.assert_not_called()
|
||||
Reference in New Issue
Block a user