mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 15:37:55 +08:00
63 lines
2.4 KiB
Python
63 lines
2.4 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import importlib
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from plyngent.tools.truncate_token import TruncateToken, decode_truncate_token, encode_truncate_token
|
||
|
|
from tests.test_tools.helpers import call_async
|
||
|
|
|
||
|
|
get_truncated_module = importlib.import_module("plyngent.tools.file.get_truncated")
|
||
|
|
get_truncated = get_truncated_module.get_truncated
|
||
|
|
|
||
|
|
|
||
|
|
def _extract_token(out: str) -> str | None:
|
||
|
|
for line in out.splitlines():
|
||
|
|
if "TRUNCATE_TOKEN:" in line:
|
||
|
|
return line.split("[TRUNCATE_TOKEN: ", 1)[1].rstrip("]")
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
async def test_get_truncated_invalid_token() -> None:
|
||
|
|
out = await call_async(get_truncated, "garbage-token")
|
||
|
|
assert out.startswith("error:")
|
||
|
|
assert "invalid" in out
|
||
|
|
|
||
|
|
|
||
|
|
async def test_get_truncated_file_chains_without_gap(workspace: Path) -> None:
|
||
|
|
(workspace / "big.txt").write_text("a" * 1000, encoding="utf-8")
|
||
|
|
token = encode_truncate_token(TruncateToken(kind="file", location="big.txt", offset=0, limit=200))
|
||
|
|
total_read = 0
|
||
|
|
for _ in range(20):
|
||
|
|
out = await call_async(get_truncated, token, max_chars=200)
|
||
|
|
content, _, _ = out.partition("\n...[")
|
||
|
|
assert content == "a" * len(content)
|
||
|
|
total_read += len(content)
|
||
|
|
next_token = _extract_token(out)
|
||
|
|
if next_token is None:
|
||
|
|
break
|
||
|
|
parsed = decode_truncate_token(next_token)
|
||
|
|
assert parsed is not None
|
||
|
|
assert parsed.offset == total_read # no gap, no overlap
|
||
|
|
token = next_token
|
||
|
|
assert total_read == 1000
|
||
|
|
|
||
|
|
|
||
|
|
async def test_get_truncated_file_missing(workspace: Path) -> None:
|
||
|
|
token = encode_truncate_token(TruncateToken(kind="file", location="nope.txt", offset=0, limit=200))
|
||
|
|
out = await call_async(get_truncated, token, max_chars=200)
|
||
|
|
assert out.startswith("error:")
|
||
|
|
|
||
|
|
|
||
|
|
async def test_get_truncated_http_passthrough(monkeypatch) -> None:
|
||
|
|
calls: list[tuple[str, int, int]] = []
|
||
|
|
|
||
|
|
async def fake_fetch(url: str, *, offset: int, max_chars: int) -> str:
|
||
|
|
calls.append((url, offset, max_chars))
|
||
|
|
return f"body-from-{offset}"
|
||
|
|
|
||
|
|
monkeypatch.setattr(get_truncated_module.fetch, "handler", fake_fetch)
|
||
|
|
token = encode_truncate_token(TruncateToken(kind="http", location="https://example.com/x", offset=500, limit=400))
|
||
|
|
out = await call_async(get_truncated, token, max_chars=200)
|
||
|
|
assert out == "body-from-500"
|
||
|
|
assert calls == [("https://example.com/x", 500, 400)]
|