From 6cacfced09fd33bbaaf7735c3c97c0e93d364446 Mon Sep 17 00:00:00 2001
From: Gabriel Almeida <gabriel@logspace.ai>
Date: Tue, 30 May 2023 21:24:23 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=A7=20chore(conftest.py):=20add=20fixt?=
 =?UTF-8?q?ures=20for=20database=20testing=20=E2=9C=85=20test(database.py)?=
 =?UTF-8?q?:=20add=20tests=20for=20CRUD=20operations=20on=20Flow=20model?=
 =?UTF-8?q?=20The=20new=20fixtures=20added=20to=20conftest.py=20are=20sess?=
 =?UTF-8?q?ion=20and=20client=20fixtures.=20These=20fixtures=20are=20used?=
 =?UTF-8?q?=20to=20create=20a=20test=20database=20and=20a=20test=20client?=
 =?UTF-8?q?=20for=20testing=20the=20database.=20The=20tests=20added=20to?=
 =?UTF-8?q?=20test=5Fdatabase.py=20test=20the=20CRUD=20operations=20on=20t?=
 =?UTF-8?q?he=20Flow=20model.=20The=20tests=20include=20creating=20a=20flo?=
 =?UTF-8?q?w,=20reading=20all=20flows,=20reading=20a=20single=20flow,=20up?=
 =?UTF-8?q?dating=20a=20flow,=20and=20deleting=20a=20flow.=20These=20tests?=
 =?UTF-8?q?=20ensure=20that=20the=20database=20is=20working=20as=20expecte?=
 =?UTF-8?q?d=20and=20that=20the=20API=20endpoints=20for=20the=20Flow=20mod?=
 =?UTF-8?q?el=20are=20functioning=20correctly.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/conftest.py      | 35 +++++++++++++++++++++++++
 tests/test_database.py | 58 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 tests/test_database.py

diff --git a/tests/conftest.py b/tests/conftest.py
index 870c48a32..763adf077 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,10 +1,13 @@
 import json
 from pathlib import Path
 from typing import AsyncGenerator
+from langflow.api.database import get_session
 
 import pytest
 from fastapi.testclient import TestClient
 from httpx import AsyncClient
+from sqlmodel import SQLModel, Session, create_engine
+from sqlmodel.pool import StaticPool
 
 
 def pytest_configure():
@@ -76,3 +79,35 @@ def complex_graph():
 @pytest.fixture
 def openapi_graph():
     return get_graph("openapi")
+
+
+@pytest.fixture
+def json_flow():
+    with open(pytest.BASIC_EXAMPLE_PATH, "r") as f:
+        return f.read()
+
+
+@pytest.fixture(name="session")  #
+def session_fixture():  #
+    engine = create_engine(
+        "sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool
+    )
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as session:
+        yield session
+
+
+@pytest.fixture(name="client")  #
+def client_fixture(session: Session):  #
+    def get_session_override():  #
+        return session
+
+    from langflow.main import create_app
+
+    app = create_app()
+
+    app.dependency_overrides[get_session] = get_session_override  #
+
+    client = TestClient(app)  #
+    yield client  #
+    app.dependency_overrides.clear()  #
diff --git a/tests/test_database.py b/tests/test_database.py
new file mode 100644
index 000000000..247d55e12
--- /dev/null
+++ b/tests/test_database.py
@@ -0,0 +1,58 @@
+from langflow.database.models.flow import FlowCreate
+
+
+def test_create_flow(client, json_flow):
+    flow = FlowCreate(name="Test Flow", flow=json_flow)
+    response = client.post("/flows/", json=flow.dict())
+    assert response.status_code == 200
+    assert response.json()["name"] == flow.name
+    assert response.json()["flow"] == flow.flow
+
+
+def test_read_flows(client, json_flow):
+    flow = FlowCreate(name="Test Flow", flow=json_flow)
+    response = client.post("/flows/", json=flow.dict())
+    assert response.status_code == 200
+    assert response.json()["name"] == flow.name
+    assert response.json()["flow"] == flow.flow
+    flow = FlowCreate(name="Test Flow", flow=json_flow)
+    response = client.post("/flows/", json=flow.dict())
+    assert response.status_code == 200
+    assert response.json()["name"] == flow.name
+    assert response.json()["flow"] == flow.flow
+
+    response = client.get("/flows/")
+    assert response.status_code == 200
+    assert len(response.json()) > 0
+
+
+def test_read_flow(client, json_flow):
+    flow = FlowCreate(name="Test Flow", flow=json_flow)
+    response = client.post("/flows/", json=flow.dict())
+    flow_id = response.json()["id"]
+    response = client.get(f"/flows/{flow_id}")
+    assert response.status_code == 200
+    assert response.json()["name"] == flow.name
+    assert response.json()["flow"] == flow.flow
+
+
+def test_update_flow(client, json_flow):
+    flow = FlowCreate(name="Test Flow", flow=json_flow)
+    response = client.post("/flows/", json=flow.dict())
+    flow_id = response.json()["id"]
+    updated_flow = FlowCreate(
+        name="Updated Flow", flow=json_flow.replace("BasicExample", "Updated Flow")
+    )
+    response = client.put(f"/flows/{flow_id}", json=updated_flow.dict())
+    assert response.status_code == 200
+    assert response.json()["name"] == updated_flow.name
+    assert response.json()["flow"] == updated_flow.flow
+
+
+def test_delete_flow(client, json_flow):
+    flow = FlowCreate(name="Test Flow", flow=json_flow)
+    response = client.post("/flows/", json=flow.dict())
+    flow_id = response.json()["id"]
+    response = client.delete(f"/flows/{flow_id}")
+    assert response.status_code == 200
+    assert response.json()["message"] == "Flow deleted successfully"