71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
"""Tests for codex.parsing.nougat."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from codex.parsing.nougat import pdf_to_markdown
|
|
|
|
|
|
class _FakeResponse:
|
|
"""Minimal stand-in for httpx.Response."""
|
|
|
|
def __init__(self, text: str, status_code: int = 200) -> None:
|
|
self.text = text
|
|
self.status_code = status_code
|
|
|
|
def raise_for_status(self) -> None:
|
|
if self.status_code >= 400:
|
|
raise httpx.HTTPStatusError(
|
|
"error",
|
|
request=MagicMock(),
|
|
response=MagicMock(status_code=self.status_code),
|
|
)
|
|
|
|
|
|
class TestPdfToMarkdown:
|
|
def test_success_returns_response_text(self, tmp_path: pytest.TempdirFactory) -> None:
|
|
pdf = tmp_path / "paper.pdf" # type: ignore[operator]
|
|
pdf.write_bytes(b"%PDF fake")
|
|
|
|
expected = "# Title\nContent"
|
|
fake_response = _FakeResponse(expected)
|
|
|
|
with patch("httpx.Client") as mock_client_cls:
|
|
mock_client = MagicMock()
|
|
mock_client_cls.return_value.__enter__.return_value = mock_client
|
|
mock_client.post.return_value = fake_response
|
|
|
|
result = pdf_to_markdown(str(pdf), nougat_url="http://localhost:8080")
|
|
|
|
assert result == expected
|
|
mock_client.post.assert_called_once()
|
|
call_kwargs = mock_client.post.call_args
|
|
assert "/predict" in call_kwargs[0][0]
|
|
|
|
def test_connect_error_retried_exactly_twice(self, tmp_path: pytest.TempdirFactory) -> None:
|
|
"""ConnectError must trigger 2 retries (3 total attempts)."""
|
|
pdf = tmp_path / "paper.pdf" # type: ignore[operator]
|
|
pdf.write_bytes(b"%PDF fake")
|
|
|
|
call_count = 0
|
|
|
|
def _raising_post(*args: object, **kwargs: object) -> _FakeResponse:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
raise httpx.ConnectError("refused")
|
|
|
|
with patch("httpx.Client") as mock_client_cls:
|
|
mock_client = MagicMock()
|
|
mock_client_cls.return_value.__enter__.return_value = mock_client
|
|
mock_client.post.side_effect = _raising_post
|
|
|
|
with pytest.raises(httpx.ConnectError):
|
|
pdf_to_markdown(str(pdf), nougat_url="http://localhost:8080")
|
|
|
|
# tenacity: stop_after_attempt(3) → 1 original + 2 retries = 3 calls
|
|
assert call_count == 3
|