mrleo1nidandClaude Sonnet 5.5
f49ee90477
Includes earlier uncommitted work: Tomie and Unity adapters, benchmark, structured output, call metrics. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
347 lines
20 KiB
Python
347 lines
20 KiB
Python
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import httpx
|
|
from pydantic import ValidationError
|
|
|
|
from local_translator.api import create_app
|
|
from local_translator.config import ProviderConfig, Settings, load_settings
|
|
from local_translator.models import Project, ProtectionRule, TranslationRequest
|
|
from local_translator.protection import protect
|
|
from local_translator.provider import Provider, ProviderError, autostart
|
|
from local_translator.storage import Storage
|
|
|
|
|
|
class ProtectionTests(unittest.TestCase):
|
|
def test_preserve_common_game_syntax_and_unicode(self):
|
|
source = '<color=red>Hello {name}! ${quest} {{value}} %1$s %02d %NAME% [wait] &\r\n\\n\tПривет'
|
|
masked = protect(source, Project(name="demo"))
|
|
self.assertNotIn("{name}", masked.text)
|
|
self.assertEqual(masked.restore(masked.text.replace("Hello", "Привет")), source.replace("Hello", "Привет"))
|
|
|
|
def test_token_damage_and_invented_markup_rejected(self):
|
|
masked = protect("Hello {name}, %s", Project(name="demo"))
|
|
tokens = list(masked.tokens)
|
|
bad_values = [masked.text.replace(tokens[0], ""), masked.text + tokens[0], masked.text + " <b>",
|
|
masked.text.replace(tokens[0], "TEMP").replace(tokens[1], tokens[0]).replace("TEMP", tokens[1])]
|
|
for value in bad_values:
|
|
with self.subTest(value=value), self.assertRaises(ValueError):
|
|
masked.restore(value)
|
|
|
|
def test_custom_rule_overrides_builtin_and_disable(self):
|
|
project = Project(name="demo", disabled_rules=["square"], protection_rules=[ProtectionRule(name="engine", pattern=r"@[A-Z]+@")])
|
|
masked = protect("Hello @PLAYER@ [translate me]", project)
|
|
self.assertIn("[translate me]", masked.text)
|
|
self.assertNotIn("@PLAYER@", masked.text)
|
|
self.assertEqual(masked.restore(masked.text.replace("[translate me]", "[переведи меня]")), "Hello @PLAYER@ [переведи меня]")
|
|
|
|
def test_source_sentinel_is_preserved(self):
|
|
source = "Hello __LT_abc_0__ {name}"
|
|
masked = protect(source, Project(name="demo"))
|
|
self.assertEqual(masked.restore(masked.text), source)
|
|
|
|
def test_invalid_rules_rejected(self):
|
|
for pattern in ("[", ".*"):
|
|
with self.assertRaises(ValidationError):
|
|
ProtectionRule(name="bad", pattern=pattern)
|
|
with self.assertRaises(ValidationError):
|
|
Project(name="demo", disabled_rules=["typo"])
|
|
|
|
|
|
class APIEngineTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.settings = Settings(database=str(Path(self.temp.name) / "test.sqlite3"))
|
|
self.calls = []
|
|
self.requests = []
|
|
self.fail_ids = set()
|
|
self.fail_once = set()
|
|
|
|
def handler(request):
|
|
self.requests.append(request)
|
|
if request.url.path.endswith("/models"):
|
|
return httpx.Response(200, json={"data": [{"id": "test-hy"}]})
|
|
payload = json.loads(request.content)
|
|
content = payload["messages"][-1]["content"]
|
|
inputs = json.loads(content.split("[Source JSON]\n")[-1])["items"]
|
|
self.calls.append(inputs)
|
|
output = []
|
|
for item in inputs:
|
|
if item["text"] in self.fail_ids:
|
|
continue
|
|
if item["text"] in self.fail_once:
|
|
self.fail_once.remove(item["text"])
|
|
continue
|
|
output.append({"id": item["id"], "text": item["text"].replace("Hello", "Привет")})
|
|
return httpx.Response(200, json={"usage": {"prompt_tokens": 10, "completion_tokens": 4}, "choices": [{"finish_reason": "stop", "message": {"content": json.dumps({"items": output}, ensure_ascii=False)}}]})
|
|
|
|
self.transport = httpx.MockTransport(handler)
|
|
self.factory = lambda config: Provider(config, transport=self.transport)
|
|
self.app = create_app(self.settings, self.factory)
|
|
self.lifespan = self.app.router.lifespan_context(self.app)
|
|
await self.lifespan.__aenter__()
|
|
self.client = httpx.AsyncClient(transport=httpx.ASGITransport(app=self.app), base_url="http://test")
|
|
response = await self.client.post("/projects", json={"name": "demo", "context": "Fantasy game", "glossary": {"Moon": "Луна"}})
|
|
self.assertEqual(response.status_code, 201)
|
|
self.project = response.json()
|
|
|
|
async def asyncTearDown(self):
|
|
await self.client.aclose()
|
|
await self.lifespan.__aexit__(None, None, None)
|
|
self.temp.cleanup()
|
|
|
|
async def wait_job(self, identity):
|
|
for _ in range(250):
|
|
response = await self.client.get("/jobs/" + identity)
|
|
job = response.json()
|
|
if job["status"] not in {"queued", "running"}:
|
|
return job
|
|
await asyncio.sleep(0.02)
|
|
self.fail("Job did not finish")
|
|
|
|
async def submit(self, items, **options):
|
|
response = await self.client.post(f"/projects/{self.project['id']}/jobs", json={"items": items, **options})
|
|
self.assertEqual(response.status_code, 202, response.text)
|
|
return await self.wait_job(response.json()["id"])
|
|
|
|
async def test_batch_masking_cache_and_project_revision(self):
|
|
items = [{"id": str(i), "text": f"Hello {{player}} {i}"} for i in range(12)]
|
|
job = await self.submit(items)
|
|
self.assertEqual(job["status"], "completed")
|
|
self.assertEqual([len(batch) for batch in self.calls], [10, 2])
|
|
self.assertEqual(job["metrics"]["provider_calls"], 2)
|
|
self.assertEqual(job["metrics"]["prompt_tokens"], 20)
|
|
self.assertEqual(job["metrics"]["completion_tokens"], 8)
|
|
self.assertEqual(job["metrics"]["calls_with_completion_usage"], 2)
|
|
self.assertGreater(job["metrics"]["provider_seconds"], 0)
|
|
self.assertEqual(job["items"][0]["translation"], "Привет {player} 0")
|
|
self.assertNotIn("{player}", self.calls[0][0]["text"])
|
|
payload = json.loads(self.requests[1].content)
|
|
self.assertEqual([x["role"] for x in payload["messages"]], ["user"])
|
|
self.assertIn("Fantasy game", payload["messages"][0]["content"])
|
|
self.assertEqual(payload["top_p"], 0.6)
|
|
self.assertEqual(payload["response_format"]["type"], "json_schema")
|
|
self.assertEqual([item["id"] for item in self.calls[0]], [str(i) for i in range(10)])
|
|
cached = await self.submit(items)
|
|
self.assertTrue(all(item["cached"] for item in cached["items"]))
|
|
self.assertEqual(len(self.calls), 2)
|
|
updated = {key: value for key, value in self.project.items() if key not in {"id", "revision"}}
|
|
updated["context"] = "A new setting"
|
|
response = await self.client.put(f"/projects/{self.project['id']}", json=updated)
|
|
self.assertEqual(response.json()["revision"], 2)
|
|
new = await self.submit(items)
|
|
self.assertEqual(len(self.calls), 4)
|
|
self.assertTrue(all(not item["cached"] for item in new["items"]))
|
|
old = (await self.client.get("/jobs/" + job["id"])).json()
|
|
self.assertEqual(old["revision"], 1)
|
|
self.assertEqual(old["items"], job["items"])
|
|
|
|
async def test_only_missing_items_are_retried(self):
|
|
self.fail_once.add("Hello again")
|
|
job = await self.submit([{"id": "good", "text": "Hello"}, {"id": "bad", "text": "Hello again"}], retries=1)
|
|
self.assertEqual(job["status"], "completed")
|
|
self.assertEqual([[item["text"] for item in batch] for batch in self.calls], [["Hello", "Hello again"], ["Hello again"]])
|
|
self.assertEqual([item["attempts"] for item in job["items"]], [1, 2])
|
|
|
|
async def test_failure_is_explicit_and_retry_preserves_successes(self):
|
|
self.fail_ids.add("Hello again")
|
|
job = await self.submit([{"id": "ok", "text": "Hello"}, {"id": "bad", "text": "Hello again"}], retries=0)
|
|
self.assertEqual(job["status"], "completed_with_errors")
|
|
self.assertIsNone(job["items"][1]["translation"])
|
|
self.fail_ids.clear()
|
|
response = await self.client.post(f"/jobs/{job['id']}/retry")
|
|
self.assertEqual(response.status_code, 202)
|
|
retried = await self.wait_job(job["id"])
|
|
self.assertEqual(retried["status"], "completed")
|
|
self.assertEqual([item["text"] for item in self.calls[-1]], ["Hello again"])
|
|
self.assertEqual((await self.client.post(f"/jobs/{job['id']}/retry")).status_code, 409)
|
|
|
|
async def test_crash_recovery_keeps_completed_items_and_snapshot(self):
|
|
await self.app.state.engine.stop()
|
|
storage = self.app.state.storage
|
|
request = TranslationRequest(items=[{"id": "done", "text": "Hello"}, {"id": "pending", "text": "Hello again"}])
|
|
job = storage.create_job(self.project, self.settings.providers["lmstudio"].model_dump(), request.model_dump())
|
|
storage.item_result(job["id"], 0, "completed", "Сохранённый перевод")
|
|
storage.job_status(job["id"], "running", model="original-model")
|
|
self.app.state.engine.start()
|
|
recovered = await self.wait_job(job["id"])
|
|
self.assertEqual(recovered["status"], "completed")
|
|
self.assertEqual(recovered["model"], "original-model")
|
|
self.assertEqual(recovered["items"][0]["translation"], "Сохранённый перевод")
|
|
self.assertEqual([item["text"] for item in self.calls[0]], ["Hello again"])
|
|
|
|
async def test_cache_bypass_and_whitespace(self):
|
|
items = [{"id": "1", "text": "Hello"}, {"id": "2", "text": "\r\n "}]
|
|
await self.submit(items)
|
|
result = await self.submit(items, use_cache=False)
|
|
self.assertEqual(len(self.calls), 2)
|
|
self.assertEqual(result["items"][1]["translation"], "\r\n ")
|
|
|
|
async def test_exact_glossary_is_authoritative_and_preserves_syntax(self):
|
|
updated = {key: value for key, value in self.project.items() if key not in {"id", "revision"}}
|
|
updated["glossary"] = {"Q.LOAD": "БЫСТРАЯ ЗАГРУЗКА", "Hello {name}": "Привет"}
|
|
await self.client.put(f"/projects/{self.project['id']}", json=updated)
|
|
result = await self.submit([{"id": "label", "text": "Q.LOAD"}, {"id": "invalid", "text": "Hello {name}"}], retries=0)
|
|
self.assertEqual(result["items"][0]["translation"], "БЫСТРАЯ ЗАГРУЗКА")
|
|
self.assertEqual(result["items"][0]["attempts"], 0)
|
|
self.assertEqual(result["items"][1]["status"], "failed")
|
|
self.assertEqual(self.calls, [])
|
|
|
|
async def test_invalid_requests(self):
|
|
path = f"/projects/{self.project['id']}/jobs"
|
|
for payload in ({"items": []}, {"items": [{"id": "x", "text": "a"}, {"id": "x", "text": "b"}]}, {"items": [{"id": "x", "text": "a"}], "retries": 999}):
|
|
self.assertEqual((await self.client.post(path, json=payload)).status_code, 422)
|
|
self.assertEqual((await self.client.get("/jobs/missing")).status_code, 404)
|
|
self.assertEqual((await self.client.post("/projects", json={"name": "bad", "provider": "missing"})).status_code, 422)
|
|
|
|
async def test_cancel_queued_job_keeps_it_cancelled(self):
|
|
request = {"items": [{"id": "1", "text": "Hello"}], "batch_size": 1, "max_batch_chars": 1000, "retries": 0, "use_cache": False}
|
|
job = self.app.state.storage.create_job(self.project, self.settings.providers["lmstudio"].model_dump(), request)
|
|
response = await self.client.post(f"/jobs/{job['id']}/cancel")
|
|
self.assertEqual(response.json()["status"], "cancelled")
|
|
self.assertEqual(response.json()["items"][0]["error"], "Cancelled")
|
|
self.app.state.engine.wakeup.set()
|
|
await asyncio.sleep(0.3)
|
|
self.assertEqual((await self.client.get(f"/jobs/{job['id']}")).json()["status"], "cancelled")
|
|
self.assertEqual(self.calls, [])
|
|
self.assertEqual((await self.client.post(f"/jobs/{job['id']}/cancel")).json()["status"], "cancelled")
|
|
self.assertEqual((await self.client.get("/health")).json()["api_version"], 2)
|
|
|
|
async def test_synchronous_translation_and_openapi(self):
|
|
result = await self.client.post(f"/projects/{self.project['id']}/translate", json={"items": [{"id": "one", "text": "Hello"}]})
|
|
self.assertEqual(result.status_code, 200)
|
|
self.assertEqual(result.json()["items"][0]["translation"], "Привет")
|
|
self.assertEqual((await self.client.get("/openapi.json")).status_code, 200)
|
|
|
|
|
|
class ProviderTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_short_wire_ids_roundtrip_and_schema_are_explicit(self):
|
|
seen = []
|
|
def handler(request):
|
|
payload = json.loads(request.content)
|
|
seen.append(payload)
|
|
return httpx.Response(200, json={"choices": [{"message": {"content": json.dumps({"items": [
|
|
{"id": "1", "text": "Второй"}, {"id": "0", "text": "Первый"},
|
|
]})}}]})
|
|
provider = Provider(ProviderConfig(structured_output=True), transport=httpx.MockTransport(handler))
|
|
try:
|
|
result = await provider.translate("test", Project(name="demo"), [
|
|
{"id": "Naninovel.UI.Unusually.Long.Key.First", "text": "First"},
|
|
{"id": "arbitrary-client-key:second", "text": "Second"},
|
|
])
|
|
finally:
|
|
await provider.close()
|
|
self.assertEqual(result, {"arbitrary-client-key:second": "Второй", "Naninovel.UI.Unusually.Long.Key.First": "Первый"})
|
|
self.assertNotIn("Unusually.Long.Key", json.dumps(seen))
|
|
schema = seen[0]["response_format"]["json_schema"]["schema"]["properties"]["items"]
|
|
self.assertEqual(schema["minItems"], 2)
|
|
self.assertEqual(schema["maxItems"], 2)
|
|
self.assertEqual(schema["items"]["properties"]["id"]["enum"], ["0", "1"])
|
|
|
|
async def test_generic_provider_does_not_require_structured_output(self):
|
|
seen = []
|
|
def handler(request):
|
|
seen.append(json.loads(request.content))
|
|
return httpx.Response(200, json={"choices": [{"message": {"content": '{"items":[{"id":"0","text":"Привет"}]}'}}]})
|
|
provider = Provider(ProviderConfig(), transport=httpx.MockTransport(handler))
|
|
try:
|
|
result = await provider.translate("test", Project(name="demo"), [{"id": "original", "text": "Hello"}])
|
|
finally:
|
|
await provider.close()
|
|
self.assertEqual(result, {"original": "Привет"})
|
|
self.assertNotIn("response_format", seen[0])
|
|
|
|
async def test_autostart_reuses_already_available_identifier(self):
|
|
transport = httpx.MockTransport(lambda request: httpx.Response(200, json={"data": [{"id": "already-loaded"}]}))
|
|
factory = lambda config: Provider(config, transport=transport)
|
|
with patch("local_translator.provider.Provider", factory), patch("local_translator.provider.shutil.which") as which:
|
|
await autostart(ProviderConfig(model="already-loaded", load_model="local-file"))
|
|
which.assert_not_called()
|
|
|
|
async def test_malformed_duplicate_missing_and_truncated_responses(self):
|
|
cases = [
|
|
("not json", None, True),
|
|
(json.dumps({"items": [{"id": "0", "text": "x"}, {"id": "0", "text": "y"}, {"id": "1", "text": "z"}]}), None, False),
|
|
(json.dumps({"items": []}), "length", True),
|
|
]
|
|
for text, finish, should_raise in cases:
|
|
with self.subTest(text=text):
|
|
transport = httpx.MockTransport(lambda request: httpx.Response(200, json={"choices": [{"finish_reason": finish, "message": {"content": text}}]}))
|
|
provider = Provider(ProviderConfig(), transport=transport)
|
|
try:
|
|
if should_raise:
|
|
with self.assertRaises(ProviderError):
|
|
await provider.translate("test", Project(name="demo"), [{"id": "a", "text": "hello"}])
|
|
else:
|
|
self.assertEqual(await provider.translate("test", Project(name="demo"), [{"id": "a", "text": "hello"}, {"id": "b", "text": "world"}]), {"b": "z"})
|
|
finally:
|
|
await provider.close()
|
|
|
|
async def test_provider_errors_and_ambiguous_model(self):
|
|
for response in (httpx.Response(401, text="SECRET"), httpx.Response(200, json={"data": [{"id": "a"}, {"id": "b"}]})):
|
|
provider = Provider(ProviderConfig(), transport=httpx.MockTransport(lambda request: response))
|
|
try:
|
|
with self.assertRaises(ProviderError) as caught:
|
|
await provider.resolve_model()
|
|
self.assertNotIn("SECRET", str(caught.exception))
|
|
finally:
|
|
await provider.close()
|
|
|
|
|
|
class StorageConfigTests(unittest.TestCase):
|
|
def test_ui_dictionary_preserves_keys_and_value_colons(self):
|
|
from examples.extract_tomie_ui import parse_dictionary
|
|
rows = parse_dictionary('; comment\r\nSettings.Time: Time: 10:00\r\nTip: <color=red>Help</color>')
|
|
self.assertEqual(rows, [
|
|
{"id": "Settings.Time", "text": "Time: 10:00", "line": 2},
|
|
{"id": "Tip", "text": "<color=red>Help</color>", "line": 3},
|
|
])
|
|
with self.assertRaises(ValueError):
|
|
parse_dictionary("Key: A\nKey: B")
|
|
|
|
def test_config_paths_and_single_worker_lock(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "settings.toml"
|
|
path.write_text('database = "db/test.sqlite3"\n', encoding="utf-8")
|
|
settings = load_settings(str(path))
|
|
self.assertEqual(Path(settings.database).resolve(), (Path(directory) / "db/test.sqlite3").resolve())
|
|
storage = Storage(settings.database)
|
|
with storage.exclusive_worker():
|
|
with self.assertRaises(RuntimeError):
|
|
with Storage(settings.database).exclusive_worker():
|
|
self.fail("Two workers acquired the same database")
|
|
with storage.exclusive_worker():
|
|
pass
|
|
|
|
|
|
|
|
class CorsTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_only_configured_origins_may_call_from_a_browser(self):
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
settings = Settings(database=str(Path(temp) / "t.sqlite3"), allowed_origins=["https://games.example/"])
|
|
self.assertEqual(settings.allowed_origins, ["https://games.example"])
|
|
app = create_app(settings)
|
|
async with app.router.lifespan_context(app):
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
preflight = await client.options("/projects", headers={
|
|
"Origin": "https://games.example", "Access-Control-Request-Method": "POST",
|
|
"Access-Control-Request-Headers": "content-type", "Access-Control-Request-Private-Network": "true"})
|
|
self.assertEqual(preflight.headers["access-control-allow-origin"], "https://games.example")
|
|
self.assertEqual(preflight.headers["access-control-allow-private-network"], "true")
|
|
allowed = await client.get("/health", headers={"Origin": "https://games.example"})
|
|
self.assertEqual(allowed.headers["access-control-allow-origin"], "https://games.example")
|
|
other = await client.get("/health", headers={"Origin": "https://evil.example"})
|
|
self.assertNotIn("access-control-allow-origin", other.headers)
|
|
|
|
def test_wildcard_and_paths_are_rejected(self):
|
|
for value in ("*", "https://games.example/app", "games.example"):
|
|
with self.subTest(value=value), self.assertRaises(ValidationError):
|
|
Settings(allowed_origins=[value])
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|