94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
from src.services.ollama_endpoint_circuit_breaker import (
|
|
filter_ollama_urls_with_cooldown,
|
|
get_ollama_endpoint_cooldown_remaining_seconds,
|
|
is_ollama_endpoint_blocked,
|
|
record_ollama_endpoint_failure,
|
|
record_ollama_endpoint_success,
|
|
reset_ollama_endpoint_cooldown_for_tests,
|
|
resolve_ollama_order_with_cooldown,
|
|
)
|
|
from src.services.ollama_endpoint_resolver import resolve_ollama_order
|
|
|
|
|
|
def setup_function() -> None:
|
|
reset_ollama_endpoint_cooldown_for_tests()
|
|
|
|
|
|
def teardown_function() -> None:
|
|
reset_ollama_endpoint_cooldown_for_tests()
|
|
|
|
|
|
def test_cooldown_filters_failed_urls_without_reordering() -> None:
|
|
urls = ("http://gcp-a:11434", "http://gcp-b:11434", "http://local-111:11434")
|
|
|
|
record_ollama_endpoint_failure("http://gcp-a:11434", now=100.0)
|
|
record_ollama_endpoint_failure("http://gcp-b:11434", now=100.0)
|
|
|
|
assert filter_ollama_urls_with_cooldown(urls, now=101.0) == ("http://local-111:11434",)
|
|
|
|
|
|
def test_cooldown_expires_and_success_clears_block() -> None:
|
|
record_ollama_endpoint_failure("http://gcp-a:11434", cooldown_seconds=10.0, now=100.0)
|
|
|
|
assert is_ollama_endpoint_blocked("http://gcp-a:11434", now=109.9)
|
|
assert not is_ollama_endpoint_blocked("http://gcp-a:11434", now=110.1)
|
|
|
|
record_ollama_endpoint_failure("http://gcp-a:11434", now=200.0)
|
|
record_ollama_endpoint_success("http://gcp-a:11434")
|
|
assert not is_ollama_endpoint_blocked("http://gcp-a:11434", now=201.0)
|
|
|
|
|
|
def test_cooldown_remaining_seconds_expires_cleanly() -> None:
|
|
record_ollama_endpoint_failure("http://gcp-a:11434", cooldown_seconds=10.0, now=100.0)
|
|
|
|
assert (
|
|
get_ollama_endpoint_cooldown_remaining_seconds(
|
|
"http://gcp-a:11434",
|
|
now=103.0,
|
|
)
|
|
== 7.0
|
|
)
|
|
assert (
|
|
get_ollama_endpoint_cooldown_remaining_seconds(
|
|
"http://gcp-a:11434",
|
|
now=111.0,
|
|
)
|
|
== 0.0
|
|
)
|
|
assert not is_ollama_endpoint_blocked("http://gcp-a:11434", now=112.0)
|
|
|
|
|
|
def test_all_blocked_returns_full_order_for_recovery_probe() -> None:
|
|
urls = ("http://gcp-a:11434", "http://gcp-b:11434")
|
|
|
|
record_ollama_endpoint_failure("http://gcp-a:11434", now=100.0)
|
|
record_ollama_endpoint_failure("http://gcp-b:11434", now=100.0)
|
|
|
|
assert filter_ollama_urls_with_cooldown(urls, now=101.0) == urls
|
|
|
|
|
|
def test_resolver_order_itself_remains_global_policy(monkeypatch) -> None:
|
|
cfg = SimpleNamespace(
|
|
OLLAMA_URL="http://gcp-a:11434",
|
|
OLLAMA_SECONDARY_URL="http://gcp-b:11434",
|
|
OLLAMA_FALLBACK_URL="http://local-111:11434",
|
|
)
|
|
import src.services.ollama_endpoint_resolver as resolver
|
|
|
|
monkeypatch.setattr(resolver, "settings", cfg)
|
|
|
|
record_ollama_endpoint_failure("http://gcp-a:11434", now=100.0)
|
|
cooled = resolve_ollama_order_with_cooldown("embedding", now=101.0)
|
|
policy = resolve_ollama_order("embedding", config=cfg)
|
|
|
|
assert [endpoint.provider_name for endpoint in cooled] == ["ollama_gcp_b", "ollama_local"]
|
|
assert [endpoint.provider_name for endpoint in policy] == [
|
|
"ollama_gcp_a",
|
|
"ollama_gcp_b",
|
|
"ollama_local",
|
|
]
|