pdf2zh/test/test_translator.py

import unittest
from pdf2zh import cache
import threading
import multiprocessing
import random
import string


class TestCache(unittest.TestCase):
    def setUp(self):
        self.test_db = cache.init_test_db()

    def tearDown(self):
        # Clean up
        cache.clean_test_db(self.test_db)

    def test_basic_set_get(self):
        """Test basic set and get operations"""
        cache_instance = cache.TranslationCache("test_engine")

        # Test get with non-existent entry
        result = cache_instance.get("hello")
        self.assertIsNone(result)

        # Test set and get
        cache_instance.set("hello", "你好")
        result = cache_instance.get("hello")
        self.assertEqual(result, "你好")

    def test_cache_overwrite(self):
        """Test that cache entries can be overwritten"""
        cache_instance = cache.TranslationCache("test_engine")

        # Set initial translation
        cache_instance.set("hello", "你好")

        # Overwrite with new translation
        cache_instance.set("hello", "您好")

        # Verify the new translation is returned
        result = cache_instance.get("hello")
        self.assertEqual(result, "您好")

    def test_non_string_params(self):
        """Test that non-string parameters are automatically converted to JSON"""
        params = {"model": "gpt-3.5", "temperature": 0.7}
        cache_instance = cache.TranslationCache("test_engine", params)

        # Test that params are converted to JSON string internally
        cache_instance.set("hello", "你好")
        result = cache_instance.get("hello")
        self.assertEqual(result, "你好")

        # Test with different param types
        array_params = ["param1", "param2"]
        cache_instance2 = cache.TranslationCache("test_engine", array_params)
        cache_instance2.set("hello", "你好2")
        self.assertEqual(cache_instance2.get("hello"), "你好2")

        # Test with nested structures
        nested_params = {"options": {"temp": 0.8, "models": ["a", "b"]}}
        cache_instance3 = cache.TranslationCache("test_engine", nested_params)
        cache_instance3.set("hello", "你好3")
        self.assertEqual(cache_instance3.get("hello"), "你好3")

    def test_engine_distinction(self):
        """Test that cache distinguishes between different translation engines"""
        cache1 = cache.TranslationCache("engine1")
        cache2 = cache.TranslationCache("engine2")

        # Set same text with different engines
        cache1.set("hello", "你好 1")
        cache2.set("hello", "你好 2")

        # Verify each engine gets its own translation
        self.assertEqual(cache1.get("hello"), "你好 1")
        self.assertEqual(cache2.get("hello"), "你好 2")

    def test_params_distinction(self):
        """Test that cache distinguishes between different engine parameters"""
        params1 = {"param": "value1"}
        params2 = {"param": "value2"}
        cache1 = cache.TranslationCache("test_engine", params1)
        cache2 = cache.TranslationCache("test_engine", params2)

        # Set same text with different parameters
        cache1.set("hello", "你好 1")
        cache2.set("hello", "你好 2")

        # Verify each parameter set gets its own translation
        self.assertEqual(cache1.get("hello"), "你好 1")
        self.assertEqual(cache2.get("hello"), "你好 2")

    def test_consistent_param_serialization(self):
        """Test that dictionary parameters are consistently serialized regardless of key order"""
        # Test simple dictionary
        params1 = {"b": 1, "a": 2}
        params2 = {"a": 2, "b": 1}
        cache1 = cache.TranslationCache("test_engine", params1)
        cache2 = cache.TranslationCache("test_engine", params2)
        self.assertEqual(cache1.translate_engine_params, cache2.translate_engine_params)

        # Test nested dictionary
        params1 = {"outer2": {"inner2": 2, "inner1": 1}, "outer1": 3}
        params2 = {"outer1": 3, "outer2": {"inner1": 1, "inner2": 2}}
        cache1 = cache.TranslationCache("test_engine", params1)
        cache2 = cache.TranslationCache("test_engine", params2)
        self.assertEqual(cache1.translate_engine_params, cache2.translate_engine_params)

        # Test dictionary with list of dictionaries
        params1 = {"b": [{"y": 1, "x": 2}], "a": 3}
        params2 = {"a": 3, "b": [{"x": 2, "y": 1}]}
        cache1 = cache.TranslationCache("test_engine", params1)
        cache2 = cache.TranslationCache("test_engine", params2)
        self.assertEqual(cache1.translate_engine_params, cache2.translate_engine_params)

        # Test that different values still produce different results
        params1 = {"a": 1, "b": 2}
        params2 = {"a": 2, "b": 1}
        cache1 = cache.TranslationCache("test_engine", params1)
        cache2 = cache.TranslationCache("test_engine", params2)
        self.assertNotEqual(
            cache1.translate_engine_params, cache2.translate_engine_params
        )

    def test_cache_with_sorted_params(self):
        """Test that cache works correctly with sorted parameters"""
        params1 = {"b": [{"y": 1, "x": 2}], "a": 3}
        params2 = {"a": 3, "b": [{"x": 2, "y": 1}]}

        # Both caches should work with the same key
        cache1 = cache.TranslationCache("test_engine", params1)
        cache1.set("hello", "你好")

        cache2 = cache.TranslationCache("test_engine", params2)
        self.assertEqual(cache2.get("hello"), "你好")

    def test_append_params(self):
        """Test the append_params method"""
        cache_instance = cache.TranslationCache("test_engine", {"initial": "value"})

        # Test appending new parameter
        cache_instance.add_params("new_param", "new_value")
        self.assertEqual(
            cache_instance.params, {"initial": "value", "new_param": "new_value"}
        )

        # Test that cache with appended params works correctly
        cache_instance.set("hello", "你好")
        self.assertEqual(cache_instance.get("hello"), "你好")

        # Test overwriting existing parameter
        cache_instance.add_params("initial", "new_value")
        self.assertEqual(
            cache_instance.params, {"initial": "new_value", "new_param": "new_value"}
        )

        # Cache should work with updated params
        cache_instance.set("hello2", "你好2")
        self.assertEqual(cache_instance.get("hello2"), "你好2")

    def test_thread_safety(self):
        """Test thread safety of cache operations"""
        cache_instance = cache.TranslationCache("test_engine")
        lock = threading.Lock()
        results = []
        num_threads = multiprocessing.cpu_count()
        items_per_thread = 100

        def generate_random_text(length=10):
            return "".join(
                random.choices(string.ascii_letters + string.digits, k=length)
            )

        def worker():
            thread_results = []  # 线程本地存储结果
            for _ in range(items_per_thread):
                text = generate_random_text()
                translation = f"翻译_{text}"

                # Write operation
                cache_instance.set(text, translation)

                # Read operation - verify our own write
                result = cache_instance.get(text)
                thread_results.append((text, result))

            # 所有操作完成后，一次性加锁并追加结果
            with lock:
                results.extend(thread_results)

        # Create threads equal to CPU core count
        threads = []
        for _ in range(num_threads):
            thread = threading.Thread(target=worker)
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        # Verify all operations were successful
        expected_total = num_threads * items_per_thread
        self.assertEqual(len(results), expected_total)

        # Verify each thread got its correct value
        for text, result in results:
            expected = f"翻译_{text}"
            self.assertEqual(result, expected)


if __name__ == "__main__":
    unittest.main()