267 lines
14 KiB
Python
267 lines
14 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["id"] in self.fail_ids:
|
|
continue
|
|
if item["id"] in self.fail_once:
|
|
self.fail_once.remove(item["id"])
|
|
continue
|
|
output.append({"id": item["id"], "text": item["text"].replace("Hello", "Привет")})
|
|
return httpx.Response(200, json={"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["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)
|
|
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("bad")
|
|
job = await self.submit([{"id": "good", "text": "Hello"}, {"id": "bad", "text": "Hello again"}], retries=1)
|
|
self.assertEqual(job["status"], "completed")
|
|
self.assertEqual([[item["id"] for item in batch] for batch in self.calls], [["good", "bad"], ["bad"]])
|
|
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("bad")
|
|
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["id"] for item in self.calls[-1]], ["bad"])
|
|
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["id"] for item in self.calls[0]], ["pending"])
|
|
|
|
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_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_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": "a", "text": "x"}, {"id": "a", "text": "y"}, {"id": "b", "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
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|