228 lines
9.5 KiB
Python
228 lines
9.5 KiB
Python
|
# This file is auto-generated by alef — DO NOT EDIT.
|
||
|
|
# alef:hash:4e15143f4af1ae8bafbdb1506ef057da924484c66a19483966333558ad437e75
|
||
|
|
# To regenerate: alef generate
|
||
|
|
# To verify freshness: alef verify --exit-code
|
||
|
|
# Issues & docs: https://github.com/kreuzberg-dev/alef
|
||
|
|
|
||
|
|
"""E2e tests for category: contract."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
import pytest # noqa: F401
|
||
|
|
from kreuzberg import extract_file, extract_file_sync, extract_bytes_sync, ExtractionConfig
|
||
|
|
|
||
|
|
|
||
|
|
def _alef_e2e_text(value: object) -> str:
|
||
|
|
return "" if value is None else str(value)
|
||
|
|
|
||
|
|
|
||
|
|
def _alef_e2e_item_texts(item: object) -> tuple[str, ...]:
|
||
|
|
raw_items = getattr(item, "items", None)
|
||
|
|
items_text = " ".join(str(value) for value in raw_items) if isinstance(raw_items, list) else ""
|
||
|
|
return (
|
||
|
|
_alef_e2e_text(item),
|
||
|
|
_alef_e2e_text(getattr(item, "kind", None)),
|
||
|
|
_alef_e2e_text(getattr(item, "name", None)),
|
||
|
|
_alef_e2e_text(getattr(item, "source", None)),
|
||
|
|
_alef_e2e_text(getattr(item, "alias", None)),
|
||
|
|
_alef_e2e_text(getattr(item, "text", None)),
|
||
|
|
_alef_e2e_text(getattr(item, "signature", None)),
|
||
|
|
items_text,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
|
||
|
|
async def test_api_batch_bytes_async() -> None:
|
||
|
|
"""Tests async batch bytes extraction API (batch_extract_bytes)."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
|
||
|
|
result = await extract_file(path, None, None)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
assert any(v in result.content for v in ["May 5, 2023", "Mallori"]) # noqa: S101
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
|
||
|
|
async def test_api_batch_bytes_with_configs_async() -> None:
|
||
|
|
"""Tests async batch bytes extraction with per-file configs (batch_extract_bytes with file_configs parameter)."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
config = ExtractionConfig(output_format="markdown")
|
||
|
|
|
||
|
|
result = await extract_file(path, None, config)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
# skipped: field 'metadata.output_format' not available on result type
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
|
||
|
|
async def test_api_batch_file_async() -> None:
|
||
|
|
"""Tests async batch file extraction API (batch_extract_file)."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
|
||
|
|
result = await extract_file(path, None, None)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
assert any(v in result.content for v in ["May 5, 2023", "Mallori"]) # noqa: S101
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
|
||
|
|
async def test_api_batch_file_with_configs_async() -> None:
|
||
|
|
"""Tests async batch file extraction with per-file configs (batch_extract_files with file_configs parameter)."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
config = ExtractionConfig(output_format="markdown")
|
||
|
|
|
||
|
|
result = await extract_file(path, None, config)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
# skipped: field 'metadata.output_format' not available on result type
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
|
||
|
|
async def test_api_extract_bytes_async() -> None:
|
||
|
|
"""Tests async bytes extraction API (extract_bytes)."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
|
||
|
|
result = await extract_file(path, None, None)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
assert any(v in result.content for v in ["May 5, 2023", "Mallori"]) # noqa: S101
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
|
||
|
|
async def test_api_extract_file_async() -> None:
|
||
|
|
"""Tests async file extraction API (extract_file)."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
|
||
|
|
result = await extract_file(path, None, None)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
assert any(v in result.content for v in ["May 5, 2023", "Mallori"]) # noqa: S101
|
||
|
|
|
||
|
|
|
||
|
|
def test_config_chunking_prepend_heading_context() -> None:
|
||
|
|
"""Tests markdown chunker prepends heading hierarchy to chunk content."""
|
||
|
|
path = "markdown/extraction_test.md"
|
||
|
|
config = ExtractionConfig(chunking={"chunker_type": "markdown", "max_chars": 300, "max_overlap": 50, "prepend_heading_context": True})
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
# skipped: field 'chunks' not available on result type
|
||
|
|
assert all(c.content for c in (result.chunks or [])) # noqa: S101
|
||
|
|
assert all(c.metadata and c.metadata.heading_context is not None for c in (result.chunks or [])) # noqa: S101
|
||
|
|
assert bool((result.chunks or []) and (result.chunks[0].metadata and result.chunks[0].metadata.heading_context)) # noqa: S101
|
||
|
|
|
||
|
|
|
||
|
|
def test_config_document_structure_with_headings() -> None:
|
||
|
|
"""Tests document structure with DOCX heading-driven nesting."""
|
||
|
|
path = "docx/fake.docx"
|
||
|
|
config = ExtractionConfig(include_document_structure=True)
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert result.mime_type.strip() == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" # noqa: S101
|
||
|
|
# skipped: field 'document' not available on result type
|
||
|
|
# skipped: field 'document.nodes' not available on result type
|
||
|
|
|
||
|
|
|
||
|
|
def test_config_element_types() -> None:
|
||
|
|
"""Tests element-based result format with element type assertions on DOCX."""
|
||
|
|
path = "docx/unit_test_headers.docx"
|
||
|
|
config = ExtractionConfig(result_format="element_based")
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert any(v in result.mime_type for v in ["application/vnd.openxmlformats-officedocument.wordprocessingml.document"]) # noqa: S101
|
||
|
|
# skipped: field 'elements' not available on result type
|
||
|
|
|
||
|
|
|
||
|
|
def test_config_extraction_timeout() -> None:
|
||
|
|
"""Tests that extraction_timeout_secs config field is accepted and does not affect fast extractions."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
config = ExtractionConfig(extraction_timeout_secs=300)
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
|
||
|
|
|
||
|
|
def test_config_keywords() -> None:
|
||
|
|
"""Tests keyword extraction via YAKE algorithm."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
config = ExtractionConfig(keywords={"algorithm": "yake", "max_keywords": 10})
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
# skipped: field 'keywords' not available on Python ExtractionResult
|
||
|
|
# skipped: field 'keywords' not available on Python ExtractionResult
|
||
|
|
|
||
|
|
|
||
|
|
def test_config_pages() -> None:
|
||
|
|
"""Tests page extraction and page marker configuration."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
config = ExtractionConfig(pages={"extract_pages": True, "insert_page_markers": True})
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
assert any(v in result.content for v in ["PAGE"]) # noqa: S101
|
||
|
|
|
||
|
|
|
||
|
|
def test_config_quality_enabled() -> None:
|
||
|
|
"""Tests quality scoring produces a score value in [0.0, 1.0]."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
config = ExtractionConfig(enable_quality_processing=True)
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
# skipped: field 'quality_score' not available on result type
|
||
|
|
# skipped: field 'quality_score' not available on result type
|
||
|
|
# skipped: field 'quality_score' not available on result type
|
||
|
|
|
||
|
|
|
||
|
|
def test_config_security_limits() -> None:
|
||
|
|
"""Tests archive extraction with custom security limits."""
|
||
|
|
path = "archives/documents.zip"
|
||
|
|
config = ExtractionConfig(security_limits={"max_archive_size": 104857600, "max_compression_ratio": 50, "max_files_in_archive": 100})
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert any(v in result.mime_type for v in ["application/zip", "application/x-zip-compressed"]) # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
|
||
|
|
|
||
|
|
def test_config_tree_sitter() -> None:
|
||
|
|
"""Tests tree-sitter configuration round-trip."""
|
||
|
|
path = "code/hello.py"
|
||
|
|
config = ExtractionConfig(tree_sitter={"groups": ["web"], "languages": ["python", "rust"], "process": {"comments": False, "diagnostics": False, "docstrings": False, "exports": True, "imports": True, "structure": True, "symbols": False}})
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert result.mime_type.strip() == "text/x-source-code" # noqa: S101
|
||
|
|
assert len(result.content) >= 5 # noqa: S101
|
||
|
|
|
||
|
|
|
||
|
|
def test_output_format_bytes_markdown() -> None:
|
||
|
|
"""Tests markdown output format via bytes extraction API."""
|
||
|
|
content = Path("pdf/fake_memo.pdf").read_bytes()
|
||
|
|
mime_type = "application/pdf"
|
||
|
|
config = ExtractionConfig(output_format="markdown")
|
||
|
|
|
||
|
|
result = extract_bytes_sync(content, mime_type, config)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
# skipped: field 'metadata.output_format' not available on result type
|
||
|
|
|
||
|
|
|
||
|
|
def test_output_format_markdown() -> None:
|
||
|
|
"""Tests Markdown output format."""
|
||
|
|
path = "pdf/fake_memo.pdf"
|
||
|
|
config = ExtractionConfig(output_format="markdown")
|
||
|
|
|
||
|
|
result = extract_file_sync(path, None, config)
|
||
|
|
assert result.mime_type.strip() == "application/pdf" # noqa: S101
|
||
|
|
assert len(result.content) >= 10 # noqa: S101
|
||
|
|
# skipped: field 'metadata.output_format' not available on result type
|