forked from linkong/planet
558 lines
17 KiB
Python
558 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi import FastAPI, HTTPException, Response
|
|
|
|
from app.api.v1 import visualization
|
|
from app.services.earth_layer_cache import (
|
|
EarthLayerCachePolicy,
|
|
apply_payload_budget,
|
|
earth_layer_cache,
|
|
format_bbox_key,
|
|
quantize_bbox,
|
|
resolve_layer_payload,
|
|
)
|
|
|
|
|
|
class FakeRedis:
|
|
def __init__(self, *, fail: bool = False) -> None:
|
|
self.store: dict[str, str] = {}
|
|
self.fail = fail
|
|
self.lock_claimed = False
|
|
|
|
def _maybe_fail(self) -> None:
|
|
if self.fail:
|
|
raise RuntimeError("redis unavailable")
|
|
|
|
def get(self, key: str):
|
|
self._maybe_fail()
|
|
return self.store.get(key)
|
|
|
|
def set(self, key: str, value: str, nx: bool = False, ex: int | None = None):
|
|
self._maybe_fail()
|
|
if nx and key in self.store:
|
|
return False
|
|
self.store[key] = value
|
|
return True
|
|
|
|
def setex(self, key: str, _seconds: int, value: str):
|
|
self._maybe_fail()
|
|
self.store[key] = value
|
|
return True
|
|
|
|
def delete(self, *keys: str):
|
|
self._maybe_fail()
|
|
deleted = 0
|
|
for key in keys:
|
|
deleted += 1 if self.store.pop(key, None) is not None else 0
|
|
return deleted
|
|
|
|
def eval(self, script: str, numkeys: int, *args):
|
|
self._maybe_fail()
|
|
keys, values = args[:numkeys], args[numkeys:]
|
|
if self.store.get(keys[0]) != values[0]:
|
|
return 0
|
|
if numkeys == 3:
|
|
self.setex(keys[1], values[1], values[3])
|
|
self.setex(keys[2], values[2], values[3])
|
|
return 1
|
|
if "'expire'" in script:
|
|
return 1
|
|
return self.delete(keys[0])
|
|
|
|
def scan_iter(self, match: str):
|
|
self._maybe_fail()
|
|
prefix = match.rstrip("*")
|
|
for key in list(self.store):
|
|
if key.startswith(prefix):
|
|
yield key
|
|
|
|
def memory_usage(self, key: str):
|
|
value = self.store.get(key, "")
|
|
return len(value.encode("utf-8"))
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def fake_cache_client():
|
|
previous = earth_layer_cache._client
|
|
fake = FakeRedis()
|
|
earth_layer_cache._client = fake
|
|
try:
|
|
yield fake
|
|
finally:
|
|
earth_layer_cache._client = previous
|
|
|
|
|
|
def test_quantized_bbox_key_is_stable_for_small_movements():
|
|
first = format_bbox_key(quantize_bbox((10.01, 59.04, 10.96, 60.02)))
|
|
second = format_bbox_key(quantize_bbox((10.04, 59.01, 10.99, 60.04)))
|
|
|
|
assert first == second
|
|
assert first == "10.0,59.0,11.0,60.0"
|
|
|
|
|
|
def test_payload_budget_truncates_features():
|
|
payload = {
|
|
"type": "FeatureCollection",
|
|
"features": [{"id": index} for index in range(5)],
|
|
}
|
|
policy = EarthLayerCachePolicy(60, 120, max_features=2, max_bytes=1024)
|
|
|
|
result = apply_payload_budget(payload, policy)
|
|
|
|
assert len(result["features"]) == 2
|
|
assert result["diagnostics"]["truncated"] is True
|
|
assert result["diagnostics"]["limit_reason"] == "feature_budget"
|
|
assert result["diagnostics"]["original_feature_count"] == 5
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_layer_payload_writes_fresh_and_stale(fake_cache_client):
|
|
calls = 0
|
|
|
|
async def builder():
|
|
nonlocal calls
|
|
calls += 1
|
|
return {"type": "FeatureCollection", "features": [{"id": "a"}]}
|
|
|
|
key = earth_layer_cache.key("satellites", limit="all")
|
|
policy = EarthLayerCachePolicy(60, 120)
|
|
|
|
first = await resolve_layer_payload(key=key, policy=policy, builder=builder)
|
|
second = await resolve_layer_payload(key=key, policy=policy, builder=builder)
|
|
|
|
assert first.state == "refresh"
|
|
assert second.state == "hit"
|
|
assert calls == 1
|
|
assert key in fake_cache_client.store
|
|
assert f"{key}:stale" in fake_cache_client.store
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_layer_payload_returns_stale_when_builder_fails(fake_cache_client):
|
|
key = earth_layer_cache.key("bgp-incidents", status="active")
|
|
fake_cache_client.store[f"{key}:stale"] = json.dumps(
|
|
{"type": "FeatureCollection", "features": []}
|
|
)
|
|
|
|
async def builder():
|
|
raise RuntimeError("db exploded")
|
|
|
|
result = await resolve_layer_payload(
|
|
key=key,
|
|
policy=EarthLayerCachePolicy(60, 120),
|
|
builder=builder,
|
|
)
|
|
|
|
assert result.state == "stale"
|
|
assert result.payload["features"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_layer_payload_uses_stale_during_lock_contention(fake_cache_client):
|
|
key = earth_layer_cache.key("cables")
|
|
fake_cache_client.store[earth_layer_cache.lock_key(key)] = "1"
|
|
fake_cache_client.store[f"{key}:stale"] = json.dumps(
|
|
{"type": "FeatureCollection", "features": [{"id": "stale-cable"}]}
|
|
)
|
|
calls = 0
|
|
|
|
async def builder():
|
|
nonlocal calls
|
|
calls += 1
|
|
return {"type": "FeatureCollection", "features": [{"id": "fresh-cable"}]}
|
|
|
|
result = await resolve_layer_payload(
|
|
key=key,
|
|
policy=EarthLayerCachePolicy(60, 120),
|
|
builder=builder,
|
|
)
|
|
|
|
assert result.state == "stale"
|
|
assert result.payload["features"][0]["id"] == "stale-cable"
|
|
assert calls == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_layer_payload_bypasses_redis_failure():
|
|
previous = earth_layer_cache._client
|
|
earth_layer_cache._client = FakeRedis(fail=True)
|
|
try:
|
|
|
|
async def builder():
|
|
return {"type": "FeatureCollection", "features": [{"id": "safe"}]}
|
|
|
|
result = await resolve_layer_payload(
|
|
key=earth_layer_cache.key("cables"),
|
|
policy=EarthLayerCachePolicy(60, 120),
|
|
builder=builder,
|
|
)
|
|
|
|
assert result.state == "bypass"
|
|
assert result.payload["features"][0]["id"] == "safe"
|
|
finally:
|
|
earth_layer_cache._client = previous
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_visualization_endpoint_sets_cache_headers(fake_cache_client, monkeypatch):
|
|
calls = 0
|
|
|
|
async def fake_build_satellites_geojson(*, limit, db):
|
|
nonlocal calls
|
|
calls += 1
|
|
return {"type": "FeatureCollection", "features": [{"id": f"sat-{limit}"}], "count": 1}
|
|
|
|
monkeypatch.setattr(visualization, "_build_satellites_geojson", fake_build_satellites_geojson)
|
|
|
|
first_response = Response()
|
|
first = await visualization.get_satellites_geojson(
|
|
limit=25, db=object(), response=first_response
|
|
)
|
|
second_response = Response()
|
|
second = await visualization.get_satellites_geojson(
|
|
limit=25, db=object(), response=second_response
|
|
)
|
|
|
|
assert first == second
|
|
assert calls == 1
|
|
assert first_response.headers["X-Planet-Cache"] == "refresh"
|
|
assert second_response.headers["X-Planet-Cache"] == "hit"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vessel_snapshot_uses_short_cache(fake_cache_client, monkeypatch):
|
|
calls = 0
|
|
|
|
async def fake_load_raw_vessel_snapshot_features(db, *, bbox, limit, observed_since):
|
|
nonlocal calls
|
|
calls += 1
|
|
return (
|
|
[
|
|
{
|
|
"type": "Feature",
|
|
"geometry": {"type": "Point", "coordinates": [10.1, 59.1]},
|
|
"properties": {"mmsi": 123, "vessel_type_name": "Cargo"},
|
|
}
|
|
],
|
|
{"raw_feature_count": 1},
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
visualization,
|
|
"_load_raw_vessel_snapshot_features",
|
|
fake_load_raw_vessel_snapshot_features,
|
|
)
|
|
|
|
first_response = Response()
|
|
first = await visualization.build_vessel_snapshot_response(
|
|
object(),
|
|
bbox=(10.01, 59.04, 10.96, 60.02),
|
|
zoom=12,
|
|
type_filter=None,
|
|
limit=1000,
|
|
since_minutes=60,
|
|
response=first_response,
|
|
)
|
|
second_response = Response()
|
|
second = await visualization.build_vessel_snapshot_response(
|
|
object(),
|
|
bbox=(10.04, 59.01, 10.99, 60.04),
|
|
zoom=12,
|
|
type_filter=None,
|
|
limit=1000,
|
|
since_minutes=60,
|
|
response=second_response,
|
|
)
|
|
|
|
assert first["count"] == 1
|
|
assert second == first
|
|
assert calls == 1
|
|
assert first_response.headers["X-Planet-Cache"] == "refresh"
|
|
assert second_response.headers["X-Planet-Cache"] == "hit"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_summary_requests_share_one_slow_build(monkeypatch):
|
|
started = asyncio.Event()
|
|
finish = asyncio.Event()
|
|
calls = 0
|
|
|
|
async def build_summary(db):
|
|
nonlocal calls
|
|
calls += 1
|
|
started.set()
|
|
await finish.wait()
|
|
return {"stats": {"vessel_count": 42}}
|
|
|
|
monkeypatch.setattr(visualization, "_build_visualization_geo_summary", build_summary)
|
|
app = FastAPI()
|
|
app.include_router(visualization.router, prefix="/api/v1/visualization")
|
|
app.dependency_overrides[visualization.get_db] = lambda: object()
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app=app), base_url="http://test"
|
|
) as client:
|
|
requests = [
|
|
asyncio.create_task(client.get("/api/v1/visualization/geo/summary")) for _ in range(6)
|
|
]
|
|
await asyncio.wait_for(started.wait(), 1)
|
|
# Keep the build open beyond the old 200 ms fallback to duplicate SQL.
|
|
await asyncio.sleep(0.25)
|
|
assert calls == 1
|
|
finish.set()
|
|
results = await asyncio.gather(*requests)
|
|
assert all(response.status_code == 200 for response in results)
|
|
assert all(response.json() == {"stats": {"vessel_count": 42}} for response in results)
|
|
assert sorted(response.headers["X-Planet-Cache"] for response in results) == [
|
|
"hit",
|
|
"hit",
|
|
"hit",
|
|
"hit",
|
|
"hit",
|
|
"refresh",
|
|
]
|
|
assert not earth_layer_cache._inflight
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("error_type", [RuntimeError, TimeoutError])
|
|
async def test_shared_builder_error_is_not_retried(error_type):
|
|
started = asyncio.Event()
|
|
finish = asyncio.Event()
|
|
calls = 0
|
|
|
|
async def builder():
|
|
nonlocal calls
|
|
calls += 1
|
|
started.set()
|
|
await finish.wait()
|
|
raise error_type("db failed")
|
|
|
|
requests = [
|
|
asyncio.create_task(
|
|
resolve_layer_payload(
|
|
key="shared-error", policy=EarthLayerCachePolicy(60, 120), builder=builder
|
|
)
|
|
)
|
|
for _ in range(4)
|
|
]
|
|
await started.wait()
|
|
await asyncio.sleep(0)
|
|
finish.set()
|
|
results = await asyncio.gather(*requests, return_exceptions=True)
|
|
assert calls == 1
|
|
assert all(isinstance(result, error_type) for result in results)
|
|
assert not earth_layer_cache._inflight
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_redis_failure_shares_bypass_result(fake_cache_client):
|
|
fake_cache_client.fail = True
|
|
started = asyncio.Event()
|
|
finish = asyncio.Event()
|
|
calls = 0
|
|
|
|
async def builder():
|
|
nonlocal calls
|
|
calls += 1
|
|
started.set()
|
|
await finish.wait()
|
|
return {"count": 42}
|
|
|
|
requests = [
|
|
asyncio.create_task(
|
|
resolve_layer_payload(
|
|
key="shared-bypass", policy=EarthLayerCachePolicy(60, 120), builder=builder
|
|
)
|
|
)
|
|
for _ in range(4)
|
|
]
|
|
await started.wait()
|
|
await asyncio.sleep(0)
|
|
finish.set()
|
|
results = await asyncio.gather(*requests)
|
|
assert calls == 1
|
|
assert all(result.state == "bypass" and result.payload == {"count": 42} for result in results)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_waiter_does_not_cancel_owner():
|
|
started = asyncio.Event()
|
|
finish = asyncio.Event()
|
|
policy = EarthLayerCachePolicy(60, 120)
|
|
|
|
async def builder():
|
|
started.set()
|
|
await finish.wait()
|
|
return {"count": 1}
|
|
|
|
owner = asyncio.create_task(
|
|
resolve_layer_payload(key="cancel-waiter", policy=policy, builder=builder)
|
|
)
|
|
await started.wait()
|
|
waiter = asyncio.create_task(
|
|
resolve_layer_payload(key="cancel-waiter", policy=policy, builder=builder)
|
|
)
|
|
await asyncio.sleep(0)
|
|
waiter.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiter
|
|
finish.set()
|
|
assert (await owner).payload == {"count": 1}
|
|
assert not earth_layer_cache._inflight
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_owner_allows_waiter_to_use_its_own_builder():
|
|
started = asyncio.Event()
|
|
policy = EarthLayerCachePolicy(60, 120)
|
|
|
|
async def owner_builder():
|
|
started.set()
|
|
await asyncio.Event().wait()
|
|
|
|
async def waiter_builder():
|
|
return {"count": 2}
|
|
|
|
owner = asyncio.create_task(
|
|
resolve_layer_payload(key="cancel-owner", policy=policy, builder=owner_builder)
|
|
)
|
|
await started.wait()
|
|
waiter = asyncio.create_task(
|
|
resolve_layer_payload(key="cancel-owner", policy=policy, builder=waiter_builder)
|
|
)
|
|
await asyncio.sleep(0)
|
|
owner.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await owner
|
|
assert (await asyncio.wait_for(waiter, 1)).payload == {"count": 2}
|
|
assert not earth_layer_cache._inflight
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cross_process_lock_waits_for_payload_instead_of_building(fake_cache_client):
|
|
key = "remote-owner"
|
|
fake_cache_client.store[earth_layer_cache.lock_key(key)] = "another-owner"
|
|
policy = EarthLayerCachePolicy(60, 120, lock_wait_seconds=0.005)
|
|
calls = 0
|
|
|
|
async def builder():
|
|
nonlocal calls
|
|
calls += 1
|
|
return {"count": -1}
|
|
|
|
waiter = asyncio.create_task(resolve_layer_payload(key=key, policy=policy, builder=builder))
|
|
await asyncio.sleep(0.03)
|
|
assert calls == 0
|
|
fake_cache_client.store[key] = json.dumps({"count": 3})
|
|
assert (await asyncio.wait_for(waiter, 1)).payload == {"count": 3}
|
|
assert calls == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wait_timeout_does_not_bypass_other_owner(fake_cache_client):
|
|
key = "remote-timeout"
|
|
fake_cache_client.store[earth_layer_cache.lock_key(key)] = "another-owner"
|
|
calls = 0
|
|
|
|
async def builder():
|
|
nonlocal calls
|
|
calls += 1
|
|
return {}
|
|
|
|
with pytest.raises(HTTPException) as error:
|
|
await resolve_layer_payload(
|
|
key=key,
|
|
policy=EarthLayerCachePolicy(60, 120, lock_wait_seconds=0.005, build_wait_seconds=0.02),
|
|
builder=builder,
|
|
)
|
|
assert error.value.status_code == 503
|
|
assert calls == 0
|
|
assert not earth_layer_cache._inflight
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lost_owner_cannot_publish_or_release_new_lock(fake_cache_client):
|
|
key = "lock-replaced"
|
|
policy = EarthLayerCachePolicy(60, 120)
|
|
|
|
async def builder():
|
|
fake_cache_client.store[earth_layer_cache.lock_key(key)] = "new-owner"
|
|
fake_cache_client.store[key] = json.dumps({"count": 99})
|
|
return {"count": 1}
|
|
|
|
result = await resolve_layer_payload(key=key, policy=policy, builder=builder)
|
|
assert result.payload == {"count": 1}
|
|
assert result.state == "bypass"
|
|
assert fake_cache_client.store[earth_layer_cache.lock_key(key)] == "new-owner"
|
|
assert json.loads(fake_cache_client.store[key]) == {"count": 99}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_publish_failure_returns_built_payload_without_retry(fake_cache_client):
|
|
calls = 0
|
|
|
|
async def builder():
|
|
nonlocal calls
|
|
calls += 1
|
|
fake_cache_client.fail = True
|
|
return {"count": 7}
|
|
|
|
result = await resolve_layer_payload(
|
|
key="publish-error", policy=EarthLayerCachePolicy(60, 120), builder=builder
|
|
)
|
|
assert result.state == "bypass"
|
|
assert result.payload == {"count": 7}
|
|
assert calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_wait_timeout_leaves_owner_running():
|
|
started = asyncio.Event()
|
|
finish = asyncio.Event()
|
|
policy = EarthLayerCachePolicy(60, 120, build_wait_seconds=0.02)
|
|
|
|
async def builder():
|
|
started.set()
|
|
await finish.wait()
|
|
return {"count": 1}
|
|
|
|
owner = asyncio.create_task(
|
|
resolve_layer_payload(key="local-timeout", policy=policy, builder=builder)
|
|
)
|
|
await started.wait()
|
|
with pytest.raises(HTTPException) as error:
|
|
await resolve_layer_payload(key="local-timeout", policy=policy, builder=builder)
|
|
assert error.value.status_code == 503
|
|
assert not owner.done()
|
|
finish.set()
|
|
assert (await owner).payload == {"count": 1}
|
|
assert not earth_layer_cache._inflight
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_different_keys_can_build_concurrently():
|
|
started = set()
|
|
both_started = asyncio.Event()
|
|
|
|
async def build(key):
|
|
started.add(key)
|
|
if len(started) == 2:
|
|
both_started.set()
|
|
await asyncio.wait_for(both_started.wait(), 1)
|
|
return {"key": key}
|
|
|
|
results = await asyncio.gather(
|
|
*(
|
|
resolve_layer_payload(
|
|
key=key, policy=EarthLayerCachePolicy(60, 120), builder=lambda key=key: build(key)
|
|
)
|
|
for key in ("one", "two")
|
|
)
|
|
)
|
|
assert {result.payload["key"] for result in results} == {"one", "two"}
|