Подготовить AI Lead Processor и проверить обработку заявок

tests/test_storage.py

@@ -0,0 +1,72 @@
+import unittest
+from concurrent.futures import ThreadPoolExecutor
+from io import BytesIO
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest.mock import patch
+from uuid import uuid4
+
+from openpyxl import load_workbook
+
+from src import storage
+
+
+class StorageTests(unittest.TestCase):
+    def setUp(self):
+        directory = TemporaryDirectory()
+        self.addCleanup(directory.cleanup)
+        self.path = Path(directory.name) / "входящие" / "leads.xlsx"
+
+    def test_empty_snapshot_creates_downloadable_workbook(self):
+        rows, data = storage.snapshot(self.path)
+        self.assertEqual(rows, [])
+        self.assertEqual(data, self.path.read_bytes())
+        sheet = load_workbook(BytesIO(data)).active
+        self.assertEqual(sheet["C1"].value, "Текст заявки")
+        self.assertEqual(sheet["L1"].value, "Оценочная стоимость API-запроса, USD")
+
+    def test_retry_keeps_one_row_with_literal_text(self):
+        row = {"ID": str(uuid4()), "Текст заявки": "=1+2\nПривет!", "Входные токены": 0}
+        self.assertTrue(storage.append(self.path, row))
+        self.assertFalse(storage.append(self.path, row))
+        rows, data = storage.snapshot(self.path)
+        self.assertEqual(len(rows), 1)
+        self.assertEqual({key: rows[0][key] for key in row}, row)
+        self.assertEqual(load_workbook(BytesIO(data)).active["C2"].data_type, "s")
+
+    def test_parallel_writers_preserve_all_rows(self):
+        ids = [str(uuid4()) for _ in range(6)]
+        with ThreadPoolExecutor(max_workers=4) as executor:
+            results = executor.map(lambda key: storage.append(
+                self.path, {"ID": key, "Текст заявки": "Здравствуйте!"}
+            ), ids)
+            self.assertTrue(all(results))
+        rows, _ = storage.snapshot(self.path)
+        self.assertCountEqual([row["ID"] for row in rows], ids)
+
+    def test_bad_input_preserves_original_bytes(self):
+        _, original = storage.snapshot(self.path)
+        for value in ["   ", "x\x00y", "ю" * 32768]:
+            with self.subTest(length=len(value)), self.assertRaises(ValueError):
+                storage.append(self.path, {"ID": str(uuid4()), "Текст заявки": value})
+        self.assertEqual(original, self.path.read_bytes())
+
+    def test_corrupt_file_is_preserved(self):
+        self.path.parent.mkdir()
+        self.path.write_bytes(b"corrupt")
+        with self.assertRaises(ValueError):
+            storage.append(self.path, {"ID": str(uuid4()), "Текст заявки": "тест"})
+        self.assertEqual(self.path.read_bytes(), b"corrupt")
+
+    def test_locked_excel_file_can_be_retried_without_losing_rows(self):
+        first = {"ID": str(uuid4()), "Текст заявки": "Первая заявка"}
+        second = {"ID": str(uuid4()), "Текст заявки": "Вторая заявка"}
+        storage.append(self.path, first)
+        original = self.path.read_bytes()
+        with patch("src.storage.os.replace", side_effect=PermissionError):
+            with self.assertRaises(PermissionError):
+                storage.append(self.path, second)
+        self.assertEqual(self.path.read_bytes(), original)
+        self.assertTrue(storage.append(self.path, second))
+        self.assertEqual([row["ID"] for row in storage.snapshot(self.path)[0]],
+                         [first["ID"], second["ID"]])