ruwex 0.1.0

//! End-to-end tests: the same synthetic dump packaged as plain XML,
//! single-stream bz2, and multistream bz2 (with index) must produce
//! identical, ordered output.

use std::io::{self, Write};
use std::path::{Path, PathBuf};

use bzip2::Compression;
use bzip2::write::BzEncoder;
use ruwex::{DocSink, ExtractorConfig, OutputFormat, PageSource, ShardedWriter};

const HEADER: &str = r#"<mediawiki xmlns="http://www.mediawiki.org/xml/export-0.10/" version="0.10" xml:lang="en">
  <siteinfo>
    <sitename>Wikipedia</sitename>
    <base>https://en.wikipedia.org/wiki/Main_Page</base>
    <namespaces>
      <namespace key="0" />
      <namespace key="1">Talk</namespace>
      <namespace key="10">Template</namespace>
    </namespaces>
  </siteinfo>
"#;

fn page_xml(id: u64, ns: i32, title: &str, text: &str, redirect: Option<&str>) -> String {
    let redirect = redirect
        .map(|target| format!("    <redirect title=\"{target}\" />\n"))
        .unwrap_or_default();
    format!(
        "  <page>\n    <title>{title}</title>\n    <ns>{ns}</ns>\n    <id>{id}</id>\n{redirect}\
             <revision>\n      <id>{revid}</id>\n      <contributor><id>4242</id></contributor>\n\
         \x20     <text>{text}</text>\n    </revision>\n  </page>\n",
        revid = id * 100,
    )
}

/// Four pages: two articles, one redirect, one talk page.
fn pages() -> [String; 4] {
    [
        page_xml(1, 0, "Alpha", "Alpha is a letter. &amp; so on.", None),
        page_xml(2, 0, "Alpha (letter)", "#REDIRECT [[Alpha]]", Some("Alpha")),
        page_xml(3, 1, "Talk:Alpha", "Discussion about Alpha.", None),
        page_xml(4, 0, "Beta", "Beta follows Alpha — naturally.", None),
    ]
}

fn full_dump_xml() -> String {
    format!("{HEADER}{}</mediawiki>\n", pages().join(""))
}

fn bz2(bytes: &[u8]) -> Vec<u8> {
    let mut encoder = BzEncoder::new(Vec::new(), Compression::best());
    encoder.write_all(bytes).unwrap();
    encoder.finish().unwrap()
}

/// Builds a multistream dump (header, two 2-page streams, footer) plus its
/// index file, and returns the dump path.
fn write_multistream_dump(dir: &Path) -> PathBuf {
    let pages = pages();
    let header = bz2(HEADER.as_bytes());
    let stream_a = bz2(format!("{}{}", pages[0], pages[1]).as_bytes());
    let stream_b = bz2(format!("{}{}", pages[2], pages[3]).as_bytes());
    let footer = bz2(b"</mediawiki>\n");

    let offset_a = header.len() as u64;
    let offset_b = offset_a + stream_a.len() as u64;

    let dump_path = dir.join("test-pages-articles-multistream.xml.bz2");
    let dump = [header, stream_a, stream_b, footer].concat();
    std::fs::write(&dump_path, dump).unwrap();

    let index = format!(
        "{offset_a}:1:Alpha\n{offset_a}:2:Alpha (letter)\n{offset_b}:3:Talk:Alpha\n{offset_b}:4:Beta\n"
    );
    let index_path = dir.join("test-pages-articles-multistream-index.txt.bz2");
    std::fs::write(&index_path, bz2(index.as_bytes())).unwrap();

    dump_path
}

/// Collects all documents into a string, in write order.
#[derive(Default)]
struct StringSink(String);

impl DocSink for StringSink {
    fn write_doc(&mut self, doc: &str) -> io::Result<()> {
        self.0.push_str(doc);
        Ok(())
    }

    fn finish(&mut self) -> io::Result<()> {
        Ok(())
    }
}

fn extract_to_string(source: PageSource, config: &ExtractorConfig) -> (String, ruwex::Stats) {
    let mut sink = StringSink::default();
    let stats = ruwex::run(source, config, &mut sink).unwrap();
    (sink.0, stats)
}

fn config_with_workers(workers: usize) -> ExtractorConfig {
    ExtractorConfig {
        workers,
        ..Default::default()
    }
}

#[test]
fn plain_xml_extracts_articles_in_order() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("dump.xml");
    std::fs::write(&path, full_dump_xml()).unwrap();

    let (out, stats) = extract_to_string(PageSource::open(&path).unwrap(), &config_with_workers(3));

    assert_eq!(stats.pages, 4);
    assert_eq!(stats.docs, 2, "redirect and talk page are filtered out");
    let alpha = out
        .find("<doc id=\"1\" url=\"https://en.wikipedia.org/wiki?curid=1\" title=\"Alpha\">")
        .unwrap();
    let beta = out
        .find("<doc id=\"4\" url=\"https://en.wikipedia.org/wiki?curid=4\" title=\"Beta\">")
        .unwrap();
    assert!(alpha < beta, "dump order must be preserved");
    assert!(
        out.contains("Alpha is a letter. &amp; so on."),
        "entities are decoded"
    );
    assert!(!out.contains("Discussion"), "talk page excluded by default");
    assert!(!out.contains("#REDIRECT"), "redirects excluded");
}

#[test]
fn namespaces_flag_includes_talk_pages() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("dump.xml");
    std::fs::write(&path, full_dump_xml()).unwrap();

    let config = ExtractorConfig {
        namespaces: vec!["Talk".to_string()],
        ..config_with_workers(2)
    };
    let (out, stats) = extract_to_string(PageSource::open(&path).unwrap(), &config);

    assert_eq!(stats.docs, 3);
    let talk = out.find("title=\"Talk:Alpha\"").unwrap();
    let beta = out.find("title=\"Beta\"").unwrap();
    assert!(talk < beta, "dump order must be preserved");
}

#[test]
fn bz2_and_multistream_match_plain_xml_output() {
    let dir = tempfile::tempdir().unwrap();
    let config = config_with_workers(3);

    let plain_path = dir.path().join("dump.xml");
    std::fs::write(&plain_path, full_dump_xml()).unwrap();
    let (expected, _) = extract_to_string(PageSource::open(&plain_path).unwrap(), &config);

    let bz2_path = dir.path().join("dump.xml.bz2");
    std::fs::write(&bz2_path, bz2(full_dump_xml().as_bytes())).unwrap();
    let source = PageSource::open(&bz2_path).unwrap();
    assert!(matches!(source, PageSource::Sequential(_)));
    let (from_bz2, _) = extract_to_string(source, &config);
    assert_eq!(from_bz2, expected);

    let ms_path = write_multistream_dump(dir.path());
    let source = PageSource::open(&ms_path).unwrap();
    assert!(
        matches!(source, PageSource::Multistream(_)),
        "index file must trigger the multistream fast path"
    );
    let (from_multistream, stats) = extract_to_string(source, &config);
    assert_eq!(from_multistream, expected);
    assert_eq!(stats.pages, 4);
}

#[test]
fn json_lines_have_wikiextractor_fields() {
    let dir = tempfile::tempdir().unwrap();
    let path = write_multistream_dump(dir.path());

    let config = ExtractorConfig {
        format: OutputFormat::Json,
        ..config_with_workers(2)
    };
    let (out, _) = extract_to_string(PageSource::open(&path).unwrap(), &config);

    let lines: Vec<serde_json::Value> = out
        .lines()
        .map(|line| serde_json::from_str(line).unwrap())
        .collect();
    assert_eq!(lines.len(), 2);
    assert_eq!(lines[0]["id"], "1");
    assert_eq!(lines[0]["revid"], "100");
    assert_eq!(lines[0]["url"], "https://en.wikipedia.org/wiki?curid=1");
    assert_eq!(lines[0]["title"], "Alpha");
    assert_eq!(lines[0]["text"], "Alpha is a letter. &amp; so on.");
    assert_eq!(lines[1]["id"], "4");
}

/// Templates must load (in parallel) from a multistream dump and expand
/// identically to the plain-XML path.
#[test]
fn templates_load_and_expand_from_all_sources() {
    let dir = tempfile::tempdir().unwrap();

    let article = page_xml(1, 0, "Article", "Greeting: {{Hi|World}} end.", None);
    let template = page_xml(101, 10, "Template:Hi", "hello '''{{{1|nobody}}}'''", None);
    let body = format!("{article}{template}");
    let full = format!("{HEADER}{body}</mediawiki>\n");

    let plain_path = dir.path().join("dump.xml");
    std::fs::write(&plain_path, &full).unwrap();

    let ms_path = dir.path().join("t-pages-articles-multistream.xml.bz2");
    let header = bz2(HEADER.as_bytes());
    let stream_a = bz2(article.as_bytes());
    let stream_b = bz2(template.as_bytes());
    let offset_a = header.len() as u64;
    let offset_b = offset_a + stream_a.len() as u64;
    std::fs::write(
        &ms_path,
        [header, stream_a, stream_b, bz2(b"</mediawiki>\n")].concat(),
    )
    .unwrap();
    std::fs::write(
        dir.path()
            .join("t-pages-articles-multistream-index.txt.bz2"),
        bz2(format!("{offset_a}:1:Article\n{offset_b}:101:Template:Hi\n").as_bytes()),
    )
    .unwrap();

    let config = config_with_workers(2);
    let mut outputs = Vec::new();
    for path in [&plain_path, &ms_path] {
        let templates =
            std::sync::Arc::new(ruwex::expand::templates::load(Some(path), None, 2).unwrap());
        assert_eq!(templates.len(), 1, "template collected from {path:?}");
        let mut sink = StringSink::default();
        ruwex::run_with_templates(
            PageSource::open(path).unwrap(),
            &templates,
            &config,
            &mut sink,
        )
        .unwrap();
        outputs.push(sink.0);
    }
    assert!(
        outputs[0].contains("Greeting: hello World end."),
        "template did not expand: {}",
        outputs[0]
    );
    assert_eq!(
        outputs[0], outputs[1],
        "multistream and plain XML must agree"
    );

    // Lazy expansion (fetching the template through the title index, no bulk
    // database) must produce the same document as the bulk path.
    let index = ruwex::TitleIndex::open_or_build(&ms_path).unwrap();
    let article = index.find_page("Article").unwrap().unwrap();
    let site = index.site_info().unwrap();
    let lazy = ruwex::LazyTemplateSource::new(index).unwrap();
    let bulk =
        std::sync::Arc::new(ruwex::expand::templates::load(Some(&ms_path), None, 2).unwrap());
    let lazy_doc = ruwex::render_page(&article, &site, &config, &lazy);
    let bulk_doc = ruwex::render_page(&article, &site, &config, bulk.as_ref());
    assert!(
        lazy_doc.contains("Greeting: hello World end."),
        "lazy: {lazy_doc}"
    );
    assert_eq!(lazy_doc, bulk_doc, "lazy and bulk expansion must agree");
}

/// The title index must build beside the dump, be reused on reopen, and
/// resolve titles to the right page (including titles in later streams).
#[test]
fn title_index_builds_and_looks_up_pages() {
    use ruwex::TitleIndex;

    let dir = tempfile::tempdir().unwrap();
    let dump_path = write_multistream_dump(dir.path());

    let fst_path = TitleIndex::index_path_for(&dump_path);
    assert!(
        !fst_path.exists(),
        "index should not exist before first use"
    );

    let index = TitleIndex::open_or_build(&dump_path).unwrap();
    assert!(fst_path.exists(), "index built beside the dump");
    // redirects and talk pages are indexed too (the index is title->offset,
    // filtering happens at extraction time, not here)
    assert_eq!(index.len(), 4);

    // "Beta" lives in the second stream, exercising the seek path
    let beta = index.find_page("Beta").unwrap().expect("Beta present");
    assert_eq!(beta.id, 4);
    assert!(beta.text.contains("Beta follows Alpha"));

    let alpha = index.find_page("Alpha").unwrap().expect("Alpha present");
    assert_eq!(alpha.id, 1);

    // titles are normalized: underscores→spaces and surrounding whitespace
    // trimmed (internal spacing preserved, not collapsed)
    assert_eq!(index.find_page("_Alpha_").unwrap().unwrap().id, 1);
    assert_eq!(index.find_page("Alpha_(letter)").unwrap().unwrap().id, 2);
    assert_eq!(index.find_page("  Talk:Alpha ").unwrap().unwrap().id, 3);

    assert!(index.find_page("Nonexistent").unwrap().is_none());

    // reopening reuses the existing index (no rebuild)
    let reopened = TitleIndex::open_or_build(&dump_path).unwrap();
    assert_eq!(reopened.find_page("Talk:Alpha").unwrap().unwrap().id, 3);

    // correct URL comes from the dump's <siteinfo> header
    let site = index.site_info().unwrap();
    assert_eq!(site.page_url(1), "https://en.wikipedia.org/wiki?curid=1");
}

#[test]
fn sharded_writer_splits_output_files() {
    let dir = tempfile::tempdir().unwrap();
    let dump_path = dir.path().join("dump.xml");
    std::fs::write(&dump_path, full_dump_xml()).unwrap();
    let out_dir = dir.path().join("out");

    let config = config_with_workers(2);
    let mut sink = ShardedWriter::create(&out_dir, 120, false).unwrap();
    ruwex::run(PageSource::open(&dump_path).unwrap(), &config, &mut sink).unwrap();
    sink.finish().unwrap();

    let wiki_00 = std::fs::read_to_string(out_dir.join("AA").join("wiki_00")).unwrap();
    let wiki_01 = std::fs::read_to_string(out_dir.join("AA").join("wiki_01")).unwrap();
    assert!(wiki_00.contains("title=\"Alpha\""));
    assert!(wiki_01.contains("title=\"Beta\""));

    // Concatenated shards must equal the unsplit extraction.
    let (expected, _) = extract_to_string(PageSource::open(&dump_path).unwrap(), &config);
    assert_eq!(format!("{wiki_00}{wiki_01}"), expected);
}