"""Tests for codex.parsing.grobid."""
from __future__ import annotations
from pathlib import Path
from unittest.mock import MagicMock, patch
from codex.parsing.grobid import extract_references, extract_structure
# Real GROBID TEI output uses type="arXiv" (camel-case), not "arxiv".
_TEI_XML = """\
Deep Learning for NLP
Jane
Doe
10.1234/nlp.2021
2101.00001
Attention Is All You Need
Ashish
Vaswani
Noam
Shazeer
10.5678/attention
"""
class _FakeResponse:
text = _TEI_XML
status_code = 200
def raise_for_status(self) -> None:
pass
def _mock_post(xml: str = _TEI_XML) -> MagicMock:
fake = MagicMock()
fake.text = xml
fake.status_code = 200
fake.raise_for_status = MagicMock()
return fake
class TestExtractReferences:
def test_returns_two_dicts(self, tmp_path: Path) -> None:
pdf = tmp_path / "paper.pdf"
pdf.write_bytes(b"%PDF fake")
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 = _mock_post()
refs = extract_references(str(pdf), grobid_url="http://localhost:8070")
assert len(refs) == 2
def test_all_keys_present(self, tmp_path: Path) -> None:
pdf = tmp_path / "paper.pdf"
pdf.write_bytes(b"%PDF fake")
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 = _mock_post()
refs = extract_references(str(pdf), grobid_url="http://localhost:8070")
required_keys = {"title", "authors", "year", "doi", "arxiv_id"}
for ref in refs:
assert required_keys == set(ref.keys()), f"Missing keys in {ref}"
def test_first_ref_values(self, tmp_path: Path) -> None:
pdf = tmp_path / "paper.pdf"
pdf.write_bytes(b"%PDF fake")
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 = _mock_post()
refs = extract_references(str(pdf), grobid_url="http://localhost:8070")
first = refs[0]
assert first["title"] == "Deep Learning for NLP"
assert "Jane Doe" in first["authors"]
assert first["year"] == "2021"
assert first["doi"] == "10.1234/nlp.2021"
assert first["arxiv_id"] == "2101.00001"
def test_second_ref_missing_arxiv_is_empty_string(self, tmp_path: Path) -> None:
pdf = tmp_path / "paper.pdf"
pdf.write_bytes(b"%PDF fake")
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 = _mock_post()
refs = extract_references(str(pdf), grobid_url="http://localhost:8070")
second = refs[1]
assert second["arxiv_id"] == ""
assert second["title"] == "Attention Is All You Need"
assert second["year"] == "2017"
class TestExtractStructure:
def test_returns_xml_string(self, tmp_path: Path) -> None:
pdf = tmp_path / "paper.pdf"
pdf.write_bytes(b"%PDF fake")
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 = _mock_post("raw xml")
result = extract_structure(str(pdf), grobid_url="http://localhost:8070")
assert result == "raw xml"
mock_client.post.assert_called_once()
assert "processFulltextDocument" in mock_client.post.call_args[0][0]