Spaces:
Running
Running
| import json | |
| import unittest | |
| from unittest.mock import Mock, patch | |
| import requests | |
| import translator | |
| class TranslatorTests(unittest.TestCase): | |
| def test_index_preserves_order_and_does_not_call_legacy_providers(self): | |
| response = Mock() | |
| response.json.return_value = { | |
| "choices": [{ | |
| "finish_reason": "stop", | |
| "message": {"content": '{"2":"่็ผ็","0":"้ฟๅ","1":"้ฟๅ"}'}, | |
| }] | |
| } | |
| with ( | |
| patch.object(translator.requests, "post", return_value=response) as post, | |
| patch.object(translator, "INDEX_TRANSLATE_API_KEY", ""), | |
| patch.object(translator, "_translate_with_tencent_pool") as tencent, | |
| patch.object(translator, "_translate_with_baidu_pool") as baidu, | |
| ): | |
| result = translator.translate_texts(["long_hair", "long_hair", "blue_eyes"]) | |
| self.assertEqual(result, ["้ฟๅ", "้ฟๅ", "่็ผ็"]) | |
| tencent.assert_not_called() | |
| baidu.assert_not_called() | |
| self.assertEqual(post.call_args.args, (translator.INDEX_TRANSLATE_URL,)) | |
| payload = post.call_args.kwargs["json"] | |
| self.assertEqual(payload["model"], translator.INDEX_TRANSLATE_MODEL) | |
| self.assertEqual(payload["chat_template_kwargs"], {"enable_thinking": False}) | |
| self.assertEqual(payload["temperature"], 0) | |
| self.assertEqual( | |
| post.call_args.kwargs["headers"], | |
| {"User-Agent": "Index-Translate-Client/1.0"}, | |
| ) | |
| def test_index_batches_by_item_count_and_text_length(self): | |
| for texts, batch_sizes in [ | |
| ([f"tag_{index}" for index in range(33)], [32, 1]), | |
| (["a" * 1000, "b" * 1000, "c" * 1000], [2, 1]), | |
| ]: | |
| with self.subTest(batch_sizes=batch_sizes): | |
| responses = [] | |
| expected = [] | |
| for batch_number, size in enumerate(batch_sizes): | |
| translations = {str(index): f"่ฏๆ{batch_number}-{index}" for index in range(size)} | |
| expected.extend(translations.values()) | |
| response = Mock() | |
| response.json.return_value = { | |
| "choices": [{ | |
| "finish_reason": "stop", | |
| "message": {"content": json.dumps(translations)}, | |
| }] | |
| } | |
| responses.append(response) | |
| with ( | |
| patch.object(translator.requests, "post", side_effect=responses) as post, | |
| patch.object(translator, "INDEX_TRANSLATE_API_KEY", "test-key"), | |
| ): | |
| self.assertEqual(translator.translate_texts(texts), expected) | |
| self.assertEqual(post.call_count, 2) | |
| self.assertEqual(post.call_args.kwargs["headers"]["Authorization"], "Bearer test-key") | |
| sources = [ | |
| json.loads(call.kwargs["json"]["messages"][0]["content"].split("ใๆบๆใ\n")[1].split("\n\n")[0]) | |
| for call in post.call_args_list | |
| ] | |
| self.assertEqual([len(source) for source in sources], batch_sizes) | |
| self.assertEqual([value for source in sources for value in source.values()], texts) | |
| def test_invalid_index_responses_fall_back_without_misaligning_tags(self): | |
| cases = [ | |
| ("not json", "stop"), | |
| ('{"0":"้ฟๅ"}', "stop"), | |
| ('{"0":"้ฟๅ","0":"่็ผ็"}', "stop"), | |
| ('{"0":"้ฟๅ","1":"่็ผ็","1":"็ผ็"}', "stop"), | |
| ('{"0":"้ฟๅ","2":"่็ผ็"}', "stop"), | |
| ('{"0":"้ฟๅ","1":""}', "stop"), | |
| ('{"0":"้ฟๅ","1":null}', "stop"), | |
| ('{"0":"้ฟๅ","1":"่็ผ็"}', "length"), | |
| (None, "stop"), | |
| ] | |
| for content, finish_reason in cases: | |
| with self.subTest(content=content, finish_reason=finish_reason): | |
| response = Mock() | |
| response.json.return_value = { | |
| "choices": [{"finish_reason": finish_reason, "message": {"content": content}}] | |
| } | |
| with ( | |
| patch.object(translator.requests, "post", return_value=response), | |
| patch.object(translator, "_translate_with_tencent_pool", return_value=["่ พ่ฎฏ้ฟๅ", "่ พ่ฎฏ่็ผ็"]) as tencent, | |
| patch.object(translator, "_translate_with_baidu_pool") as baidu, | |
| ): | |
| result = translator.translate_texts(["long_hair", "blue_eyes"], "en", "zh") | |
| self.assertEqual(result, ["่ พ่ฎฏ้ฟๅ", "่ พ่ฎฏ่็ผ็"]) | |
| tencent.assert_called_once_with(["long_hair", "blue_eyes"], "en", "zh") | |
| baidu.assert_not_called() | |
| def test_index_http_and_transport_errors_use_existing_fallback_chain(self): | |
| for error in [requests.Timeout(), requests.HTTPError("429")]: | |
| with self.subTest(error=type(error).__name__): | |
| with ( | |
| patch.object(translator.requests, "post", side_effect=error), | |
| patch.object(translator, "_translate_with_tencent_pool", return_value=None) as tencent, | |
| patch.object(translator, "_translate_with_baidu_pool", return_value=["็พๅบฆ้ฟๅ"]) as baidu, | |
| ): | |
| self.assertEqual(translator.translate_texts(["long_hair"]), ["็พๅบฆ้ฟๅ"]) | |
| tencent.assert_called_once() | |
| baidu.assert_called_once() | |
| def test_later_index_batch_failure_retranslates_all_inputs_in_order(self): | |
| texts = [f"tag_{index}" for index in range(33)] | |
| response = Mock() | |
| response.json.return_value = { | |
| "choices": [{ | |
| "finish_reason": "stop", | |
| "message": {"content": json.dumps({str(index): "Index่ฏๆ" for index in range(32)})}, | |
| }] | |
| } | |
| expected = [f"่ พ่ฎฏ่ฏๆ{index}" for index in range(33)] | |
| with ( | |
| patch.object(translator.requests, "post", side_effect=[response, requests.Timeout()]), | |
| patch.object(translator, "_translate_with_tencent_pool", return_value=expected) as tencent, | |
| ): | |
| self.assertEqual(translator.translate_texts(texts), expected) | |
| tencent.assert_called_once_with(texts, "auto", "zh") | |
| def test_incomplete_legacy_results_return_originals_and_empty_inputs_do_not_call_api(self): | |
| for result in (None, [], ["็ผบๅฐไธ้กน"]): | |
| with self.subTest(result=result): | |
| with ( | |
| patch.object(translator.requests, "post", side_effect=requests.Timeout()), | |
| patch.object(translator, "_translate_with_tencent_pool", return_value=result), | |
| patch.object(translator, "_translate_with_baidu_pool", return_value=result), | |
| ): | |
| self.assertEqual(translator.translate_texts(["long_hair", "blue_eyes"]), ["long_hair", "blue_eyes"]) | |
| with patch.object(translator.requests, "post") as post: | |
| self.assertEqual(translator.translate_texts([]), []) | |
| post.assert_not_called() | |
| if __name__ == "__main__": | |
| unittest.main() | |