by-ae
17858e1963
A letter plus a combining accent was rejected, and when it was already in the file it did not match the composed name.
360 lines
14 KiB
Python
360 lines
14 KiB
Python
"""Unit tests for the store module."""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
|
|
|
# conftest.py installs the folder_paths mock before store is imported.
|
|
# Replacing sys.modules['folder_paths'] here binds a second mock that the
|
|
# route tests then configure, while store keeps writing through the original.
|
|
import store
|
|
from store import _validate_name, load_bucket, save_entry, delete_entry
|
|
|
|
|
|
class TestStore(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.patcher = patch('store.folder_paths')
|
|
self.mock_folder_paths = self.patcher.start()
|
|
self.mock_folder_paths.get_user_directory.return_value = self.temp_dir
|
|
|
|
def tearDown(self):
|
|
self.patcher.stop()
|
|
import shutil
|
|
if os.path.exists(self.temp_dir):
|
|
shutil.rmtree(self.temp_dir)
|
|
|
|
def test_validate_name_valid(self):
|
|
self.assertEqual(_validate_name("MyCharacter"), "MyCharacter")
|
|
self.assertEqual(_validate_name("char_123"), "char_123")
|
|
self.assertEqual(_validate_name("test-name"), "test-name")
|
|
self.assertEqual(_validate_name("name with spaces"), "name with spaces")
|
|
|
|
def test_validate_name_invalid(self):
|
|
with self.assertRaises(ValueError):
|
|
_validate_name("")
|
|
with self.assertRaises(ValueError):
|
|
_validate_name(" ")
|
|
with self.assertRaises(ValueError):
|
|
_validate_name("-starts-with-dash")
|
|
with self.assertRaises(ValueError):
|
|
_validate_name("has/slash")
|
|
with self.assertRaises(ValueError):
|
|
_validate_name("has\\backslash")
|
|
with self.assertRaises(ValueError):
|
|
_validate_name("a" * 65)
|
|
|
|
def test_save_and_load_entry(self):
|
|
test_data = {"desc": "A test character", "palette": ["#FF0000", "#00FF00"]}
|
|
save_entry("characters", "TestChar", test_data)
|
|
|
|
bucket = load_bucket("characters")
|
|
self.assertIn("TestChar", bucket)
|
|
self.assertEqual(bucket["TestChar"], test_data)
|
|
|
|
def test_save_multiple_entries(self):
|
|
save_entry("characters", "Char1", {"data": 1})
|
|
save_entry("characters", "Char2", {"data": 2})
|
|
|
|
bucket = load_bucket("characters")
|
|
self.assertEqual(len(bucket), 2)
|
|
self.assertEqual(bucket["Char1"]["data"], 1)
|
|
self.assertEqual(bucket["Char2"]["data"], 2)
|
|
|
|
def test_overwrite_entry(self):
|
|
save_entry("objects", "Object1", {"version": 1})
|
|
save_entry("objects", "Object1", {"version": 2})
|
|
|
|
bucket = load_bucket("objects")
|
|
self.assertEqual(bucket["Object1"]["version"], 2)
|
|
|
|
def test_delete_entry(self):
|
|
save_entry("scenes", "Scene1", {"data": "test"})
|
|
self.assertTrue(delete_entry("scenes", "Scene1"))
|
|
|
|
bucket = load_bucket("scenes")
|
|
self.assertNotIn("Scene1", bucket)
|
|
|
|
def test_delete_nonexistent_entry(self):
|
|
self.assertFalse(delete_entry("characters", "NonExistent"))
|
|
|
|
def test_load_empty_bucket(self):
|
|
bucket = load_bucket("characters")
|
|
self.assertEqual(bucket, {})
|
|
|
|
def test_invalid_bucket(self):
|
|
with self.assertRaises(ValueError):
|
|
load_bucket("invalid_bucket")
|
|
with self.assertRaises(ValueError):
|
|
save_entry("invalid_bucket", "name", {})
|
|
with self.assertRaises(ValueError):
|
|
delete_entry("invalid_bucket", "name")
|
|
|
|
def test_persistence_across_loads(self):
|
|
save_entry("objects", "Persistent", {"value": "data"})
|
|
|
|
first_load = load_bucket("objects")
|
|
second_load = load_bucket("objects")
|
|
|
|
self.assertEqual(first_load["Persistent"], second_load["Persistent"])
|
|
|
|
def test_corrupt_json_creates_backup_and_returns_empty(self):
|
|
# Create a corrupt JSON file
|
|
bucket_path = os.path.join(self.temp_dir, "byae-ideogram-tools", "characters.json")
|
|
os.makedirs(os.path.dirname(bucket_path), exist_ok=True)
|
|
with open(bucket_path, "w") as f:
|
|
f.write("{invalid json content")
|
|
|
|
# Loading should return empty dict and create backup
|
|
bucket = load_bucket("characters")
|
|
self.assertEqual(bucket, {})
|
|
|
|
# Check that backup file was created
|
|
backup_files = [f for f in os.listdir(os.path.dirname(bucket_path))
|
|
if f.startswith("characters.json.corrupt-")]
|
|
self.assertEqual(len(backup_files), 1)
|
|
|
|
# Verify backup contains original corrupt content
|
|
backup_path = os.path.join(os.path.dirname(bucket_path), backup_files[0])
|
|
with open(backup_path, "r") as f:
|
|
self.assertEqual(f.read(), "{invalid json content")
|
|
|
|
# Verify original file no longer exists
|
|
self.assertFalse(os.path.exists(bucket_path))
|
|
|
|
def test_corrupt_file_does_not_affect_other_buckets(self):
|
|
# Save valid entry in objects bucket
|
|
save_entry("objects", "ValidObject", {"data": "test"})
|
|
|
|
# Create corrupt characters file
|
|
bucket_path = os.path.join(self.temp_dir, "byae-ideogram-tools", "characters.json")
|
|
os.makedirs(os.path.dirname(bucket_path), exist_ok=True)
|
|
with open(bucket_path, "w") as f:
|
|
f.write("corrupted")
|
|
|
|
# Load all should still work, with characters empty but objects intact
|
|
from store import load_all
|
|
stores = load_all()
|
|
|
|
self.assertEqual(stores["characters"], {})
|
|
self.assertIn("ValidObject", stores["objects"])
|
|
self.assertEqual(stores["objects"]["ValidObject"]["data"], "test")
|
|
|
|
def test_non_object_json_is_backed_up_before_overwrite(self):
|
|
bucket_path = os.path.join(self.temp_dir, "byae-ideogram-tools", "characters.json")
|
|
os.makedirs(os.path.dirname(bucket_path), exist_ok=True)
|
|
original = '[{"desc": "keep me"}]'
|
|
with open(bucket_path, "w", encoding="utf-8") as f:
|
|
f.write(original)
|
|
|
|
self.assertEqual(load_bucket("characters"), {})
|
|
self.assertFalse(os.path.exists(bucket_path))
|
|
backup_files = [
|
|
name for name in os.listdir(os.path.dirname(bucket_path))
|
|
if name.startswith("characters.json.corrupt-")
|
|
]
|
|
self.assertEqual(len(backup_files), 1)
|
|
with open(os.path.join(os.path.dirname(bucket_path), backup_files[0]), encoding="utf-8") as f:
|
|
self.assertEqual(f.read(), original)
|
|
|
|
save_entry("characters", "New", {"desc": "n"})
|
|
bucket = load_bucket("characters")
|
|
self.assertEqual(bucket["New"]["desc"], "n")
|
|
with open(os.path.join(os.path.dirname(bucket_path), backup_files[0]), encoding="utf-8") as f:
|
|
self.assertEqual(f.read(), original)
|
|
|
|
def test_second_corrupt_backup_in_the_same_second_keeps_the_first(self):
|
|
store_dir = os.path.join(self.temp_dir, "byae-ideogram-tools")
|
|
os.makedirs(store_dir, exist_ok=True)
|
|
bucket_path = os.path.join(store_dir, "objects.json")
|
|
stamp = 1_700_000_000
|
|
existing = os.path.join(store_dir, f"objects.json.corrupt-{stamp}")
|
|
with open(existing, "w", encoding="utf-8") as f:
|
|
f.write("first")
|
|
with open(bucket_path, "w", encoding="utf-8") as f:
|
|
f.write("{not json")
|
|
|
|
with patch("store.time.time", return_value=stamp):
|
|
self.assertEqual(load_bucket("objects"), {})
|
|
|
|
self.assertFalse(os.path.exists(bucket_path))
|
|
with open(existing, encoding="utf-8") as f:
|
|
self.assertEqual(f.read(), "first")
|
|
extras = [
|
|
name for name in os.listdir(store_dir)
|
|
if name.startswith(f"objects.json.corrupt-{stamp}-")
|
|
]
|
|
self.assertEqual(len(extras), 1)
|
|
with open(os.path.join(store_dir, extras[0]), encoding="utf-8") as f:
|
|
self.assertEqual(f.read(), "{not json")
|
|
|
|
def test_utf8_bom_is_loaded_without_backup(self):
|
|
# Editors sometimes write a leading byte-order mark. json.load rejects
|
|
# that as corrupt, which used to move the whole bucket aside.
|
|
bucket_path = os.path.join(self.temp_dir, "byae-ideogram-tools", "characters.json")
|
|
os.makedirs(os.path.dirname(bucket_path), exist_ok=True)
|
|
payload = '{"Hero": {"desc": "kept"}, "角色": {"desc": "unicode"}}'
|
|
with open(bucket_path, "wb") as f:
|
|
f.write(b"\xef\xbb\xbf" + payload.encode("utf-8"))
|
|
|
|
with self.assertNoLogs(store.logger, level="WARNING"):
|
|
bucket = load_bucket("characters")
|
|
|
|
self.assertEqual(bucket["Hero"]["desc"], "kept")
|
|
self.assertEqual(bucket["角色"]["desc"], "unicode")
|
|
self.assertTrue(os.path.exists(bucket_path))
|
|
backups = [
|
|
name for name in os.listdir(os.path.dirname(bucket_path))
|
|
if name.startswith("characters.json.corrupt-")
|
|
]
|
|
self.assertEqual(backups, [])
|
|
|
|
save_entry("characters", "Next", {"desc": "added"})
|
|
reloaded = load_bucket("characters")
|
|
self.assertEqual(reloaded["Hero"]["desc"], "kept")
|
|
self.assertEqual(reloaded["Next"]["desc"], "added")
|
|
with open(bucket_path, "rb") as f:
|
|
saved = f.read()
|
|
self.assertFalse(saved.startswith(b"\xef\xbb\xbf"))
|
|
self.assertIn("Hero".encode("utf-8"), saved)
|
|
|
|
def test_save_works_after_corrupt_file_recovery(self):
|
|
# Create corrupt file
|
|
bucket_path = os.path.join(self.temp_dir, "byae-ideogram-tools", "scenes.json")
|
|
os.makedirs(os.path.dirname(bucket_path), exist_ok=True)
|
|
with open(bucket_path, "w") as f:
|
|
f.write("[not a dict]")
|
|
|
|
# Load should recover
|
|
load_bucket("scenes")
|
|
|
|
# Save should now work
|
|
save_entry("scenes", "NewScene", {"data": "value"})
|
|
|
|
# Verify save succeeded
|
|
bucket = load_bucket("scenes")
|
|
self.assertIn("NewScene", bucket)
|
|
self.assertEqual(bucket["NewScene"]["data"], "value")
|
|
|
|
def test_unicode_names_that_look_the_same_are_one_preset(self):
|
|
import unicodedata
|
|
nfd = unicodedata.normalize("NFD", "Émilie")
|
|
nfc = unicodedata.normalize("NFC", "Émilie")
|
|
self.assertNotEqual(nfd, nfc)
|
|
|
|
save_entry("characters", nfd, {"desc": "first"})
|
|
save_entry("characters", nfc, {"desc": "second"})
|
|
bucket = load_bucket("characters")
|
|
self.assertEqual(list(bucket), [nfc])
|
|
self.assertEqual(bucket[nfc]["desc"], "second")
|
|
|
|
self.assertTrue(delete_entry("characters", nfd))
|
|
self.assertEqual(load_bucket("characters"), {})
|
|
|
|
def test_composed_and_decomposed_keys_in_one_file_collapse(self):
|
|
import unicodedata
|
|
nfd = unicodedata.normalize("NFD", "Émilie")
|
|
nfc = unicodedata.normalize("NFC", "Émilie")
|
|
bucket_path = os.path.join(self.temp_dir, "byae-ideogram-tools", "characters.json")
|
|
os.makedirs(os.path.dirname(bucket_path), exist_ok=True)
|
|
with open(bucket_path, "w", encoding="utf-8") as f:
|
|
json.dump({nfd: {"desc": "old"}, nfc: {"desc": "new"}, " Hero ": {"desc": "spaced"}}, f, ensure_ascii=False)
|
|
|
|
with self.assertLogs(store.logger, level="WARNING") as cm:
|
|
bucket = load_bucket("characters")
|
|
self.assertEqual(set(bucket), {nfc, "Hero"})
|
|
self.assertEqual(bucket[nfc]["desc"], "new")
|
|
self.assertEqual(bucket["Hero"]["desc"], "spaced")
|
|
self.assertTrue(any("Émilie" in line or nfd in line for line in cm.output))
|
|
|
|
def test_long_name_says_how_long_a_name_can_be(self):
|
|
with self.assertRaises(ValueError) as caught:
|
|
_validate_name("a" * 65)
|
|
self.assertIn("64", str(caught.exception))
|
|
|
|
def test_unicode_character_names(self):
|
|
# Test that Unicode names work correctly
|
|
save_entry("characters", "角色123", {"desc": "Japanese character"})
|
|
save_entry("characters", "Персонаж", {"desc": "Russian character"})
|
|
save_entry("characters", "Émilie", {"desc": "French character"})
|
|
|
|
bucket = load_bucket("characters")
|
|
self.assertIn("角色123", bucket)
|
|
self.assertIn("Персонаж", bucket)
|
|
self.assertIn("Émilie", bucket)
|
|
self.assertEqual(bucket["角色123"]["desc"], "Japanese character")
|
|
self.assertEqual(bucket["Персонаж"]["desc"], "Russian character")
|
|
self.assertEqual(bucket["Émilie"]["desc"], "French character")
|
|
|
|
def test_concurrent_write_safety(self):
|
|
# Test that concurrent writes don't corrupt the file
|
|
import threading
|
|
import time
|
|
|
|
errors = []
|
|
|
|
def writer(name, value):
|
|
try:
|
|
for i in range(10):
|
|
save_entry("characters", name, {"iteration": i, "value": value})
|
|
time.sleep(0.001)
|
|
except Exception as e:
|
|
errors.append((name, str(e)))
|
|
|
|
threads = [
|
|
threading.Thread(target=writer, args=(f"char{i}", i))
|
|
for i in range(5)
|
|
]
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# Check no errors occurred
|
|
self.assertEqual(errors, [], f"Concurrent write errors: {errors}")
|
|
|
|
# Verify all characters were saved
|
|
bucket = load_bucket("characters")
|
|
for i in range(5):
|
|
self.assertIn(f"char{i}", bucket)
|
|
self.assertIn("iteration", bucket[f"char{i}"])
|
|
self.assertEqual(bucket[f"char{i}"]["value"], i)
|
|
|
|
def test_first_bucket_lock_excludes_other_writers(self):
|
|
import threading
|
|
import time
|
|
from store import _BUCKET_LOCKS, _BUCKET_LOCKS_GUARD, _bucket_lock
|
|
|
|
with _BUCKET_LOCKS_GUARD:
|
|
_BUCKET_LOCKS.pop("objects", None)
|
|
|
|
state = {"current": 0, "max": 0}
|
|
started = threading.Barrier(12)
|
|
|
|
def worker():
|
|
started.wait()
|
|
with _bucket_lock("objects"):
|
|
state["current"] += 1
|
|
state["max"] = max(state["max"], state["current"])
|
|
time.sleep(0.01)
|
|
state["current"] -= 1
|
|
|
|
threads = [threading.Thread(target=worker) for _ in range(12)]
|
|
for thread in threads:
|
|
thread.start()
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
self.assertEqual(state["max"], 1)
|
|
self.assertEqual(state["current"], 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|