194 lines
6.1 KiB
Python
194 lines
6.1 KiB
Python
"""The admission gate a rename closes (section 2.1 of
|
|
docs/private/id_refac_rename.md)."""
|
|
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from starlette.applications import Starlette
|
|
from starlette.requests import Request
|
|
from starlette.responses import PlainTextResponse
|
|
from starlette.routing import Route, WebSocketRoute
|
|
from starlette.websockets import WebSocket, WebSocketDisconnect
|
|
|
|
from fluksio.core.config import settings
|
|
from fluksio.flow.rename import RenameGate
|
|
from fluksio.main import AdmissionGate, app
|
|
from fluksio.mcp import server
|
|
|
|
API = settings.API_V1_STR
|
|
|
|
|
|
def _gated(
|
|
handler_gate: threading.Event | None = None,
|
|
) -> tuple[TestClient, RenameGate]:
|
|
"""A bare app behind a gate of its own. A mutation waits for
|
|
*handler_gate*, when given, so it can be caught while admitted."""
|
|
|
|
async def answer(request: Request) -> PlainTextResponse:
|
|
if handler_gate is not None and request.method == "POST":
|
|
await asyncio.to_thread(handler_gate.wait)
|
|
return PlainTextResponse("ok")
|
|
|
|
async def socket(websocket: WebSocket) -> None:
|
|
await websocket.accept()
|
|
await websocket.send_text("ok")
|
|
await websocket.close()
|
|
|
|
methods = ["GET", "POST", "PUT", "DELETE"]
|
|
inner = Starlette(
|
|
routes=[
|
|
Route("/{path:path}", answer, methods=methods),
|
|
WebSocketRoute("/hooks/{flow}/ws", socket),
|
|
]
|
|
)
|
|
gate = RenameGate()
|
|
return TestClient(AdmissionGate(inner, gate)), gate
|
|
|
|
|
|
def _close(gate: RenameGate, affected: list[str] | None = None) -> None:
|
|
assert gate.close("op-1", affected or [], threading.Event())
|
|
|
|
|
|
def test_a_mutation_is_turned_away_while_closed_and_reads_pass():
|
|
client, gate = _gated()
|
|
_close(gate)
|
|
|
|
refused = client.post(f"{API}/flows/f_x/publish")
|
|
assert refused.status_code == 503
|
|
assert refused.headers["retry-after"] == "5"
|
|
assert refused.json()["code"] == "renaming"
|
|
assert refused.json()["op"] == "op-1"
|
|
for method in ("put", "delete"):
|
|
assert getattr(client, method)(f"{API}/dashboards/d_x").status_code == 503
|
|
assert client.post(f"{API}/messages/heat.temp").status_code == 503
|
|
assert client.post(f"{API}/runs/flows/f_x").status_code == 503
|
|
|
|
assert client.get(f"{API}/flows/f_x").status_code == 200
|
|
for path in (
|
|
"/login/access-token",
|
|
"/users/me",
|
|
"/flows/ws-ticket",
|
|
"/ai/chat",
|
|
"/alerts/webpush/subscriptions",
|
|
"/artifacts",
|
|
"/flows/f_x/rename/plan",
|
|
"/dashboards/d_x/rename",
|
|
):
|
|
assert client.post(f"{API}{path}").status_code == 200, path
|
|
|
|
gate.open()
|
|
assert client.post(f"{API}/flows/f_x/publish").status_code == 200
|
|
|
|
|
|
def test_an_admitted_request_is_drained_before_close_returns():
|
|
release = threading.Event()
|
|
client, gate = _gated(release)
|
|
answered: list[int] = []
|
|
request = threading.Thread(
|
|
target=lambda: answered.append(client.post(f"{API}/flows/").status_code)
|
|
)
|
|
request.start()
|
|
while not gate._admitted:
|
|
time.sleep(0.01)
|
|
|
|
heard: list[list[str]] = []
|
|
closed = threading.Thread(
|
|
target=lambda: gate.close("op-1", [], threading.Event(), heard.append)
|
|
)
|
|
closed.start()
|
|
closed.join(0.3)
|
|
assert closed.is_alive()
|
|
assert heard == [[f"POST {API}/flows/"]]
|
|
# Turned away meanwhile, while the one inside finishes.
|
|
assert client.post(f"{API}/flows/").status_code == 503
|
|
|
|
release.set()
|
|
closed.join(5)
|
|
request.join(5)
|
|
assert not closed.is_alive()
|
|
assert answered == [200]
|
|
assert gate.closed
|
|
|
|
|
|
def test_an_aborted_drain_reopens_the_gate():
|
|
release = threading.Event()
|
|
client, gate = _gated(release)
|
|
request = threading.Thread(target=lambda: client.post(f"{API}/flows/"))
|
|
request.start()
|
|
while not gate._admitted:
|
|
time.sleep(0.01)
|
|
|
|
abort = threading.Event()
|
|
result: list[bool] = []
|
|
closed = threading.Thread(
|
|
target=lambda: result.append(gate.close("op-1", [], abort))
|
|
)
|
|
closed.start()
|
|
abort.set()
|
|
closed.join(5)
|
|
assert result == [False]
|
|
assert not gate.closed
|
|
release.set()
|
|
request.join(5)
|
|
assert client.post(f"{API}/flows/").status_code == 200
|
|
|
|
|
|
def test_only_the_hooks_of_affected_flows_are_turned_away():
|
|
client, gate = _gated()
|
|
_close(gate, ["f_a"])
|
|
|
|
assert client.get("/hooks/f_a/in").status_code == 503
|
|
assert client.post("/hooks/f_a/in").status_code == 503
|
|
assert client.post("/hooks/f_b/in").status_code == 200
|
|
with pytest.raises(WebSocketDisconnect) as closed:
|
|
with client.websocket_connect("/hooks/f_a/ws") as socket:
|
|
socket.receive_text()
|
|
assert closed.value.code == 1013
|
|
with client.websocket_connect("/hooks/f_b/ws") as socket:
|
|
assert socket.receive_text() == "ok"
|
|
|
|
gate.open()
|
|
assert client.post("/hooks/f_a/in").status_code == 200
|
|
|
|
|
|
def test_the_app_is_gated(client: TestClient, superuser_token_headers):
|
|
gate: RenameGate = app.state.rename_gate
|
|
_close(gate, ["f_x"])
|
|
try:
|
|
assert client.get(f"{API}/flows/", headers=superuser_token_headers).is_success
|
|
made = client.post(
|
|
f"{API}/flows/", headers=superuser_token_headers, json={"name": "gated"}
|
|
)
|
|
assert made.status_code == 503
|
|
assert client.post(
|
|
f"{API}/login/access-token",
|
|
data={
|
|
"username": settings.FIRST_SUPERUSER,
|
|
"password": settings.FIRST_SUPERUSER_PASSWORD,
|
|
},
|
|
).is_success
|
|
|
|
async def tool_call() -> object:
|
|
previous = server._client
|
|
server.set_client(None)
|
|
server.connect(app)
|
|
token = server.caller_authorization.set(
|
|
superuser_token_headers["Authorization"]
|
|
)
|
|
try:
|
|
return await server.create_flow({"name": "gated"})
|
|
finally:
|
|
server.caller_authorization.reset(token)
|
|
await server.aclose()
|
|
server.set_client(previous)
|
|
|
|
assert asyncio.run(tool_call()) == {
|
|
"error": "A rename is in progress; try again shortly",
|
|
"status": 503,
|
|
}
|
|
finally:
|
|
gate.open()
|