134 lines
4.7 KiB
Python
134 lines
4.7 KiB
Python
"""Tests for cassette transport fidelity and redaction."""
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
|
|
import httpx
|
|
import httpx2
|
|
from recording import Httpx2RecordingTransport, HttpxRecordingTransport, redact_cassette
|
|
|
|
|
|
class HttpxChunks(httpx.AsyncByteStream):
|
|
"""Yield fixed httpx byte chunks."""
|
|
|
|
async def __aiter__(self) -> AsyncIterator[bytes]:
|
|
yield b"first-"
|
|
yield b"second"
|
|
|
|
|
|
class Httpx2Chunks(httpx2.AsyncByteStream):
|
|
"""Yield fixed httpx2 byte chunks."""
|
|
|
|
async def __aiter__(self) -> AsyncIterator[bytes]:
|
|
yield b"first-"
|
|
yield b"second"
|
|
|
|
|
|
def test_httpx_recording_transport_preserves_chunk_boundaries() -> None:
|
|
async def run() -> None:
|
|
async def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(200, stream=HttpxChunks())
|
|
|
|
transport = HttpxRecordingTransport(httpx.MockTransport(handler))
|
|
async with httpx.AsyncClient(transport=transport) as client:
|
|
response = await client.get("https://api.anthropic.com/v1/messages")
|
|
|
|
assert response.content == b"first-second"
|
|
assert transport.interactions[0].response_chunks == (b"first-", b"second")
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_httpx2_recording_transport_preserves_chunk_boundaries() -> None:
|
|
async def run() -> None:
|
|
async def handler(request: httpx2.Request) -> httpx2.Response:
|
|
return httpx2.Response(200, stream=Httpx2Chunks())
|
|
|
|
transport = Httpx2RecordingTransport(httpx2.MockTransport(handler))
|
|
async with httpx2.AsyncClient(transport=transport) as client:
|
|
response = await client.get("https://api.spotify.com/v1/search?q=test")
|
|
|
|
assert response.content == b"first-second"
|
|
assert transport.interactions[0].response_chunks == (b"first-", b"second")
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_redaction_removes_headers_tokens_identity_and_taste() -> None:
|
|
interactions = [
|
|
{
|
|
"method": "POST",
|
|
"url": "https://accounts.spotify.com/api/token",
|
|
"request_headers": {"Authorization": "Bearer raw-token"},
|
|
"request_body_base64": _encode(b"client_secret=secret"),
|
|
"response_body_base64": _encode(b'{"access_token":"raw-token"}'),
|
|
"response_chunks_base64": [],
|
|
"status": 200,
|
|
},
|
|
{
|
|
"method": "GET",
|
|
"url": "https://api.spotify.com/v1/me",
|
|
"response_headers": {"Set-Cookie": "secret"},
|
|
"request_body_base64": "",
|
|
"response_body_base64": _encode(
|
|
b'{"id":"real-account-id","display_name":"Real Listener"}'
|
|
),
|
|
"response_chunks_base64": [],
|
|
"status": 200,
|
|
},
|
|
{
|
|
"method": "GET",
|
|
"url": "https://api.spotify.com/v1/me/top/artists?limit=50",
|
|
"headers": {"Authorization": "Bearer raw-token"},
|
|
"request_body_base64": "",
|
|
"response_body_base64": _encode(b'{"items":[{"name":"Private Artist"}]}'),
|
|
"response_chunks_base64": [],
|
|
"status": 200,
|
|
},
|
|
{
|
|
"method": "POST",
|
|
"url": "https://api.anthropic.com/v1/messages",
|
|
"request_body_base64": _encode(
|
|
json.dumps(
|
|
{
|
|
"messages": [
|
|
{
|
|
"content": json.dumps(
|
|
{
|
|
"taste_profile": "Private Artist",
|
|
"account": "real-account-id",
|
|
}
|
|
)
|
|
}
|
|
]
|
|
}
|
|
).encode()
|
|
),
|
|
"response_body_base64": _encode(b"Bearer raw-token"),
|
|
"response_chunks_base64": [_encode(b"Bearer raw-token")],
|
|
"status": 200,
|
|
},
|
|
]
|
|
|
|
redacted = redact_cassette(interactions)
|
|
rendered = json.dumps(redacted)
|
|
decoded_bodies = " ".join(
|
|
base64.b64decode(str(item.get(field, ""))).decode(errors="replace")
|
|
for item in redacted
|
|
for field in ("request_body_base64", "response_body_base64")
|
|
)
|
|
|
|
assert len(redacted) == 3
|
|
assert "headers" not in rendered.casefold()
|
|
assert "raw-token" not in rendered + decoded_bodies
|
|
assert "real-account-id" not in rendered + decoded_bodies
|
|
assert "Real Listener" not in rendered + decoded_bodies
|
|
assert "Private Artist" not in rendered + decoded_bodies
|
|
assert "Synthetic Focus Artist" in decoded_bodies
|
|
assert "accounts.spotify.com" not in rendered
|
|
|
|
|
|
def _encode(value: bytes) -> str:
|
|
return base64.b64encode(value).decode("ascii")
|