Spaces:
Running
Running
File size: 7,271 Bytes
5dd4a4e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | 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()
|