268 lines
9.3 KiB
Python
268 lines
9.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Router tests — registry semantics + health loop behavior.
|
|
|
|
Avoids real WebSocket proxying (the bridge requires the ``websockets``
|
|
package and two running uvicorn processes); test_server.py already
|
|
covers the end-to-end WS protocol against a direct backend.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from fastvideo.entrypoints.streaming.router.config import (
|
|
ReplicaEndpoint,
|
|
RouterConfig,
|
|
)
|
|
from fastvideo.entrypoints.streaming.router.registry import (
|
|
ReplicaRegistry,
|
|
ReplicaStatus,
|
|
run_health_check_loop,
|
|
)
|
|
|
|
|
|
def _registry(
|
|
*,
|
|
num_primary: int = 1,
|
|
num_secondary: int = 1,
|
|
) -> ReplicaRegistry:
|
|
replicas = [ReplicaEndpoint(
|
|
url=f"http://primary-{i}:8000",
|
|
primary=True,
|
|
) for i in range(num_primary)
|
|
] + [ReplicaEndpoint(
|
|
url=f"http://secondary-{i}:8000",
|
|
primary=False,
|
|
) for i in range(num_secondary)]
|
|
return ReplicaRegistry(replicas)
|
|
|
|
|
|
class TestReplicaRegistry:
|
|
|
|
def test_requires_replicas(self):
|
|
with pytest.raises(ValueError):
|
|
ReplicaRegistry([])
|
|
|
|
def test_select_none_when_all_unknown(self):
|
|
assert _registry().select() is None
|
|
|
|
def test_select_prefers_healthy_primary(self):
|
|
reg = _registry()
|
|
|
|
async def promote():
|
|
primary = reg.primaries()[0]
|
|
await reg.record_success(primary, recovery_threshold=1, latency_ms=1.0)
|
|
# Also mark secondary healthy — primary should still win.
|
|
sec = next(r for r in reg.all() if not r.primary)
|
|
await reg.record_success(sec, recovery_threshold=1, latency_ms=1.0)
|
|
return reg.select()
|
|
|
|
pick = asyncio.run(promote())
|
|
assert pick is not None
|
|
assert pick.primary
|
|
|
|
def test_falls_back_to_secondary_when_primary_unhealthy(self):
|
|
reg = _registry()
|
|
|
|
async def run():
|
|
primary = reg.primaries()[0]
|
|
sec = next(r for r in reg.all() if not r.primary)
|
|
# Fail primary past threshold, succeed secondary.
|
|
for _ in range(3):
|
|
await reg.record_failure(primary, failure_threshold=3, reason="mock")
|
|
await reg.record_success(sec, recovery_threshold=1, latency_ms=1.0)
|
|
return reg.select()
|
|
|
|
pick = asyncio.run(run())
|
|
assert pick is not None
|
|
assert not pick.primary
|
|
|
|
def test_failure_threshold_transitions_to_unhealthy(self):
|
|
reg = _registry()
|
|
primary = reg.primaries()[0]
|
|
|
|
async def run():
|
|
for _ in range(2):
|
|
await reg.record_failure(primary, failure_threshold=3, reason="x")
|
|
assert primary.health.status is not ReplicaStatus.UNHEALTHY
|
|
await reg.record_failure(primary, failure_threshold=3, reason="x")
|
|
return primary
|
|
|
|
result = asyncio.run(run())
|
|
assert result.health.status is ReplicaStatus.UNHEALTHY
|
|
assert result.health.consecutive_failures == 3
|
|
|
|
def test_recovery_threshold_returns_to_healthy(self):
|
|
reg = _registry()
|
|
primary = reg.primaries()[0]
|
|
|
|
async def run():
|
|
for _ in range(3):
|
|
await reg.record_failure(primary, failure_threshold=3, reason="x")
|
|
assert primary.health.status is ReplicaStatus.UNHEALTHY
|
|
for _ in range(2):
|
|
await reg.record_success(primary, recovery_threshold=2, latency_ms=5.0)
|
|
return primary
|
|
|
|
result = asyncio.run(run())
|
|
assert result.health.status is ReplicaStatus.HEALTHY
|
|
assert result.health.last_latency_ms == 5.0
|
|
|
|
def test_record_success_resets_failure_counter(self):
|
|
reg = _registry()
|
|
primary = reg.primaries()[0]
|
|
|
|
async def run():
|
|
await reg.record_failure(primary, failure_threshold=10, reason="x")
|
|
await reg.record_success(primary, recovery_threshold=1, latency_ms=1.0)
|
|
return primary
|
|
|
|
result = asyncio.run(run())
|
|
assert result.health.consecutive_failures == 0
|
|
|
|
|
|
class TestHealthCheckLoop:
|
|
|
|
def test_loop_transitions_replicas_on_probe_results(self):
|
|
config = RouterConfig(
|
|
replicas=[
|
|
ReplicaEndpoint(url="http://a", primary=True),
|
|
ReplicaEndpoint(url="http://b"),
|
|
],
|
|
health_check_interval_seconds=0.01,
|
|
failure_threshold=1,
|
|
recovery_threshold=1,
|
|
)
|
|
reg = ReplicaRegistry(config.replicas)
|
|
stop_event = asyncio.Event()
|
|
|
|
calls: list[str] = []
|
|
|
|
async def probe(url, *, timeout):
|
|
calls.append(url)
|
|
if "http://a" in url:
|
|
return 1.0, None
|
|
return 0.0, "mock failure"
|
|
|
|
async def run() -> None:
|
|
task = asyncio.create_task(
|
|
run_health_check_loop(
|
|
registry=reg,
|
|
config=config,
|
|
stop_event=stop_event,
|
|
http_get=probe,
|
|
))
|
|
await asyncio.sleep(0.05)
|
|
stop_event.set()
|
|
await task
|
|
|
|
asyncio.run(run())
|
|
a = reg.get("http://a")
|
|
b = reg.get("http://b")
|
|
assert a is not None
|
|
assert b is not None
|
|
assert a.health.status is ReplicaStatus.HEALTHY
|
|
assert b.health.status is ReplicaStatus.UNHEALTHY
|
|
assert any("/health" in c for c in calls)
|
|
|
|
|
|
class TestRouterApp:
|
|
|
|
def test_status_endpoint_lists_replicas(self):
|
|
from starlette.testclient import TestClient
|
|
|
|
from fastvideo.entrypoints.streaming.router.main import build_router_app
|
|
|
|
config = RouterConfig(
|
|
replicas=[
|
|
ReplicaEndpoint(url="http://a", primary=True),
|
|
ReplicaEndpoint(url="http://b"),
|
|
],
|
|
health_check_interval_seconds=60, # don't actually poll
|
|
)
|
|
reg = ReplicaRegistry(config.replicas)
|
|
app = build_router_app(config, registry=reg)
|
|
client = TestClient(app)
|
|
response = client.get("/status")
|
|
body = response.json()
|
|
urls = {r["url"] for r in body["replicas"]}
|
|
assert urls == {"http://a", "http://b"}
|
|
# Initial status is UNKNOWN.
|
|
assert all(r["status"] == "unknown" for r in body["replicas"])
|
|
|
|
def test_ws_rejects_when_no_healthy_replica(self):
|
|
from starlette.testclient import TestClient
|
|
|
|
from fastvideo.entrypoints.streaming.router.main import build_router_app
|
|
|
|
config = RouterConfig(
|
|
replicas=[ReplicaEndpoint(url="http://a", primary=True)],
|
|
health_check_interval_seconds=60,
|
|
)
|
|
reg = ReplicaRegistry(config.replicas)
|
|
app = build_router_app(config, registry=reg)
|
|
client = TestClient(app)
|
|
with client.websocket_connect("/v1/stream") as ws:
|
|
err = ws.receive_json()
|
|
assert err["type"] == "error"
|
|
assert err["code"] == "gpu_unavailable"
|
|
|
|
|
|
class TestUnknownToHealthyImmediate:
|
|
"""Initial probe must promote UNKNOWN -> HEALTHY without waiting for recovery_threshold."""
|
|
|
|
def test_first_success_promotes_unknown(self):
|
|
reg = _registry(num_primary=1, num_secondary=0)
|
|
primary = reg.primaries()[0]
|
|
assert primary.health.status is ReplicaStatus.UNKNOWN
|
|
|
|
async def run():
|
|
await reg.record_success(primary, recovery_threshold=10, latency_ms=1.0)
|
|
|
|
asyncio.run(run())
|
|
assert primary.health.status is ReplicaStatus.HEALTHY
|
|
|
|
def test_unhealthy_recovery_still_gated_by_threshold(self):
|
|
reg = _registry(num_primary=1, num_secondary=0)
|
|
primary = reg.primaries()[0]
|
|
|
|
async def run():
|
|
for _ in range(3):
|
|
await reg.record_failure(primary, failure_threshold=3, reason="x")
|
|
assert primary.health.status is ReplicaStatus.UNHEALTHY
|
|
await reg.record_success(primary, recovery_threshold=2, latency_ms=1.0)
|
|
assert primary.health.status is ReplicaStatus.UNHEALTHY # 1/2 successes
|
|
await reg.record_success(primary, recovery_threshold=2, latency_ms=1.0)
|
|
assert primary.health.status is ReplicaStatus.HEALTHY # 2/2 successes
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
class TestConfigValidation:
|
|
"""RouterConfig.__post_init__ rejects malformed configs."""
|
|
|
|
def test_rejects_path_in_url(self):
|
|
with pytest.raises(ValueError, match="without a path"):
|
|
RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000/api")])
|
|
|
|
def test_rejects_query_in_url(self):
|
|
with pytest.raises(ValueError, match="query/fragment"):
|
|
RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000?x=1")])
|
|
|
|
def test_rejects_fragment_in_url(self):
|
|
with pytest.raises(ValueError, match="query/fragment"):
|
|
RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000#frag")])
|
|
|
|
def test_rejects_duplicate_urls(self):
|
|
with pytest.raises(ValueError, match="Duplicate"):
|
|
RouterConfig(replicas=[
|
|
ReplicaEndpoint(url="http://host:8000"),
|
|
ReplicaEndpoint(url="http://host:8000"),
|
|
])
|
|
|
|
def test_accepts_trailing_slash(self):
|
|
# parsed.path == "/" should be allowed
|
|
cfg = RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000/")])
|
|
assert cfg.replicas[0].url == "http://host:8000/"
|