from __future__ import annotations import io import zipfile import pytest from astrbot.core.knowledge_base.parsers.epub_parser import EpubParser from astrbot.core.knowledge_base.parsers.markitdown_parser import MarkitdownParser from astrbot.core.knowledge_base.parsers.util import select_parser def _make_epub_bytes() -> bytes: buffer = io.BytesIO() with zipfile.ZipFile(buffer, mode="w") as archive: archive.writestr( "mimetype", "application/epub+zip", compress_type=zipfile.ZIP_STORED, ) archive.writestr( "META-INF/container.xml", """ """, ) archive.writestr( "OEBPS/nav.xhtml", """ Navigation """, ) archive.writestr( "OEBPS/chapter2.xhtml", """ Chapter 2

Second

Beta paragraph.

""", ) archive.writestr( "OEBPS/chapter1.xhtml", """ Chapter 1

First

Alpha paragraph.

""", ) archive.writestr( "OEBPS/content.opf", """ test-book Test Book en """, ) return buffer.getvalue() def _make_epub_bytes_with_generic_content() -> bytes: buffer = io.BytesIO() with zipfile.ZipFile(buffer, mode="w") as archive: archive.writestr( "mimetype", "application/epub+zip", compress_type=zipfile.ZIP_STORED, ) archive.writestr( "META-INF/container.xml", """ """, ) archive.writestr( "OEBPS/chapter1.xhtml", """ Chapter 1

First

Lead text

Piura*5, continued.

Inside div
Inside section
Cell A Cell B
""", ) archive.writestr( "OEBPS/content.opf", """ test-book Test Book en """, ) return buffer.getvalue() @pytest.mark.asyncio async def test_select_parser_supports_epub(): parser = await select_parser(".epub") assert isinstance(parser, EpubParser) @pytest.mark.asyncio @pytest.mark.parametrize("ext", [".rst", ".adoc"]) async def test_select_parser_supports_text_markup_formats(ext): parser = await select_parser(ext) assert isinstance(parser, MarkitdownParser) @pytest.mark.asyncio async def test_epub_parser_reads_spine_order_as_text(): result = await EpubParser().parse(_make_epub_bytes(), "book.epub") assert result.media == [] assert "**Title:**" not in result.text assert "[Chapter 2](chapter2.xhtml)" not in result.text assert result.text.startswith("1. Chapter 2") assert "2. Chapter 1" in result.text assert "Beta paragraph." in result.text assert "# First" in result.text assert "* Point A" in result.text assert result.text.index("1. Chapter 2") < result.text.index("## Second") assert result.text.index("## Second") < result.text.index("# First") @pytest.mark.asyncio async def test_epub_parser_preserves_generic_container_text(): result = await EpubParser().parse( _make_epub_bytes_with_generic_content(), "book.epub", ) assert "**Title:**" not in result.text assert "# First" in result.text assert "Lead text" in result.text assert r"Piura\*5, continued." in result.text assert "filepos" not in result.text assert r"[\*5]" not in result.text assert "Image00000.jpg" not in result.text assert "![](" not in result.text assert "\n\n\n" not in result.text assert "Inside div" in result.text assert "Inside section" in result.text assert "| Cell A | Cell B |" in result.text