Refactor test_initial_setup.py to use RunOutputs instead of ResultData

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-04-01 08:58:55 -03:00
commit 3c42c6e37d

View file

@ -1,14 +1,13 @@
from itertools import chain
import pytest import pytest
from langflow.graph.graph.base import Graph from langflow.graph.graph.base import Graph
from langflow.graph.schema import ResultData from langflow.graph.schema import RunOutputs
from langflow.initial_setup.setup import ( from langflow.initial_setup.setup import (
STARTER_FOLDER_NAME, STARTER_FOLDER_NAME,
create_or_update_starter_projects, create_or_update_starter_projects,
get_project_data, get_project_data,
load_starter_projects, load_starter_projects,
) )
from langflow.memory import delete_messages
from langflow.services.database.models.flow.model import Flow from langflow.services.database.models.flow.model import Flow
from langflow.services.deps import session_scope from langflow.services.deps import session_scope
from sqlalchemy import func from sqlalchemy import func
@ -61,16 +60,15 @@ async def test_starter_project_can_run_successfully(client):
# Get all the starter projects # Get all the starter projects
projects = session.exec(select(Flow).where(Flow.folder == STARTER_FOLDER_NAME)).all() projects = session.exec(select(Flow).where(Flow.folder == STARTER_FOLDER_NAME)).all()
graphs: list[Graph] = [ graphs: list[tuple[str, Graph]] = [
(project.name, Graph.from_payload(project.data, flow_id=project.id)) for project in projects (project.name, Graph.from_payload(project.data, flow_id=project.id)) for project in projects
] ]
assert len(graphs) == len(projects) assert len(graphs) == len(projects)
for name, graph in graphs: for name, graph in graphs:
outputs = await graph.run( outputs = await graph.arun(
inputs={}, inputs={},
outputs=[], outputs=[],
session_id="test", session_id="test",
) )
assert all( assert all(isinstance(output, RunOutputs) for output in outputs), f"Project {name} error: {outputs}"
isinstance(output, ResultData) for output in chain.from_iterable(outputs) delete_messages(session_id="test")
), f"Project {name} error: {outputs}"