diff --git a/src/kernel/__init__.py b/src/kernel/__init__.py index 333d7018..af9462fb 100644 --- a/src/kernel/__init__.py +++ b/src/kernel/__init__.py @@ -51,6 +51,11 @@ app_registry, export_registry, ) +from .lib.audit_log_download import ( + AuditLogDownloadError, + AuditLogDownloadResult, + AuditLogDownloadProgress, +) __all__ = [ "types", @@ -105,6 +110,9 @@ "App", "app_registry", "export_registry", + "AuditLogDownloadError", + "AuditLogDownloadProgress", + "AuditLogDownloadResult", ] if not _t.TYPE_CHECKING: diff --git a/src/kernel/lib/audit_log_download.py b/src/kernel/lib/audit_log_download.py new file mode 100644 index 00000000..f334d6ae --- /dev/null +++ b/src/kernel/lib/audit_log_download.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +import time +import hashlib +import inspect +from typing import BinaryIO, Callable, Optional, Protocol, Awaitable, ContextManager, AsyncContextManager +from dataclasses import dataclass + +import anyio +import httpx +from anyio import to_thread + +from .._exceptions import KernelError + +_DEFAULT_MAX_TRANSFER_RETRIES = 6 +_MAX_CHUNK_ROWS = 50_000 +_MAX_RETRY_DELAY = 8.0 + + +class AuditLogDownloadError(KernelError): + pass + + +@dataclass(frozen=True) +class AuditLogDownloadResult: + bytes_written: int + chunks: int + rows: int + + +@dataclass(frozen=True) +class AuditLogDownloadProgress(AuditLogDownloadResult): + chunk_rows: int + + +class _SyncChunkResponse(Protocol): + @property + def headers(self) -> httpx.Headers: ... + + def read(self) -> bytes: ... + + def close(self) -> None: ... + + +class _AsyncChunkResponse(Protocol): + @property + def headers(self) -> httpx.Headers: ... + + async def read(self) -> bytes: ... + + async def close(self) -> None: ... + + +SyncFetchChunk = Callable[[Optional[str]], ContextManager[_SyncChunkResponse]] +AsyncFetchChunk = Callable[[Optional[str]], AsyncContextManager[_AsyncChunkResponse]] +ProgressCallback = Callable[[AuditLogDownloadProgress], None] +AsyncProgressCallback = Callable[[AuditLogDownloadProgress], Optional[Awaitable[None]]] + + +def download_audit_logs( + fetch_chunk: SyncFetchChunk, + destination: BinaryIO, + *, + on_progress: ProgressCallback | None = None, + max_transfer_retries: int = _DEFAULT_MAX_TRANSFER_RETRIES, +) -> AuditLogDownloadResult: + if not callable(getattr(destination, "write", None)): + raise TypeError("audit log download destination must provide write()") + _validate_max_transfer_retries(max_transfer_retries) + + cursor: str | None = None + result = AuditLogDownloadResult(bytes_written=0, chunks=0, rows=0) + seen_cursors: set[str] = set() + while True: + body, headers = _fetch_verified_chunk(fetch_chunk, cursor, max_transfer_retries) + chunk_rows, next_cursor, has_more = _parse_chunk_headers(headers, cursor) + if has_more and next_cursor is not None: + if next_cursor in seen_cursors: + raise AuditLogDownloadError("response repeated X-Next-Cursor header") + seen_cursors.add(next_cursor) + _write_chunk(destination, body) + + cursor = next_cursor + result = AuditLogDownloadResult( + bytes_written=result.bytes_written + len(body), + chunks=result.chunks + 1, + rows=result.rows + chunk_rows, + ) + if on_progress is not None: + on_progress( + AuditLogDownloadProgress( + bytes_written=result.bytes_written, + chunks=result.chunks, + rows=result.rows, + chunk_rows=chunk_rows, + ) + ) + if not has_more: + return result + + +async def async_download_audit_logs( + fetch_chunk: AsyncFetchChunk, + destination: BinaryIO, + *, + on_progress: AsyncProgressCallback | None = None, + max_transfer_retries: int = _DEFAULT_MAX_TRANSFER_RETRIES, +) -> AuditLogDownloadResult: + if not callable(getattr(destination, "write", None)): + raise TypeError("audit log download destination must provide write()") + _validate_max_transfer_retries(max_transfer_retries) + + cursor: str | None = None + result = AuditLogDownloadResult(bytes_written=0, chunks=0, rows=0) + seen_cursors: set[str] = set() + while True: + body, headers = await _async_fetch_verified_chunk(fetch_chunk, cursor, max_transfer_retries) + chunk_rows, next_cursor, has_more = _parse_chunk_headers(headers, cursor) + if has_more and next_cursor is not None: + if next_cursor in seen_cursors: + raise AuditLogDownloadError("response repeated X-Next-Cursor header") + seen_cursors.add(next_cursor) + await to_thread.run_sync(_write_chunk, destination, body) + + cursor = next_cursor + result = AuditLogDownloadResult( + bytes_written=result.bytes_written + len(body), + chunks=result.chunks + 1, + rows=result.rows + chunk_rows, + ) + if on_progress is not None: + callback_result = on_progress( + AuditLogDownloadProgress( + bytes_written=result.bytes_written, + chunks=result.chunks, + rows=result.rows, + chunk_rows=chunk_rows, + ) + ) + if inspect.isawaitable(callback_result): + await callback_result + if not has_more: + return result + + +def _fetch_verified_chunk( + fetch_chunk: SyncFetchChunk, cursor: str | None, max_transfer_retries: int +) -> tuple[bytes, httpx.Headers]: + for retries in range(max_transfer_retries + 1): + transfer_error: Exception | None = None + with fetch_chunk(cursor) as response: + try: + body = response.read() + headers = httpx.Headers(response.headers) + _verify_checksum(body, headers) + except Exception as error: + transfer_error = error + else: + return body, headers + assert transfer_error is not None + if retries == max_transfer_retries: + raise transfer_error + time.sleep(min(2**retries, _MAX_RETRY_DELAY)) + raise AssertionError("unreachable") + + +async def _async_fetch_verified_chunk( + fetch_chunk: AsyncFetchChunk, cursor: str | None, max_transfer_retries: int +) -> tuple[bytes, httpx.Headers]: + for retries in range(max_transfer_retries + 1): + transfer_error: Exception | None = None + async with fetch_chunk(cursor) as response: + try: + body = await response.read() + headers = httpx.Headers(response.headers) + await to_thread.run_sync(_verify_checksum, body, headers) + except Exception as error: + transfer_error = error + else: + return body, headers + assert transfer_error is not None + if retries == max_transfer_retries: + raise transfer_error + await anyio.sleep(min(2**retries, _MAX_RETRY_DELAY)) + raise AssertionError("unreachable") + + +def _verify_checksum(body: bytes, headers: httpx.Headers) -> None: + expected = headers.get("x-content-sha256") + if not expected: + raise AuditLogDownloadError("response missing X-Content-Sha256 header") + actual = hashlib.sha256(body).hexdigest() + if actual != expected: + raise AuditLogDownloadError(f"audit log chunk checksum mismatch (got {actual}, want {expected})") + + +def _parse_chunk_headers(headers: httpx.Headers, current_cursor: str | None) -> tuple[int, str | None, bool]: + has_more_value = headers.get("x-has-more") + if has_more_value not in {"true", "false"}: + raise AuditLogDownloadError("response missing or invalid X-Has-More header") + has_more = has_more_value == "true" + + row_count = headers.get("x-row-count") + if row_count is None or not row_count.isascii() or not row_count.isdecimal(): + raise AuditLogDownloadError("response missing or invalid X-Row-Count header") + normalized_row_count = row_count.lstrip("0") or "0" + if len(normalized_row_count) > len(str(_MAX_CHUNK_ROWS)): + raise AuditLogDownloadError("response missing or invalid X-Row-Count header") + rows = int(normalized_row_count) + if rows > _MAX_CHUNK_ROWS: + raise AuditLogDownloadError("response missing or invalid X-Row-Count header") + + next_cursor = headers.get("x-next-cursor") or None + if has_more and (not next_cursor or next_cursor == current_cursor): + raise AuditLogDownloadError("response has invalid X-Next-Cursor header") + if not has_more and next_cursor: + raise AuditLogDownloadError("response returned a cursor after the final chunk") + return rows, next_cursor, has_more + + +def _write_chunk(destination: BinaryIO, body: bytes) -> None: + remaining = memoryview(body) + while remaining: + written = destination.write(remaining) + if type(written) is not int or written <= 0 or written > len(remaining): + raise AuditLogDownloadError("audit log download destination performed a short write") + remaining = remaining[written:] + + +def _validate_max_transfer_retries(max_transfer_retries: int) -> None: + if type(max_transfer_retries) is not int or max_transfer_retries < 0: + raise ValueError("max_transfer_retries must be a non-negative integer") diff --git a/src/kernel/resources/audit_logs.py b/src/kernel/resources/audit_logs.py index 2970e153..aac5d71c 100644 --- a/src/kernel/resources/audit_logs.py +++ b/src/kernel/resources/audit_logs.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Union +from typing import Union, BinaryIO, ContextManager, AsyncContextManager from datetime import datetime from typing_extensions import Literal @@ -30,6 +30,13 @@ from ..pagination import SyncPageTokenPagination, AsyncPageTokenPagination from .._base_client import AsyncPaginator, make_request_options from ..types.audit_log_entry import AuditLogEntry +from ..lib.audit_log_download import ( + ProgressCallback, + AsyncProgressCallback, + AuditLogDownloadResult, + download_audit_logs, + async_download_audit_logs, +) __all__ = ["AuditLogsResource", "AsyncAuditLogsResource"] @@ -222,6 +229,60 @@ def export_chunk( cast_to=BinaryAPIResponse, ) + def download( + self, + *, + to: BinaryIO, + end: Union[str, datetime], + start: Union[str, datetime], + auth_strategy: str | Omit = omit, + exclude_method: SequenceNotStr[str] | Omit = omit, + limit: int | Omit = omit, + method: str | Omit = omit, + search: str | Omit = omit, + search_user_id: SequenceNotStr[str] | Omit = omit, + service: str | Omit = omit, + on_progress: ProgressCallback | None = None, + max_transfer_retries: int = 6, + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> AuditLogDownloadResult: + """Download a complete gzip-compressed JSON Lines audit log export. + + The SDK writes the export to a writable binary destination, verifies + every chunk, and retries transient transfer failures. It does not close + the destination. If the download fails, the destination may contain a + partial export; use a temporary file and atomic rename when the completed + export must be published atomically. + """ + + def fetch_chunk(cursor: str | None) -> ContextManager[StreamedBinaryAPIResponse]: + return self.with_streaming_response.export_chunk( + end=end, + start=start, + auth_strategy=auth_strategy, + cursor=cursor if cursor is not None else omit, + exclude_method=exclude_method, + limit=limit, + method=method, + search=search, + search_user_id=search_user_id, + service=service, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + return download_audit_logs( + fetch_chunk, + to, + on_progress=on_progress, + max_transfer_retries=max_transfer_retries, + ) + class AsyncAuditLogsResource(AsyncAPIResource): """Read audit log records for the authenticated organization.""" @@ -411,6 +472,60 @@ async def export_chunk( cast_to=AsyncBinaryAPIResponse, ) + async def download( + self, + *, + to: BinaryIO, + end: Union[str, datetime], + start: Union[str, datetime], + auth_strategy: str | Omit = omit, + exclude_method: SequenceNotStr[str] | Omit = omit, + limit: int | Omit = omit, + method: str | Omit = omit, + search: str | Omit = omit, + search_user_id: SequenceNotStr[str] | Omit = omit, + service: str | Omit = omit, + on_progress: AsyncProgressCallback | None = None, + max_transfer_retries: int = 6, + extra_headers: Headers | None = None, + extra_query: Query | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = not_given, + ) -> AuditLogDownloadResult: + """Download a complete gzip-compressed JSON Lines audit log export. + + The SDK writes the export to a writable binary destination, verifies + every chunk, and retries transient transfer failures. It does not close + the destination. If the download fails, the destination may contain a + partial export; use a temporary file and atomic rename when the completed + export must be published atomically. + """ + + def fetch_chunk(cursor: str | None) -> AsyncContextManager[AsyncStreamedBinaryAPIResponse]: + return self.with_streaming_response.export_chunk( + end=end, + start=start, + auth_strategy=auth_strategy, + cursor=cursor if cursor is not None else omit, + exclude_method=exclude_method, + limit=limit, + method=method, + search=search, + search_user_id=search_user_id, + service=service, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + + return await async_download_audit_logs( + fetch_chunk, + to, + on_progress=on_progress, + max_transfer_retries=max_transfer_retries, + ) + class AuditLogsResourceWithRawResponse: def __init__(self, audit_logs: AuditLogsResource) -> None: diff --git a/tests/test_audit_log_download.py b/tests/test_audit_log_download.py new file mode 100644 index 00000000..e254066e --- /dev/null +++ b/tests/test_audit_log_download.py @@ -0,0 +1,409 @@ +from __future__ import annotations + +import hashlib +import threading +from io import BytesIO +from typing import BinaryIO, Iterator, AsyncIterator, cast +from typing_extensions import override + +import httpx +import pytest +from respx import MockRouter + +from kernel import Kernel, AsyncKernel, InternalServerError, AuditLogDownloadError +from kernel.lib import audit_log_download + + +def chunk_headers(body: bytes, *, rows: int, has_more: bool, next_cursor: str | None = None) -> dict[str, str]: + headers = { + "x-content-sha256": hashlib.sha256(body).hexdigest(), + "x-has-more": str(has_more).lower(), + "x-row-count": str(rows), + } + if next_cursor is not None: + headers["x-next-cursor"] = next_cursor + return headers + + +def chunk_response(body: bytes, *, rows: int, has_more: bool, next_cursor: str | None = None) -> httpx.Response: + return httpx.Response( + 200, + content=body, + headers=chunk_headers(body, rows=rows, has_more=has_more, next_cursor=next_cursor), + ) + + +class BrokenSyncByteStream(httpx.SyncByteStream): + @override + def __iter__(self) -> Iterator[bytes]: + yield b"partial" + raise httpx.ReadError("truncated response body") + + +class BrokenAsyncByteStream(httpx.AsyncByteStream): + @override + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"partial" + raise httpx.ReadError("truncated response body") + + +def test_download_writes_verified_chunks(client: Kernel, respx_mock: MockRouter) -> None: + responses = iter( + [ + chunk_response(b"first", rows=2, has_more=True, next_cursor="next"), + chunk_response(b"second", rows=1, has_more=False), + ] + ) + cursors: list[str | None] = [] + formats: list[str | None] = [] + + def respond(request: httpx.Request) -> httpx.Response: + cursors.append(request.url.params.get("cursor")) + formats.append(request.url.params.get("format")) + return next(responses) + + respx_mock.get("/audit-logs/export/chunk").mock(side_effect=respond) + destination = BytesIO() + progress: list[tuple[int, int, int, int]] = [] + + result = client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + on_progress=lambda update: progress.append( + (update.bytes_written, update.chunks, update.rows, update.chunk_rows) + ), + ) + + assert destination.getvalue() == b"firstsecond" + assert result.bytes_written == 11 + assert result.chunks == 2 + assert result.rows == 3 + assert progress == [(5, 1, 2, 2), (11, 2, 3, 1)] + assert cursors == [None, "next"] + assert formats == [None, None] + + +async def test_async_download_writes_verified_chunks( + async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + responses = iter( + [ + chunk_response(b"first", rows=2, has_more=True, next_cursor="next"), + chunk_response(b"second", rows=1, has_more=False), + ] + ) + formats: list[str | None] = [] + + def respond(request: httpx.Request) -> httpx.Response: + formats.append(request.url.params.get("format")) + return next(responses) + + respx_mock.get("/audit-logs/export/chunk").mock(side_effect=respond) + thread_ids: list[int] = [] + write_chunk = audit_log_download._write_chunk + + def record_write(destination: BinaryIO, body: bytes) -> None: + thread_ids.append(threading.get_ident()) + write_chunk(destination, body) + + monkeypatch.setattr(audit_log_download, "_write_chunk", record_write) + destination = BytesIO() + progress_called = False + + async def on_progress(_: object) -> None: + nonlocal progress_called + progress_called = True + + result = await async_client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + on_progress=on_progress, + ) + + assert destination.getvalue() == b"firstsecond" + assert result.bytes_written == 11 + assert result.chunks == 2 + assert result.rows == 3 + assert progress_called + assert thread_ids and all(thread_id != threading.get_ident() for thread_id in thread_ids) + assert formats == [None, None] + + +async def test_async_download_retries_body_read_failure( + async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + async def no_sleep(_: float) -> None: + pass + + monkeypatch.setattr(audit_log_download.anyio, "sleep", no_sleep) + headers = chunk_headers(b"good", rows=1, has_more=False) + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + httpx.Response(200, headers=headers, stream=BrokenAsyncByteStream()), + chunk_response(b"good", rows=1, has_more=False), + ] + ) + destination = BytesIO() + + await async_client.with_options(max_retries=0).audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + max_transfer_retries=1, + ) + + assert destination.getvalue() == b"good" + assert route.call_count == 2 + + +async def test_async_download_retries_checksum_mismatch( + async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + async def no_sleep(_: float) -> None: + pass + + monkeypatch.setattr(audit_log_download.anyio, "sleep", no_sleep) + bad = chunk_response(b"bad", rows=1, has_more=False) + bad.headers["x-content-sha256"] = hashlib.sha256(b"good").hexdigest() + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[bad, chunk_response(b"good", rows=1, has_more=False)] + ) + destination = BytesIO() + + await async_client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"good" + assert route.call_count == 2 + + +async def test_async_download_uses_client_http_retries(async_client: AsyncKernel, respx_mock: MockRouter) -> None: + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + httpx.Response(500, json={"message": "temporary failure"}), + chunk_response(b"good", rows=1, has_more=False), + ] + ) + + await async_client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert route.call_count == 2 + + +async def test_async_download_respects_disabled_http_retries(async_client: AsyncKernel, respx_mock: MockRouter) -> None: + route = respx_mock.get("/audit-logs/export/chunk").mock( + return_value=httpx.Response(500, json={"message": "temporary failure"}) + ) + + with pytest.raises(InternalServerError, match="500"): + await async_client.with_options(max_retries=0).audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert route.call_count == 1 + + +@pytest.mark.parametrize("row_count", ["", "1.0", "50001", "9" * 5000]) +def test_download_rejects_invalid_row_count(row_count: str, client: Kernel, respx_mock: MockRouter) -> None: + response = chunk_response(b"chunk", rows=1, has_more=False) + response.headers["x-row-count"] = row_count + respx_mock.get("/audit-logs/export/chunk").mock(return_value=response) + destination = BytesIO() + + with pytest.raises(AuditLogDownloadError, match="response missing or invalid X-Row-Count header"): + client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"" + + +def test_download_accepts_zero_padded_row_count(client: Kernel, respx_mock: MockRouter) -> None: + response = chunk_response(b"chunk", rows=1, has_more=False) + response.headers["x-row-count"] = "0" * 5000 + "1" + respx_mock.get("/audit-logs/export/chunk").mock(return_value=response) + + result = client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert result.rows == 1 + + +class InvalidWriteDestination: + def __init__(self, result: object) -> None: + self.result = result + + def write(self, _: object) -> object: + return self.result + + +@pytest.mark.parametrize("write_result", [float("nan"), 0.5, True]) +def test_download_rejects_invalid_write_result(write_result: object, client: Kernel, respx_mock: MockRouter) -> None: + respx_mock.get("/audit-logs/export/chunk").mock(return_value=chunk_response(b"chunk", rows=1, has_more=False)) + + with pytest.raises(AuditLogDownloadError, match="audit log download destination performed a short write"): + client.audit_logs.download( + to=cast(BinaryIO, InvalidWriteDestination(write_result)), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + +@pytest.mark.parametrize("max_transfer_retries", [True, 1.5]) +def test_download_rejects_invalid_transfer_retry_count(max_transfer_retries: object, client: Kernel) -> None: + with pytest.raises(ValueError, match="max_transfer_retries must be a non-negative integer"): + client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + max_transfer_retries=cast(int, max_transfer_retries), + ) + + +def test_download_retries_body_read_failure( + client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + def no_sleep(_: float) -> None: + pass + + monkeypatch.setattr(audit_log_download.time, "sleep", no_sleep) + headers = chunk_headers(b"good", rows=1, has_more=False) + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + httpx.Response(200, headers=headers, stream=BrokenSyncByteStream()), + chunk_response(b"good", rows=1, has_more=False), + ] + ) + destination = BytesIO() + + client.with_options(max_retries=0).audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + max_transfer_retries=1, + ) + + assert destination.getvalue() == b"good" + assert route.call_count == 2 + + +def test_download_retries_checksum_mismatch( + client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + def no_sleep(_: float) -> None: + pass + + monkeypatch.setattr(audit_log_download.time, "sleep", no_sleep) + bad = chunk_response(b"bad", rows=1, has_more=False) + bad.headers["x-content-sha256"] = hashlib.sha256(b"good").hexdigest() + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[bad, chunk_response(b"good", rows=1, has_more=False)] + ) + destination = BytesIO() + + client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"good" + assert route.call_count == 2 + + +def test_download_uses_client_http_retries(client: Kernel, respx_mock: MockRouter) -> None: + route = respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + httpx.Response(500, json={"message": "temporary failure"}), + chunk_response(b"good", rows=1, has_more=False), + ] + ) + + client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert route.call_count == 2 + + +def test_download_respects_disabled_http_retries(client: Kernel, respx_mock: MockRouter) -> None: + route = respx_mock.get("/audit-logs/export/chunk").mock( + return_value=httpx.Response(500, json={"message": "temporary failure"}) + ) + + with pytest.raises(InternalServerError, match="500"): + client.with_options(max_retries=0).audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert route.call_count == 1 + + +def test_download_respects_transfer_retry_limit(client: Kernel, respx_mock: MockRouter) -> None: + bad = chunk_response(b"bad", rows=1, has_more=False) + bad.headers["x-content-sha256"] = hashlib.sha256(b"good").hexdigest() + route = respx_mock.get("/audit-logs/export/chunk").mock(return_value=bad) + + with pytest.raises(AuditLogDownloadError, match="checksum mismatch"): + client.audit_logs.download( + to=BytesIO(), + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + max_transfer_retries=0, + ) + + assert route.call_count == 1 + + +def test_download_rejects_cursor_cycle(client: Kernel, respx_mock: MockRouter) -> None: + respx_mock.get("/audit-logs/export/chunk").mock( + side_effect=[ + chunk_response(b"first", rows=1, has_more=True, next_cursor="a"), + chunk_response(b"second", rows=1, has_more=True, next_cursor="b"), + chunk_response(b"duplicate", rows=1, has_more=True, next_cursor="a"), + ] + ) + destination = BytesIO() + + with pytest.raises(AuditLogDownloadError, match="response repeated X-Next-Cursor header"): + client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b"firstsecond" + + +def test_download_rejects_invalid_cursor_before_writing(client: Kernel, respx_mock: MockRouter) -> None: + respx_mock.get("/audit-logs/export/chunk").mock(return_value=chunk_response(b"chunk", rows=1, has_more=True)) + destination = BytesIO() + + with pytest.raises(AuditLogDownloadError, match="response has invalid X-Next-Cursor header"): + client.audit_logs.download( + to=destination, + start="2026-06-01T00:00:00Z", + end="2026-06-02T00:00:00Z", + ) + + assert destination.getvalue() == b""