Merge remote-tracking branch 'origin/master' into fix/chunk-context-indexed-lookup
# Conflicts: # nextcloud_mcp_server/api/visualization.py # nextcloud_mcp_server/auth/viz_routes.py
This commit is contained in:
@@ -0,0 +1,269 @@
|
||||
"""Unit tests for Mistral provider."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from mistralai.client.errors import SDKError
|
||||
|
||||
from nextcloud_mcp_server.providers.mistral import (
|
||||
BATCH_SIZE,
|
||||
MISTRAL_EMBEDDING_DIMENSIONS,
|
||||
MistralProvider,
|
||||
_is_rate_limit,
|
||||
)
|
||||
|
||||
|
||||
def _make_data(embedding: list[float], index: int) -> MagicMock:
|
||||
"""Build a mock EmbeddingResponseData entry."""
|
||||
item = MagicMock()
|
||||
item.embedding = embedding
|
||||
item.index = index
|
||||
return item
|
||||
|
||||
|
||||
def _make_response(embeddings: list[list[float]]) -> MagicMock:
|
||||
"""Build a mock EmbeddingResponse with `embeddings` indexed in order."""
|
||||
response = MagicMock()
|
||||
response.data = [_make_data(emb, i) for i, emb in enumerate(embeddings)]
|
||||
return response
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_mistral_client(mocker):
|
||||
"""Mock the Mistral SDK constructor."""
|
||||
mock_client = MagicMock()
|
||||
mock_client.embeddings = MagicMock()
|
||||
mocker.patch(
|
||||
"nextcloud_mcp_server.providers.mistral.Mistral", return_value=mock_client
|
||||
)
|
||||
return mock_client
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_embedding_single(mock_mistral_client):
|
||||
"""Single text embed: round-trip through SDK with correct kwargs."""
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(
|
||||
return_value=_make_response([[0.1, 0.2, 0.3]])
|
||||
)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
embedding = await provider.embed("hello world")
|
||||
|
||||
assert embedding == [0.1, 0.2, 0.3]
|
||||
mock_mistral_client.embeddings.create_async.assert_awaited_once_with(
|
||||
model="mistral-embed",
|
||||
inputs=["hello world"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_embedding_batch_single_call(mock_mistral_client):
|
||||
"""Batch smaller than BATCH_SIZE issues a single API call."""
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(
|
||||
return_value=_make_response([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
|
||||
)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
embeddings = await provider.embed_batch(["a", "b", "c"])
|
||||
|
||||
assert embeddings == [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]
|
||||
assert mock_mistral_client.embeddings.create_async.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_embedding_batch_chunking(mock_mistral_client):
|
||||
"""Batches exceeding BATCH_SIZE are split into multiple API calls."""
|
||||
|
||||
# Each call returns one embedding per input it received; capture by side
|
||||
# effect so we can inspect lengths per chunk.
|
||||
def _side_effect(*, model, inputs, **_kwargs):
|
||||
return _make_response([[float(i)] for i in range(len(inputs))])
|
||||
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(side_effect=_side_effect)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
total = BATCH_SIZE * 2 + 5 # forces three chunks: 64, 64, 5 (with default)
|
||||
embeddings = await provider.embed_batch([f"text-{i}" for i in range(total)])
|
||||
|
||||
assert len(embeddings) == total
|
||||
assert mock_mistral_client.embeddings.create_async.await_count == 3
|
||||
# Verify the chunk sizes the SDK was actually called with.
|
||||
chunk_sizes = [
|
||||
len(call.kwargs["inputs"])
|
||||
for call in mock_mistral_client.embeddings.create_async.await_args_list
|
||||
]
|
||||
assert chunk_sizes == [BATCH_SIZE, BATCH_SIZE, 5]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_embedding_batch_order_preserved(mock_mistral_client):
|
||||
"""Out-of-order index in response data is sorted before returning."""
|
||||
response = MagicMock()
|
||||
response.data = [
|
||||
_make_data([0.3, 0.3], 2),
|
||||
_make_data([0.1, 0.1], 0),
|
||||
_make_data([0.2, 0.2], 1),
|
||||
]
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(return_value=response)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
embeddings = await provider.embed_batch(["x", "y", "z"])
|
||||
|
||||
assert embeddings == [[0.1, 0.1], [0.2, 0.2], [0.3, 0.3]]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_supports_capabilities(mock_mistral_client):
|
||||
"""Mistral provider advertises embeddings only."""
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
assert provider.supports_embeddings is True
|
||||
assert provider.supports_generation is False
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_generate_not_implemented(mock_mistral_client):
|
||||
"""generate() always raises NotImplementedError."""
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
with pytest.raises(NotImplementedError, match="does not support generation"):
|
||||
await provider.generate("test prompt")
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_get_dimension_known_model(mock_mistral_client):
|
||||
"""Known model: dimension available without an API call."""
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
assert provider.get_dimension() == MISTRAL_EMBEDDING_DIMENSIONS["mistral-embed"]
|
||||
mock_mistral_client.embeddings.create_async.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_get_dimension_unknown_model_detected(mock_mistral_client):
|
||||
"""Unknown model: dimension detected on first embed() call."""
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(
|
||||
return_value=_make_response([[0.1] * 768])
|
||||
)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="custom-mistral")
|
||||
|
||||
with pytest.raises(RuntimeError, match="not detected yet"):
|
||||
provider.get_dimension()
|
||||
|
||||
await provider.embed("test")
|
||||
assert provider.get_dimension() == 768
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_no_embeddings_disabled(mock_mistral_client):
|
||||
"""Setting embedding_model=None disables the embedding capability."""
|
||||
provider = MistralProvider(api_key="test-key", embedding_model=None)
|
||||
assert provider.supports_embeddings is False
|
||||
|
||||
with pytest.raises(NotImplementedError, match="no embedding_model configured"):
|
||||
await provider.embed("test")
|
||||
with pytest.raises(NotImplementedError, match="no embedding_model configured"):
|
||||
await provider.embed_batch(["test"])
|
||||
with pytest.raises(NotImplementedError, match="no embedding_model configured"):
|
||||
provider.get_dimension()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_empty_batch(mock_mistral_client):
|
||||
"""An empty batch returns [] without calling the API."""
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
assert await provider.embed_batch([]) == []
|
||||
mock_mistral_client.embeddings.create_async.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_close_no_error(mock_mistral_client):
|
||||
"""close() is best-effort and does not raise."""
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
# No __aexit__ on the mock by default → close() should silently no-op.
|
||||
await provider.close()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_base_url_passed_to_sdk(mocker):
|
||||
"""base_url is forwarded as server_url to the Mistral SDK constructor."""
|
||||
mock_ctor = mocker.patch(
|
||||
"nextcloud_mcp_server.providers.mistral.Mistral", return_value=MagicMock()
|
||||
)
|
||||
|
||||
MistralProvider(
|
||||
api_key="test-key",
|
||||
embedding_model="mistral-embed",
|
||||
base_url="https://example.com/mistral",
|
||||
)
|
||||
mock_ctor.assert_called_once_with(
|
||||
api_key="test-key",
|
||||
server_url="https://example.com/mistral",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_embed_raises_on_empty_response_data(mock_mistral_client):
|
||||
"""embed(): empty response.data triggers the defensive RuntimeError guard."""
|
||||
empty_response = MagicMock()
|
||||
empty_response.data = []
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(return_value=empty_response)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
with pytest.raises(RuntimeError, match="returned no embedding"):
|
||||
await provider.embed("test")
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_embed_raises_on_null_embedding(mock_mistral_client):
|
||||
"""embed(): a single response item with embedding=None is rejected."""
|
||||
null_item = MagicMock()
|
||||
null_item.embedding = None
|
||||
null_item.index = 0
|
||||
null_response = MagicMock()
|
||||
null_response.data = [null_item]
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(return_value=null_response)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
with pytest.raises(RuntimeError, match="returned no embedding"):
|
||||
await provider.embed("test")
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_batch_raises_on_null_embedding(mock_mistral_client):
|
||||
"""_embed_batch_request: a null embedding inside a batch raises explicitly."""
|
||||
good = _make_data([0.1, 0.2], 0)
|
||||
bad = MagicMock()
|
||||
bad.embedding = None
|
||||
bad.index = 1
|
||||
response = MagicMock()
|
||||
response.data = [good, bad]
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(return_value=response)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
with pytest.raises(RuntimeError, match="null embedding"):
|
||||
await provider.embed_batch(["a", "b"])
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_batch_raises_on_count_mismatch(mock_mistral_client):
|
||||
"""_embed_batch_request: fewer embeddings returned than inputs sent."""
|
||||
# Two inputs sent, one embedding returned.
|
||||
response = _make_response([[0.1, 0.2]])
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(return_value=response)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
with pytest.raises(RuntimeError, match="returned 1 embeddings for 2 inputs"):
|
||||
await provider.embed_batch(["a", "b"])
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_mistral_is_rate_limit_predicate():
|
||||
"""_is_rate_limit returns True only for SDKErrors with status_code == 429."""
|
||||
err_429 = MagicMock(spec=SDKError)
|
||||
err_429.status_code = 429
|
||||
err_500 = MagicMock(spec=SDKError)
|
||||
err_500.status_code = 500
|
||||
|
||||
assert _is_rate_limit(err_429) is True
|
||||
assert _is_rate_limit(err_500) is False
|
||||
# ValueError has no status_code attr → getattr returns None → False.
|
||||
assert _is_rate_limit(ValueError()) is False
|
||||
@@ -0,0 +1,136 @@
|
||||
"""Unit tests for ProviderRegistry — dynaconf-driven auto-detection."""
|
||||
|
||||
import pytest
|
||||
|
||||
from nextcloud_mcp_server.config import _reload_config
|
||||
from nextcloud_mcp_server.providers import (
|
||||
BedrockProvider,
|
||||
MistralProvider,
|
||||
OllamaProvider,
|
||||
OpenAIProvider,
|
||||
SimpleProvider,
|
||||
get_provider,
|
||||
reset_provider,
|
||||
)
|
||||
from nextcloud_mcp_server.providers.bedrock import BOTO3_AVAILABLE
|
||||
|
||||
|
||||
def _clear_provider_envs(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Strip every provider-selection env var so each test starts clean."""
|
||||
for name in (
|
||||
"AWS_REGION",
|
||||
"AWS_ACCESS_KEY_ID",
|
||||
"AWS_SECRET_ACCESS_KEY",
|
||||
"BEDROCK_EMBEDDING_MODEL",
|
||||
"BEDROCK_GENERATION_MODEL",
|
||||
"OPENAI_API_KEY",
|
||||
"OPENAI_BASE_URL",
|
||||
"OPENAI_EMBEDDING_MODEL",
|
||||
"OPENAI_GENERATION_MODEL",
|
||||
"MISTRAL_API_KEY",
|
||||
"MISTRAL_BASE_URL",
|
||||
"MISTRAL_EMBEDDING_MODEL",
|
||||
"OLLAMA_BASE_URL",
|
||||
"OLLAMA_EMBEDDING_MODEL",
|
||||
"OLLAMA_GENERATION_MODEL",
|
||||
"OLLAMA_VERIFY_SSL",
|
||||
"SIMPLE_EMBEDDING_DIMENSION",
|
||||
):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clean_provider_env(monkeypatch):
|
||||
"""Reset provider singleton + dynaconf cache around each test."""
|
||||
_clear_provider_envs(monkeypatch)
|
||||
reset_provider()
|
||||
_reload_config()
|
||||
yield monkeypatch
|
||||
reset_provider()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_registry_falls_back_to_simple(clean_provider_env):
|
||||
"""No provider env set → SimpleProvider (with default dimension)."""
|
||||
provider = get_provider()
|
||||
assert isinstance(provider, SimpleProvider)
|
||||
assert provider.get_dimension() == 384
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_registry_picks_simple_with_custom_dimension(clean_provider_env):
|
||||
"""SIMPLE_EMBEDDING_DIMENSION flows through dynaconf to SimpleProvider."""
|
||||
clean_provider_env.setenv("SIMPLE_EMBEDDING_DIMENSION", "512")
|
||||
_reload_config()
|
||||
|
||||
provider = get_provider()
|
||||
assert isinstance(provider, SimpleProvider)
|
||||
assert provider.get_dimension() == 512
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_registry_picks_mistral_when_api_key_set(clean_provider_env, mocker):
|
||||
"""MISTRAL_API_KEY alone is enough to select MistralProvider."""
|
||||
# MistralProvider eagerly constructs the SDK client in __init__; stub it
|
||||
# so the test doesn't depend on the SDK accepting arbitrary keys.
|
||||
mocker.patch("nextcloud_mcp_server.providers.mistral.Mistral")
|
||||
clean_provider_env.setenv("MISTRAL_API_KEY", "test-key")
|
||||
_reload_config()
|
||||
|
||||
provider = get_provider()
|
||||
assert isinstance(provider, MistralProvider)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_registry_picks_ollama_when_base_url_set(clean_provider_env, mocker):
|
||||
"""OLLAMA_BASE_URL selects OllamaProvider."""
|
||||
# OllamaProvider eagerly probes /api/tags in __init__; stub it out.
|
||||
mocker.patch(
|
||||
"nextcloud_mcp_server.providers.ollama.OllamaProvider._check_model_is_loaded"
|
||||
)
|
||||
clean_provider_env.setenv("OLLAMA_BASE_URL", "http://localhost:11434")
|
||||
_reload_config()
|
||||
|
||||
provider = get_provider()
|
||||
assert isinstance(provider, OllamaProvider)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_registry_openai_wins_over_mistral_and_ollama(clean_provider_env):
|
||||
"""OpenAI takes priority when multiple provider env vars are set."""
|
||||
clean_provider_env.setenv("OPENAI_API_KEY", "openai-key")
|
||||
clean_provider_env.setenv("MISTRAL_API_KEY", "mistral-key")
|
||||
clean_provider_env.setenv("OLLAMA_BASE_URL", "http://localhost:11434")
|
||||
_reload_config()
|
||||
|
||||
provider = get_provider()
|
||||
assert isinstance(provider, OpenAIProvider)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_registry_mistral_wins_over_ollama(clean_provider_env, mocker):
|
||||
"""Mistral takes priority over Ollama when both are configured."""
|
||||
# Stub the Mistral SDK constructor for the same reason as the sibling
|
||||
# picker test — keeps the registry test independent of SDK key validation.
|
||||
mocker.patch("nextcloud_mcp_server.providers.mistral.Mistral")
|
||||
clean_provider_env.setenv("MISTRAL_API_KEY", "mistral-key")
|
||||
clean_provider_env.setenv("OLLAMA_BASE_URL", "http://localhost:11434")
|
||||
_reload_config()
|
||||
|
||||
provider = get_provider()
|
||||
assert isinstance(provider, MistralProvider)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_registry_bedrock_wins_when_aws_region_set(clean_provider_env):
|
||||
"""AWS_REGION alone routes to Bedrock, even with other providers configured."""
|
||||
if not BOTO3_AVAILABLE:
|
||||
pytest.skip("boto3 not installed")
|
||||
|
||||
clean_provider_env.setenv("AWS_REGION", "us-east-1")
|
||||
clean_provider_env.setenv("OPENAI_API_KEY", "openai-key")
|
||||
clean_provider_env.setenv("MISTRAL_API_KEY", "mistral-key")
|
||||
_reload_config()
|
||||
|
||||
provider = get_provider()
|
||||
assert isinstance(provider, BedrockProvider)
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Unit tests for the shared rate-limit retry decorator."""
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from nextcloud_mcp_server.providers import _retry
|
||||
|
||||
|
||||
class _FakeError(Exception):
|
||||
"""Stand-in for an SDK exception with an HTTP status code attached."""
|
||||
|
||||
def __init__(self, status_code: int):
|
||||
super().__init__(f"status {status_code}")
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_real_sleep(monkeypatch):
|
||||
"""Replace anyio.sleep with an awaitable no-op so retries don't waste time."""
|
||||
monkeypatch.setattr(_retry.anyio, "sleep", AsyncMock(return_value=None))
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_retry_succeeds_after_429():
|
||||
"""A 429 followed by success returns the success value."""
|
||||
calls = {"n": 0}
|
||||
|
||||
@_retry.retry_on_rate_limit(
|
||||
_FakeError, is_rate_limit=lambda e: e.status_code == 429
|
||||
)
|
||||
async def flaky():
|
||||
calls["n"] += 1
|
||||
if calls["n"] < 3:
|
||||
raise _FakeError(429)
|
||||
return "ok"
|
||||
|
||||
result = await flaky()
|
||||
assert result == "ok"
|
||||
assert calls["n"] == 3
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_retry_reraises_non_rate_limit_immediately():
|
||||
"""A non-rate-limit error of the same class is re-raised on first hit."""
|
||||
calls = {"n": 0}
|
||||
|
||||
@_retry.retry_on_rate_limit(
|
||||
_FakeError, is_rate_limit=lambda e: e.status_code == 429
|
||||
)
|
||||
async def boom():
|
||||
calls["n"] += 1
|
||||
raise _FakeError(500)
|
||||
|
||||
with pytest.raises(_FakeError, match="status 500"):
|
||||
await boom()
|
||||
assert calls["n"] == 1 # No retries on non-429.
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_retry_gives_up_after_max_retries():
|
||||
"""After MAX_RETRIES failed attempts the last error is re-raised."""
|
||||
calls = {"n": 0}
|
||||
|
||||
@_retry.retry_on_rate_limit(
|
||||
_FakeError, is_rate_limit=lambda e: e.status_code == 429
|
||||
)
|
||||
async def always_429():
|
||||
calls["n"] += 1
|
||||
raise _FakeError(429)
|
||||
|
||||
with pytest.raises(_FakeError, match="status 429"):
|
||||
await always_429()
|
||||
assert calls["n"] == _retry.MAX_RETRIES
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_retry_default_predicate_treats_all_as_rate_limit():
|
||||
"""Default predicate (`lambda _: True`) retries every caught exception."""
|
||||
calls = {"n": 0}
|
||||
|
||||
@_retry.retry_on_rate_limit(_FakeError)
|
||||
async def fail_once():
|
||||
calls["n"] += 1
|
||||
if calls["n"] < 2:
|
||||
raise _FakeError(503)
|
||||
return "recovered"
|
||||
|
||||
result = await fail_once()
|
||||
assert result == "recovered"
|
||||
assert calls["n"] == 2
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_retry_does_not_catch_unrelated_exceptions():
|
||||
"""Exceptions of a different class bypass the decorator entirely."""
|
||||
|
||||
@_retry.retry_on_rate_limit(_FakeError)
|
||||
async def value_error():
|
||||
raise ValueError("nope")
|
||||
|
||||
with pytest.raises(ValueError, match="nope"):
|
||||
await value_error()
|
||||
@@ -0,0 +1,232 @@
|
||||
"""Unit tests for PDFHighlighter.compute_chunk_bboxes_batch (Deck #76).
|
||||
|
||||
Replaces the legacy `highlight_chunks_batch`-+-base64 pipeline that inflated
|
||||
Qdrant payloads with per-chunk PNG screenshots. The new path returns
|
||||
normalized bounding boxes only.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pymupdf
|
||||
import pytest
|
||||
|
||||
from nextcloud_mcp_server.search.pdf_highlighter import PDFHighlighter
|
||||
|
||||
|
||||
def _make_pdf(pages: list[str]) -> bytes:
|
||||
"""Build an in-memory PDF whose pages contain the given text."""
|
||||
doc = pymupdf.open()
|
||||
for body in pages:
|
||||
page = doc.new_page(width=595, height=842) # A4
|
||||
page.insert_text((50, 50), body)
|
||||
pdf_bytes = doc.tobytes()
|
||||
doc.close()
|
||||
return pdf_bytes
|
||||
|
||||
|
||||
def _page_boundaries(pages: list[str]) -> tuple[list[dict], str]:
|
||||
"""Build (page_boundaries, full_text) compatible with the highlighter API."""
|
||||
boundaries: list[dict] = []
|
||||
cursor = 0
|
||||
parts: list[str] = []
|
||||
for i, body in enumerate(pages, start=1):
|
||||
end = cursor + len(body)
|
||||
boundaries.append({"page": i, "start_offset": cursor, "end_offset": end})
|
||||
parts.append(body)
|
||||
cursor = end
|
||||
return boundaries, "".join(parts)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_compute_chunk_bboxes_returns_normalized_rects():
|
||||
"""Each returned bbox should be 4 floats in [0, 1] tagged with the page."""
|
||||
pages = [
|
||||
"Chapter 1: Introduction. Nextcloud is a self-hosted collaboration platform "
|
||||
"covering installation, configuration and maintenance topics.",
|
||||
"Chapter 2: Installation. Download the package, extract it to the web "
|
||||
"server directory, and configure the database connection.",
|
||||
]
|
||||
pdf_bytes = _make_pdf(pages)
|
||||
boundaries, full_text = _page_boundaries(pages)
|
||||
|
||||
chunks = [
|
||||
(
|
||||
0,
|
||||
0,
|
||||
len(pages[0]),
|
||||
1,
|
||||
"Chapter 1: Introduction. Nextcloud is a self-hosted collaboration platform.",
|
||||
),
|
||||
(
|
||||
1,
|
||||
len(pages[0]),
|
||||
len(pages[0]) + len(pages[1]),
|
||||
2,
|
||||
"Chapter 2: Installation. Download the package.",
|
||||
),
|
||||
]
|
||||
|
||||
results = PDFHighlighter.compute_chunk_bboxes_batch(
|
||||
pdf_bytes=pdf_bytes,
|
||||
chunks=chunks,
|
||||
page_boundaries=boundaries,
|
||||
full_text=full_text,
|
||||
)
|
||||
|
||||
assert set(results) == {0, 1}
|
||||
|
||||
bboxes_p1, page_p1 = results[0]
|
||||
bboxes_p2, page_p2 = results[1]
|
||||
|
||||
assert page_p1 == 1
|
||||
assert page_p2 == 2
|
||||
|
||||
for rects in (bboxes_p1, bboxes_p2):
|
||||
assert len(rects) >= 1
|
||||
for rect in rects:
|
||||
assert len(rect) == 4
|
||||
x0, y0, x1, y1 = rect
|
||||
assert 0.0 <= x0 < x1 <= 1.0
|
||||
assert 0.0 <= y0 < y1 <= 1.0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_compute_chunk_bboxes_empty_input():
|
||||
assert (
|
||||
PDFHighlighter.compute_chunk_bboxes_batch(
|
||||
pdf_bytes=b"",
|
||||
chunks=[],
|
||||
page_boundaries=[],
|
||||
full_text="",
|
||||
)
|
||||
== {}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_compute_chunk_bboxes_omits_when_offsets_out_of_range():
|
||||
"""Chunks whose offsets fall outside every page boundary are omitted.
|
||||
|
||||
Verifies the docstring contract: *"Chunks whose bbox cannot be located
|
||||
are omitted from the result."* (path: ``find_chunk_page`` returns None).
|
||||
"""
|
||||
pages = ["Page one body text content here for the test."]
|
||||
pdf_bytes = _make_pdf(pages)
|
||||
boundaries, full_text = _page_boundaries(pages)
|
||||
|
||||
# Offsets way beyond the document end — no page boundary matches.
|
||||
out_of_range_start = len(full_text) + 1000
|
||||
out_of_range_end = out_of_range_start + 50
|
||||
chunks = [(0, out_of_range_start, out_of_range_end, 1, "irrelevant")]
|
||||
|
||||
results = PDFHighlighter.compute_chunk_bboxes_batch(
|
||||
pdf_bytes=pdf_bytes,
|
||||
chunks=chunks,
|
||||
page_boundaries=boundaries,
|
||||
full_text=full_text,
|
||||
)
|
||||
|
||||
assert results == {}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_compute_chunk_bboxes_omits_when_text_not_in_pdf():
|
||||
"""Chunks whose page-relative text isn't on the page are omitted.
|
||||
|
||||
Verifies the second omission path: ``_find_chunk_bbox`` returns None
|
||||
when the supplied text cannot be located on the rendered page.
|
||||
"""
|
||||
pages = ["Hello world."]
|
||||
pdf_bytes = _make_pdf(pages)
|
||||
# Build boundaries from the real text but pass a *different* full_text
|
||||
# so the page-relative slice is content that does not exist in the PDF.
|
||||
boundaries, _ = _page_boundaries(pages)
|
||||
bogus_full_text = "Z" * len(pages[0])
|
||||
|
||||
chunks = [(0, 0, len(pages[0]), 1, "ignored")]
|
||||
|
||||
results = PDFHighlighter.compute_chunk_bboxes_batch(
|
||||
pdf_bytes=pdf_bytes,
|
||||
chunks=chunks,
|
||||
page_boundaries=boundaries,
|
||||
full_text=bogus_full_text,
|
||||
)
|
||||
|
||||
assert results == {}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.parametrize("page_index", [0, 1])
|
||||
def test_compute_chunk_bboxes_assigns_correct_page(page_index: int):
|
||||
"""Verify the page number returned matches the page the chunk lives on."""
|
||||
pages = [
|
||||
"Page one talks about apples and oranges in detail.",
|
||||
"Page two discusses bananas and grapes thoroughly.",
|
||||
]
|
||||
pdf_bytes = _make_pdf(pages)
|
||||
boundaries, full_text = _page_boundaries(pages)
|
||||
|
||||
if page_index == 0:
|
||||
chunk_text = "apples and oranges"
|
||||
offsets = (0, len(pages[0]))
|
||||
else:
|
||||
chunk_text = "bananas and grapes"
|
||||
offsets = (len(pages[0]), len(pages[0]) + len(pages[1]))
|
||||
|
||||
chunks = [(0, offsets[0], offsets[1], page_index + 1, chunk_text)]
|
||||
|
||||
results = PDFHighlighter.compute_chunk_bboxes_batch(
|
||||
pdf_bytes=pdf_bytes,
|
||||
chunks=chunks,
|
||||
page_boundaries=boundaries,
|
||||
full_text=full_text,
|
||||
)
|
||||
|
||||
assert results, "expected a bbox for the chunk"
|
||||
_, page_num = results[0]
|
||||
assert page_num == page_index + 1
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_compute_chunk_bboxes_handles_unordered_page_boundaries():
|
||||
"""Page lookup must match by ``page`` key, not by list position.
|
||||
|
||||
Regression guard: an earlier implementation indexed
|
||||
``page_boundaries[page_num - 1]``, which silently produces a wrong
|
||||
bbox if boundaries are passed out of order. Reverse the boundaries
|
||||
and assert the result is identical to the in-order case.
|
||||
"""
|
||||
pages = [
|
||||
"Page one talks about apples and oranges in detail.",
|
||||
"Page two discusses bananas and grapes thoroughly.",
|
||||
]
|
||||
pdf_bytes = _make_pdf(pages)
|
||||
boundaries, full_text = _page_boundaries(pages)
|
||||
|
||||
chunks = [
|
||||
(0, 0, len(pages[0]), 1, "apples and oranges"),
|
||||
(
|
||||
1,
|
||||
len(pages[0]),
|
||||
len(pages[0]) + len(pages[1]),
|
||||
2,
|
||||
"bananas and grapes",
|
||||
),
|
||||
]
|
||||
|
||||
in_order = PDFHighlighter.compute_chunk_bboxes_batch(
|
||||
pdf_bytes=pdf_bytes,
|
||||
chunks=chunks,
|
||||
page_boundaries=boundaries,
|
||||
full_text=full_text,
|
||||
)
|
||||
reversed_order = PDFHighlighter.compute_chunk_bboxes_batch(
|
||||
pdf_bytes=pdf_bytes,
|
||||
chunks=chunks,
|
||||
page_boundaries=list(reversed(boundaries)),
|
||||
full_text=full_text,
|
||||
)
|
||||
|
||||
assert in_order == reversed_order
|
||||
assert reversed_order[0][1] == 1
|
||||
assert reversed_order[1][1] == 2
|
||||
Reference in New Issue
Block a user