mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
feat(item_storage): implement item_storage_memory max_size
Implemented with unordered dict and set.
This commit is contained in:
111
tests/test_item_storage_memory.py
Normal file
111
tests/test_item_storage_memory.py
Normal file
@ -0,0 +1,111 @@
|
||||
import re
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from invokeai.app.services.item_storage.item_storage_memory import ItemStorageMemory
|
||||
|
||||
|
||||
class MockItemModel(BaseModel):
|
||||
id: str
|
||||
value: int
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def item_storage_memory():
|
||||
return ItemStorageMemory[MockItemModel]()
|
||||
|
||||
|
||||
def test_item_storage_memory_initializes():
|
||||
item_storage_memory = ItemStorageMemory()
|
||||
assert item_storage_memory._items == {}
|
||||
assert item_storage_memory._item_ids == set()
|
||||
assert item_storage_memory._id_field == "id"
|
||||
assert item_storage_memory._max_items == 10
|
||||
|
||||
item_storage_memory = ItemStorageMemory(id_field="bananas", max_items=20)
|
||||
assert item_storage_memory._id_field == "bananas"
|
||||
assert item_storage_memory._max_items == 20
|
||||
|
||||
with pytest.raises(ValueError, match=re.escape("max_items must be at least 1")):
|
||||
item_storage_memory = ItemStorageMemory(max_items=0)
|
||||
with pytest.raises(ValueError, match=re.escape("id_field must not be empty")):
|
||||
item_storage_memory = ItemStorageMemory(id_field="")
|
||||
|
||||
|
||||
def test_item_storage_memory_sets(item_storage_memory: ItemStorageMemory[MockItemModel]):
|
||||
item_1 = MockItemModel(id="1", value=1)
|
||||
item_storage_memory.set(item_1)
|
||||
assert item_storage_memory._items == {"1": item_1}
|
||||
assert item_storage_memory._item_ids == {"1"}
|
||||
|
||||
item_2 = MockItemModel(id="2", value=2)
|
||||
item_storage_memory.set(item_2)
|
||||
assert item_storage_memory._items == {"1": item_1, "2": item_2}
|
||||
assert item_storage_memory._item_ids == {"1", "2"}
|
||||
|
||||
# Updating value of existing item
|
||||
item_2_updated = MockItemModel(id="2", value=9001)
|
||||
item_storage_memory.set(item_2_updated)
|
||||
assert item_storage_memory._items == {"1": item_1, "2": item_2_updated}
|
||||
assert item_storage_memory._item_ids == {"1", "2"}
|
||||
|
||||
|
||||
def test_item_storage_memory_gets(item_storage_memory: ItemStorageMemory[MockItemModel]):
|
||||
item_1 = MockItemModel(id="1", value=1)
|
||||
item_storage_memory.set(item_1)
|
||||
item = item_storage_memory.get("1")
|
||||
assert item == item_1
|
||||
|
||||
item_2 = MockItemModel(id="2", value=2)
|
||||
item_storage_memory.set(item_2)
|
||||
item = item_storage_memory.get("2")
|
||||
assert item == item_2
|
||||
|
||||
item = item_storage_memory.get("3")
|
||||
assert item is None
|
||||
|
||||
|
||||
def test_item_storage_memory_deletes(item_storage_memory: ItemStorageMemory[MockItemModel]):
|
||||
item_1 = MockItemModel(id="1", value=1)
|
||||
item_2 = MockItemModel(id="2", value=2)
|
||||
item_storage_memory.set(item_1)
|
||||
item_storage_memory.set(item_2)
|
||||
|
||||
item_storage_memory.delete("2")
|
||||
assert item_storage_memory._items == {"1": item_1}
|
||||
assert item_storage_memory._item_ids == {"1"}
|
||||
|
||||
|
||||
def test_item_storage_memory_respects_max():
|
||||
item_storage_memory = ItemStorageMemory(max_items=3)
|
||||
for i in range(10):
|
||||
item_storage_memory.set(MockItemModel(id=str(i), value=i))
|
||||
assert len(item_storage_memory._items) == 3
|
||||
|
||||
|
||||
def test_item_storage_memory_calls_set_callback(item_storage_memory: ItemStorageMemory[MockItemModel]):
|
||||
called_item = None
|
||||
item = MockItemModel(id="1", value=1)
|
||||
|
||||
def on_changed(item: MockItemModel):
|
||||
nonlocal called_item
|
||||
called_item = item
|
||||
|
||||
item_storage_memory.on_changed(on_changed)
|
||||
item_storage_memory.set(item)
|
||||
assert called_item == item
|
||||
|
||||
|
||||
def test_item_storage_memory_calls_delete_callback(item_storage_memory: ItemStorageMemory[MockItemModel]):
|
||||
called_item_id = None
|
||||
item = MockItemModel(id="1", value=1)
|
||||
|
||||
def on_deleted(item_id: str):
|
||||
nonlocal called_item_id
|
||||
called_item_id = item_id
|
||||
|
||||
item_storage_memory.on_deleted(on_deleted)
|
||||
item_storage_memory.set(item)
|
||||
item_storage_memory.delete("1")
|
||||
assert called_item_id == "1"
|
Reference in New Issue
Block a user