ss-tools/agent/tests/test_agent/test_document_parser.py

# #region TestAgentChat.DocumentParser [C:2] [TYPE Module] [SEMANTICS test,agent,document,parser]
# @BRIEF Tests for document parser — PDF, XLSX, unsupported formats, empty files, errors.
# @RELATION BINDS_TO -> [AgentChat.Document.Parser]
import sys
from pathlib import Path

sys.path.append(str(Path(__file__).parent.parent.parent / "src"))

import os
import tempfile
from unittest.mock import MagicMock, patch

import pytest

from ss_tools.agent.document_parser import parse_pdf, parse_xlsx, parse_upload, ParseError

# #region TestAgentChat.DocumentParser.PDF [C:2] [TYPE Function] [SEMANTICS test,parser,pdf]
# @BRIEF PDF parsing: extracts text from valid PDF, handles empty and encrypted PDFs gracefully.


def test_parse_pdf_valid():
    """Parse a valid PDF and verify text extraction."""
    # Create a minimal valid PDF
    pdf_content = (
        b"%PDF-1.4\n"
        b"1 0 obj<</Type/Catalog/Pages 2 0 R>>endobj\n"
        b"2 0 obj<</Type/Pages/Kids[3 0 R]/Count 1>>endobj\n"
        b"3 0 obj<</Type/Page/Parent 2 0 R/MediaBox[0 0 612 792]"
        b"/Contents 4 0 R/Resources<</Font<</F1 5 0 R>>>>>>endobj\n"
        b"4 0 obj<</Length 44>>stream\n"
        b"BT /F1 12 Tf 100 700 Td (Hello PDF) Tj ET\n"
        b"endstream\nendobj\n"
        b"5 0 obj<</Type/Font/Subtype/Type1/BaseFont/Helvetica>>endobj\n"
        b"xref\n"
        b"0 6\n"
        b"trailer<</Size 6/Root 1 0 R>>\n"
        b"startxref\n"
        b"169\n"
        b"%%EOF"
    )
    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
        f.write(pdf_content)
        tmp_path = f.name

    try:
        result = parse_pdf(tmp_path)
        assert isinstance(result, str)
        assert "Hello" in result
    finally:
        os.unlink(tmp_path)


def test_parse_pdf_empty():
    """Empty/non-existent PDF file raises ParseError."""
    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
        f.write(b"")
        tmp_path = f.name

    try:
        with pytest.raises(ParseError):
            parse_pdf(tmp_path)
    finally:
        os.unlink(tmp_path)


def test_parse_pdf_nonexistent():
    """Non-existent PDF file raises ParseError."""
    with pytest.raises((ParseError, FileNotFoundError)):
        parse_pdf("/tmp/nonexistent_file_12345.pdf")
# #endregion TestAgentChat.DocumentParser.PDF


# #region TestAgentChat.DocumentParser.XLSX [C:2] [TYPE Function] [SEMANTICS test,parser,xlsx]
# @BRIEF XLSX parsing: extracts sheet names and cell data.


def test_parse_xlsx_valid():
    """Parse a valid XLSX and verify sheet+cell extraction."""
    import openpyxl

    wb = openpyxl.Workbook()
    ws = wb.active
    ws.title = "Sheet1"
    ws["A1"] = "Name"
    ws["B1"] = "Value"
    ws["A2"] = "Test"
    ws["B2"] = 42

    with tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False) as f:
        wb.save(f.name)
        tmp_path = f.name

    try:
        result = parse_xlsx(tmp_path)
        assert isinstance(result, str)
        assert "Sheet1" in result
        assert "Name" in result
        assert "Value" in result
        assert "Test" in result
        assert "42" in result
    finally:
        os.unlink(tmp_path)


def test_parse_xlsx_empty_sheet():
    """XLSX with empty sheet returns headers but no data rows."""
    import openpyxl

    wb = openpyxl.Workbook()
    ws = wb.active
    ws.title = "EmptySheet"

    with tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False) as f:
        wb.save(f.name)
        tmp_path = f.name

    try:
        result = parse_xlsx(tmp_path)
        assert isinstance(result, str)
        assert "EmptySheet" in result
    finally:
        os.unlink(tmp_path)


def test_parse_xlsx_not_excel():
    """Non-XLSX file raises ParseError."""
    with tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False) as f:
        f.write(b"not an excel file")
        tmp_path = f.name

    try:
        with pytest.raises(ParseError):
            parse_xlsx(tmp_path)
    finally:
        os.unlink(tmp_path)
# #endregion TestAgentChat.DocumentParser.XLSX


# #region TestAgentChat.DocumentParser.ParseUpload [C:2] [TYPE Function] [SEMANTICS test,parser,upload]
# @BRIEF parse_upload dispatches to correct parser based on extension. Unsupported → error.


def test_parse_upload_txt():
    """Parse a .txt file returns its content."""
    with tempfile.NamedTemporaryFile(suffix=".txt", mode="w", delete=False) as f:
        f.write("Hello, world!")
        tmp_path = f.name

    try:
        result = parse_upload(tmp_path)
        assert "Hello" in result
    finally:
        os.unlink(tmp_path)


def test_parse_upload_json():
    """Parse a .json file returns its text."""
    with tempfile.NamedTemporaryFile(suffix=".json", mode="w", delete=False) as f:
        f.write('{"key": "value"}')
        tmp_path = f.name

    try:
        result = parse_upload(tmp_path)
        assert "key" in result
    finally:
        os.unlink(tmp_path)


def test_parse_upload_unsupported():
    """Unsupported format raises ParseError."""
    with tempfile.NamedTemporaryFile(suffix=".exe", delete=False) as f:
        f.write(b"binary")
        tmp_path = f.name

    try:
        with pytest.raises(ParseError, match="Unsupported format"):
            parse_upload(tmp_path)
    finally:
        os.unlink(tmp_path)


def test_parse_upload_dict():
    """parse_upload accepts dict with name+path (Gradio file format)."""
    with tempfile.NamedTemporaryFile(suffix=".txt", mode="w", delete=False) as f:
        f.write("Dict test")
        tmp_path = f.name

    try:
        result = parse_upload({"name": "test.txt", "path": tmp_path})
        assert "Dict test" in result
    finally:
        os.unlink(tmp_path)


def test_parse_upload_file_path_fallback():
    """parse_upload falls back to file_path key if path is missing."""
    with tempfile.NamedTemporaryFile(suffix=".txt", mode="w", delete=False) as f:
        f.write("Fallback key test")
        tmp_path = f.name

    try:
        result = parse_upload({"name": "test.txt", "file_path": tmp_path})
        assert "Fallback" in result
    finally:
        os.unlink(tmp_path)
# #endregion TestAgentChat.DocumentParser.ParseUpload


# #region TestAgentChat.DocumentParser.ImportErrors [C:2] [TYPE Function] [SEMANTICS test,parser,import,error]
# @BRIEF Test ImportError paths in parse_pdf and parse_xlsx.

def test_parse_pdf_import_error():
    """pdfplumber import fails → ParseError."""
    import builtins
    real_import = builtins.__import__

    def raise_on_pdfplumber(name, globals=None, locals=None, fromlist=(), level=0):
        if name == 'pdfplumber':
            raise ImportError(f"No module named pdfplumber", name=name)
        return real_import(name, globals, locals, fromlist, level)

    with patch('builtins.__import__', side_effect=raise_on_pdfplumber):
        with pytest.raises(ParseError, match="pdfplumber not installed"):
            parse_pdf("/tmp/dummy.pdf")


def test_parse_xlsx_import_error():
    """openpyxl import fails → ParseError."""
    import builtins
    real_import = builtins.__import__

    def raise_on_openpyxl(name, globals=None, locals=None, fromlist=(), level=0):
        if name == 'openpyxl':
            raise ImportError(f"No module named openpyxl", name=name)
        return real_import(name, globals, locals, fromlist, level)

    with patch('builtins.__import__', side_effect=raise_on_openpyxl):
        with pytest.raises(ParseError, match="openpyxl not installed"):
            parse_xlsx("/tmp/dummy.xlsx")
# #endregion TestAgentChat.DocumentParser.ImportErrors


# #region TestAgentChat.DocumentParser.PyPDF2Fallback [C:2] [TYPE Function] [SEMANTICS test,parser,pypdf2,fallback]
# @BRIEF When pdfplumber fails, PyPDF2 is used as fallback path.

def test_parse_pdf_pypdf2_fallback():
    """pdfplumber failure triggers PyPDF2 fallback."""
    # Build mock PyPDF2 module in sys.modules
    mock_pypdf2 = MagicMock()
    mock_reader = MagicMock()
    mock_page = MagicMock()
    mock_page.extract_text.return_value = "Fallback PyPDF2 text"
    mock_reader.pages = [mock_page]
    mock_pypdf2.PdfReader = MagicMock(return_value=mock_reader)

    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
        f.write(b"dummy pdf content")
        tmp_path = f.name

    try:
        # Patch pdfplumber.open directly (the real module) — NOT document_parser.pdfplumber
        # because pdfplumber is imported inside the function body, not at module level.
        with patch('pdfplumber.open', side_effect=Exception("pdfplumber error")), \
             patch.dict('sys.modules', {'PyPDF2': mock_pypdf2}):
            result = parse_pdf(tmp_path)
            assert isinstance(result, str)
            assert "PyPDF2" in result
            mock_pypdf2.PdfReader.assert_called_once()
    finally:
        os.unlink(tmp_path)
# #endregion TestAgentChat.DocumentParser.PyPDF2Fallback


# #region TestAgentChat.DocumentParser.ParseUploadBranches [C:2] [TYPE Function] [SEMANTICS test,parser,upload,branches]
# @BRIEF Test parse_upload dispatches to correct parser for PDF, XLSX, XLS.

def test_parse_upload_pdf_branch():
    """parse_upload with .pdf dispatches to parse_pdf."""
    import ss_tools.agent.document_parser as dp
    with patch.object(dp, 'parse_pdf', return_value="pdf result") as mock_parse:
        result = parse_upload({"name": "report.pdf", "path": "/tmp/report.pdf"})
        assert result == "pdf result"
        mock_parse.assert_called_once_with("/tmp/report.pdf")


def test_parse_upload_xlsx_branch():
    """parse_upload with .xlsx dispatches to parse_xlsx."""
    import ss_tools.agent.document_parser as dp
    with patch.object(dp, 'parse_xlsx', return_value="xlsx result") as mock_parse:
        result = parse_upload({"name": "data.xlsx", "path": "/tmp/data.xlsx"})
        assert result == "xlsx result"
        mock_parse.assert_called_once_with("/tmp/data.xlsx")


def test_parse_upload_xls_branch():
    """parse_upload with .xls (legacy) also dispatches to parse_xlsx."""
    import ss_tools.agent.document_parser as dp
    with patch.object(dp, 'parse_xlsx', return_value="xls result") as mock_parse:
        result = parse_upload({"name": "legacy.xls", "path": "/tmp/legacy.xls"})
        assert result == "xls result"
        mock_parse.assert_called_once_with("/tmp/legacy.xls")
# #endregion TestAgentChat.DocumentParser.ParseUploadBranches
# #endregion TestAgentChat.DocumentParser