mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 12:05:07 +00:00
imSp4rky
b10fd2d3a4
A tunnel closed with code 1000 is freed without the 60s grace, a shared: true API key counts tunnels per X-Client-Id, and --remote waits out Retry-After and a full tunnel cap with a countdown; responses carry X-RateLimit-* headers.
338 lines
12 KiB
Python
338 lines
12 KiB
Python
"""Unit tests for unshackle.core.remote_service.RemoteClient."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import click
|
|
import pytest
|
|
import responses
|
|
|
|
from unshackle.core.remote_service import RemoteClient
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
|
|
@pytest.fixture
|
|
def client() -> RemoteClient:
|
|
return RemoteClient(server_url="http://srv:8786", api_key="secret-xyz")
|
|
|
|
|
|
def test_session_sets_secret_key_header(client: RemoteClient) -> None:
|
|
s = client.session
|
|
assert s.headers.get("X-Secret-Key") == "secret-xyz"
|
|
assert s.headers["User-Agent"].startswith("unshackle/")
|
|
|
|
|
|
def test_session_omits_secret_key_when_empty() -> None:
|
|
c = RemoteClient(server_url="http://srv:8786", api_key="")
|
|
assert "X-Secret-Key" not in c.session.headers
|
|
|
|
|
|
def test_server_url_trailing_slash_stripped() -> None:
|
|
c = RemoteClient(server_url="http://srv:8786/", api_key="")
|
|
assert c.server_url == "http://srv:8786"
|
|
|
|
|
|
@responses.activate
|
|
def test_get_returns_json(client: RemoteClient) -> None:
|
|
responses.add(
|
|
responses.GET,
|
|
"http://srv:8786/api/health",
|
|
json={"status": "ok"},
|
|
status=200,
|
|
)
|
|
assert client.get("/api/health") == {"status": "ok"}
|
|
|
|
|
|
@responses.activate
|
|
def test_post_sends_json_body(client: RemoteClient) -> None:
|
|
captured = {}
|
|
|
|
def cb(request):
|
|
captured["body"] = json.loads(request.body)
|
|
return (200, {}, json.dumps({"session_id": "abc"}))
|
|
|
|
responses.add_callback(
|
|
responses.POST, "http://srv:8786/api/session/create", callback=cb, content_type="application/json"
|
|
)
|
|
result = client.post("/api/session/create", {"service": "EXAMPLE"})
|
|
assert result == {"session_id": "abc"}
|
|
assert captured["body"] == {"service": "EXAMPLE"}
|
|
|
|
|
|
@responses.activate
|
|
def test_delete_returns_json(client: RemoteClient) -> None:
|
|
responses.add(
|
|
responses.DELETE,
|
|
"http://srv:8786/api/session/abc",
|
|
json={"status": "deleted"},
|
|
status=200,
|
|
)
|
|
assert client.delete("/api/session/abc") == {"status": "deleted"}
|
|
|
|
|
|
@responses.activate
|
|
def test_4xx_raises_systemexit_with_logged_error(client: RemoteClient, caplog: pytest.LogCaptureFixture) -> None:
|
|
responses.add(
|
|
responses.GET,
|
|
"http://srv:8786/api/session/none",
|
|
json={"error_code": "SESSION_NOT_FOUND", "message": "no such session"},
|
|
status=404,
|
|
)
|
|
with caplog.at_level("ERROR"), pytest.raises(SystemExit):
|
|
client.get("/api/session/none")
|
|
assert "SESSION_NOT_FOUND" in caplog.text
|
|
assert "no such session" in caplog.text
|
|
|
|
|
|
PROXY_502 = "<html><head><title>unshackle.dev | 502: Bad gateway</title></head></html>"
|
|
|
|
|
|
@responses.activate
|
|
def test_proxy_502_page_is_retried_once(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr("unshackle.core.remote_service.time.sleep", lambda _s: None)
|
|
url = "http://srv:8786/api/session/abc/license"
|
|
responses.add(responses.POST, url, body=PROXY_502, status=502, content_type="text/html")
|
|
responses.add(responses.POST, url, json={"keys": {}}, status=200)
|
|
assert client.post("/api/session/abc/license", {}) == {"keys": {}}
|
|
assert len(responses.calls) == 2
|
|
|
|
|
|
@responses.activate
|
|
def test_second_proxy_502_page_exits(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr("unshackle.core.remote_service.time.sleep", lambda _s: None)
|
|
url = "http://srv:8786/api/session/abc/license"
|
|
responses.add(responses.POST, url, body=PROXY_502, status=502, content_type="text/html")
|
|
with pytest.raises(SystemExit):
|
|
client.post("/api/session/abc/license", {})
|
|
assert len(responses.calls) == 2
|
|
|
|
|
|
@responses.activate
|
|
def test_server_json_503_is_not_retried(client: RemoteClient) -> None:
|
|
responses.add(
|
|
responses.GET,
|
|
"http://srv:8786/api/services",
|
|
json={"error_code": "SERVICE_UNAVAILABLE", "message": "busy"},
|
|
status=503,
|
|
)
|
|
with pytest.raises(SystemExit):
|
|
client.get("/api/services")
|
|
assert len(responses.calls) == 1
|
|
|
|
|
|
@responses.activate
|
|
def test_401_falls_back_to_next_auth_header(client: RemoteClient) -> None:
|
|
seen: list[tuple[str | None, str | None]] = []
|
|
|
|
def cb(request):
|
|
seen.append((request.headers.get("X-Secret-Key"), request.headers.get("X-Api-Key")))
|
|
if request.headers.get("X-Api-Key") == "secret-xyz":
|
|
return (200, {}, json.dumps({"status": "ok"}))
|
|
return (401, {}, json.dumps({"message": "Secret Key is Invalid."}))
|
|
|
|
responses.add_callback(responses.GET, "http://srv:8786/api/services", callback=cb)
|
|
|
|
assert client.get("/api/services") == {"status": "ok"}
|
|
assert seen == [("secret-xyz", None), (None, "secret-xyz")]
|
|
|
|
# the winning header is sticky: no second attempt on the next request
|
|
seen.clear()
|
|
assert client.get("/api/services") == {"status": "ok"}
|
|
assert seen == [(None, "secret-xyz")]
|
|
|
|
|
|
@responses.activate
|
|
def test_non_401_error_does_not_fall_back(client: RemoteClient, caplog: pytest.LogCaptureFixture) -> None:
|
|
responses.add(
|
|
responses.GET,
|
|
"http://srv:8786/api/services",
|
|
json={"error_code": "FORBIDDEN", "message": "nope"},
|
|
status=403,
|
|
)
|
|
with caplog.at_level("ERROR"), pytest.raises(SystemExit):
|
|
client.get("/api/services")
|
|
assert len(responses.calls) == 1
|
|
assert client.session.headers.get("X-Secret-Key") == "secret-xyz"
|
|
|
|
|
|
@responses.activate
|
|
def test_401_on_every_header_exits(client: RemoteClient, caplog: pytest.LogCaptureFixture) -> None:
|
|
responses.add(
|
|
responses.GET,
|
|
"http://srv:8786/api/services",
|
|
json={"error_code": "UNAUTHORIZED", "message": "Secret Key is Invalid."},
|
|
status=401,
|
|
)
|
|
with caplog.at_level("ERROR"), pytest.raises(SystemExit):
|
|
client.get("/api/services")
|
|
assert len(responses.calls) == 2
|
|
assert "Secret Key is Invalid." in caplog.text
|
|
|
|
|
|
@responses.activate
|
|
def test_single_configured_auth_header_does_not_retry() -> None:
|
|
c = RemoteClient(server_url="http://srv:8786", api_key="k", auth_headers=["X-Api-Key"])
|
|
responses.add(responses.GET, "http://srv:8786/api/services", json={"message": "nope"}, status=401)
|
|
with pytest.raises(SystemExit):
|
|
c.get("/api/services")
|
|
assert len(responses.calls) == 1
|
|
assert c.session.headers.get("X-Api-Key") == "k"
|
|
|
|
|
|
@responses.activate
|
|
def test_connection_error_raises_systemexit(client: RemoteClient) -> None:
|
|
import requests
|
|
|
|
responses.add(
|
|
responses.GET,
|
|
"http://srv:8786/api/health",
|
|
body=requests.ConnectionError("boom"),
|
|
)
|
|
with pytest.raises(SystemExit):
|
|
client.get("/api/health")
|
|
|
|
|
|
@responses.activate
|
|
def test_timeout_raises_systemexit(client: RemoteClient) -> None:
|
|
import requests
|
|
|
|
responses.add(
|
|
responses.GET,
|
|
"http://srv:8786/api/health",
|
|
body=requests.Timeout("slow"),
|
|
)
|
|
with pytest.raises(SystemExit):
|
|
client.get("/api/health")
|
|
|
|
|
|
@responses.activate
|
|
def test_late_error_after_heartbeat_is_fatal(client: RemoteClient) -> None:
|
|
"""A heartbeat response carries a late failure as a 200 error body after leading newlines."""
|
|
body = "\n\n" + json.dumps({"status": "error", "error_code": "SERVICE_ERROR", "message": "late"})
|
|
responses.add(responses.GET, "http://srv:8786/api/session/s/titles", body=body, status=200)
|
|
with pytest.raises(SystemExit):
|
|
client.get("/api/session/s/titles")
|
|
assert client.request("get", "/api/session/s/titles", optional=True) == {}
|
|
|
|
|
|
@responses.activate
|
|
def test_heartbeat_body_parses_past_leading_newlines(client: RemoteClient) -> None:
|
|
responses.add(responses.GET, "http://srv:8786/api/session/s/titles", body='\n\n{"titles": []}', status=200)
|
|
assert client.get("/api/session/s/titles") == {"titles": []}
|
|
|
|
|
|
def test_read_timeout_in_body_reports_timeout(client: RemoteClient, caplog: pytest.LogCaptureFixture) -> None:
|
|
"""requests re-raises a read timeout in the body as a ConnectionError; the client still names the timeout."""
|
|
import requests
|
|
import urllib3
|
|
|
|
def stalled(*args: object, **kwargs: object) -> None:
|
|
raise requests.ConnectionError(urllib3.exceptions.ReadTimeoutError(None, "", "Read timed out."))
|
|
|
|
client.session.get = stalled # type: ignore[method-assign]
|
|
with pytest.raises(SystemExit):
|
|
client.get("/api/session/s/titles")
|
|
assert "timed out after 120" in caplog.text
|
|
|
|
|
|
@responses.activate
|
|
def test_proxy_error_page_logs_its_title(client: RemoteClient, caplog: pytest.LogCaptureFixture) -> None:
|
|
page = "<!DOCTYPE html><html><head><title>example.com | 524: A timeout occurred</title></head><body>x</body></html>"
|
|
responses.add(responses.GET, "http://srv:8786/api/session/s/titles", body=page, status=524)
|
|
with pytest.raises(SystemExit):
|
|
client.get("/api/session/s/titles")
|
|
assert "HTTP 524: example.com | 524: A timeout occurred" in caplog.text
|
|
assert "DOCTYPE" not in caplog.text
|
|
|
|
|
|
@responses.activate
|
|
def test_a_429_with_retry_after_waits_then_retries(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
waits: list[int] = []
|
|
monkeypatch.setattr("unshackle.core.remote_service.wait_retry_after", waits.append)
|
|
url = "http://srv:8786/api/services"
|
|
responses.add(responses.GET, url, status=429, json={"status": 429}, headers={"Retry-After": "7"})
|
|
responses.add(responses.GET, url, json={"status": "ok"})
|
|
|
|
assert client.get("/api/services") == {"status": "ok"}
|
|
assert waits == [7]
|
|
|
|
|
|
@responses.activate
|
|
def test_a_429_without_retry_after_fails_at_once(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr("unshackle.core.remote_service.wait_retry_after", pytest.fail)
|
|
responses.add(responses.GET, "http://srv:8786/api/services", status=429, json={"error_code": "RATE_LIMITED"})
|
|
|
|
with pytest.raises(SystemExit):
|
|
client.get("/api/services")
|
|
|
|
|
|
def test_wait_retry_after_waits_out_its_own_deadline(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from unshackle.core import remote_service
|
|
|
|
clock = [0.0]
|
|
monkeypatch.setattr(remote_service.time, "monotonic", lambda: clock[0])
|
|
monkeypatch.setattr(remote_service.time, "sleep", lambda s: clock.__setitem__(0, clock[0] + s))
|
|
remote_service.wait_retry_after(3)
|
|
assert clock[0] == 3.0
|
|
remote_service.wait_retry_after(99999)
|
|
assert clock[0] == 3.0 + remote_service.RETRY_AFTER_MAX
|
|
|
|
|
|
def test_a_rate_limited_tunnel_waits_then_opens(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from unshackle.core import remote_service
|
|
|
|
attempts: list[object] = []
|
|
|
|
class Refused:
|
|
def __init__(self, *_: object) -> None:
|
|
self.retry_after = 5 if not attempts else None
|
|
self.at_limit = len(attempts) == 1
|
|
attempts.append(self)
|
|
|
|
def start(self) -> str:
|
|
if self.retry_after or self.at_limit:
|
|
raise ConnectionError("HTTP 429")
|
|
return "id"
|
|
|
|
waits: list[int] = []
|
|
monkeypatch.setattr(remote_service, "TunnelClient", Refused)
|
|
monkeypatch.setattr(remote_service, "wait_retry_after", lambda s, *_: waits.append(s))
|
|
assert remote_service.start_tunnel(client) is attempts[2]
|
|
assert waits == [5, remote_service.TUNNEL_SLOT_POLL]
|
|
|
|
|
|
def test_a_tunnel_refused_for_another_reason_fails_at_once(
|
|
client: RemoteClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
from unshackle.core import remote_service
|
|
|
|
class Refused:
|
|
retry_after = None
|
|
at_limit = False
|
|
|
|
def __init__(self, *_: object) -> None:
|
|
pass
|
|
|
|
def start(self) -> str:
|
|
raise ConnectionError("HTTP 403")
|
|
|
|
monkeypatch.setattr(remote_service, "TunnelClient", Refused)
|
|
monkeypatch.setattr(remote_service, "wait_retry_after", pytest.fail)
|
|
with pytest.raises(click.ClickException, match="HTTP 403"):
|
|
remote_service.start_tunnel(client)
|
|
|
|
|
|
@responses.activate
|
|
def test_the_429_wait_gives_up_after_its_tries(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from unshackle.core import remote_service
|
|
|
|
waits: list[int] = []
|
|
monkeypatch.setattr(remote_service, "wait_retry_after", waits.append)
|
|
responses.add(responses.GET, "http://srv:8786/api/services", status=429, json={}, headers={"Retry-After": "0"})
|
|
|
|
with pytest.raises(SystemExit):
|
|
client.get("/api/services")
|
|
assert waits == [1] * remote_service.RATE_LIMIT_TRIES
|