cornflow 2.0.0a11__py3-none-any.whl → 2.0.0a12__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (52) hide show
  1. airflow_config/airflow_local_settings.py +1 -1
  2. cornflow/app.py +8 -3
  3. cornflow/cli/migrations.py +23 -3
  4. cornflow/cli/service.py +18 -18
  5. cornflow/cli/utils.py +16 -1
  6. cornflow/commands/dag.py +1 -1
  7. cornflow/config.py +13 -8
  8. cornflow/endpoints/__init__.py +8 -2
  9. cornflow/endpoints/alarms.py +66 -2
  10. cornflow/endpoints/data_check.py +53 -26
  11. cornflow/endpoints/execution.py +387 -132
  12. cornflow/endpoints/login.py +81 -63
  13. cornflow/endpoints/meta_resource.py +11 -3
  14. cornflow/migrations/versions/999b98e24225.py +34 -0
  15. cornflow/models/base_data_model.py +4 -32
  16. cornflow/models/execution.py +2 -3
  17. cornflow/models/meta_models.py +28 -22
  18. cornflow/models/user.py +7 -10
  19. cornflow/schemas/alarms.py +8 -0
  20. cornflow/schemas/execution.py +1 -1
  21. cornflow/schemas/query.py +2 -1
  22. cornflow/schemas/user.py +5 -20
  23. cornflow/shared/authentication/auth.py +201 -264
  24. cornflow/shared/const.py +3 -14
  25. cornflow/shared/databricks.py +5 -1
  26. cornflow/tests/const.py +1 -0
  27. cornflow/tests/custom_test_case.py +77 -26
  28. cornflow/tests/unit/test_actions.py +2 -2
  29. cornflow/tests/unit/test_alarms.py +55 -1
  30. cornflow/tests/unit/test_apiview.py +108 -3
  31. cornflow/tests/unit/test_cases.py +20 -29
  32. cornflow/tests/unit/test_cli.py +6 -5
  33. cornflow/tests/unit/test_commands.py +3 -3
  34. cornflow/tests/unit/test_dags.py +5 -6
  35. cornflow/tests/unit/test_executions.py +443 -123
  36. cornflow/tests/unit/test_instances.py +14 -2
  37. cornflow/tests/unit/test_instances_file.py +1 -1
  38. cornflow/tests/unit/test_licenses.py +1 -1
  39. cornflow/tests/unit/test_log_in.py +230 -207
  40. cornflow/tests/unit/test_permissions.py +8 -8
  41. cornflow/tests/unit/test_roles.py +48 -10
  42. cornflow/tests/unit/test_schemas.py +1 -1
  43. cornflow/tests/unit/test_tables.py +7 -7
  44. cornflow/tests/unit/test_token.py +19 -5
  45. cornflow/tests/unit/test_users.py +22 -6
  46. cornflow/tests/unit/tools.py +75 -10
  47. {cornflow-2.0.0a11.dist-info → cornflow-2.0.0a12.dist-info}/METADATA +16 -15
  48. {cornflow-2.0.0a11.dist-info → cornflow-2.0.0a12.dist-info}/RECORD +51 -51
  49. {cornflow-2.0.0a11.dist-info → cornflow-2.0.0a12.dist-info}/WHEEL +1 -1
  50. cornflow/endpoints/execution_databricks.py +0 -808
  51. {cornflow-2.0.0a11.dist-info → cornflow-2.0.0a12.dist-info}/entry_points.txt +0 -0
  52. {cornflow-2.0.0a11.dist-info → cornflow-2.0.0a12.dist-info}/top_level.txt +0 -0
@@ -6,6 +6,8 @@ Unit test for the executions endpoints
6
6
  import json
7
7
  from unittest.mock import patch
8
8
 
9
+ from cornflow.app import create_app
10
+
9
11
  # Import from internal modules
10
12
  from cornflow.models import ExecutionModel, InstanceModel
11
13
  from cornflow.tests.const import (
@@ -20,7 +22,7 @@ from cornflow.tests.const import (
20
22
  EXECUTION_SOLUTION_PATH,
21
23
  )
22
24
  from cornflow.tests.custom_test_case import CustomTestCase, BaseTestCases
23
- from cornflow.tests.unit.tools import patch_af_client
25
+ from cornflow.tests.unit.tools import patch_af_client, patch_db_client
24
26
 
25
27
 
26
28
  class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
@@ -62,13 +64,13 @@ class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
62
64
  def test_new_execution(self):
63
65
  self.create_new_row(self.url, self.model, payload=self.payload)
64
66
 
65
- @patch("cornflow.endpoints.execution_databricks.Airflow")
67
+ @patch("cornflow.endpoints.execution.Airflow")
66
68
  def test_new_execution_run(self, af_client_class):
67
69
  patch_af_client(af_client_class)
68
70
 
69
71
  self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
70
72
 
71
- @patch("cornflow.endpoints.execution_databricks.Airflow")
73
+ @patch("cornflow.endpoints.execution.Airflow")
72
74
  def test_new_execution_bad_config(self, af_client_class):
73
75
  patch_af_client(af_client_class)
74
76
  response = self.create_new_row(
@@ -81,7 +83,7 @@ class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
81
83
  self.assertIn("error", response)
82
84
  self.assertIn("jsonschema_errors", response)
83
85
 
84
- @patch("cornflow.endpoints.execution_databricks.Airflow")
86
+ @patch("cornflow.endpoints.execution.Airflow")
85
87
  def test_new_execution_partial_config(self, af_client_class):
86
88
  patch_af_client(af_client_class)
87
89
  self.payload["config"].pop("solver")
@@ -91,7 +93,7 @@ class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
91
93
  self.assertIn("solver", response["config"])
92
94
  self.assertEqual(response["config"]["solver"], "cbc")
93
95
 
94
- @patch("cornflow.endpoints.execution_databricks.Airflow")
96
+ @patch("cornflow.endpoints.execution.Airflow")
95
97
  def test_new_execution_with_solution(self, af_client_class):
96
98
  patch_af_client(af_client_class)
97
99
  self.payload["data"] = self.solution
@@ -102,7 +104,7 @@ class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
102
104
  check_payload=False,
103
105
  )
104
106
 
105
- @patch("cornflow.endpoints.execution_databricks.Airflow")
107
+ @patch("cornflow.endpoints.execution.Airflow")
106
108
  def test_new_execution_with_solution_bad(self, af_client_class):
107
109
  patch_af_client(af_client_class)
108
110
  patch_af_client(af_client_class)
@@ -144,6 +146,131 @@ class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
144
146
  self.assertEqual(len(rows.json), len(self.payloads))
145
147
 
146
148
 
149
+ class TestExecutionsListEndpointDatabricks(BaseTestCases.ListFilters):
150
+ def setUp(self):
151
+ super().setUp()
152
+
153
+ with open(INSTANCE_PATH) as f:
154
+ payload = json.load(f)
155
+ fk_id = self.create_new_row(INSTANCE_URL, InstanceModel, payload)
156
+ self.url = EXECUTION_URL_NORUN
157
+ self.model = ExecutionModel
158
+
159
+ def load_file_fk(_file):
160
+ with open(_file) as f:
161
+ temp = json.load(f)
162
+ temp["instance_id"] = fk_id
163
+ return temp
164
+
165
+ self.payload = load_file_fk(EXECUTION_PATH)
166
+ self.bad_payload = load_file_fk(BAD_EXECUTION_PATH)
167
+ self.payloads = [load_file_fk(f) for f in EXECUTIONS_LIST]
168
+ self.solution = load_file_fk(EXECUTION_SOLUTION_PATH)
169
+ self.keys_to_check = [
170
+ "data_hash",
171
+ "created_at",
172
+ "config",
173
+ "state",
174
+ "message",
175
+ "schema",
176
+ "description",
177
+ "id",
178
+ "user_id",
179
+ "log",
180
+ "instance_id",
181
+ "name",
182
+ "indicators",
183
+ ]
184
+
185
+ def create_app(self):
186
+ app = create_app("testing-databricks")
187
+ return app
188
+
189
+ def test_new_execution(self):
190
+ self.create_new_row(self.url, self.model, payload=self.payload)
191
+
192
+ @patch("cornflow.endpoints.execution.Databricks")
193
+ def test_new_execution_run_databricks(self, db_client_class):
194
+ patch_db_client(db_client_class)
195
+
196
+ self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
197
+
198
+ @patch("cornflow.endpoints.execution.Databricks")
199
+ def test_new_execution_bad_config(self, db_client_class):
200
+ patch_db_client(db_client_class)
201
+ response = self.create_new_row(
202
+ EXECUTION_URL,
203
+ self.model,
204
+ payload=self.bad_payload,
205
+ expected_status=400,
206
+ check_payload=False,
207
+ )
208
+ self.assertIn("error", response)
209
+ self.assertIn("jsonschema_errors", response)
210
+
211
+ @patch("cornflow.endpoints.execution.Databricks")
212
+ def test_new_execution_partial_config(self, db_client_class):
213
+ patch_db_client(db_client_class)
214
+ self.payload["config"].pop("solver")
215
+ response = self.create_new_row(
216
+ EXECUTION_URL, self.model, payload=self.payload, check_payload=False
217
+ )
218
+ self.assertIn("solver", response["config"])
219
+ self.assertEqual(response["config"]["solver"], "cbc")
220
+
221
+ @patch("cornflow.endpoints.execution.Databricks")
222
+ def test_new_execution_with_solution(self, db_client_class):
223
+ patch_db_client(db_client_class)
224
+ self.payload["data"] = self.solution
225
+ response = self.create_new_row(
226
+ EXECUTION_URL,
227
+ self.model,
228
+ payload=self.payload,
229
+ check_payload=False,
230
+ )
231
+
232
+ @patch("cornflow.endpoints.execution.Databricks")
233
+ def test_new_execution_with_solution_bad(self, db_client_class):
234
+ patch_db_client(db_client_class)
235
+ patch_db_client(db_client_class)
236
+ self.payload["data"] = {"message": "THIS IS NOT A VALID SOLUTION"}
237
+ response = self.create_new_row(
238
+ EXECUTION_URL,
239
+ self.model,
240
+ payload=self.payload,
241
+ check_payload=False,
242
+ expected_status=400,
243
+ )
244
+ self.assertIn("error", response)
245
+ self.assertIn("jsonschema_errors", response)
246
+
247
+ def test_new_execution_no_instance(self):
248
+ payload = dict(self.payload)
249
+ payload["instance_id"] = "bad_id"
250
+ response = self.client.post(
251
+ self.url,
252
+ data=json.dumps(payload),
253
+ follow_redirects=True,
254
+ headers=self.get_header_with_auth(self.token),
255
+ )
256
+ self.assertEqual(404, response.status_code)
257
+ self.assertTrue("error" in response.json)
258
+
259
+ def test_get_executions(self):
260
+ self.get_rows(self.url, self.payloads, keys_to_check=self.keys_to_check)
261
+
262
+ def test_get_no_executions(self):
263
+ self.get_no_rows(self.url)
264
+
265
+ def test_get_executions_superadmin(self):
266
+ self.get_rows(self.url, self.payloads, keys_to_check=self.keys_to_check)
267
+ token = self.create_service_user()
268
+ rows = self.client.get(
269
+ self.url, follow_redirects=True, headers=self.get_header_with_auth(token)
270
+ )
271
+ self.assertEqual(len(rows.json), len(self.payloads))
272
+
273
+
147
274
  class TestExecutionRelaunchEndpoint(CustomTestCase):
148
275
  def setUp(self):
149
276
  super().setUp()
@@ -193,7 +320,7 @@ class TestExecutionRelaunchEndpoint(CustomTestCase):
193
320
  self.assertEqual(row["config"], self.payload["config"])
194
321
  self.assertIsNone(row["checks"])
195
322
 
196
- @patch("cornflow.endpoints.execution_databricks.Airflow")
323
+ @patch("cornflow.endpoints.execution.Airflow")
197
324
  def test_relaunch_execution_run(self, af_client_class):
198
325
  patch_af_client(af_client_class)
199
326
 
@@ -240,6 +367,106 @@ class TestExecutionRelaunchEndpoint(CustomTestCase):
240
367
  self.assertEqual(404, response.status_code)
241
368
 
242
369
 
370
+ class TestExecutionRelaunchEndpointDatabricks(CustomTestCase):
371
+ def setUp(self):
372
+ super().setUp()
373
+
374
+ with open(INSTANCE_PATH) as f:
375
+ payload = json.load(f)
376
+ fk_id = self.create_new_row(INSTANCE_URL, InstanceModel, payload)
377
+ self.url = EXECUTION_URL_NORUN
378
+ self.model = ExecutionModel
379
+
380
+ def load_file_fk(_file):
381
+ with open(_file) as f:
382
+ temp = json.load(f)
383
+ temp["instance_id"] = fk_id
384
+ return temp
385
+
386
+ self.payload = load_file_fk(EXECUTION_PATH)
387
+
388
+ def create_app(self):
389
+ app = create_app("testing-databricks")
390
+ return app
391
+
392
+ def test_relaunch_execution(self):
393
+ idx = self.create_new_row(self.url, self.model, payload=self.payload)
394
+
395
+ # Add solution checks to see if they are deleted correctly
396
+ token = self.create_service_user()
397
+ self.update_row(
398
+ url=DAG_URL + idx + "/",
399
+ payload_to_check=dict(),
400
+ change=dict(solution_schema="_data_checks", checks=dict(check_1=[])),
401
+ token=token,
402
+ check_payload=False,
403
+ )
404
+
405
+ url = EXECUTION_URL + idx + "/relaunch/?run=0"
406
+ self.payload["config"]["warmStart"] = False
407
+ response = self.client.post(
408
+ url,
409
+ data=json.dumps({"config": self.payload["config"]}),
410
+ follow_redirects=True,
411
+ headers=self.get_header_with_auth(self.token),
412
+ )
413
+ self.assertEqual(201, response.status_code)
414
+
415
+ url = EXECUTION_URL + idx + "/data"
416
+ row = self.client.get(
417
+ url, follow_redirects=True, headers=self.get_header_with_auth(self.token)
418
+ ).json
419
+
420
+ self.assertEqual(row["config"], self.payload["config"])
421
+ self.assertIsNone(row["checks"])
422
+
423
+ @patch("cornflow.endpoints.execution.Databricks")
424
+ def test_relaunch_execution_run(self, db_client_class):
425
+ patch_db_client(db_client_class)
426
+
427
+ idx = self.create_new_row(self.url, self.model, payload=self.payload)
428
+
429
+ # Add solution checks to see if they are deleted correctly
430
+ token = self.create_service_user()
431
+ self.update_row(
432
+ url=DAG_URL + idx + "/",
433
+ payload_to_check=dict(),
434
+ change=dict(solution_schema="_data_checks", checks=dict(check_1=[])),
435
+ token=token,
436
+ check_payload=False,
437
+ )
438
+
439
+ url = EXECUTION_URL + idx + "/relaunch/"
440
+ self.payload["config"]["warmStart"] = False
441
+ response = self.client.post(
442
+ url,
443
+ data=json.dumps({"config": self.payload["config"]}),
444
+ follow_redirects=True,
445
+ headers=self.get_header_with_auth(self.token),
446
+ )
447
+ self.assertEqual(201, response.status_code)
448
+
449
+ url = EXECUTION_URL + idx + "/data"
450
+ row = self.client.get(
451
+ url, follow_redirects=True, headers=self.get_header_with_auth(self.token)
452
+ ).json
453
+
454
+ self.assertEqual(row["config"], self.payload["config"])
455
+ self.assertIsNone(row["checks"])
456
+
457
+ def test_relaunch_invalid_execution(self):
458
+ idx = "thisIsAnInvalidExecutionId"
459
+ url = EXECUTION_URL + idx + "/relaunch/?run=0"
460
+ self.payload["config"]["warmStart"] = False
461
+ response = self.client.post(
462
+ url,
463
+ data=json.dumps({"config": self.payload["config"]}),
464
+ follow_redirects=True,
465
+ headers=self.get_header_with_auth(self.token),
466
+ )
467
+ self.assertEqual(404, response.status_code)
468
+
469
+
243
470
  class TestExecutionsDetailEndpointMock(CustomTestCase):
244
471
  def setUp(self):
245
472
  super().setUp()
@@ -269,7 +496,7 @@ class TestExecutionsDetailEndpointMock(CustomTestCase):
269
496
  self.payload["instance_id"] = fk_id
270
497
 
271
498
 
272
- class TestExecutionsDetailEndpoint(
499
+ class TestExecutionsDetailEndpointAirflow(
273
500
  TestExecutionsDetailEndpointMock, BaseTestCases.DetailEndpoint
274
501
  ):
275
502
  def setUp(self):
@@ -277,115 +504,141 @@ class TestExecutionsDetailEndpoint(
277
504
  self.url = self.url
278
505
  self.query_arguments = {"run": 0}
279
506
 
280
- # TODO: this test should be moved as it is not using the detail endpoint
281
- def test_incomplete_payload2(self):
282
- payload = {"description": "arg", "instance_id": self.payload["instance_id"]}
283
- response = self.create_new_row(
284
- self.url + "?run=0",
285
- self.model,
286
- payload,
287
- expected_status=400,
288
- check_payload=False,
289
- )
507
+ def create_app(self):
508
+ app = create_app("testing") # configuración para Airflow
509
+ return app
290
510
 
291
- def test_create_delete_instance_load(self):
292
- idx = self.create_new_row(self.url + "?run=0", self.model, self.payload)
293
- keys_to_check = [
294
- "message",
295
- "id",
296
- "schema",
297
- "data_hash",
298
- "config",
299
- "instance_id",
300
- "user_id",
301
- "indicators",
302
- "description",
303
- "name",
304
- "created_at",
305
- "state",
306
- ]
307
- execution = self.get_one_row(
308
- self.url + idx,
309
- payload={**self.payload, **dict(id=idx)},
310
- keys_to_check=keys_to_check,
311
- )
312
- self.delete_row(self.url + idx + "/")
313
- keys_to_check = [
314
- "id",
315
- "schema",
316
- "description",
317
- "name",
318
- "user_id",
319
- "executions",
320
- "created_at",
321
- "data_hash",
322
- ]
323
- instance = self.get_one_row(
324
- INSTANCE_URL + execution["instance_id"] + "/",
325
- payload={},
326
- expected_status=200,
327
- check_payload=False,
328
- keys_to_check=keys_to_check,
511
+ @patch("cornflow.endpoints.execution.Airflow")
512
+ def test_stop_execution(self, af_client_class):
513
+ patch_af_client(af_client_class)
514
+
515
+ idx = self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
516
+
517
+ response = self.client.post(
518
+ self.url + str(idx) + "/",
519
+ follow_redirects=True,
520
+ headers=self.get_header_with_auth(self.token),
329
521
  )
330
- executions = [execution["id"] for execution in instance["executions"]]
331
- self.assertFalse(idx in executions)
332
522
 
333
- def test_delete_instance_deletes_execution(self):
334
- # we create a new instance
335
- with open(INSTANCE_PATH) as f:
336
- payload = json.load(f)
337
- fk_id = self.create_new_row(INSTANCE_URL, InstanceModel, payload)
338
- payload = {**self.payload, **dict(instance_id=fk_id)}
339
- # we create an execution for that instance
340
- idx = self.create_new_row(self.url + "?run=0", self.model, payload)
341
- self.get_one_row(self.url + idx, payload={**self.payload, **dict(id=idx)})
342
- # we delete the new instance
343
- self.delete_row(INSTANCE_URL + fk_id + "/")
344
- # we check the execution does not exist
345
- self.get_one_row(
346
- self.url + idx, payload={}, expected_status=404, check_payload=False
523
+ self.assertEqual(200, response.status_code)
524
+ self.assertEqual(response.json["message"], "The execution has been stopped")
525
+
526
+
527
+ class TestExecutionsDetailEndpointDatabricks(
528
+ TestExecutionsDetailEndpointMock, BaseTestCases.DetailEndpoint
529
+ ):
530
+ def setUp(self):
531
+ super().setUp()
532
+ self.url = self.url
533
+ self.query_arguments = {"run": 0}
534
+
535
+ def create_app(self):
536
+ app = create_app("testing-databricks")
537
+ return app
538
+
539
+ @patch("cornflow.endpoints.execution.Databricks")
540
+ def test_stop_execution(self, db_client_class):
541
+ patch_db_client(db_client_class)
542
+
543
+ idx = self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
544
+
545
+ response = self.client.post(
546
+ self.url + str(idx) + "/",
547
+ follow_redirects=True,
548
+ headers=self.get_header_with_auth(self.token),
347
549
  )
348
550
 
349
- def test_update_one_row_data(self):
350
- idx = self.create_new_row(
351
- self.url_with_query_arguments(), self.model, self.payload
551
+ self.assertEqual(200, response.status_code)
552
+ self.assertEqual(
553
+ response.json["message"], "This feature is not available for Databricks"
352
554
  )
353
- with open(INSTANCE_PATH) as f:
354
- payload = json.load(f)
355
- payload["data"]["parameters"]["name"] = "NewName"
356
555
 
357
- url = self.url + str(idx) + "/"
358
- payload = {
359
- **self.payload,
360
- **dict(id=idx, name="new_name", data=payload["data"]),
361
- }
362
- self.update_row(
363
- url,
364
- dict(name="new_name", data=payload["data"]),
556
+
557
+ class TestExecutionsStatusEndpointAirflow(TestExecutionsDetailEndpointMock):
558
+ def setUp(self):
559
+ super().setUp()
560
+ self.response_items = {"id", "name", "status"}
561
+ self.items_to_check = []
562
+
563
+ def create_app(self):
564
+ app = create_app("testing")
565
+ return app
566
+
567
+ @patch("cornflow.endpoints.execution.Airflow")
568
+ def test_get_one_status(self, af_client_class):
569
+ patch_af_client(af_client_class)
570
+
571
+ idx = self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
572
+ payload = dict(self.payload)
573
+ payload["id"] = idx
574
+ keys_to_check = ["state", "message", "id", "data_hash"]
575
+ data = self.get_one_row(
576
+ EXECUTION_URL + idx + "/status/",
365
577
  payload,
578
+ check_payload=False,
579
+ keys_to_check=keys_to_check,
366
580
  )
581
+ self.assertEqual(data["state"], 1)
367
582
 
368
- url += "data/"
369
- row = self.client.get(
370
- url, follow_redirects=True, headers=self.get_header_with_auth(self.token)
583
+ @patch("cornflow.endpoints.execution.Airflow")
584
+ def test_put_one_status(self, af_client_class):
585
+ patch_af_client(af_client_class)
586
+
587
+ idx = self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
588
+ payload = dict(self.payload)
589
+ payload["id"] = idx
590
+ response = self.client.put(
591
+ EXECUTION_URL + idx + "/status/",
592
+ data=json.dumps({"status": 0}),
593
+ follow_redirects=True,
594
+ headers=self.get_header_with_auth(self.token),
371
595
  )
372
596
 
373
- self.assertEqual(row.json["checks"], None)
597
+ self.assertEqual(200, response.status_code)
598
+ self.assertEqual(f"execution {idx} updated correctly", response.json["message"])
374
599
 
375
- @patch("cornflow.endpoints.execution_databricks.Airflow")
376
- def test_stop_execution(self, af_client_class):
377
- patch_af_client(af_client_class)
378
600
 
601
+ class TestExecutionsStatusEndpointDatabricks(TestExecutionsDetailEndpointMock):
602
+ def setUp(self):
603
+ super().setUp()
604
+ self.response_items = {"id", "name", "status"}
605
+ self.items_to_check = []
606
+
607
+ def create_app(self):
608
+ app = create_app("testing-databricks")
609
+ return app
610
+
611
+ @patch("cornflow.endpoints.execution.Databricks")
612
+ def test_get_one_status(self, db_client_class):
613
+ patch_db_client(db_client_class)
379
614
  idx = self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
615
+ payload = dict(self.payload)
616
+ payload["id"] = idx
617
+ keys_to_check = ["state", "message", "id", "data_hash"]
618
+ data = self.get_one_row(
619
+ EXECUTION_URL + idx + "/status/",
620
+ payload,
621
+ check_payload=False,
622
+ keys_to_check=keys_to_check,
623
+ )
624
+ self.assertEqual(data["state"], -1)
380
625
 
381
- response = self.client.post(
382
- self.url + str(idx) + "/",
626
+ @patch("cornflow.endpoints.execution.Databricks")
627
+ def test_put_one_status(self, db_client_class):
628
+ patch_db_client(db_client_class)
629
+
630
+ idx = self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
631
+ payload = dict(self.payload)
632
+ payload["id"] = idx
633
+ response = self.client.put(
634
+ EXECUTION_URL + idx + "/status/",
635
+ data=json.dumps({"status": 0}),
383
636
  follow_redirects=True,
384
637
  headers=self.get_header_with_auth(self.token),
385
638
  )
386
639
 
387
640
  self.assertEqual(200, response.status_code)
388
- self.assertEqual(response.json["message"], "The execution has been stopped")
641
+ self.assertEqual(f"execution {idx} updated correctly", response.json["message"])
389
642
 
390
643
 
391
644
  class TestExecutionsDataEndpoint(TestExecutionsDetailEndpointMock):
@@ -430,6 +683,52 @@ class TestExecutionsDataEndpoint(TestExecutionsDetailEndpointMock):
430
683
  )
431
684
 
432
685
 
686
+ class TestExecutionsDataEndpointDatabricks(TestExecutionsDetailEndpointMock):
687
+ def setUp(self):
688
+ super().setUp()
689
+ self.response_items = {"id", "name", "data"}
690
+ self.items_to_check = ["name"]
691
+ self.keys_to_check = [
692
+ "created_at",
693
+ "checks",
694
+ "instance_id",
695
+ "schema",
696
+ "data",
697
+ "user_id",
698
+ "message",
699
+ "data_hash",
700
+ "log",
701
+ "config",
702
+ "description",
703
+ "state",
704
+ "name",
705
+ "id",
706
+ ]
707
+
708
+ def create_app(self):
709
+ app = create_app("testing-databricks")
710
+ return app
711
+
712
+ def test_get_one_execution(self):
713
+ idx = self.create_new_row(EXECUTION_URL_NORUN, self.model, self.payload)
714
+ self.url = EXECUTION_URL + idx + "/data/"
715
+ payload = dict(self.payload)
716
+ payload["id"] = idx
717
+ self.get_one_row(self.url, payload, keys_to_check=self.keys_to_check)
718
+
719
+ def test_get_one_execution_superadmin(self):
720
+ idx = self.create_new_row(EXECUTION_URL_NORUN, self.model, self.payload)
721
+ payload = dict(self.payload)
722
+ payload["id"] = idx
723
+ token = self.create_service_user()
724
+ self.get_one_row(
725
+ EXECUTION_URL + idx + "/data/",
726
+ payload,
727
+ token=token,
728
+ keys_to_check=self.keys_to_check,
729
+ )
730
+
731
+
433
732
  class TestExecutionsLogEndpoint(TestExecutionsDetailEndpointMock):
434
733
  def setUp(self):
435
734
  super().setUp()
@@ -473,45 +772,52 @@ class TestExecutionsLogEndpoint(TestExecutionsDetailEndpointMock):
473
772
  )
474
773
 
475
774
 
476
- class TestExecutionsStatusEndpoint(TestExecutionsDetailEndpointMock):
775
+ class TestExecutionsLogEndpointDatabricks(TestExecutionsDetailEndpointMock):
477
776
  def setUp(self):
478
777
  super().setUp()
479
- self.response_items = {"id", "name", "status"}
480
- self.items_to_check = []
778
+ self.response_items = {"id", "name", "log", "indicators"}
779
+ self.items_to_check = ["name"]
780
+ self.keys_to_check = [
781
+ "created_at",
782
+ "id",
783
+ "log_text",
784
+ "instance_id",
785
+ "state",
786
+ "message",
787
+ "description",
788
+ "data_hash",
789
+ "name",
790
+ "log",
791
+ "schema",
792
+ "user_id",
793
+ "config",
794
+ "indicators",
795
+ ]
481
796
 
482
- @patch("cornflow.endpoints.execution_databricks.Airflow")
483
- def test_get_one_status(self, af_client_class):
484
- patch_af_client(af_client_class)
797
+ def create_app(self):
798
+ app = create_app("testing-databricks")
799
+ return app
485
800
 
486
- idx = self.create_new_row(EXECUTION_URL, self.model, self.payload)
801
+ def test_get_one_execution(self):
802
+ idx = self.create_new_row(EXECUTION_URL_NORUN, self.model, self.payload)
487
803
  payload = dict(self.payload)
488
804
  payload["id"] = idx
489
- keys_to_check = ["state", "message", "id", "data_hash"]
490
- data = self.get_one_row(
491
- EXECUTION_URL + idx + "/status/",
492
- payload,
493
- check_payload=False,
494
- keys_to_check=keys_to_check,
805
+ self.get_one_row(
806
+ EXECUTION_URL + idx + "/log/", payload, keys_to_check=self.keys_to_check
495
807
  )
496
- self.assertEqual(data["state"], 1)
497
-
498
- @patch("cornflow.endpoints.execution_databricks.Airflow")
499
- def test_put_one_status(self, af_client_class):
500
- patch_af_client(af_client_class)
501
808
 
502
- idx = self.create_new_row(EXECUTION_URL, self.model, self.payload)
809
+ def test_get_one_execution_superadmin(self):
810
+ idx = self.create_new_row(EXECUTION_URL_NORUN, self.model, self.payload)
503
811
  payload = dict(self.payload)
504
812
  payload["id"] = idx
505
- response = self.client.put(
506
- EXECUTION_URL + idx + "/status/",
507
- data=json.dumps({"status": 0}),
508
- follow_redirects=True,
509
- headers=self.get_header_with_auth(self.token),
813
+ token = self.create_service_user()
814
+ self.get_one_row(
815
+ EXECUTION_URL + idx + "/log/",
816
+ payload,
817
+ token=token,
818
+ keys_to_check=self.keys_to_check,
510
819
  )
511
820
 
512
- self.assertEqual(200, response.status_code)
513
- self.assertEqual(f"execution {idx} updated correctly", response.json["message"])
514
-
515
821
 
516
822
  class TestExecutionsModel(TestExecutionsDetailEndpointMock):
517
823
  def test_repr_method(self):
@@ -521,3 +827,17 @@ class TestExecutionsModel(TestExecutionsDetailEndpointMock):
521
827
  def test_str_method(self):
522
828
  idx = self.create_new_row(self.url + "?run=0", self.model, self.payload)
523
829
  self.str_method(idx, f"<Execution {idx}>")
830
+
831
+
832
+ class TestExecutionsModelDatabricks(TestExecutionsDetailEndpointMock):
833
+ def create_app(self):
834
+ app = create_app("testing-databricks")
835
+ return app
836
+
837
+ def test_repr_method(self):
838
+ idx = self.create_new_row(self.url + "?run=0", self.model, self.payload)
839
+ self.repr_method(idx, f"<Execution {idx}>")
840
+
841
+ def test_str_method(self):
842
+ idx = self.create_new_row(self.url + "?run=0", self.model, self.payload)
843
+ self.str_method(idx, f"<Execution {idx}>")