80 lines
3.4 KiB
Python
80 lines
3.4 KiB
Python
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),
|
|
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"]],
|
|
)
|