portfolio-optimiser/tests/test_retrieval.py

86 lines
2.9 KiB
Python

"""Step 5 tests — local-folder retriever core (exact locators + EscapeRoute path-security).
Citations are exact by construction (locator slices the snippet byte-for-byte) and
deterministic; the path-security checks are behavioural (traversal / symlink-escape /
prefix-collision sibling all fail-closed). Pattern: tests/test_reference_domain.py +
tests/test_backends.py (raises).
"""
import os
import pytest
from portfolio_optimiser.retrieval import (
PathSecurityError,
is_within_dir,
retrieve,
safe_resolve,
)
@pytest.fixture()
def docs(tmp_path):
d = tmp_path / "docs"
d.mkdir()
(d / "asphalt.txt").write_text(
"Asphalt Ab11 unit rate renegotiation reduced the paving cost on the school stretch.",
encoding="utf-8",
)
(d / "drainage.txt").write_text(
"Drainage length was reduced after a survey on the low-traffic section.",
encoding="utf-8",
)
return d
def test_locator_exactly_slices_snippet(docs) -> None:
hits = retrieve("asphalt paving cost", str(docs), top_k=5)
assert hits
for h in hits:
text = (docs / h.file).read_text(encoding="utf-8")
assert text[h.locator.start_index : h.locator.end_index] == h.snippet
def test_retrieve_is_deterministic(docs) -> None:
a = retrieve("drainage survey", str(docs), top_k=3)
b = retrieve("drainage survey", str(docs), top_k=3)
assert [(h.file, h.locator) for h in a] == [(h.file, h.locator) for h in b]
def test_top_k_respected(docs) -> None:
assert len(retrieve("the", str(docs), top_k=1)) <= 1
def test_non_positive_top_k_rejected(docs) -> None:
with pytest.raises(ValueError):
retrieve("x", str(docs), top_k=0)
def test_parent_traversal_rejected(docs) -> None:
# A `..` traversal escaping the docs folder is rejected fail-closed.
with pytest.raises(PathSecurityError):
safe_resolve(str(docs), "../outside.txt")
def test_symlink_escape_is_not_read(tmp_path, docs) -> None:
# A symlink INSIDE docs_dir pointing OUTSIDE must never be read (fail-closed).
secret = tmp_path / "secret.txt"
secret.write_text("TOPSECRET asphalt asphalt asphalt exfiltrated", encoding="utf-8")
os.symlink(str(secret), str(docs / "link.txt"))
# safe_resolve refuses the escaping symlink...
with pytest.raises(PathSecurityError):
safe_resolve(str(docs), "link.txt")
# ...and retrieve never surfaces the outside content even on a matching query.
hits = retrieve("asphalt", str(docs), top_k=10)
assert all("TOPSECRET" not in h.snippet for h in hits)
def test_prefix_collision_sibling_rejected(tmp_path) -> None:
# A sibling dir sharing a name prefix (docs vs docs-evil) is NOT inside docs.
root = tmp_path / "docs"
root.mkdir()
sibling = tmp_path / "docs-evil"
sibling.mkdir()
intruder = sibling / "secret.txt"
intruder.write_text("x", encoding="utf-8")
assert is_within_dir(str(intruder), str(root)) is False