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 = '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 + " ", 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: Help') self.assertEqual(rows, [ {"id": "Settings.Time", "text": "Time: 10:00", "line": 2}, {"id": "Tip", "text": "Help", "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()