From 1165afbdd17cf072c434020376d653b5f411fdd6 Mon Sep 17 00:00:00 2001
From: Jiajie Zhong <zhongjiajie955@hotmail.com>
Date: Thu, 4 Nov 2021 22:25:45 +0800
Subject: [PATCH] [python] Add param workgroup to process definition (#6667)

---
 .../src/pydolphinscheduler/core/process_definition.py       | 6 +++++-
 .../tests/core/test_process_definition.py                   | 3 +++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/dolphinscheduler-python/pydolphinscheduler/src/pydolphinscheduler/core/process_definition.py b/dolphinscheduler-python/pydolphinscheduler/src/pydolphinscheduler/core/process_definition.py
index 8d295d0809..fa6ad97a9a 100644
--- a/dolphinscheduler-python/pydolphinscheduler/src/pydolphinscheduler/core/process_definition.py
+++ b/dolphinscheduler-python/pydolphinscheduler/src/pydolphinscheduler/core/process_definition.py
@@ -60,6 +60,7 @@ class ProcessDefinition(Base):
         "description",
         "_project",
         "_tenant",
+        "worker_group",
         "timeout",
         "release_state",
         "param",
@@ -76,6 +77,7 @@ class ProcessDefinition(Base):
             project: Optional[str] = ProcessDefinitionDefault.PROJECT,
             tenant: Optional[str] = ProcessDefinitionDefault.TENANT,
             queue: Optional[str] = ProcessDefinitionDefault.QUEUE,
+            worker_group: Optional[str] = ProcessDefinitionDefault.WORKER_GROUP,
             timeout: Optional[int] = 0,
             release_state: Optional[str] = ProcessDefinitionReleaseState.ONLINE,
             param: Optional[List] = None
@@ -85,6 +87,7 @@ class ProcessDefinition(Base):
         self._project = project
         self._tenant = tenant
         self._queue = queue
+        self.worker_group = worker_group
         self.timeout = timeout
         self.release_state = release_state
         self.param = param
@@ -226,6 +229,7 @@ class ProcessDefinition(Base):
             str(self.param) if self.param else None,
             json.dumps(self.task_location),
             self.timeout,
+            self.worker_group,
             self._tenant,
             # TODO add serialization function
             json.dumps(self.task_relation_json),
@@ -244,6 +248,6 @@ class ProcessDefinition(Base):
             self._project,
             self.name,
             "",
-            "default",
+            self.worker_group,
             24 * 3600,
         )
diff --git a/dolphinscheduler-python/pydolphinscheduler/tests/core/test_process_definition.py b/dolphinscheduler-python/pydolphinscheduler/tests/core/test_process_definition.py
index 83a4c2b965..f4b6b1c89c 100644
--- a/dolphinscheduler-python/pydolphinscheduler/tests/core/test_process_definition.py
+++ b/dolphinscheduler-python/pydolphinscheduler/tests/core/test_process_definition.py
@@ -49,6 +49,7 @@ def test_process_definition_key_attr(func):
                       ProcessDefinitionDefault.TENANT,
                       ProcessDefinitionDefault.QUEUE,
                       ProcessDefinitionDefault.USER_STATE)),
+        ("worker_group", ProcessDefinitionDefault.WORKER_GROUP),
         ("release_state", ProcessDefinitionReleaseState.ONLINE),
     ],
 )
@@ -63,6 +64,7 @@ def test_process_definition_default_value(name, value):
     [
         ("project", Project, "project"),
         ("tenant", Tenant, "tenant"),
+        ("worker_group", str, "worker_group"),
     ],
 )
 def test_process_definition_set_attr(name, cls, expect):
@@ -78,6 +80,7 @@ def test_process_definition_to_dict_without_task():
         "description": None,
         "project": ProcessDefinitionDefault.PROJECT,
         "tenant": ProcessDefinitionDefault.TENANT,
+        "workerGroup": ProcessDefinitionDefault.WORKER_GROUP,
         "timeout": 0,
         "releaseState": ProcessDefinitionReleaseState.ONLINE,
         "param": None,