From b81ea11edd3aef4b344e31f7d46c31e2b367d074 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 25 Jun 2024 16:54:05 -0300 Subject: [PATCH] refactor: Add test for adding messages --- tests/test_messages_endpoints.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/test_messages_endpoints.py b/tests/test_messages_endpoints.py index 62fd0c941..163ce610b 100644 --- a/tests/test_messages_endpoints.py +++ b/tests/test_messages_endpoints.py @@ -1,9 +1,12 @@ +import uuid from uuid import UUID import pytest from fastapi.testclient import TestClient +from sqlmodel import select -from langflow.memory import add_messagetables +from langflow.memory import add_messages, add_messagetables +from langflow.schema.message import Message # Assuming you have these imports available from langflow.services.database.models.message import MessageCreate, MessageRead, MessageUpdate @@ -37,6 +40,21 @@ def created_messages(session): return messages_read +def test_add_message(session, flow): + session_id = str(uuid.uuid4()) + message = Message(text="Test message", sender="User", sender_name="User", session_id=session_id) + messages = add_messages([message], flow.id) + + with session_scope() as session: + message = session.exec(select(MessageTable).where(MessageTable.session_id == session_id)).first() + assert message is not None + assert len(messages) == 1 + assert message.text == "Test message" + assert message.sender == "User" + assert message.sender_name == "User" + assert message.session_id == session_id + + def test_delete_messages(client: TestClient, created_messages, logged_in_headers): response = client.request( "DELETE", "api/v1/monitor/messages", json=[str(msg.id) for msg in created_messages], headers=logged_in_headers