307 lines
11 KiB
Python
307 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextvars
|
|
import os
|
|
import threading
|
|
from concurrent.futures import Future, ThreadPoolExecutor
|
|
from concurrent.futures.process import BrokenProcessPool
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from deerflow.utils import file_conversion
|
|
from deerflow.utils.file_io import FileIOPool
|
|
|
|
|
|
async def _wait_for_event(event: threading.Event) -> None:
|
|
async def wait():
|
|
while not event.is_set():
|
|
await asyncio.sleep(0.005)
|
|
|
|
await asyncio.wait_for(wait(), timeout=3)
|
|
|
|
|
|
@pytest.fixture
|
|
def conversion_pool(monkeypatch):
|
|
pool = FileIOPool(2, 2, "test-document-conversion")
|
|
monkeypatch.setattr(file_conversion, "_CONVERSION_IO", pool)
|
|
monkeypatch.setattr(file_conversion, "_get_pdf_converter", lambda: "auto")
|
|
# These tests coordinate parser entry/exit with in-process threading events.
|
|
# Spawn isolation has its own real-worker regression below.
|
|
monkeypatch.setattr(
|
|
file_conversion,
|
|
"_run_conversion_in_process",
|
|
lambda path, converter: file_conversion._convert_document_worker(str(path), converter),
|
|
)
|
|
yield pool
|
|
pool._pool.shutdown(wait=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parser_and_sidecar_write_run_off_event_loop(monkeypatch, tmp_path: Path, conversion_pool):
|
|
source = tmp_path / "report.docx"
|
|
source.write_bytes(b"source")
|
|
main_thread = threading.get_ident()
|
|
user_context = contextvars.ContextVar("conversion-test-user", default="missing")
|
|
token = user_context.set("alice")
|
|
observed = []
|
|
write_text = Path.write_text
|
|
|
|
def parse(_path, _converter):
|
|
observed.append(("parse", threading.get_ident(), user_context.get()))
|
|
return "# Complete report"
|
|
|
|
def write(path, text, **kwargs):
|
|
observed.append(("write", threading.get_ident(), user_context.get()))
|
|
return write_text(path, text, **kwargs)
|
|
|
|
monkeypatch.setattr(file_conversion, "_do_convert_unbounded", parse)
|
|
monkeypatch.setattr(Path, "write_text", write)
|
|
try:
|
|
result = await file_conversion.convert_file_to_markdown(source)
|
|
finally:
|
|
user_context.reset(token)
|
|
|
|
assert result == source.with_name("report.docx.md")
|
|
assert result.read_text(encoding="utf-8") == "# Complete report"
|
|
assert [operation for operation, _, _ in observed] == ["parse", "write"]
|
|
assert all(worker != main_thread and user == "alice" for _, worker, user in observed)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_saturated_conversion_does_not_use_default_executor(monkeypatch, tmp_path: Path, conversion_pool):
|
|
release = threading.Event()
|
|
saturated = threading.Event()
|
|
lock = threading.Lock()
|
|
active = peak = started = 0
|
|
|
|
def parse(_path, _converter):
|
|
nonlocal active, peak, started
|
|
with lock:
|
|
started += 1
|
|
active += 1
|
|
peak = max(peak, active)
|
|
if active == 2:
|
|
saturated.set()
|
|
try:
|
|
assert release.wait(5), "test did not release conversion workers"
|
|
return "converted"
|
|
finally:
|
|
with lock:
|
|
active -= 1
|
|
|
|
monkeypatch.setattr(file_conversion, "_do_convert_unbounded", parse)
|
|
# One general worker makes starvation observable without starting dozens of
|
|
# threads. Conversion must not consume it, even when its own pool is full.
|
|
default_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="test-default")
|
|
asyncio.get_running_loop().set_default_executor(default_pool)
|
|
jobs = [asyncio.create_task(file_conversion.extract_file_to_markdown(tmp_path / f"{index}.pdf")) for index in range(6)]
|
|
try:
|
|
await _wait_for_event(saturated)
|
|
assert await asyncio.wait_for(asyncio.to_thread(lambda: "ordinary request"), timeout=1) == "ordinary request"
|
|
await asyncio.sleep(0.02)
|
|
assert started == 2
|
|
finally:
|
|
release.set()
|
|
results = await asyncio.gather(*jobs)
|
|
default_pool.shutdown(wait=True)
|
|
|
|
assert results == ["converted"] * 6
|
|
assert peak == 2
|
|
assert not list(tmp_path.iterdir()), "extraction must not create sidecars"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("extract_only", [False, True])
|
|
async def test_cancelled_conversion_finishes_before_input_cleanup(monkeypatch, tmp_path: Path, conversion_pool, extract_only):
|
|
source = tmp_path / "report.pdf"
|
|
source.write_bytes(b"input must remain available")
|
|
started = threading.Event()
|
|
release = threading.Event()
|
|
finished = threading.Event()
|
|
|
|
def parse(path, _converter):
|
|
started.set()
|
|
assert release.wait(5), "test did not release conversion worker"
|
|
assert path.read_bytes() == b"input must remain available"
|
|
finished.set()
|
|
return "complete output"
|
|
|
|
monkeypatch.setattr(file_conversion, "_do_convert_unbounded", parse)
|
|
convert = file_conversion.extract_file_to_markdown if extract_only else file_conversion.convert_file_to_markdown
|
|
job = asyncio.create_task(convert(source))
|
|
try:
|
|
await _wait_for_event(started)
|
|
job.cancel()
|
|
await asyncio.sleep(0.01)
|
|
assert not job.done()
|
|
job.cancel() # A second cancellation must not release input ownership.
|
|
await asyncio.sleep(0.01)
|
|
assert not job.done()
|
|
finally:
|
|
release.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await job
|
|
|
|
assert finished.is_set()
|
|
sidecar = file_conversion.markdown_sidecar_path(source)
|
|
if extract_only:
|
|
assert not sidecar.exists()
|
|
else:
|
|
assert sidecar.read_text(encoding="utf-8") == "complete output"
|
|
source.unlink() # The caller may now remove the source safely.
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_queued_file_work_never_runs():
|
|
pool = FileIOPool(1, 1, "test-cancelled-queue")
|
|
started = threading.Event()
|
|
release = threading.Event()
|
|
cancelled = threading.Event()
|
|
ran_queued = threading.Event()
|
|
|
|
def blocked():
|
|
started.set()
|
|
assert release.wait(5)
|
|
|
|
first = asyncio.create_task(pool.run_until_complete(blocked))
|
|
queued = None
|
|
try:
|
|
await _wait_for_event(started)
|
|
queued = asyncio.create_task(pool.run_until_complete(ran_queued.set, on_cancel=cancelled.set))
|
|
await asyncio.sleep(0.01)
|
|
queued.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await queued
|
|
finally:
|
|
release.set()
|
|
await first
|
|
pool._pool.shutdown(wait=True)
|
|
|
|
assert not ran_queued.is_set()
|
|
assert not cancelled.is_set(), "a cancellation callback is only needed for submitted work"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_callback_drains_failed_worker_and_keeps_admission():
|
|
pool = FileIOPool(1, 1, "test-cooperative-cancel")
|
|
started = threading.Event()
|
|
release = threading.Event()
|
|
cancelled = threading.Event()
|
|
queued_ran = threading.Event()
|
|
|
|
def blocked():
|
|
started.set()
|
|
assert release.wait(5)
|
|
raise RuntimeError("worker failed while cancellation was pending")
|
|
|
|
first = asyncio.create_task(pool.run_until_complete(blocked, on_cancel=cancelled.set))
|
|
queued = None
|
|
try:
|
|
await _wait_for_event(started)
|
|
first.cancel()
|
|
await _wait_for_event(cancelled)
|
|
queued = asyncio.create_task(pool.run_until_complete(queued_ran.set))
|
|
await asyncio.sleep(0.01)
|
|
assert not queued_ran.is_set()
|
|
assert not first.done()
|
|
release.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await first
|
|
await queued
|
|
finally:
|
|
release.set()
|
|
await asyncio.gather(first, *([queued] if queued else []), return_exceptions=True)
|
|
pool._pool.shutdown(wait=True)
|
|
|
|
assert queued_ran.is_set()
|
|
|
|
|
|
def test_sync_conversion_produces_complete_sidecar(monkeypatch, tmp_path: Path, conversion_pool):
|
|
source = tmp_path / "slides.pptx"
|
|
source.write_bytes(b"slides")
|
|
monkeypatch.setattr(file_conversion, "_get_pdf_converter", lambda: "auto")
|
|
monkeypatch.setattr(file_conversion, "_do_convert_unbounded", lambda *_args: "# Presentation")
|
|
|
|
result = file_conversion.convert_file_to_markdown_sync(source)
|
|
|
|
assert result == source.with_name("slides.pptx.md")
|
|
assert result.read_text(encoding="utf-8") == "# Presentation"
|
|
|
|
|
|
def test_crashed_conversion_pool_is_replaced_without_retrying_document(monkeypatch, tmp_path: Path):
|
|
source = tmp_path / "report.pdf"
|
|
source.write_bytes(b"source")
|
|
pools = []
|
|
|
|
class FakePool:
|
|
def __init__(self, **kwargs):
|
|
assert kwargs["max_workers"] == 2
|
|
assert kwargs["mp_context"].get_start_method() == "spawn"
|
|
self.failed = not pools
|
|
self.calls = 0
|
|
self.closed = False
|
|
pools.append(self)
|
|
|
|
def submit(self, *_args):
|
|
self.calls += 1
|
|
future = Future()
|
|
if self.failed:
|
|
future.set_exception(BrokenProcessPool("native reader crashed"))
|
|
else:
|
|
future.set_result("recovered")
|
|
return future
|
|
|
|
def shutdown(self, **_kwargs):
|
|
self.closed = True
|
|
|
|
monkeypatch.setattr(file_conversion, "ProcessPoolExecutor", FakePool)
|
|
monkeypatch.setattr(file_conversion, "_PROCESS_POOL", None)
|
|
monkeypatch.setattr(file_conversion, "_get_pdf_converter", lambda: "auto")
|
|
|
|
assert file_conversion.convert_file_to_markdown_sync(source) is None
|
|
assert pools[0].calls == 1
|
|
assert pools[0].closed
|
|
assert file_conversion._PROCESS_POOL is None
|
|
assert not file_conversion.markdown_sidecar_path(source).exists()
|
|
|
|
result = file_conversion.convert_file_to_markdown_sync(source)
|
|
assert len(pools) == 2
|
|
assert result.read_text(encoding="utf-8") == "recovered"
|
|
file_conversion.shutdown_conversion_workers()
|
|
assert pools[1].closed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_spawn_conversion_is_isolated_and_bounds_returned_text(monkeypatch, tmp_path: Path):
|
|
source = tmp_path / "large.txt"
|
|
source.write_text("word " * (file_conversion.MAX_CONVERTED_MARKDOWN_CHARS // 5 + 100), encoding="utf-8")
|
|
monkeypatch.setattr(file_conversion, "_get_pdf_converter", lambda: "markitdown")
|
|
|
|
def reject_parent_parser(*_args):
|
|
raise AssertionError("document parsing must execute in a fresh process")
|
|
|
|
# Spawn imports a clean module; a parser executed in the gateway would hit
|
|
# this guard. This test intentionally uses the real installed text reader.
|
|
monkeypatch.setattr(file_conversion, "_do_convert_unbounded", reject_parent_parser)
|
|
pool = file_conversion._get_conversion_process_pool()
|
|
job = None
|
|
try:
|
|
child_pid = await asyncio.to_thread(lambda: pool.submit(os.getpid).result(timeout=15))
|
|
assert child_pid != os.getpid()
|
|
assert pool._mp_context.get_start_method() == "spawn"
|
|
job = asyncio.create_task(file_conversion.convert_file_to_markdown(source))
|
|
# The gateway can execute unrelated coroutine/default-pool work while
|
|
# the process imports and runs the real parser.
|
|
assert await asyncio.wait_for(asyncio.to_thread(lambda: "available"), timeout=2) == "available"
|
|
result = await asyncio.wait_for(job, timeout=30)
|
|
assert result == file_conversion.markdown_sidecar_path(source)
|
|
converted = result.read_text(encoding="utf-8")
|
|
assert len(converted) <= file_conversion.MAX_CONVERTED_MARKDOWN_CHARS
|
|
assert "deerflow-conversion-truncated" in converted
|
|
finally:
|
|
if job is not None:
|
|
await asyncio.gather(job, return_exceptions=True)
|
|
await asyncio.to_thread(file_conversion.shutdown_conversion_workers)
|