"""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()