qalmsw / tests /test_concurrency.py
pebaryan
Fan out LLM calls in grammar and reviewer checkers
7c804e7
Raw
History Blame Contribute Delete
2.64 kB
import time
from pathlib import Path
from qalmsw._concurrency import ordered_parallel_map
from qalmsw.checkers import GrammarChecker, ReviewerChecker
from qalmsw.document import Document
from qalmsw.parse import Paragraph
def test_ordered_parallel_map_preserves_order_with_jitter():
def slow(x: int) -> int:
time.sleep(0.02 if x % 2 == 0 else 0.0)
return x * 10
result = ordered_parallel_map(slow, [1, 2, 3, 4, 5], concurrency=4)
assert result == [10, 20, 30, 40, 50]
def test_ordered_parallel_map_serial_path_for_concurrency_one():
result = ordered_parallel_map(lambda x: x + 1, [1, 2, 3], concurrency=1)
assert result == [2, 3, 4]
def test_ordered_parallel_map_empty_input():
assert ordered_parallel_map(lambda x: x, [], concurrency=4) == []
class _ConcurrentLLM:
"""FakeLLM that returns a per-paragraph-indexed response so we can verify
that a parallel checker still maps each result back to the right paragraph."""
def __init__(self) -> None:
self.calls: list[str] = []
def complete_json(self, system: str, user: str) -> dict:
self.calls.append(user)
# A short sleep magnifies the probability that out-of-order completion
# would corrupt the result-to-paragraph mapping if the checker had a bug.
time.sleep(0.01)
token = user.strip().split()[-1]
return {
"issues": [{"excerpt": token, "message": f"flag-{token}", "severity": "info"}],
"comments": [{"message": f"comment-{token}", "severity": "info"}],
}
def _doc(paragraphs: list[Paragraph], source: str = "") -> Document:
return Document(path=Path("test.tex"), source=source, paragraphs=paragraphs)
def test_grammar_parallel_run_preserves_paragraph_association():
paras = [
Paragraph(text=f"The paragraph number is alpha{i}", start_line=i + 1, end_line=i + 1)
for i in range(6)
]
findings = GrammarChecker(_ConcurrentLLM(), concurrency=4).check(_doc(paras))
assert [f.message for f in findings] == [f"flag-alpha{i}" for i in range(6)]
assert [f.line for f in findings] == [i + 1 for i in range(6)]
def test_reviewer_parallel_run_preserves_section_association():
source = (
"\\begin{document}\n"
"\\section{One}\ntoken-one\n"
"\\section{Two}\ntoken-two\n"
"\\section{Three}\ntoken-three\n"
"\\end{document}\n"
)
findings = ReviewerChecker(_ConcurrentLLM(), concurrency=4).check(_doc([], source))
messages = [f.message for f in findings]
assert messages == ["comment-token-one", "comment-token-two", "comment-token-three"]