73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
"""A socket's credential comes as a header or a one-time ticket, not in its URL."""
|
|
|
|
from time import monotonic
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from starlette.websockets import WebSocketDisconnect
|
|
|
|
from fluksio.core.config import settings
|
|
from fluksio.flow.remote import PROTOCOL
|
|
|
|
FLOWS = f"{settings.API_V1_STR}/flows"
|
|
ATTACH = f"{settings.API_V1_STR}/workers/attach"
|
|
|
|
|
|
def ticket(client: TestClient, headers: dict[str, str]) -> str:
|
|
answer = client.post(f"{FLOWS}/ws-ticket", headers=headers)
|
|
assert answer.status_code == 200, answer.text
|
|
return str(answer.json()["ticket"])
|
|
|
|
|
|
def opens(client: TestClient, url: str, headers: dict[str, str] | None = None) -> bool:
|
|
try:
|
|
with client.websocket_connect(url, headers=headers or {}) as ws:
|
|
assert ws.receive_json()["type"] == "snapshot"
|
|
except WebSocketDisconnect:
|
|
return False
|
|
return True
|
|
|
|
|
|
def test_the_live_socket_takes_a_header_and_still_a_token(
|
|
client: TestClient, superuser_token_headers: dict[str, str]
|
|
) -> None:
|
|
assert opens(client, f"{FLOWS}/ws", superuser_token_headers)
|
|
token = superuser_token_headers["Authorization"][7:]
|
|
assert opens(client, f"{FLOWS}/ws?token={token}")
|
|
assert not opens(client, f"{FLOWS}/ws")
|
|
|
|
|
|
def test_a_ticket_opens_the_socket_once(
|
|
client: TestClient, superuser_token_headers: dict[str, str]
|
|
) -> None:
|
|
once = ticket(client, superuser_token_headers)
|
|
assert opens(client, f"{FLOWS}/ws?ticket={once}")
|
|
assert not opens(client, f"{FLOWS}/ws?ticket={once}")
|
|
|
|
|
|
def test_a_ticket_lapses_after_thirty_seconds(
|
|
client: TestClient,
|
|
superuser_token_headers: dict[str, str],
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
stale = ticket(client, superuser_token_headers)
|
|
later = monotonic() + 31
|
|
monkeypatch.setattr("fluksio.api.socket_auth.monotonic", lambda: later)
|
|
assert not opens(client, f"{FLOWS}/ws?ticket={stale}")
|
|
|
|
|
|
def test_a_worker_attaches_with_a_header_and_still_a_token(
|
|
client: TestClient, superuser_token_headers: dict[str, str]
|
|
) -> None:
|
|
token = client.post(
|
|
f"{settings.API_V1_STR}/workers/tokens",
|
|
headers=superuser_token_headers,
|
|
json={"name": "header-box"},
|
|
).json()["token"]
|
|
hello = {"op": "hello", "protocol": PROTOCOL}
|
|
bearer = {"Authorization": f"Bearer {token}"}
|
|
for url, headers in ((ATTACH, bearer), (f"{ATTACH}?token={token}", {})):
|
|
with client.websocket_connect(url, headers=headers) as ws:
|
|
ws.send_json(hello)
|
|
assert ws.receive_json()["op"] == "welcome"
|