This commit is contained in:
takatost
2024-07-17 11:26:33 +08:00
parent 16e2d00157
commit cc96acdae3
4 changed files with 199 additions and 116 deletions

View File

@ -1,9 +1,20 @@
from unittest.mock import patch
from flask import Flask
from core.app.entities.app_invoke_entities import InvokeFrom
from core.workflow.entities.node_entities import SystemVariable, UserFrom
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.graph_engine.entities.event import (
BaseNodeEvent,
GraphRunFailedEvent,
GraphRunStartedEvent,
NodeRunFailedEvent,
NodeRunStartedEvent,
NodeRunSucceededEvent,
)
from core.workflow.graph_engine.entities.graph import Graph
from core.workflow.graph_engine.entities.runtime_route_state import RouteNodeState
from core.workflow.graph_engine.graph_engine import GraphEngine
from models.workflow import WorkflowType
@ -14,31 +25,29 @@ def test_run(mock_close, mock_remove):
graph_config = {
"edges": [
{
"id": "llm-source-answer-target",
"source": "llm",
"target": "answer",
},
{
"id": "start-source-qc-target",
"id": "1",
"source": "start",
"target": "qc",
"target": "answer1",
},
{
"id": "qc-1-llm-target",
"source": "qc",
"sourceHandle": "1",
"target": "llm",
},
{
"id": "qc-2-http-target",
"source": "qc",
"sourceHandle": "2",
"target": "http",
},
{
"id": "http-source-answer2-target",
"source": "http",
"id": "2",
"source": "answer1",
"target": "answer2",
},
{
"id": "3",
"source": "answer1",
"target": "answer3",
},
{
"id": "4",
"source": "answer2",
"target": "answer4",
},
{
"id": "5",
"source": "answer3",
"target": "answer5",
}
],
"nodes": [
@ -51,38 +60,43 @@ def test_run(mock_close, mock_remove):
},
{
"data": {
"type": "llm",
"title": "llm"
"type": "answer",
"title": "answer1",
"answer": "1"
},
"id": "llm"
"id": "answer1"
},
{
"data": {
"type": "answer",
"title": "answer"
},
"id": "answer",
},
{
"data": {
"type": "question-classifier",
"title": "qc"
},
"id": "qc",
},
{
"data": {
"type": "http-request",
"title": "http"
},
"id": "http",
},
{
"data": {
"type": "answer",
"title": "answer2"
"title": "answer2",
"answer": "2"
},
"id": "answer2",
},
{
"data": {
"type": "answer",
"title": "answer3",
"answer": "3"
},
"id": "answer3",
},
{
"data": {
"type": "answer",
"title": "answer4",
"answer": "4"
},
"id": "answer4",
},
{
"data": {
"type": "answer",
"title": "answer5",
"answer": "5"
},
"id": "answer5",
}
],
}
@ -115,6 +129,28 @@ def test_run(mock_close, mock_remove):
print("")
generator = graph_engine.run()
for item in generator:
print(type(item), item)
app = Flask('test')
items = []
with app.app_context():
generator = graph_engine.run()
for item in generator:
print(type(item), item)
items.append(item)
if isinstance(item, NodeRunSucceededEvent):
assert item.route_node_state.status == RouteNodeState.Status.SUCCESS
assert not isinstance(item, NodeRunFailedEvent)
assert not isinstance(item, GraphRunFailedEvent)
if isinstance(item, BaseNodeEvent) and item.route_node_state.node_id in ['answer2', 'answer3']:
assert item.parallel_id is not None
assert len(items) == 12
assert isinstance(items[0], GraphRunStartedEvent)
assert isinstance(items[1], NodeRunStartedEvent)
assert items[1].route_node_state.node_id == 'start'
assert isinstance(items[2], NodeRunSucceededEvent)
assert items[2].route_node_state.node_id == 'start'
print(graph_engine.graph_runtime_state)