tests/test_ingestion.py

import json
import os
import tempfile
from pathlib import Path
import sys

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

import pytest
from ingestion import (
    normalize_text,
    extract_text_pypdf,
    extract_text_pdfplumber,
    chunk_text,
    make_chunks,
    persist_chunks,
    ingest_pdf,
    Chunk
)


class TestTextNormalization:
    def test_removes_control_characters(self):
        text = "Hello\x00World\x0C\x1FTest"
        result = normalize_text(text)
        assert result == "Hello World Test"
    
    def test_normalizes_whitespace(self):
        text = "Multiple   spaces\n\n\nand\t\ttabs"
        result = normalize_text(text)
        assert result == "Multiple spaces and tabs"
    
    def test_handles_unicode_errors(self):
        text = "Valid text with émojis 😀"
        result = normalize_text(text)
        assert "Valid text" in result


class TestChunking:
    def test_basic_chunking(self):
        text = "a" * 1000
        chunks = chunk_text(text, chunk_size=100, overlap=10)
        assert len(chunks) == 11  # 1000 chars / ~90 effective chars per chunk
        assert all(len(c) <= 100 for c in chunks)
    
    def test_overlap_functionality(self):
        text = "0123456789" * 10  # 100 chars
        chunks = chunk_text(text, chunk_size=30, overlap=10)
        # Verify overlap exists
        for i in range(len(chunks) - 1):
            assert chunks[i][-10:] == chunks[i+1][:10]
    
    def test_empty_text(self):
        chunks = chunk_text("", chunk_size=100, overlap=10)
        assert chunks == []  # Empty text returns empty list
    
    def test_text_shorter_than_chunk_size(self):
        text = "Short text"
        chunks = chunk_text(text, chunk_size=100, overlap=10)
        assert len(chunks) == 1
        assert chunks[0] == text


class TestChunkPersistence:
    def test_persist_and_load_chunks(self):
        chunks = [
            Chunk(content="Test chunk 1", source="test.pdf", page=1),
            Chunk(content="Test chunk 2", source="test.pdf", page=2)
        ]
        session_id = "test_session_123"
        
        # Persist chunks
        output_path = persist_chunks(session_id, chunks)
        assert os.path.exists(output_path)
        
        # Load and verify
        with open(output_path, "r") as f:
            loaded = json.load(f)
        
        assert len(loaded) == 2
        assert loaded[0]["content"] == "Test chunk 1"
        assert loaded[1]["page"] == 2
        
        # Cleanup
        os.remove(output_path)


class TestPDFGeneration:
    @pytest.fixture
    def sample_pdf_path(self):
        """Generate a simple PDF for testing"""
        from reportlab.pdfgen import canvas
        from reportlab.lib.pagesizes import letter
        
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
            pdf_path = tmp.name
            
        c = canvas.Canvas(pdf_path, pagesize=letter)
        
        # Page 1 - Short story sample
        c.drawString(100, 750, "The Test Story")
        c.drawString(100, 700, "Once upon a time in a unit test,")
        c.drawString(100, 680, "there was a PDF that needed parsing.")
        c.drawString(100, 660, "It contained multiple pages and text.")
        c.showPage()
        
        # Page 2 - Technical content
        c.drawString(100, 750, "Technical Documentation")
        c.drawString(100, 700, "API Reference: parse_pdf(file_path: str)")
        c.drawString(100, 680, "Returns: List[Tuple[int, str]]")
        c.drawString(100, 660, "Raises: ValueError if file not found")
        c.showPage()
        
        c.save()
        
        yield pdf_path
        
        # Cleanup
        os.remove(pdf_path)
    
    def test_pypdf_extraction(self, sample_pdf_path):
        pages = extract_text_pypdf(sample_pdf_path)
        assert len(pages) == 2
        assert "Test Story" in pages[0][1]
        assert "Technical Documentation" in pages[1][1]
    
    @pytest.mark.skipif(not os.system("which pdfplumber > /dev/null 2>&1") == 0, 
                        reason="pdfplumber not installed")
    def test_pdfplumber_extraction(self, sample_pdf_path):
        pages = extract_text_pdfplumber(sample_pdf_path)
        assert len(pages) == 2
        assert "Test Story" in pages[0][1]
        assert "Technical Documentation" in pages[1][1]
    
    def test_full_ingestion_pipeline(self, sample_pdf_path):
        session_id = "test_full_pipeline"
        chunks, output_path = ingest_pdf(
            sample_pdf_path,
            session_id,
            parser="pypdf",
            chunk_size=200,
            overlap=20
        )
        
        # Verify chunks created
        assert len(chunks) > 0
        assert all(isinstance(c, Chunk) for c in chunks)
        
        # Verify persistence
        assert os.path.exists(output_path)
        with open(output_path, "r") as f:
            persisted = json.load(f)
        assert len(persisted) == len(chunks)
        
        # Cleanup
        os.remove(output_path)


class TestErrorHandling:
    def test_invalid_parser(self):
        with pytest.raises(ValueError, match="Unknown parser"):
            ingest_pdf("dummy.pdf", "session", parser="invalid_parser")
    
    def test_nonexistent_file(self):
        with pytest.raises(Exception):  # FileNotFoundError or similar
            extract_text_pypdf("/nonexistent/file.pdf")