From 5b5e6e31bf2549ca7e1fa057f74725cfefe598fb Mon Sep 17 00:00:00 2001
From: takatost <takatost@gmail.com>
Date: Fri, 16 Aug 2024 01:44:00 +0800
Subject: [PATCH] fix: answer node unit tests

---
 .../core/workflow/nodes/test_answer.py        | 77 ++++++++++++++++---
 1 file changed, 65 insertions(+), 12 deletions(-)

diff --git a/api/tests/unit_tests/core/workflow/nodes/test_answer.py b/api/tests/unit_tests/core/workflow/nodes/test_answer.py
index 4617b6a42f..71660229de 100644
--- a/api/tests/unit_tests/core/workflow/nodes/test_answer.py
+++ b/api/tests/unit_tests/core/workflow/nodes/test_answer.py
@@ -1,22 +1,83 @@
+import time
+import uuid
 from unittest.mock import MagicMock
 
 from core.app.entities.app_invoke_entities import InvokeFrom
+from core.workflow.entities.node_entities import UserFrom
 from core.workflow.entities.variable_pool import VariablePool
 from core.workflow.enums import SystemVariable
+from core.workflow.graph_engine.entities.graph import Graph
+from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams
+from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
 from core.workflow.nodes.answer.answer_node import AnswerNode
-from core.workflow.nodes.base_node import UserFrom
 from extensions.ext_database import db
-from models.workflow import WorkflowNodeExecutionStatus
+from models.workflow import WorkflowNodeExecutionStatus, WorkflowType
 
 
 def test_execute_answer():
-    node = AnswerNode(
+    graph_config = {
+        "edges": [
+            {
+                "id": "start-source-answer-target",
+                "source": "start",
+                "target": "answer",
+            },
+        ],
+        "nodes": [
+            {
+                "data": {
+                    "type": "start"
+                },
+                "id": "start"
+            },
+            {
+                "data": {
+                    'title': '123',
+                    'type': 'answer',
+                    'answer': 'Today\'s weather is {{#start.weather#}}\n{{#llm.text#}}\n{{img}}\nFin.'
+                },
+                "id": "answer"
+            },
+        ]
+    }
+
+    graph = Graph.init(
+        graph_config=graph_config
+    )
+
+    init_params = GraphInitParams(
         tenant_id='1',
         app_id='1',
+        workflow_type=WorkflowType.WORKFLOW,
         workflow_id='1',
+        graph_config=graph_config,
         user_id='1',
         user_from=UserFrom.ACCOUNT,
         invoke_from=InvokeFrom.DEBUGGER,
+        call_depth=0
+    )
+
+    # construct variable pool
+    variable_pool = VariablePool(
+        system_variables={
+            SystemVariable.FILES: [],
+            SystemVariable.USER_ID: 'aaa'
+        },
+        user_inputs={},
+        environment_variables=[],
+        conversation_variables=[],
+    )
+    variable_pool.add(['start', 'weather'], 'sunny')
+    variable_pool.add(['llm', 'text'], 'You are a helpful AI.')
+
+    node = AnswerNode(
+        id=str(uuid.uuid4()),
+        graph_init_params=init_params,
+        graph=graph,
+        graph_runtime_state=GraphRuntimeState(
+            variable_pool=variable_pool,
+            start_at=time.perf_counter()
+        ),
         config={
             'id': 'answer',
             'data': {
@@ -27,19 +88,11 @@ def test_execute_answer():
         }
     )
 
-    # construct variable pool
-    pool = VariablePool(system_variables={
-        SystemVariable.FILES: [],
-        SystemVariable.USER_ID: 'aaa'
-    }, user_inputs={}, environment_variables=[])
-    pool.add(['start', 'weather'], 'sunny')
-    pool.add(['llm', 'text'], 'You are a helpful AI.')
-
     # Mock db.session.close()
     db.session.close = MagicMock()
 
     # execute node
-    result = node._run(pool)
+    result = node._run()
 
     assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
     assert result.outputs['answer'] == "Today's weather is sunny\nYou are a helpful AI.\n{{img}}\nFin."