From eae53e11e644f6f739b5890488b0b93d5f4913ad Mon Sep 17 00:00:00 2001
From: -LAN- <laipz8200@outlook.com>
Date: Tue, 20 Aug 2024 17:52:21 +0800
Subject: [PATCH] refactor(api/models/workflow.py): Add `__init__` to Workflow
 (#7443)

---
 api/models/workflow.py                       | 40 +++++++++++++-------
 api/tests/unit_tests/models/test_workflow.py | 38 ++++++++++++++++---
 2 files changed, 60 insertions(+), 18 deletions(-)

diff --git a/api/models/workflow.py b/api/models/workflow.py
index 396c023dc1..cdd5e1992d 100644
--- a/api/models/workflow.py
+++ b/api/models/workflow.py
@@ -1,5 +1,6 @@
 import json
 from collections.abc import Mapping, Sequence
+from datetime import datetime
 from enum import Enum
 from typing import Any, Optional, Union
 
@@ -110,19 +111,32 @@ class Workflow(db.Model):
         db.Index('workflow_version_idx', 'tenant_id', 'app_id', 'version'),
     )
 
-    id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
-    tenant_id = db.Column(StringUUID, nullable=False)
-    app_id = db.Column(StringUUID, nullable=False)
-    type = db.Column(db.String(255), nullable=False)
-    version = db.Column(db.String(255), nullable=False)
-    graph = db.Column(db.Text)
-    features = db.Column(db.Text)
-    created_by = db.Column(StringUUID, nullable=False)
-    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
-    updated_by = db.Column(StringUUID)
-    updated_at = db.Column(db.DateTime)
-    _environment_variables = db.Column('environment_variables', db.Text, nullable=False, server_default='{}')
-    _conversation_variables = db.Column('conversation_variables', db.Text, nullable=False, server_default='{}')
+    id: Mapped[str] = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
+    tenant_id: Mapped[str] = db.Column(StringUUID, nullable=False)
+    app_id: Mapped[str] = db.Column(StringUUID, nullable=False)
+    type: Mapped[str] = db.Column(db.String(255), nullable=False)
+    version: Mapped[str] = db.Column(db.String(255), nullable=False)
+    graph: Mapped[str] = db.Column(db.Text)
+    features: Mapped[str] = db.Column(db.Text)
+    created_by: Mapped[str] = db.Column(StringUUID, nullable=False)
+    created_at: Mapped[datetime] = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
+    updated_by: Mapped[str] = db.Column(StringUUID)
+    updated_at: Mapped[datetime] = db.Column(db.DateTime)
+    _environment_variables: Mapped[str] = db.Column('environment_variables', db.Text, nullable=False, server_default='{}')
+    _conversation_variables: Mapped[str] = db.Column('conversation_variables', db.Text, nullable=False, server_default='{}')
+
+    def __init__(self, *, tenant_id: str, app_id: str, type: str, version: str, graph: str,
+                 features: str, created_by: str, environment_variables: Sequence[Variable],
+                 conversation_variables: Sequence[Variable]):
+        self.tenant_id = tenant_id
+        self.app_id = app_id
+        self.type = type
+        self.version = version
+        self.graph = graph
+        self.features = features
+        self.created_by = created_by
+        self.environment_variables = environment_variables or []
+        self.conversation_variables = conversation_variables or []
 
     @property
     def created_by_account(self):
diff --git a/api/tests/unit_tests/models/test_workflow.py b/api/tests/unit_tests/models/test_workflow.py
index facea34b5b..bea896b83a 100644
--- a/api/tests/unit_tests/models/test_workflow.py
+++ b/api/tests/unit_tests/models/test_workflow.py
@@ -11,7 +11,17 @@ def test_environment_variables():
     contexts.tenant_id.set('tenant_id')
 
     # Create a Workflow instance
-    workflow = Workflow()
+    workflow = Workflow(
+        tenant_id='tenant_id',
+        app_id='app_id',
+        type='workflow',
+        version='draft',
+        graph='{}',
+        features='{}',
+        created_by='account_id',
+        environment_variables=[],
+        conversation_variables=[],
+    )
 
     # Create some EnvironmentVariable instances
     variable1 = StringVariable.model_validate({'name': 'var1', 'value': 'value1', 'id': str(uuid4())})
@@ -35,7 +45,17 @@ def test_update_environment_variables():
     contexts.tenant_id.set('tenant_id')
 
     # Create a Workflow instance
-    workflow = Workflow()
+    workflow = Workflow(
+        tenant_id='tenant_id',
+        app_id='app_id',
+        type='workflow',
+        version='draft',
+        graph='{}',
+        features='{}',
+        created_by='account_id',
+        environment_variables=[],
+        conversation_variables=[],
+    )
 
     # Create some EnvironmentVariable instances
     variable1 = StringVariable.model_validate({'name': 'var1', 'value': 'value1', 'id': str(uuid4())})
@@ -70,9 +90,17 @@ def test_to_dict():
     contexts.tenant_id.set('tenant_id')
 
     # Create a Workflow instance
-    workflow = Workflow()
-    workflow.graph = '{}'
-    workflow.features = '{}'
+    workflow = Workflow(
+        tenant_id='tenant_id',
+        app_id='app_id',
+        type='workflow',
+        version='draft',
+        graph='{}',
+        features='{}',
+        created_by='account_id',
+        environment_variables=[],
+        conversation_variables=[],
+    )
 
     # Create some EnvironmentVariable instances