arborist/tests/test_wikipedia_xml.py

"""Tests for the MediaWiki XML and abstract dump sources.

Synthetic XML fixtures match the Phase IV (2006+) export schema:
  pages-articles.xml: <mediawiki><page><title/><ns/><id/><revision/></page>...
  abstract.xml:       <feed><doc><title/><url/><abstract/></doc>...
No real dumps are downloaded by the test suite.
"""

from __future__ import annotations

import bz2
from pathlib import Path

from arborist.ingest import ingest_source
from arborist.sources.wikipedia_xml import (
    WikipediaAbstractDump,
    WikipediaXmlDump,
    _strip_ns,
)
from arborist.store import connect


PAGES_ARTICLES_XML = """<?xml version="1.0" encoding="UTF-8"?>
<mediawiki xmlns="http://www.mediawiki.org/xml/export-0.10/" version="0.10" xml:lang="en">
  <siteinfo>
    <sitename>Wikipedia</sitename>
    <dbname>enwiki</dbname>
    <base>http://en.wikipedia.org/wiki/Main_Page</base>
    <generator>MediaWiki 1.36</generator>
  </siteinfo>
  <page>
    <title>Anarchism</title>
    <ns>0</ns>
    <id>12</id>
    <revision>
      <id>1234567</id>
      <timestamp>2010-11-08T13:00:00Z</timestamp>
      <text xml:space="preserve">'''Anarchism''' is a [[political philosophy]] that advocates [[stateless society|stateless societies]].</text>
    </revision>
  </page>
  <page>
    <title>Foo Redirect</title>
    <ns>0</ns>
    <id>13</id>
    <redirect title="Foo Target"/>
    <revision>
      <id>1234568</id>
      <timestamp>2010-11-08T13:00:01Z</timestamp>
      <text xml:space="preserve">#REDIRECT [[Foo Target]]</text>
    </revision>
  </page>
  <page>
    <title>Talk:Anarchism</title>
    <ns>1</ns>
    <id>14</id>
    <revision>
      <id>1234569</id>
      <timestamp>2010-11-08T13:00:02Z</timestamp>
      <text xml:space="preserve">Discussion about [[anarchism]].</text>
    </revision>
  </page>
  <page>
    <title>Capitalism</title>
    <ns>0</ns>
    <id>15</id>
    <revision>
      <id>2000001</id>
      <timestamp>2010-11-08T13:00:03Z</timestamp>
      <text xml:space="preserve">First version of capitalism article.</text>
    </revision>
    <revision>
      <id>2000002</id>
      <timestamp>2010-11-08T13:00:04Z</timestamp>
      <text xml:space="preserve">Second version of capitalism article. References [[markets]].</text>
    </revision>
  </page>
  <page>
    <title>Body Redirect Without Marker</title>
    <ns>0</ns>
    <id>16</id>
    <revision>
      <id>3000001</id>
      <timestamp>2010-11-08T13:00:05Z</timestamp>
      <text xml:space="preserve">  #redirect [[Some Other Page]]</text>
    </revision>
  </page>
</mediawiki>"""


ABSTRACT_XML = """<?xml version="1.0" encoding="UTF-8"?>
<feed>
  <doc>
    <title>Wikipedia: Anarchism</title>
    <url>http://en.wikipedia.org/wiki/Anarchism</url>
    <abstract>Anarchism is a political philosophy that advocates self-governed societies.</abstract>
    <links>
      <sublink linktype="nav">
        <anchor>Origins</anchor>
        <link>http://en.wikipedia.org/wiki/Anarchism#Origins</link>
      </sublink>
      <sublink linktype="nav">
        <anchor>Schools</anchor>
        <link>http://en.wikipedia.org/wiki/Anarchism#Schools</link>
      </sublink>
    </links>
  </doc>
  <doc>
    <title>Wikipedia: Capitalism</title>
    <url>http://en.wikipedia.org/wiki/Capitalism</url>
    <abstract>Capitalism is an economic system based on private ownership.</abstract>
  </doc>
  <doc>
    <title>Wikipedia: Empty</title>
    <url>http://en.wikipedia.org/wiki/Empty</url>
    <abstract></abstract>
  </doc>
</feed>"""


def _write(path: Path, text: str) -> Path:
    path.write_text(text, encoding="utf-8")
    return path


def _write_bz2(path: Path, text: str) -> Path:
    with bz2.open(path, "wt", encoding="utf-8") as f:
        f.write(text)
    return path


def test_strip_ns_removes_xmlns_braces():
    assert _strip_ns("{http://www.mediawiki.org/xml/export-0.10/}page") == "page"
    assert _strip_ns("page") == "page"


def test_xml_dump_yields_main_namespace_pages_only(tmp_path):
    src = WikipediaXmlDump(_write(tmp_path / "wp.xml", PAGES_ARTICLES_XML))
    docs = list(src.iter_documents())
    titles = [d.title for d in docs]
    # Anarchism in, Capitalism in (last revision), Talk:Anarchism out (ns=1),
    # Foo Redirect out (<redirect/>), Body Redirect Without Marker out (#REDIRECT body).
    assert titles == ["Anarchism", "Capitalism"]
    # URI shape matches the SQL source: spaces -> underscores, base prefix.
    uris = [d.uri for d in docs]
    assert uris == [
        "https://en.wikipedia.org/wiki/Anarchism",
        "https://en.wikipedia.org/wiki/Capitalism",
    ]


def test_xml_dump_extracts_wikilinks_as_edges(tmp_path):
    src = WikipediaXmlDump(_write(tmp_path / "wp.xml", PAGES_ARTICLES_XML))
    docs = list(src.iter_documents())
    anarchism = next(d for d in docs if d.title == "Anarchism")
    edge_uris = {e.dst_uri for e in anarchism.edges}
    assert "https://en.wikipedia.org/wiki/political_philosophy" in edge_uris
    # [[stateless society|stateless societies]] -> dst is target before pipe.
    assert "https://en.wikipedia.org/wiki/stateless_society" in edge_uris
    assert all(e.edge_type == "wikilink" for e in anarchism.edges)


def test_xml_dump_default_takes_last_revision_only(tmp_path):
    src = WikipediaXmlDump(_write(tmp_path / "wp.xml", PAGES_ARTICLES_XML))
    docs = list(src.iter_documents())
    capitalism = next(d for d in docs if d.title == "Capitalism")
    assert "Second version" in capitalism.content
    assert "First version" not in capitalism.content


def test_xml_dump_multi_revision_yields_every_revision(tmp_path):
    src = WikipediaXmlDump(
        _write(tmp_path / "wp.xml", PAGES_ARTICLES_XML),
        multi_revision=True,
    )
    docs = list(src.iter_documents())
    capitalism_docs = [d for d in docs if d.title == "Capitalism"]
    assert len(capitalism_docs) == 2
    contents = [d.content for d in capitalism_docs]
    assert any("First version" in c for c in contents)
    assert any("Second version" in c for c in contents)


def test_xml_dump_handles_bz2_streaming(tmp_path):
    src = WikipediaXmlDump(_write_bz2(tmp_path / "wp.xml.bz2", PAGES_ARTICLES_XML))
    docs = list(src.iter_documents())
    assert {d.title for d in docs} == {"Anarchism", "Capitalism"}


def test_xml_dump_shard_strides_pages(tmp_path):
    src_a = WikipediaXmlDump(
        _write(tmp_path / "wp.xml", PAGES_ARTICLES_XML), shard=(0, 2)
    )
    src_b = WikipediaXmlDump(tmp_path / "wp.xml", shard=(1, 2))
    a_titles = {d.title for d in src_a.iter_documents()}
    b_titles = {d.title for d in src_b.iter_documents()}
    # Stride is at the page index level (every page seen by every shard,
    # filtered post-namespace). Together they must equal the full set.
    assert a_titles | b_titles == {"Anarchism", "Capitalism"}
    assert not (a_titles & b_titles)


def test_xml_dump_resume_skips_pages_at_or_below_start_id(tmp_path):
    # First run: bring last_id to whatever it sees.
    src = WikipediaXmlDump(_write(tmp_path / "wp.xml", PAGES_ARTICLES_XML))
    list(src.iter_documents())
    high_water = src.last_id
    assert high_water > 0
    # Second run with start_id at high_water = no docs.
    src2 = WikipediaXmlDump(tmp_path / "wp.xml", start_id=high_water)
    assert list(src2.iter_documents()) == []


def test_xml_dump_round_trip_into_sqlite(tmp_path):
    src = WikipediaXmlDump(_write(tmp_path / "wp.xml", PAGES_ARTICLES_XML))
    db_path = tmp_path / "arborist.db"
    conn = connect(db_path)
    try:
        stats = ingest_source(conn, src)
    finally:
        conn.close()
    assert stats.inserted == 2
    # Re-ingest is idempotent.
    conn = connect(db_path)
    try:
        again = ingest_source(conn, WikipediaXmlDump(tmp_path / "wp.xml"))
    finally:
        conn.close()
    assert again.inserted == 0
    assert again.skipped_duplicate == 2


def test_abstract_dump_yields_one_doc_per_nonempty_abstract(tmp_path):
    src = WikipediaAbstractDump(_write(tmp_path / "abstract.xml", ABSTRACT_XML))
    docs = list(src.iter_documents())
    assert len(docs) == 2  # "Empty" doc has no abstract -> skipped
    titles = [d.title for d in docs]
    assert titles == ["Anarchism", "Capitalism"]
    # Wikipedia: prefix stripped from titles.
    assert all(not (t or "").startswith("Wikipedia:") for t in titles)
    # URL preserved verbatim.
    assert docs[0].uri == "http://en.wikipedia.org/wiki/Anarchism"


def test_abstract_dump_extracts_sublinks_as_edges(tmp_path):
    src = WikipediaAbstractDump(_write(tmp_path / "abstract.xml", ABSTRACT_XML))
    docs = list(src.iter_documents())
    anarchism = next(d for d in docs if d.title == "Anarchism")
    edge_uris = {e.dst_uri for e in anarchism.edges}
    assert "http://en.wikipedia.org/wiki/Anarchism#Origins" in edge_uris
    assert "http://en.wikipedia.org/wiki/Anarchism#Schools" in edge_uris
    assert all(e.edge_type == "abstract_link" for e in anarchism.edges)


def test_abstract_dump_round_trip_into_sqlite(tmp_path):
    src = WikipediaAbstractDump(_write(tmp_path / "abstract.xml", ABSTRACT_XML))
    db_path = tmp_path / "arborist.db"
    conn = connect(db_path)
    try:
        stats = ingest_source(conn, src)
    finally:
        conn.close()
    assert stats.inserted == 2