Skip to content

Commit 347cac3

Browse files
pankajkotipankajkoti
committed
Refactor tests to demonstrate tasks and taskgroups as list
1 parent 75229d9 commit 347cac3

4 files changed

Lines changed: 109 additions & 269 deletions

File tree

‎tests/test_dagbuilder.py‎

Lines changed: 85 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -76,89 +76,104 @@
7676
get_schedule_key(): "0 3 * * *",
7777
"tags": ["tag1", "tag2"],
7878
"render_template_as_native_obj": True,
79-
"tasks": {
80-
"task_1": {
79+
"tasks": [
80+
{
81+
"task_id": "task_1",
8182
"operator": get_bash_operator_path(),
8283
"bash_command": "echo 1",
8384
"execution_timeout_secs": 5,
8485
},
85-
"task_2": {
86+
{
87+
"task_id": "task_2",
8688
"operator": get_bash_operator_path(),
8789
"bash_command": "echo 2",
8890
"dependencies": ["task_1"],
8991
},
90-
"task_3": {
92+
{
93+
"task_id": "task_3",
9194
"operator": get_bash_operator_path(),
9295
"bash_command": "echo 3",
9396
"dependencies": ["task_1"],
9497
},
95-
},
98+
],
9699
}
97100
DAG_CONFIG_TASK_GROUP = {
98101
"default_args": {"owner": "custom_owner"},
99102
get_schedule_key(): "0 3 * * *",
100-
"task_groups": {
101-
"task_group_1": {
103+
"task_groups": [
104+
{
105+
"group_name": "task_group_1",
102106
"tooltip": "this is a task group",
103107
"dependencies": ["task_1"],
104108
},
105-
"task_group_2": {
109+
{
110+
"group_name": "task_group_2",
106111
"dependencies": ["task_group_1"],
107112
},
108-
"task_group_3": {},
109-
},
110-
"tasks": {
111-
"task_1": {
113+
{
114+
"group_name": "task_group_3",
115+
},
116+
],
117+
"tasks": [
118+
{
119+
"task_id": "task_1",
112120
"operator": get_bash_operator_path(),
113121
"bash_command": "echo 1",
114122
},
115-
"task_2": {
123+
{
124+
"task_id": "task_2",
116125
"operator": get_bash_operator_path(),
117126
"bash_command": "echo 2",
118127
"task_group_name": "task_group_1",
119128
},
120-
"task_3": {
129+
{
130+
"task_id": "task_3",
121131
"operator": get_bash_operator_path(),
122132
"bash_command": "echo 3",
123133
"task_group_name": "task_group_1",
124134
"dependencies": ["task_2"],
125135
},
126-
"task_4": {
136+
{
137+
"task_id": "task_4",
127138
"operator": get_bash_operator_path(),
128139
"bash_command": "echo 4",
129140
"dependencies": ["task_group_1"],
130141
},
131-
"task_5": {
142+
{
143+
"task_id": "task_5",
132144
"operator": get_bash_operator_path(),
133145
"bash_command": "echo 5",
134146
"task_group_name": "task_group_2",
135147
},
136-
"task_6": {
148+
{
149+
"task_id": "task_6",
137150
"operator": get_bash_operator_path(),
138151
"bash_command": "echo 6",
139152
"task_group_name": "task_group_2",
140153
"dependencies": ["task_5"],
141154
},
142-
},
155+
],
143156
}
144157
DAG_CONFIG_DYNAMIC_TASK_MAPPING = {
145158
"default_args": {"owner": "custom_owner"},
146159
"description": "This is an example dag with dynamic task mapping",
147160
get_schedule_key(): "0 4 * * *",
148-
"tasks": {
149-
"request": {
161+
"tasks": [
162+
{
163+
"task_id": "request",
150164
"operator": get_python_operator_path(),
151165
"python_callable_name": "example_task_mapping",
152166
"python_callable_file": os.path.realpath(__file__),
153167
},
154-
"process_1": {
168+
{
169+
"task_id": "process_1",
155170
"operator": get_python_operator_path(),
156171
"python_callable_name": "expand_task",
157172
"python_callable_file": os.path.realpath(__file__),
158173
"partial": {"op_kwargs": {"test_id": "test"}},
159174
"expand": {"op_args": {"request_output": "request.output"}},
160175
},
161-
},
176+
],
162177
}
163178

164179
DAG_CONFIG_ML = {
@@ -205,15 +220,16 @@
205220
},
206221
"on_failure_callback_name": "print_context_callback",
207222
"on_failure_callback_file": __file__,
208-
"tasks": {
209-
"task_1": { # Make sure that default_args are applied to this Task
223+
"tasks": [
224+
{ # Make sure that default_args are applied to this Task
225+
"task_id": "task_1",
210226
"operator": get_bash_operator_path(),
211227
"bash_command": "echo 1",
212228
"execution_timeout_secs": 5,
213229
"on_failure_callback_name": "print_context_callback",
214230
"on_failure_callback_file": __file__,
215231
}
216-
},
232+
],
217233
}
218234

219235
DAG_CONFIG_TASK_GROUP_WITH_CALLBACKS = {
@@ -226,8 +242,9 @@
226242
},
227243
},
228244
get_schedule_key(): "0 3 * * *",
229-
"task_groups": {
230-
"task_group_1": {
245+
"task_groups": [
246+
{
247+
"group_name": "task_group_1",
231248
"tooltip": "this is a task group",
232249
"default_args": {
233250
"on_execute_callback": f"{__name__}.print_context_callback",
@@ -237,14 +254,16 @@
237254
"on_skip_callback": f"{__name__}.print_context_callback", # Throwing this in for good measure
238255
},
239256
},
240-
},
241-
"tasks": {
242-
"task_1": {
257+
],
258+
"tasks": [
259+
{
260+
"task_id": "task_1",
243261
"operator": get_bash_operator_path(),
244262
"bash_command": "echo 1",
245263
"task_group_name": "task_group_1",
246264
},
247-
"task_2": {
265+
{
266+
"task_id": "task_2",
248267
"operator": get_bash_operator_path(),
249268
"bash_command": "echo 2",
250269
"task_group_name": "task_group_1",
@@ -254,7 +273,8 @@
254273
"param_2": "value_2",
255274
},
256275
},
257-
"task_3": {
276+
{
277+
"task_id": "task_3",
258278
"operator": get_bash_operator_path(),
259279
"bash_command": "echo 3",
260280
"task_group_name": "task_group_1",
@@ -264,7 +284,8 @@
264284
# - String with no parameters
265285
# - String with parameters
266286
# - File name and path
267-
"task_4": {
287+
{
288+
"task_id": "task_4",
268289
"operator": get_bash_operator_path(),
269290
"bash_command": "echo 4",
270291
"dependencies": ["task_group_1"],
@@ -277,7 +298,7 @@
277298
"on_failure_callback_name": "print_context_callback",
278299
"on_failure_callback_file": __file__,
279300
},
280-
},
301+
],
281302
}
282303

283304

@@ -313,23 +334,26 @@ def test_get_dag_params():
313334
"dagrun_timeout": datetime.timedelta(seconds=600),
314335
"render_template_as_native_obj": True,
315336
"tags": ["tag1", "tag2"],
316-
"tasks": {
317-
"task_1": {
337+
"tasks": [
338+
{
339+
"task_id": "task_1",
318340
"operator": get_bash_operator_path(),
319341
"bash_command": "echo 1",
320342
"execution_timeout_secs": 5,
321343
},
322-
"task_2": {
344+
{
345+
"task_id": "task_2",
323346
"operator": get_bash_operator_path(),
324347
"bash_command": "echo 2",
325348
"dependencies": ["task_1"],
326349
},
327-
"task_3": {
350+
{
351+
"task_id": "task_3",
328352
"operator": get_bash_operator_path(),
329353
"bash_command": "echo 3",
330354
"dependencies": ["task_1"],
331355
},
332-
},
356+
],
333357
}
334358
actual = td.get_dag_params()
335359
assert actual == expected
@@ -572,47 +596,56 @@ def test_get_dag_params_dag_with_task_group():
572596
"retry_delay": datetime.timedelta(seconds=300),
573597
},
574598
get_schedule_key(): "0 3 * * *",
575-
"task_groups": {
576-
"task_group_1": {
599+
"task_groups": [
600+
{
601+
"group_name": "task_group_1",
577602
"tooltip": "this is a task group",
578603
"dependencies": ["task_1"],
579604
},
580-
"task_group_2": {"dependencies": ["task_group_1"]},
581-
"task_group_3": {},
582-
},
583-
"tasks": {
584-
"task_1": {
605+
{"group_name": "task_group_2", "dependencies": ["task_group_1"]},
606+
{
607+
"group_name": "task_group_3",
608+
},
609+
],
610+
"tasks": [
611+
{
612+
"task_id": "task_1",
585613
"operator": get_bash_operator_path(),
586614
"bash_command": "echo 1",
587615
},
588-
"task_2": {
616+
{
617+
"task_id": "task_2",
589618
"operator": get_bash_operator_path(),
590619
"bash_command": "echo 2",
591620
"task_group_name": "task_group_1",
592621
},
593-
"task_3": {
622+
{
623+
"task_id": "task_3",
594624
"operator": get_bash_operator_path(),
595625
"bash_command": "echo 3",
596626
"task_group_name": "task_group_1",
597627
"dependencies": ["task_2"],
598628
},
599-
"task_4": {
629+
{
630+
"task_id": "task_4",
600631
"operator": get_bash_operator_path(),
601632
"bash_command": "echo 4",
602633
"dependencies": ["task_group_1"],
603634
},
604-
"task_5": {
635+
{
636+
"task_id": "task_5",
605637
"operator": get_bash_operator_path(),
606638
"bash_command": "echo 5",
607639
"task_group_name": "task_group_2",
608640
},
609-
"task_6": {
641+
{
642+
"task_id": "task_6",
610643
"operator": get_bash_operator_path(),
611644
"bash_command": "echo 6",
612645
"task_group_name": "task_group_2",
613646
"dependencies": ["task_5"],
614647
},
615-
},
648+
],
616649
"concurrency": 1,
617650
"max_active_runs": 1,
618651
"dag_id": "test_dag",

‎tests/test_dagbuilder_httpoperator.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -175,16 +175,18 @@ def test_dag_with_http_operator():
175175
http_dag_config = {
176176
"default_args": {"owner": "test_owner", "start_date": datetime.date(2023, 1, 1)},
177177
get_schedule_key(): "0 0 * * *",
178-
"tasks": {
179-
"http_task_json": {
178+
"tasks": [
179+
{
180+
"task_id": "http_task_json",
180181
"operator": HTTP_OPERATOR_PATH,
181182
"http_conn_id": "test_conn",
182183
"method": "POST",
183184
"endpoint": "/api/test",
184185
"headers": {"Content-Type": "application/json"},
185186
"data": {"message": "test data", "value": 123},
186187
},
187-
"http_task_plain": {
188+
{
189+
"task_id": "http_task_plain",
188190
"operator": HTTP_OPERATOR_PATH,
189191
"http_conn_id": "test_conn",
190192
"method": "POST",
@@ -193,7 +195,7 @@ def test_dag_with_http_operator():
193195
"data": "plain text data",
194196
"dependencies": ["http_task_json"],
195197
},
196-
},
198+
],
197199
}
198200

199201
# Build the DAG

0 commit comments

Comments
 (0)