test(websocket.py): add tests for websocket connection, chat history and sending message

tests/conftest.py

@@ -1,4 +1,8 @@
+import json
 from pathlib import Path
+from typing import AsyncGenerator
+from httpx import AsyncClient
+
 
 import pytest
 from fastapi.testclient import TestClient
@@ -21,6 +25,15 @@ def get_text():
     """
 
 
+@pytest.fixture()
+async def async_client() -> AsyncGenerator:
+    from langflow.main import create_app
+
+    app = create_app()
+    async with AsyncClient(app=app, base_url="http://testserver") as client:
+        yield client
+
+
 # Create client fixture for FastAPI
 @pytest.fixture(scope="module")
 def client():
@@ -30,3 +43,37 @@ def client():
 
     with TestClient(app) as client:
         yield client
+
+
+def get_graph(_type="basic"):
+    """Get a graph from a json file"""
+    from langflow.graph.graph import Graph
+
+    if _type == "basic":
+        path = pytest.BASIC_EXAMPLE_PATH
+    elif _type == "complex":
+        path = pytest.COMPLEX_EXAMPLE_PATH
+    elif _type == "openapi":
+        path = pytest.OPENAPI_EXAMPLE_PATH
+
+    with open(path, "r") as f:
+        flow_graph = json.load(f)
+    data_graph = flow_graph["data"]
+    nodes = data_graph["nodes"]
+    edges = data_graph["edges"]
+    return Graph(nodes, edges)
+
+
+@pytest.fixture
+def basic_graph():
+    return get_graph()
+
+
+@pytest.fixture
+def complex_graph():
+    return get_graph("complex")
+
+
+@pytest.fixture
+def openapi_graph():
+    return get_graph("openapi")

tests/test_graph.py

@@ -15,7 +15,7 @@ from langflow.graph.nodes import (
     ToolNode,
     WrapperNode,
 )
-from langflow.interface.run import get_result_and_thought_using_graph
+from langflow.interface.run import get_result_and_steps
 from langflow.utils.payload import build_json, get_root_node
 
 # Test cases for the graph module
@@ -24,38 +24,6 @@ from langflow.utils.payload import build_json, get_root_node
 # BASIC_EXAMPLE_PATH, COMPLEX_EXAMPLE_PATH, OPENAPI_EXAMPLE_PATH
 
 
-def get_graph(_type="basic"):
-    """Get a graph from a json file"""
-    if _type == "basic":
-        path = pytest.BASIC_EXAMPLE_PATH
-    elif _type == "complex":
-        path = pytest.COMPLEX_EXAMPLE_PATH
-    elif _type == "openapi":
-        path = pytest.OPENAPI_EXAMPLE_PATH
-
-    with open(path, "r") as f:
-        flow_graph = json.load(f)
-    data_graph = flow_graph["data"]
-    nodes = data_graph["nodes"]
-    edges = data_graph["edges"]
-    return Graph(nodes, edges)
-
-
-@pytest.fixture
-def basic_graph():
-    return get_graph()
-
-
-@pytest.fixture
-def complex_graph():
-    return get_graph("complex")
-
-
-@pytest.fixture
-def openapi_graph():
-    return get_graph("openapi")
-
-
 def get_node_by_type(graph, node_type: Type[Node]) -> Union[Node, None]:
     """Get a node by type"""
     return next((node for node in graph.nodes if isinstance(node, node_type)), None)
@@ -441,7 +409,7 @@ def test_get_result_and_thought(basic_graph):
     # now build again and check if FakeListLLM was used
 
     # Get the result and thought
-    result, thought = get_result_and_thought_using_graph(langchain_object, message)
+    result, thought = get_result_and_steps(langchain_object, message)
     # The result should be a str
     assert isinstance(result, str)
     # The thought should be a Thought

tests/test_websocket.py

@@ -0,0 +1,30 @@
+import json
+
+
+def test_websocket_connection(client):
+    with client.websocket_connect("/ws") as websocket:
+        assert websocket.client == client
+        assert websocket.url.path == "/ws"
+
+
+def test_chat_history(client):
+    chat_history = ["Test message 1", "Test message 2"]
+
+    with client.websocket_connect("/ws") as websocket:
+        received_history = websocket.receive_text()
+        received_history = json.loads(received_history)
+
+        assert received_history == chat_history
+
+
+def test_send_message(client, basic_graph):
+    with client.websocket_connect("/ws") as websocket:
+        # Send the JSON payload through the WebSocket connection
+        websocket.send_text(basic_graph)
+
+        # Receive and parse the response from the server
+        response = websocket.receive_text()
+        response = json.loads(response)
+
+        # Test that the response is as expected
+        assert response == "Your response message here"