codex-py/tests/parsing/test_mathpix.py

"""Tests for codex.parsing.mathpix — formula extraction module.

All external dependencies (fitz, pix2tex, httpx, get_settings) are mocked
so the suite runs offline without real API calls or model downloads.
"""

from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock, patch

import pytest

from codex.models import FormulaChunk


def _mock_open_pdf() -> MagicMock:
    """Return a mock that satisfies ``open(path, 'rb')`` as a context manager."""
    mock_file = MagicMock()
    mock_file.__enter__ = MagicMock(return_value=mock_file)
    mock_file.__exit__ = MagicMock(return_value=False)
    mock_open = MagicMock(return_value=mock_file)
    return mock_open


class TestMathPixAPIPath:
    """Tests for the MathPix cloud API backend."""

    def test_mathpix_200_returns_formulas(self) -> None:
        """extract_formulas routes to MathPix when creds are provided and parses result."""
        mmd_content = "\\[\n\\alpha + \\beta\n\\]\n\\[\nE = mc^2\n\\]"

        mock_settings = MagicMock()
        mock_settings.mathpix_app_id = None
        mock_settings.mathpix_app_key = None
        mock_settings.pix2tex_fallback = True

        with (
            patch("codex.parsing.mathpix.get_settings", return_value=mock_settings),
            patch("builtins.open", _mock_open_pdf()),
            patch("httpx.post") as mock_post,
            patch("httpx.get") as mock_get,
            patch("time.sleep"),
        ):
            mock_post.return_value = MagicMock(
                status_code=200,
                json=lambda: {"pdf_id": "abc123"},
            )
            mock_post.return_value.raise_for_status = MagicMock()

            mock_get.return_value = MagicMock(
                status_code=200,
                text=mmd_content,
            )

            from codex.parsing.mathpix import _extract_formulas_mathpix

            results = _extract_formulas_mathpix("paper.pdf", "app_id", "app_key")

        assert len(results) == 2
        assert all(isinstance(r, FormulaChunk) for r in results)
        assert results[0].raw_latex == "\\[\n\\alpha + \\beta\n\\]"

    def test_mathpix_empty_pdf_id_returns_empty(self) -> None:
        """When MathPix returns no pdf_id, result is empty list."""
        with (
            patch("builtins.open", _mock_open_pdf()),
            patch("httpx.post") as mock_post,
            patch("time.sleep"),
        ):
            mock_post.return_value = MagicMock(
                status_code=200,
                json=lambda: {},  # No pdf_id
            )
            mock_post.return_value.raise_for_status = MagicMock()

            from codex.parsing.mathpix import _extract_formulas_mathpix

            results = _extract_formulas_mathpix("paper.pdf", "app_id", "app_key")

        assert results == []

    def test_mathpix_paper_id_from_stem(self) -> None:
        """paper_id in FormulaChunk is derived from the PDF filename stem."""
        mmd_content = "\\[\nx^2\n\\]"

        with (
            patch("builtins.open", _mock_open_pdf()),
            patch("httpx.post") as mock_post,
            patch("httpx.get") as mock_get,
            patch("time.sleep"),
        ):
            mock_post.return_value = MagicMock(
                status_code=200,
                json=lambda: {"pdf_id": "xyz"},
            )
            mock_post.return_value.raise_for_status = MagicMock()
            mock_get.return_value = MagicMock(status_code=200, text=mmd_content)

            from codex.parsing.mathpix import _extract_formulas_mathpix

            results = _extract_formulas_mathpix("/tmp/2301.07041.pdf", "id", "key")

        assert results[0].paper_id == "2301.07041"

    def test_mathpix_page_counter_increments(self) -> None:
        """Page counter increments on MathPix page markers in MMD."""
        mmd_content = "\\[\nA\n\\]\n<!-- Page 2 -->\n\\[\nB\n\\]"

        with (
            patch("builtins.open", _mock_open_pdf()),
            patch("httpx.post") as mock_post,
            patch("httpx.get") as mock_get,
            patch("time.sleep"),
        ):
            mock_post.return_value = MagicMock(
                status_code=200,
                json=lambda: {"pdf_id": "p1"},
            )
            mock_post.return_value.raise_for_status = MagicMock()
            mock_get.return_value = MagicMock(status_code=200, text=mmd_content)

            from codex.parsing.mathpix import _extract_formulas_mathpix

            results = _extract_formulas_mathpix("p.pdf", "id", "key")

        assert results[0].page == 1
        assert results[1].page == 2

    def test_mathpix_http_error_propagates(self) -> None:
        """HTTP errors from MathPix are raised (tenacity reraises after retries)."""
        import httpx

        with (
            patch("builtins.open", _mock_open_pdf()),
            patch("httpx.post") as mock_post,
        ):
            mock_post.return_value = MagicMock(status_code=401)
            mock_post.return_value.raise_for_status.side_effect = httpx.HTTPStatusError(
                "401", request=MagicMock(), response=MagicMock()
            )

            from codex.parsing.mathpix import _extract_formulas_mathpix

            with pytest.raises(httpx.HTTPStatusError):
                _extract_formulas_mathpix("p.pdf", "id", "key")


class TestPix2TexFallbackPath:
    """Tests for the local pix2tex backend."""

    def _make_fitz_page(
        self,
        blocks: list[Any],
        pixmap_samples: bytes = b"\xff" * (30 * 100 * 3),
        pixmap_w: int = 100,
        pixmap_h: int = 30,
    ) -> MagicMock:
        """Build a mock fitz page with given text blocks and a fake pixmap."""
        mock_pix = MagicMock()
        mock_pix.width = pixmap_w
        mock_pix.height = pixmap_h
        mock_pix.samples = pixmap_samples

        mock_page = MagicMock()
        mock_page.get_text.return_value = blocks
        mock_page.get_pixmap.return_value = mock_pix
        return mock_page

    def test_pix2tex_fallback_invoked_without_creds(self) -> None:
        """When no MathPix creds, pix2tex fallback is called."""
        math_block = (0.0, 0.0, 200.0, 50.0, "∑ α∫β∇γ∂δ∞ε", 0, 0)

        mock_settings = MagicMock()
        mock_settings.mathpix_app_id = None
        mock_settings.mathpix_app_key = None
        mock_settings.pix2tex_fallback = True
        mock_settings.figures_dir = "/tmp/figs"

        mock_model = MagicMock(return_value=r"\sum \alpha")

        mock_page = self._make_fitz_page([math_block])
        mock_doc = MagicMock()
        mock_doc.__len__.return_value = 1
        mock_doc.__getitem__ = MagicMock(return_value=mock_page)

        with (
            patch("codex.parsing.mathpix.get_settings", return_value=mock_settings),
            patch("codex.parsing.mathpix._get_pix2tex_model", return_value=mock_model),
            patch("fitz.open", return_value=mock_doc),
            patch("fitz.Matrix"),
            patch("fitz.Rect"),
        ):
            from PIL import Image

            with patch.object(Image, "frombytes", return_value=MagicMock()):
                from codex.parsing.mathpix import extract_formulas

                results = extract_formulas("paper.pdf")

        assert len(results) == 1
        assert results[0].raw_latex == r"\sum \alpha"

    def test_pix2tex_bbox_size_filter(self) -> None:
        """Blocks smaller than MIN_HEIGHT_PT or MIN_WIDTH_PT are skipped."""
        tiny_block = (0.0, 0.0, 10.0, 5.0, "∑∫∂∞∇ε", 0, 0)  # 10pt wide, 5pt tall

        mock_settings = MagicMock()
        mock_settings.mathpix_app_id = None
        mock_settings.mathpix_app_key = None
        mock_settings.pix2tex_fallback = True

        mock_model = MagicMock(return_value=r"\sum")
        mock_page = self._make_fitz_page([tiny_block])
        mock_doc = MagicMock()
        mock_doc.__len__.return_value = 1
        mock_doc.__getitem__ = MagicMock(return_value=mock_page)

        with (
            patch("codex.parsing.mathpix.get_settings", return_value=mock_settings),
            patch("codex.parsing.mathpix._get_pix2tex_model", return_value=mock_model),
            patch("fitz.open", return_value=mock_doc),
            patch("fitz.Matrix"),
            patch("fitz.Rect"),
        ):
            from codex.parsing.mathpix import extract_formulas

            results = extract_formulas("paper.pdf")

        assert results == []

    def test_pix2tex_math_ratio_filter(self) -> None:
        """Blocks with low math-character ratio are skipped."""
        plain_text_block = (0.0, 0.0, 200.0, 50.0, "This is just regular text with no math.", 0, 0)

        mock_settings = MagicMock()
        mock_settings.mathpix_app_id = None
        mock_settings.mathpix_app_key = None
        mock_settings.pix2tex_fallback = True

        mock_model = MagicMock(return_value=r"\alpha")
        mock_page = self._make_fitz_page([plain_text_block])
        mock_doc = MagicMock()
        mock_doc.__len__.return_value = 1
        mock_doc.__getitem__ = MagicMock(return_value=mock_page)

        with (
            patch("codex.parsing.mathpix.get_settings", return_value=mock_settings),
            patch("codex.parsing.mathpix._get_pix2tex_model", return_value=mock_model),
            patch("fitz.open", return_value=mock_doc),
            patch("fitz.Matrix"),
            patch("fitz.Rect"),
        ):
            from codex.parsing.mathpix import extract_formulas

            results = extract_formulas("paper.pdf")

        assert results == []

    def test_pix2tex_exception_swallowed(self) -> None:
        """If pix2tex raises, the block is skipped (no exception propagates)."""
        math_block = (0.0, 0.0, 200.0, 50.0, "∑ α∫β∇γ∂δ∞ε", 0, 0)

        mock_settings = MagicMock()
        mock_settings.mathpix_app_id = None
        mock_settings.mathpix_app_key = None
        mock_settings.pix2tex_fallback = True

        mock_model = MagicMock(side_effect=RuntimeError("GPU OOM"))
        mock_page = self._make_fitz_page([math_block])
        mock_doc = MagicMock()
        mock_doc.__len__.return_value = 1
        mock_doc.__getitem__ = MagicMock(return_value=mock_page)

        with (
            patch("codex.parsing.mathpix.get_settings", return_value=mock_settings),
            patch("codex.parsing.mathpix._get_pix2tex_model", return_value=mock_model),
            patch("fitz.open", return_value=mock_doc),
            patch("fitz.Matrix"),
            patch("fitz.Rect"),
        ):
            from PIL import Image

            with patch.object(Image, "frombytes", return_value=MagicMock()):
                from codex.parsing.mathpix import extract_formulas

                results = extract_formulas("paper.pdf")  # Must not raise

        assert results == []


class TestRouteSelection:
    """Tests for extract_formulas route selection logic."""

    def test_mathpix_selected_when_creds_present(self) -> None:
        """extract_formulas calls MathPix backend when both creds are set."""
        mock_settings = MagicMock()
        mock_settings.mathpix_app_id = "my_id"
        mock_settings.mathpix_app_key = "my_key"
        mock_settings.pix2tex_fallback = True

        with (
            patch("codex.parsing.mathpix.get_settings", return_value=mock_settings),
            patch(
                "codex.parsing.mathpix._extract_formulas_mathpix", return_value=[]
            ) as mock_mathpix,
            patch(
                "codex.parsing.mathpix._extract_formulas_pix2tex", return_value=[]
            ) as mock_pix2tex,
        ):
            from codex.parsing.mathpix import extract_formulas

            extract_formulas("paper.pdf")

        mock_mathpix.assert_called_once_with("paper.pdf", "my_id", "my_key")
        mock_pix2tex.assert_not_called()

    def test_pix2tex_selected_without_creds(self) -> None:
        """extract_formulas calls pix2tex when no MathPix creds are set."""
        mock_settings = MagicMock()
        mock_settings.mathpix_app_id = None
        mock_settings.mathpix_app_key = None
        mock_settings.pix2tex_fallback = True

        with (
            patch("codex.parsing.mathpix.get_settings", return_value=mock_settings),
            patch(
                "codex.parsing.mathpix._extract_formulas_mathpix", return_value=[]
            ) as mock_mathpix,
            patch(
                "codex.parsing.mathpix._extract_formulas_pix2tex", return_value=[]
            ) as mock_pix2tex,
        ):
            from codex.parsing.mathpix import extract_formulas

            extract_formulas("paper.pdf")

        mock_pix2tex.assert_called_once_with("paper.pdf")
        mock_mathpix.assert_not_called()

    def test_empty_string_creds_treated_as_absent(self) -> None:
        """Empty string credentials fall through to pix2tex (not MathPix)."""
        mock_settings = MagicMock()
        mock_settings.mathpix_app_id = ""
        mock_settings.mathpix_app_key = ""
        mock_settings.pix2tex_fallback = True

        with (
            patch("codex.parsing.mathpix.get_settings", return_value=mock_settings),
            patch(
                "codex.parsing.mathpix._extract_formulas_mathpix", return_value=[]
            ) as mock_mathpix,
            patch(
                "codex.parsing.mathpix._extract_formulas_pix2tex", return_value=[]
            ) as mock_pix2tex,
        ):
            from codex.parsing.mathpix import extract_formulas

            extract_formulas("paper.pdf")

        mock_pix2tex.assert_called_once()
        mock_mathpix.assert_not_called()

    def test_formula_extraction_disabled_when_fallback_false(self) -> None:
        """No extraction when pix2tex_fallback=False and no creds."""
        mock_settings = MagicMock()
        mock_settings.mathpix_app_id = None
        mock_settings.mathpix_app_key = None
        mock_settings.pix2tex_fallback = False

        with (
            patch("codex.parsing.mathpix.get_settings", return_value=mock_settings),
            patch(
                "codex.parsing.mathpix._extract_formulas_mathpix", return_value=[]
            ) as mock_mathpix,
            patch(
                "codex.parsing.mathpix._extract_formulas_pix2tex", return_value=[]
            ) as mock_pix2tex,
        ):
            from codex.parsing.mathpix import extract_formulas

            results = extract_formulas("paper.pdf")

        assert results == []
        mock_mathpix.assert_not_called()
        mock_pix2tex.assert_not_called()