searchi/crates/searchi-core/tests/documents.rs

use anyhow::Result;
use lopdf::{
    content::{Content, Operation},
    dictionary, Document, Object, Stream,
};
use searchi_core::SearchIndex;
use std::{fs, io::Write, path::Path};

fn pdf(path: &Path) -> Result<()> {
    pdf_with_text(
        path,
        &[
            "First introductory page",
            "Supplier cancellation clause on second page",
        ],
    )
}

fn pdf_with_text(path: &Path, texts: &[&str]) -> Result<()> {
    let mut document = Document::with_version("1.5");
    let pages_id = document.new_object_id();
    let font = document.add_object(dictionary! {"Type"=>"Font", "Subtype"=>"Type1", "BaseFont"=>"Helvetica", "Encoding"=>"WinAnsiEncoding"});
    let resources = document.add_object(dictionary! {"Font"=>dictionary! {"F1"=>font}});
    let mut pages = Vec::new();
    for text in texts {
        let content = Content {
            operations: vec![
                Operation::new("BT", vec![]),
                Operation::new("Tf", vec![Object::Name(b"F1".to_vec()), 12.into()]),
                Operation::new("Td", vec![50.into(), 700.into()]),
                Operation::new("Tj", vec![Object::string_literal(*text)]),
                Operation::new("ET", vec![]),
            ],
        };
        let stream = document.add_object(Stream::new(dictionary! {}, content.encode()?));
        let page = document
            .add_object(dictionary! {"Type"=>"Page", "Parent"=>pages_id, "Contents"=>stream});
        pages.push(Object::Reference(page));
    }
    document.objects.insert(pages_id, Object::Dictionary(dictionary! {"Type"=>"Pages", "Kids"=>pages,
        "Count"=>texts.len() as i64, "Resources"=>resources, "MediaBox"=>vec![0.into(),0.into(),595.into(),842.into()]}));
    let catalog = document.add_object(dictionary! {"Type"=>"Catalog", "Pages"=>pages_id});
    document.trailer.set("Root", catalog);
    document.compress();
    document.save(path)?;
    Ok(())
}

#[test]
fn compressed_document_limits_are_enforced() -> Result<()> {
    let root = tempfile::tempdir()?;
    let text = "a".repeat(4 * 1024 * 1024 + 100);
    pdf_with_text(&root.path().join("bomb.pdf"), &[&text])?;
    let xml = format!(
        r#"<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>{}</w:t></w:r></w:p></w:body></w:document>"#,
        "b".repeat(2 * 1024 * 1024 + 1)
    );
    docx(&root.path().join("over-limit.docx"), &xml)?;
    let mut index = SearchIndex::open(":memory:")?;
    assert_eq!(index.index_folder(root.path())?.failures, 2);
    assert_eq!(index.stats()?.content_files, 0);
    assert!(index.search("bomb.pdf", 10)?[0].extraction_error.is_some());
    assert!(index.search("over-limit.docx", 10)?[0]
        .extraction_error
        .as_ref()
        .unwrap()
        .contains("exceeds"));
    Ok(())
}

#[test]
fn encrypted_pdf_is_explicitly_unsupported() -> Result<()> {
    let root = tempfile::tempdir()?;
    let path = root.path().join("protected.pdf");
    pdf(&path)?;
    let mut document = Document::load(&path)?;
    document.trailer.set(
        "ID",
        vec![
            Object::string_literal("fixture-id"),
            Object::string_literal("fixture-id"),
        ],
    );
    let state = lopdf::EncryptionState::try_from(lopdf::EncryptionVersion::V1 {
        document: &document,
        owner_password: "owner",
        user_password: "reader",
        permissions: lopdf::Permissions::all(),
    })?;
    document.encrypt(&state)?;
    document.save(&path)?;
    let mut index = SearchIndex::open(":memory:")?;
    assert_eq!(index.index_folder(root.path())?.failures, 1);
    assert!(index.search("protected.pdf", 10)?[0]
        .extraction_error
        .as_ref()
        .unwrap()
        .contains("encrypted"));
    assert!(index.search("cancellation", 10)?.is_empty());
    Ok(())
}

fn docx(path: &Path, xml: &str) -> Result<()> {
    let mut writer = zip::ZipWriter::new(fs::File::create(path)?);
    let options = zip::write::SimpleFileOptions::default()
        .compression_method(zip::CompressionMethod::Deflated);
    for (name, content) in [
        (
            "[Content_Types].xml",
            r#"<?xml version="1.0"?><Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types"><Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/><Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/></Types>"#,
        ),
        (
            "_rels/.rels",
            r#"<?xml version="1.0"?><Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships"><Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/></Relationships>"#,
        ),
        ("word/document.xml", xml),
    ] {
        writer.start_file(name, options)?;
        writer.write_all(content.as_bytes())?;
    }
    writer.finish()?;
    Ok(())
}

#[test]
fn real_pdf_and_docx_return_matching_locations_and_persist() -> Result<()> {
    let root = tempfile::tempdir()?;
    pdf(&root.path().join("agreement.pdf"))?;
    docx(
        &root.path().join("notes.docx"),
        r#"<?xml version="1.0" encoding="UTF-8"?><w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>Introduction</w:t></w:r></w:p><w:p/><w:p><w:r><w:t>Research &amp; </w:t></w:r><w:r><w:t>development budget</w:t></w:r></w:p></w:body></w:document>"#,
    )?;
    let db = root.path().join("index.db");
    let mut index = SearchIndex::open(&db)?;
    assert_eq!(index.index_folder(root.path())?.indexed, 2);
    let result = index.search("cancellation", 10)?;
    assert_eq!(result[0].source.as_ref().unwrap().kind, "page");
    assert_eq!(result[0].source.as_ref().unwrap().number, 2);
    assert!(result[0].snippet.contains("cancellation"));
    let result = index.search("development", 10)?;
    assert_eq!(result[0].source.as_ref().unwrap().kind, "paragraph");
    assert_eq!(result[0].source.as_ref().unwrap().number, 3);
    assert!(result[0].snippet.contains("Research & development"));
    let batch = index.passage_batch(0, 2)?;
    assert_eq!(batch.len(), 2);
    let remainder = index.passage_batch(batch[1].id, 256)?;
    assert_eq!(remainder.len(), 2);
    assert!(remainder
        .iter()
        .any(|passage| passage.source.kind == "paragraph" && passage.source.number == 3));
    assert!(index.contains_path(&result[0].path)?);
    assert!(!index.contains_path("unknown")?);
    drop(index);
    let index = SearchIndex::open(&db)?;
    assert_eq!(
        index.search("budget", 10)?[0]
            .source
            .as_ref()
            .unwrap()
            .number,
        3
    );
    Ok(())
}

#[test]
fn corrupt_document_reports_failure_and_preserves_filename_search() -> Result<()> {
    let root = tempfile::tempdir()?;
    fs::write(root.path().join("notes.txt"), "persistent evidence")?;
    let mut index = SearchIndex::open(":memory:")?;
    index.index_folder(root.path())?;
    fs::write(root.path().join("broken.pdf"), b"not a PDF")?;
    assert_eq!(index.index_folder(root.path())?.failures, 1);
    let result = index.search("broken.pdf", 10)?;
    assert!(result[0].extraction_error.as_ref().unwrap().contains("PDF"));
    assert_eq!(index.search("evidence", 10)?.len(), 1);
    fs::remove_file(root.path().join("broken.pdf"))?;
    docx(
        &root.path().join("broken.docx"),
        "<w:document><w:p><w:t>broken</w:p>",
    )?;
    assert_eq!(index.index_folder(root.path())?.failures, 1);
    assert!(index.search("broken.docx", 10)?[0]
        .extraction_error
        .is_some());
    pdf(&root.path().join("agreement.pdf"))?;
    index.index_folder(root.path())?;
    assert_eq!(index.search("cancellation", 10)?.len(), 1);
    fs::write(root.path().join("agreement.pdf"), "corrupted replacement")?;
    assert_eq!(index.index_folder(root.path())?.failures, 1);
    assert!(index.search("cancellation", 10)?.is_empty());
    assert!(index.search("agreement.pdf", 10)?[0]
        .extraction_error
        .is_some());
    Ok(())
}