diff --git a/src/documents/tests/test_caching.py b/src/documents/tests/test_caching.py index dccfc9a68..4ec1b2468 100644 --- a/src/documents/tests/test_caching.py +++ b/src/documents/tests/test_caching.py @@ -2,11 +2,14 @@ from concurrent.futures import ThreadPoolExecutor from threading import Event from threading import Lock +import pytest + from documents.caching import StoredLRUCache from documents.caching import retrieve_llm_suggestions from paperless.signed_pickle import HMAC_SIZE from paperless.signed_pickle import signed_pickle_dumps from paperless.signed_pickle import signed_pickle_loads +from paperless_ai.exceptions import LLMTimeoutError def test_lru_cache_entries() -> None: @@ -117,3 +120,63 @@ def test_llm_suggestions_are_generated_once_for_concurrent_requests(mocker) -> N assert calls == 1 mock_get_classification.assert_called_once_with(document, user, None) + + +def test_llm_suggestions_waiter_does_not_rerun_a_failed_generation(mocker) -> None: + """ + A request queued behind a generation that fails should give up, not take + its turn at re-running a query that just failed. + """ + generation_started = Event() + fail_generation = Event() + waiter_started = Event() + call_lock = Lock() + calls = 0 + document = mocker.Mock(pk=43) + user = mocker.Mock() + + def generate(*args) -> dict: + nonlocal calls + with call_lock: + calls += 1 + generation_started.set() + assert fail_generation.wait(timeout=2) + raise ValueError("Unknown model") + + def wait_for_generation(_interval: float) -> None: + waiter_started.set() + assert fail_generation.wait(timeout=2) + + mocker.patch( + "paperless_ai.ai_classifier.get_ai_document_classification", + side_effect=generate, + ) + mocker.patch("documents.caching.time.sleep", side_effect=wait_for_generation) + + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit( + retrieve_llm_suggestions, + document, + user, + None, + backend="ollama:model", + lock_timeout=10, + ) + assert generation_started.wait(timeout=2) + second = executor.submit( + retrieve_llm_suggestions, + document, + user, + None, + backend="ollama:model", + lock_timeout=10, + ) + assert waiter_started.wait(timeout=2) + fail_generation.set() + + with pytest.raises(ValueError, match="Unknown model"): + first.result(timeout=2) + with pytest.raises(LLMTimeoutError): + second.result(timeout=2) + + assert calls == 1