81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
|
import os
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from llamacppctl.prompt_io import PromptSourceError, load_text_file
|
||
|
|
|
||
|
|
|
||
|
|
def test_load_utf8_file(tmp_path, policy):
|
||
|
|
p = tmp_path / "sys.txt"
|
||
|
|
p.write_text("Du bist präzise.", encoding="utf-8")
|
||
|
|
assert load_text_file(str(p), policy) == "Du bist präzise."
|
||
|
|
|
||
|
|
|
||
|
|
def test_strip_bom(tmp_path, policy):
|
||
|
|
p = tmp_path / "bom.txt"
|
||
|
|
p.write_bytes(b"\xef\xbb\xbfhello")
|
||
|
|
assert load_text_file(str(p), policy) == "hello"
|
||
|
|
|
||
|
|
|
||
|
|
def test_normalize_crlf(tmp_path, policy):
|
||
|
|
p = tmp_path / "crlf.txt"
|
||
|
|
p.write_bytes(b"line1\r\nline2\r\n")
|
||
|
|
assert load_text_file(str(p), policy) == "line1\nline2\n"
|
||
|
|
|
||
|
|
|
||
|
|
def test_reject_missing_file(tmp_path, policy):
|
||
|
|
with pytest.raises(PromptSourceError):
|
||
|
|
load_text_file(str(tmp_path / "nope.txt"), policy)
|
||
|
|
|
||
|
|
|
||
|
|
def test_reject_directory(tmp_path, policy):
|
||
|
|
d = tmp_path / "adir"
|
||
|
|
d.mkdir()
|
||
|
|
with pytest.raises(PromptSourceError):
|
||
|
|
load_text_file(str(d), policy)
|
||
|
|
|
||
|
|
|
||
|
|
def test_reject_binary_file(tmp_path, policy):
|
||
|
|
p = tmp_path / "bin.dat"
|
||
|
|
p.write_bytes(b"\x00\x01\x02binary")
|
||
|
|
with pytest.raises(PromptSourceError):
|
||
|
|
load_text_file(str(p), policy)
|
||
|
|
|
||
|
|
|
||
|
|
def test_reject_too_large_file(tmp_path, small_policy):
|
||
|
|
p = tmp_path / "big.txt"
|
||
|
|
p.write_text("x" * 100, encoding="utf-8")
|
||
|
|
with pytest.raises(PromptSourceError):
|
||
|
|
load_text_file(str(p), small_policy)
|
||
|
|
|
||
|
|
|
||
|
|
def test_reject_invalid_utf8(tmp_path, policy):
|
||
|
|
p = tmp_path / "invalid.txt"
|
||
|
|
p.write_bytes(b"\xff\xfe\xfa\xfb\x80\x81")
|
||
|
|
with pytest.raises(PromptSourceError):
|
||
|
|
load_text_file(str(p), policy)
|
||
|
|
|
||
|
|
|
||
|
|
def test_reject_symlink_when_disabled(tmp_path, policy):
|
||
|
|
target = tmp_path / "target.txt"
|
||
|
|
target.write_text("hi", encoding="utf-8")
|
||
|
|
link = tmp_path / "link.txt"
|
||
|
|
os.symlink(target, link)
|
||
|
|
|
||
|
|
from dataclasses import replace
|
||
|
|
|
||
|
|
strict_policy = replace(policy, allow_symlinks=False)
|
||
|
|
with pytest.raises(PromptSourceError):
|
||
|
|
load_text_file(str(link), strict_policy)
|
||
|
|
|
||
|
|
|
||
|
|
def test_allow_symlink_when_enabled(tmp_path, policy):
|
||
|
|
target = tmp_path / "target2.txt"
|
||
|
|
target.write_text("hi again", encoding="utf-8")
|
||
|
|
link = tmp_path / "link2.txt"
|
||
|
|
os.symlink(target, link)
|
||
|
|
|
||
|
|
from dataclasses import replace
|
||
|
|
|
||
|
|
permissive_policy = replace(policy, allow_symlinks=True)
|
||
|
|
assert load_text_file(str(link), permissive_policy) == "hi again"
|