|
|
|
|
@ -4,7 +4,7 @@ from constants import UUID_NIL
|
|
|
|
|
from core.prompt.utils.extract_thread_messages import extract_thread_messages
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class TestMessage:
|
|
|
|
|
class MockMessage:
|
|
|
|
|
def __init__(self, id, parent_message_id):
|
|
|
|
|
self.id = id
|
|
|
|
|
self.parent_message_id = parent_message_id
|
|
|
|
|
@ -14,7 +14,7 @@ class TestMessage:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_extract_thread_messages_single_message():
|
|
|
|
|
messages = [TestMessage(str(uuid4()), UUID_NIL)]
|
|
|
|
|
messages = [MockMessage(str(uuid4()), UUID_NIL)]
|
|
|
|
|
result = extract_thread_messages(messages)
|
|
|
|
|
assert len(result) == 1
|
|
|
|
|
assert result[0] == messages[0]
|
|
|
|
|
@ -23,11 +23,11 @@ def test_extract_thread_messages_single_message():
|
|
|
|
|
def test_extract_thread_messages_linear_thread():
|
|
|
|
|
id1, id2, id3, id4, id5 = str(uuid4()), str(uuid4()), str(uuid4()), str(uuid4()), str(uuid4())
|
|
|
|
|
messages = [
|
|
|
|
|
TestMessage(id5, id4),
|
|
|
|
|
TestMessage(id4, id3),
|
|
|
|
|
TestMessage(id3, id2),
|
|
|
|
|
TestMessage(id2, id1),
|
|
|
|
|
TestMessage(id1, UUID_NIL),
|
|
|
|
|
MockMessage(id5, id4),
|
|
|
|
|
MockMessage(id4, id3),
|
|
|
|
|
MockMessage(id3, id2),
|
|
|
|
|
MockMessage(id2, id1),
|
|
|
|
|
MockMessage(id1, UUID_NIL),
|
|
|
|
|
]
|
|
|
|
|
result = extract_thread_messages(messages)
|
|
|
|
|
assert len(result) == 5
|
|
|
|
|
@ -37,10 +37,10 @@ def test_extract_thread_messages_linear_thread():
|
|
|
|
|
def test_extract_thread_messages_branched_thread():
|
|
|
|
|
id1, id2, id3, id4 = str(uuid4()), str(uuid4()), str(uuid4()), str(uuid4())
|
|
|
|
|
messages = [
|
|
|
|
|
TestMessage(id4, id2),
|
|
|
|
|
TestMessage(id3, id2),
|
|
|
|
|
TestMessage(id2, id1),
|
|
|
|
|
TestMessage(id1, UUID_NIL),
|
|
|
|
|
MockMessage(id4, id2),
|
|
|
|
|
MockMessage(id3, id2),
|
|
|
|
|
MockMessage(id2, id1),
|
|
|
|
|
MockMessage(id1, UUID_NIL),
|
|
|
|
|
]
|
|
|
|
|
result = extract_thread_messages(messages)
|
|
|
|
|
assert len(result) == 3
|
|
|
|
|
@ -56,9 +56,9 @@ def test_extract_thread_messages_empty_list():
|
|
|
|
|
def test_extract_thread_messages_partially_loaded():
|
|
|
|
|
id0, id1, id2, id3 = str(uuid4()), str(uuid4()), str(uuid4()), str(uuid4())
|
|
|
|
|
messages = [
|
|
|
|
|
TestMessage(id3, id2),
|
|
|
|
|
TestMessage(id2, id1),
|
|
|
|
|
TestMessage(id1, id0),
|
|
|
|
|
MockMessage(id3, id2),
|
|
|
|
|
MockMessage(id2, id1),
|
|
|
|
|
MockMessage(id1, id0),
|
|
|
|
|
]
|
|
|
|
|
result = extract_thread_messages(messages)
|
|
|
|
|
assert len(result) == 3
|
|
|
|
|
@ -68,9 +68,9 @@ def test_extract_thread_messages_partially_loaded():
|
|
|
|
|
def test_extract_thread_messages_legacy_messages():
|
|
|
|
|
id1, id2, id3 = str(uuid4()), str(uuid4()), str(uuid4())
|
|
|
|
|
messages = [
|
|
|
|
|
TestMessage(id3, UUID_NIL),
|
|
|
|
|
TestMessage(id2, UUID_NIL),
|
|
|
|
|
TestMessage(id1, UUID_NIL),
|
|
|
|
|
MockMessage(id3, UUID_NIL),
|
|
|
|
|
MockMessage(id2, UUID_NIL),
|
|
|
|
|
MockMessage(id1, UUID_NIL),
|
|
|
|
|
]
|
|
|
|
|
result = extract_thread_messages(messages)
|
|
|
|
|
assert len(result) == 3
|
|
|
|
|
@ -80,11 +80,11 @@ def test_extract_thread_messages_legacy_messages():
|
|
|
|
|
def test_extract_thread_messages_mixed_with_legacy_messages():
|
|
|
|
|
id1, id2, id3, id4, id5 = str(uuid4()), str(uuid4()), str(uuid4()), str(uuid4()), str(uuid4())
|
|
|
|
|
messages = [
|
|
|
|
|
TestMessage(id5, id4),
|
|
|
|
|
TestMessage(id4, id2),
|
|
|
|
|
TestMessage(id3, id2),
|
|
|
|
|
TestMessage(id2, UUID_NIL),
|
|
|
|
|
TestMessage(id1, UUID_NIL),
|
|
|
|
|
MockMessage(id5, id4),
|
|
|
|
|
MockMessage(id4, id2),
|
|
|
|
|
MockMessage(id3, id2),
|
|
|
|
|
MockMessage(id2, UUID_NIL),
|
|
|
|
|
MockMessage(id1, UUID_NIL),
|
|
|
|
|
]
|
|
|
|
|
result = extract_thread_messages(messages)
|
|
|
|
|
assert len(result) == 4
|
|
|
|
|
|