146 lines
4.3 KiB
Python
146 lines
4.3 KiB
Python
import threading
|
|
import time
|
|
import uuid
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
from agentmsgs.stores.memory import InMemoryStore
|
|
|
|
|
|
def test_memory_create_and_find_thread():
|
|
s = InMemoryStore()
|
|
alice = s.get_or_create_agent("Alice")
|
|
bob = s.get_or_create_agent("Bob")
|
|
t = s.create_thread({alice, bob})
|
|
assert s.get_thread(t.id) == t
|
|
assert t in s.find_threads({alice})
|
|
assert t in s.find_threads({alice, bob})
|
|
assert s.find_threads({alice, bob, s.get_or_create_agent("Charlie")}) == []
|
|
|
|
|
|
def test_memory_soft_unique_name():
|
|
s = InMemoryStore()
|
|
a1 = s.get_or_create_agent("Alice")
|
|
a2 = s.get_or_create_agent("Alice")
|
|
assert a1.id == a2.id
|
|
|
|
|
|
def test_memory_delete_account():
|
|
s = InMemoryStore()
|
|
a = s.get_or_create_agent("Alice")
|
|
s.delete_agent(a.id)
|
|
assert s.get_agent_by_id(a.id) is None
|
|
assert s.get_agent_by_name("Alice") is None
|
|
|
|
|
|
def test_memory_append_and_list():
|
|
s = InMemoryStore()
|
|
a = s.get_or_create_agent("A")
|
|
b = s.get_or_create_agent("B")
|
|
t = s.create_thread({a, b})
|
|
m1 = s.append_message(t.id, a, "hi")
|
|
m2 = s.append_message(t.id, b, "yo")
|
|
assert m1.seq == 1 and m2.seq == 2
|
|
assert s.list_messages(t.id, after_seq=1) == [m2]
|
|
|
|
|
|
def test_memory_thread_info_and_read_state():
|
|
s = InMemoryStore()
|
|
a = s.get_or_create_agent("A")
|
|
b = s.get_or_create_agent("B")
|
|
t = s.create_thread({a, b})
|
|
s.append_message(t.id, a, "one")
|
|
s.append_message(t.id, b, "two")
|
|
m3 = s.append_message(t.id, a, "three")
|
|
|
|
info = s.get_thread_info(t.id)
|
|
assert info is not None
|
|
assert info.thread == t
|
|
assert info.message_count == 3
|
|
assert info.latest_message == m3
|
|
|
|
state = s.get_thread_read_state(t.id, a)
|
|
assert state.agent == a
|
|
assert state.cursor_seq == 0
|
|
assert state.latest_seq == 3
|
|
assert state.unread_count == 3
|
|
assert state.unread_from_others_count == 1
|
|
|
|
s.set_cursor(t.id, a, 2)
|
|
state = s.get_thread_read_state(t.id, a)
|
|
assert state.cursor_seq == 2
|
|
assert state.unread_count == 1
|
|
assert state.unread_from_others_count == 0
|
|
|
|
|
|
def test_memory_thread_info_for_empty_thread():
|
|
s = InMemoryStore()
|
|
a = s.get_or_create_agent("A")
|
|
b = s.get_or_create_agent("B")
|
|
t = s.create_thread({a, b})
|
|
|
|
info = s.get_thread_info(t.id)
|
|
state = s.get_thread_read_state(t.id, a)
|
|
|
|
assert info is not None
|
|
assert info.message_count == 0
|
|
assert info.latest_message is None
|
|
assert state.latest_seq == 0
|
|
assert state.unread_count == 0
|
|
assert state.unread_from_others_count == 0
|
|
assert s.get_thread_info(uuid.uuid4()) is None
|
|
assert s.get_thread_read_state(uuid.uuid4(), a) is None
|
|
|
|
|
|
def test_memory_concurrent_get_or_create_agent_is_idempotent():
|
|
store = InMemoryStore()
|
|
start = threading.Barrier(2)
|
|
|
|
class SlowNameIndex(dict):
|
|
def __contains__(self, key):
|
|
present = super().__contains__(key)
|
|
if not present:
|
|
time.sleep(0.01)
|
|
return present
|
|
|
|
store._name_idx = SlowNameIndex()
|
|
|
|
def create_agent():
|
|
start.wait()
|
|
return store.get_or_create_agent("Alice")
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
agents = list(executor.map(lambda _: create_agent(), range(2)))
|
|
|
|
assert len({agent.id for agent in agents}) == 1
|
|
assert len(store.list_agents()) == 1
|
|
|
|
|
|
def test_memory_concurrent_append_assigns_unique_sequences():
|
|
store = InMemoryStore()
|
|
alice = store.get_or_create_agent("Alice")
|
|
bob = store.get_or_create_agent("Bob")
|
|
thread = store.create_thread({alice, bob})
|
|
gate = threading.Barrier(2, timeout=0.2)
|
|
local = threading.local()
|
|
|
|
class SlowMessageList(list):
|
|
def __len__(self):
|
|
length = super().__len__()
|
|
if not hasattr(local, "waited"):
|
|
try:
|
|
gate.wait()
|
|
except threading.BrokenBarrierError:
|
|
pass
|
|
local.waited = True
|
|
return length
|
|
|
|
store._msgs[thread.id] = SlowMessageList()
|
|
|
|
def append_message(index):
|
|
return store.append_message(thread.id, alice, str(index))
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
messages = list(executor.map(append_message, range(2)))
|
|
|
|
assert sorted(message.seq for message in messages) == [1, 2]
|