112 lines
4.5 KiB
Python
112 lines
4.5 KiB
Python
"""熔断器测试."""
|
|
|
|
from unittest.mock import patch
|
|
|
|
from src.ai.providers.circuit_breaker import CircuitBreaker, CircuitState
|
|
|
|
|
|
class TestCircuitBreaker:
|
|
"""CircuitBreaker 状态机测试."""
|
|
|
|
def test_initial_state_closed(self) -> None:
|
|
"""初始状态为 CLOSED."""
|
|
cb = CircuitBreaker()
|
|
assert cb.get_state("openai") == CircuitState.CLOSED
|
|
|
|
def test_record_success_resets_failures(self) -> None:
|
|
"""成功重置失败计数."""
|
|
cb = CircuitBreaker()
|
|
cb.record_failure("openai")
|
|
cb.record_failure("openai")
|
|
cb.record_success("openai")
|
|
assert cb.get_state("openai") == CircuitState.CLOSED
|
|
status = cb.status()
|
|
assert status["openai"]["failures"] == 0
|
|
|
|
def test_threshold_opens_circuit(self) -> None:
|
|
"""达阈值触发 OPEN."""
|
|
cb = CircuitBreaker(failure_threshold=3)
|
|
cb.record_failure("p1")
|
|
cb.record_failure("p1")
|
|
assert cb.get_state("p1") == CircuitState.CLOSED
|
|
cb.record_failure("p1")
|
|
assert cb.get_state("p1") == CircuitState.OPEN
|
|
|
|
def test_open_blocks_calls(self) -> None:
|
|
"""OPEN 状态禁止调用."""
|
|
cb = CircuitBreaker(failure_threshold=1)
|
|
cb.record_failure("p1")
|
|
assert cb.is_closed("p1") is False
|
|
|
|
def test_closed_allows_calls(self) -> None:
|
|
"""CLOSED 状态允许调用."""
|
|
cb = CircuitBreaker()
|
|
assert cb.is_closed("p1") is True
|
|
|
|
def test_half_open_after_cooldown(self) -> None:
|
|
"""冷却后进入 HALF_OPEN."""
|
|
cb = CircuitBreaker(failure_threshold=1, cooldown_seconds=60.0)
|
|
with patch("src.ai.providers.circuit_breaker.time.monotonic", return_value=1000.0):
|
|
cb.record_failure("p1")
|
|
assert cb.get_state("p1") == CircuitState.OPEN
|
|
# 时间推进超过冷却期
|
|
with patch("src.ai.providers.circuit_breaker.time.monotonic", return_value=1061.0):
|
|
assert cb.get_state("p1") == CircuitState.HALF_OPEN
|
|
|
|
def test_half_open_allows_one_call(self) -> None:
|
|
"""HALF_OPEN 允许 1 次试探."""
|
|
cb = CircuitBreaker(failure_threshold=1, cooldown_seconds=60.0, half_open_max_calls=1)
|
|
with patch("src.ai.providers.circuit_breaker.time.monotonic", return_value=1000.0):
|
|
cb.record_failure("p1")
|
|
# 推进时间触发 HALF_OPEN
|
|
with patch("src.ai.providers.circuit_breaker.time.monotonic", return_value=1061.0):
|
|
cb.get_state("p1") # HALF_OPEN
|
|
assert cb.is_closed("p1") is True # 第 1 次允许
|
|
assert cb.is_closed("p1") is False # 第 2 次拒绝
|
|
|
|
def test_half_open_success_closes(self) -> None:
|
|
"""HALF_OPEN 成功 → CLOSED."""
|
|
cb = CircuitBreaker(failure_threshold=1, cooldown_seconds=60.0)
|
|
with patch("src.ai.providers.circuit_breaker.time.monotonic", return_value=1000.0):
|
|
cb.record_failure("p1")
|
|
with patch("src.ai.providers.circuit_breaker.time.monotonic", return_value=1061.0):
|
|
cb.get_state("p1") # HALF_OPEN
|
|
cb.record_success("p1")
|
|
assert cb.get_state("p1") == CircuitState.CLOSED
|
|
|
|
def test_half_open_failure_reopens(self) -> None:
|
|
"""HALF_OPEN 失败 → 重新 OPEN."""
|
|
cb = CircuitBreaker(failure_threshold=1, cooldown_seconds=60.0)
|
|
with patch("src.ai.providers.circuit_breaker.time.monotonic", return_value=1000.0):
|
|
cb.record_failure("p1")
|
|
with patch("src.ai.providers.circuit_breaker.time.monotonic", return_value=1061.0):
|
|
cb.get_state("p1") # HALF_OPEN
|
|
cb.record_failure("p1")
|
|
assert cb.get_state("p1") == CircuitState.OPEN
|
|
|
|
def test_reset_clears_state(self) -> None:
|
|
"""reset 清除 Provider 状态."""
|
|
cb = CircuitBreaker(failure_threshold=1)
|
|
cb.record_failure("p1")
|
|
cb.reset("p1")
|
|
assert cb.get_state("p1") == CircuitState.CLOSED
|
|
assert "p1" not in cb.status()
|
|
|
|
def test_status_snapshot(self) -> None:
|
|
"""status() 返回快照."""
|
|
cb = CircuitBreaker(failure_threshold=3)
|
|
cb.record_failure("p1")
|
|
cb.record_failure("p1")
|
|
cb.record_failure("p2")
|
|
status = cb.status()
|
|
assert "p1" in status
|
|
assert status["p1"]["failures"] == 2
|
|
assert "p2" in status
|
|
|
|
def test_different_providers_independent(self) -> None:
|
|
"""不同 Provider 状态独立."""
|
|
cb = CircuitBreaker(failure_threshold=1)
|
|
cb.record_failure("p1")
|
|
assert cb.get_state("p1") == CircuitState.OPEN
|
|
assert cb.get_state("p2") == CircuitState.CLOSED
|