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

cornflow/shared/databricks.py

@@ -48,12 +48,16 @@ class Databricks:
         oauth_token = oauth_response.json()["access_token"]
         return oauth_token
 
-    def is_alive(self):
+    def is_alive(self, config=None):
         try:
-            # TODO: this url is project specific. Either it has to be a config option or some other way has to be found
-            path = (
-                "/Workspace/Repos/nippon/nippon_production_scheduling/requirements.txt"
-            )
+            if config is None or config["DATABRICKS_HEALTH_PATH"] == "default path":
+                # We raise an error because the default path is not valid
+                raise DatabricksError(
+                    "Invalid default path. Please set DATABRICKS_HEALTH_PATH as an environment variable"
+                )
+            else:
+                path = config["DATABRICKS_HEALTH_PATH"]
+
             url = f"{self.url}/api/2.0/workspace/get-status?path={path}"
             response = self.request_headers_auth(method="GET", url=url)
             if "error_code" in response.json().keys():
@@ -75,8 +79,6 @@ class Databricks:
             raise DatabricksError("JOB not available")
         return schema_info
 
-    # TODO AGA: incluir un id de job por defecto o hacer obligatorio el uso el parámetro.
-    #   Revisar los efectos secundarios de eliminar execution_id y usar el predeterminado
     def run_workflow(
         self,
         execution_id,
@@ -87,13 +89,11 @@ class Databricks:
         """
         Run a job in Databricks
         """
-        # TODO AGA: revisar si la url esta bien/si acepta asi los parámetros
         url = f"{self.url}/api/2.1/jobs/run-now/"
-        # TODO AGA: revisar si deben ser notebook parameters o job parameters.
         #   Entender cómo se usa checks_only
         payload = dict(
             job_id=orch_name,
-            notebook_parameters=dict(
+            job_parameters=dict(
                 checks_only=checks_only,
                 execution_id=execution_id,
             ),

cornflow/shared/exceptions.py

@@ -9,6 +9,8 @@ from cornflow_client.constants import AirflowError, DatabricksError
 from werkzeug.exceptions import HTTPException
 import traceback
 
+from cornflow.shared.const import DATA_DOES_NOT_EXIST_MSG
+
 
 class InvalidUsage(Exception):
     """
@@ -48,7 +50,7 @@ class ObjectDoesNotExist(InvalidUsage):
     """
 
     status_code = 404
-    error = "The object does not exist"
+    error = DATA_DOES_NOT_EXIST_MSG
 
 
 class ObjectAlreadyExists(InvalidUsage):

cornflow/shared/utils_tables.py

@@ -2,6 +2,7 @@
 import inspect
 import os
 import sys
+import collections
 
 from importlib import import_module
 from sqlalchemy.dialects.postgresql import TEXT
@@ -31,14 +32,42 @@ def import_models():
 
 
 def all_subclasses(cls, models=None):
-    subclasses = set()
+    """Finds all direct and indirect subclasses of a given class.
+
+    Optionally filters a provided list of models first.
+
+    Args:
+        cls: The base class to find subclasses for.
+        models: An optional iterable of classes to pre-filter.
+
+    Returns:
+        A set containing all subclasses found.
+    """
+    filtered_subclasses = set()
     if models is not None:
         for val in models:
-            if issubclass(val, cls):
-                subclasses.add(val)
+            # Ensure val is a class before checking issubclass
+            if isinstance(val, type) and issubclass(val, cls):
+                filtered_subclasses.add(val)
+
+    all_descendants = set()
+    # Use a deque for efficient pop(0)
+    queue = collections.deque(cls.__subclasses__())
+    # Keep track of visited classes during the traversal to handle potential complex hierarchies
+    # (though direct subclass relationships shouldn't form cycles)
+    # Initialize with direct subclasses as they are the starting point.
+    visited_for_queue = set(cls.__subclasses__())
+
+    while queue:
+        current_sub = queue.popleft()
+        all_descendants.add(current_sub)
+
+        for grandchild in current_sub.__subclasses__():
+            if grandchild not in visited_for_queue:
+                visited_for_queue.add(grandchild)
+                queue.append(grandchild)
 
-    return subclasses.union(set(cls.__subclasses__()).union(
-        [s for c in cls.__subclasses__() for s in all_subclasses(c)]))
+    return filtered_subclasses.union(all_descendants)
 
 
 type_converter = {
@@ -76,6 +105,5 @@ def item_as_dict(item):
 
 def items_as_dict_list(ls):
     return [
-        {c.name: getattr(item, c.name) for c in item.__table__.columns}
-        for item in ls
-    ]
+        {c.name: getattr(item, c.name) for c in item.__table__.columns} for item in ls
+    ]

cornflow/shared/validators.py

@@ -1,12 +1,12 @@
 """
 This file has several validators
 """
+
 import re
 from typing import Tuple, Union
 
 from jsonschema import Draft7Validator, validators
 from disposable_email_domains import blocklist
-from jsonschema.protocols import Validator
 
 
 def is_special_character(character):

cornflow/tests/const.py

@@ -15,6 +15,7 @@ INSTANCE_URL = PREFIX + "/instance/"
 INSTANCE_MPS = _get_file("./data/test_mps.mps")
 INSTANCE_GC_20 = _get_file("./data/gc_20_7.json")
 INSTANCE_FILE_FAIL = _get_file("./unit/test_instances.py")
+EDIT_EXECUTION_SOLUTION = _get_file("./data/edit_execution_solution.json")
 
 EXECUTION_PATH = _get_file("./data/new_execution.json")
 BAD_EXECUTION_PATH = _get_file("./data/bad_execution.json")

cornflow/tests/custom_test_case.py

@@ -635,9 +635,9 @@ class BaseTestCases:
             # (we patch the request to airflow to check if the schema is valid)
             # we create 4 instances
             data_many = [self.payload for _ in range(4)]
-            data_many[-1] = {**data_many[-1], **dict(schema="timer")}
-            allrows = self.get_rows(self.url, data_many)
-            self.apply_filter(self.url, dict(schema="timer"), allrows.json[:1])
+
+            self.get_rows(self.url, data_many)
+            self.apply_filter(self.url, dict(schema="timer"), [])
 
         def test_opt_filters_date_lte(self):
             """
@@ -1117,7 +1117,7 @@ class LoginTestCases:
 
             self.assertEqual(400, response.status_code)
             self.assertEqual(
-                "Invalid token issuer. Token must be issued by a valid provider",
+                "Invalid token format or signature",
                 response.json["error"],
             )
 

cornflow/tests/unit/test_alarms.py

@@ -91,7 +91,6 @@ class TestAlarmsEndpoint(CustomTestCase):
 class TestAlarmsDetailEndpoint(TestAlarmsEndpoint, BaseTestCases.DetailEndpoint):
     def setUp(self):
         super().setUp()
-        self.url = self.url
         self.idx = 0
         self.payload = {
             "name": "Alarm 1",
@@ -138,4 +137,4 @@ class TestAlarmsDetailEndpoint(TestAlarmsEndpoint, BaseTestCases.DetailEndpoint)
                 # We check deleted at has a value
                 self.assertIsNotNone(row.deleted_at)
             else:
-                self.assertIsNone(row.deleted_at)
+                self.assertIsNone(row.deleted_at)

cornflow/tests/unit/test_cases.py

@@ -42,6 +42,7 @@ import zlib
 
 # Import from internal modules
 from cornflow.models import CaseModel, ExecutionModel, InstanceModel, UserModel
+from cornflow.shared.const import DATA_DOES_NOT_EXIST_MSG
 from cornflow.shared.utils import hash_json_256
 from cornflow.tests.const import (
     INSTANCE_URL,
@@ -227,12 +228,8 @@ class TestCasesFromInstanceExecutionEndpoint(CustomTestCase):
             "execution_id": execution_id,
             "schema": "solve_model_dag",
         }
-        self.instance = InstanceModel.get_one_object(
-            user=self.user, idx=instance_id
-        )
-        self.execution = ExecutionModel.get_one_object(
-            user=self.user, idx=execution_id
-        )
+        self.instance = InstanceModel.get_one_object(user=self.user, idx=instance_id)
+        self.execution = ExecutionModel.get_one_object(user=self.user, idx=execution_id)
 
     def test_new_case_execution(self):
         """
@@ -729,7 +726,7 @@ class TestCaseToInstanceEndpoint(CustomTestCase):
             headers=self.get_header_with_auth(self.token),
         )
         self.assertEqual(response.status_code, 404)
-        self.assertEqual(response.json["error"], "The object does not exist")
+        self.assertEqual(response.json["error"], DATA_DOES_NOT_EXIST_MSG)
 
 
 class TestCaseJsonPatch(CustomTestCase):

cornflow/tests/unit/test_executions.py

@@ -20,6 +20,7 @@ from cornflow.tests.const import (
     DAG_URL,
     BAD_EXECUTION_PATH,
     EXECUTION_SOLUTION_PATH,
+    EDIT_EXECUTION_SOLUTION,
 )
 from cornflow.tests.custom_test_case import CustomTestCase, BaseTestCases
 from cornflow.tests.unit.tools import patch_af_client, patch_db_client
@@ -186,7 +187,9 @@ class TestExecutionsListEndpointDatabricks(BaseTestCases.ListFilters):
         app = create_app("testing-databricks")
         return app
 
-    def test_new_execution(self):
+    @patch("cornflow.endpoints.execution.Databricks")
+    def test_new_execution(self, db_client_class):
+        patch_db_client(db_client_class)
         self.create_new_row(self.url, self.model, payload=self.payload)
 
     @patch("cornflow.endpoints.execution.Databricks")
@@ -231,7 +234,6 @@ class TestExecutionsListEndpointDatabricks(BaseTestCases.ListFilters):
 
     @patch("cornflow.endpoints.execution.Databricks")
     def test_new_execution_with_solution_bad(self, db_client_class):
-        patch_db_client(db_client_class)
         patch_db_client(db_client_class)
         self.payload["data"] = {"message": "THIS IS NOT A VALID SOLUTION"}
         response = self.create_new_row(
@@ -244,7 +246,9 @@ class TestExecutionsListEndpointDatabricks(BaseTestCases.ListFilters):
         self.assertIn("error", response)
         self.assertIn("jsonschema_errors", response)
 
-    def test_new_execution_no_instance(self):
+    @patch("cornflow.endpoints.execution.Databricks")
+    def test_new_execution_no_instance(self, db_client_class):
+        patch_db_client(db_client_class)
         payload = dict(self.payload)
         payload["instance_id"] = "bad_id"
         response = self.client.post(
@@ -256,13 +260,19 @@ class TestExecutionsListEndpointDatabricks(BaseTestCases.ListFilters):
         self.assertEqual(404, response.status_code)
         self.assertTrue("error" in response.json)
 
-    def test_get_executions(self):
+    @patch("cornflow.endpoints.execution.Databricks")
+    def test_get_executions(self, db_client_class):
+        patch_db_client(db_client_class)
         self.get_rows(self.url, self.payloads, keys_to_check=self.keys_to_check)
 
-    def test_get_no_executions(self):
+    @patch("cornflow.endpoints.execution.Databricks")
+    def test_get_no_executions(self, db_client_class):
+        patch_db_client(db_client_class)
         self.get_no_rows(self.url)
 
-    def test_get_executions_superadmin(self):
+    @patch("cornflow.endpoints.execution.Databricks")
+    def test_get_executions_superadmin(self, db_client_class):
+        patch_db_client(db_client_class)
         self.get_rows(self.url, self.payloads, keys_to_check=self.keys_to_check)
         token = self.create_service_user()
         rows = self.client.get(
@@ -473,6 +483,7 @@ class TestExecutionsDetailEndpointMock(CustomTestCase):
         with open(INSTANCE_PATH) as f:
             payload = json.load(f)
         fk_id = self.create_new_row(INSTANCE_URL, InstanceModel, payload)
+        self.instance_payload = payload
         self.model = ExecutionModel
         self.response_items = {
             "id",
@@ -501,7 +512,6 @@ class TestExecutionsDetailEndpointAirflow(
 ):
     def setUp(self):
         super().setUp()
-        self.url = self.url
         self.query_arguments = {"run": 0}
 
     def create_app(self):
@@ -523,51 +533,41 @@ class TestExecutionsDetailEndpointAirflow(
         self.assertEqual(200, response.status_code)
         self.assertEqual(response.json["message"], "The execution has been stopped")
 
+    def test_edit_execution(self):
 
-class TestExecutionsDetailEndpointDatabricks(
-    TestExecutionsDetailEndpointMock, BaseTestCases.DetailEndpoint
-):
-    def setUp(self):
-        super().setUp()
-        self.url = self.url
-        self.query_arguments = {"run": 0}
-
-    def create_app(self):
-        app = create_app("testing-databricks")
-        return app
-
-    @patch("cornflow.endpoints.execution.Databricks")
-    def test_stop_execution(self, db_client_class):
-        patch_db_client(db_client_class)
-
-        idx = self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
-
-        response = self.client.post(
-            self.url + str(idx) + "/",
-            follow_redirects=True,
-            headers=self.get_header_with_auth(self.token),
+        id_new_instance = self.create_new_row(
+            INSTANCE_URL, InstanceModel, self.instance_payload
         )
-
-        self.assertEqual(200, response.status_code)
-        self.assertEqual(
-            response.json["message"], "This feature is not available for Databricks"
+        idx = self.create_new_row(
+            self.url_with_query_arguments(), self.model, self.payload
         )
 
+        # Extract the data from data/edit_execution_solution.json
+        with open(EDIT_EXECUTION_SOLUTION) as f:
+            data = json.load(f)
 
-class TestExecutionsStatusEndpointAirflow(TestExecutionsDetailEndpointMock):
-    def setUp(self):
-        super().setUp()
-        self.response_items = {"id", "name", "status"}
-        self.items_to_check = []
-
-    def create_app(self):
-        app = create_app("testing")
-        return app
+        data = {
+            "name": "new_name",
+            "description": "Updated description",
+            "data": data,
+            "instance_id": id_new_instance,
+        }
+        payload_to_check = {
+            "id": idx,
+            "name": "new_name",
+            "description": "Updated description",
+            "data_hash": "74234e98afe7498fb5daf1f36ac2d78acc339464f950703b8c019892f982b90b",
+            "instance_id": "805bad3280c95e45384dc6bd91a41317f9a7858c",
+        }
+        self.update_row(
+            self.url + str(idx) + "/",
+            data,
+            payload_to_check,
+        )
 
     @patch("cornflow.endpoints.execution.Airflow")
     def test_get_one_status(self, af_client_class):
         patch_af_client(af_client_class)
-
         idx = self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
         payload = dict(self.payload)
         payload["id"] = idx
@@ -598,6 +598,68 @@ class TestExecutionsStatusEndpointAirflow(TestExecutionsDetailEndpointMock):
         self.assertEqual(f"execution {idx} updated correctly", response.json["message"])
 
 
+class TestExecutionsDetailEndpointDatabricks(
+    TestExecutionsDetailEndpointMock, BaseTestCases.DetailEndpoint
+):
+    def setUp(self):
+        super().setUp()
+        self.url = self.url
+        self.query_arguments = {"run": 0}
+
+    def create_app(self):
+        app = create_app("testing-databricks")
+        return app
+
+    @patch("cornflow.endpoints.execution.Databricks")
+    def test_stop_execution(self, db_client_class):
+        patch_db_client(db_client_class)
+
+        idx = self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
+
+        response = self.client.post(
+            self.url + str(idx) + "/",
+            follow_redirects=True,
+            headers=self.get_header_with_auth(self.token),
+        )
+
+        self.assertEqual(200, response.status_code)
+        self.assertEqual(
+            response.json["message"], "This feature is not available for Databricks"
+        )
+
+    def test_edit_execution(self):
+
+        id_new_instance = self.create_new_row(
+            INSTANCE_URL, InstanceModel, self.instance_payload
+        )
+        idx = self.create_new_row(
+            self.url_with_query_arguments(), self.model, self.payload
+        )
+
+        # Extract the data from data/edit_execution_solution.json
+        with open(EDIT_EXECUTION_SOLUTION) as f:
+            data = json.load(f)
+
+        data = {
+            "name": "new_name",
+            "description": "Updated description",
+            "data": data,
+            "instance_id": id_new_instance,
+        }
+        payload_to_check = {
+            "id": idx,
+            "name": "new_name",
+            "description": "Updated description",
+            "data_hash": "74234e98afe7498fb5daf1f36ac2d78acc339464f950703b8c019892f982b90b",
+            "instance_id": "805bad3280c95e45384dc6bd91a41317f9a7858c",
+        }
+        self.update_row(
+            self.url + str(idx) + "/",
+            data,
+            payload_to_check,
+        )
+
+
 class TestExecutionsStatusEndpointDatabricks(TestExecutionsDetailEndpointMock):
     def setUp(self):
         super().setUp()

cornflow/tests/unit/test_log_in.py

@@ -135,11 +135,20 @@ class TestLogInOpenAuth(CustomTestCase):
         Tests token validation failure when the kid is not found in public keys
         """
         # Import the real exceptions to ensure they are preserved
-        from jwt.exceptions import InvalidTokenError, ExpiredSignatureError
+        from jwt.exceptions import (
+            InvalidTokenError,
+            ExpiredSignatureError,
+            InvalidIssuerError,
+            InvalidSignatureError,
+            DecodeError,
+        )
 
         # Keep the real exception classes in the mock
         mock_jwt.ExpiredSignatureError = ExpiredSignatureError
         mock_jwt.InvalidTokenError = InvalidTokenError
+        mock_jwt.InvalidIssuerError = InvalidIssuerError
+        mock_jwt.InvalidSignatureError = InvalidSignatureError
+        mock_jwt.DecodeError = DecodeError
 
         mock_jwt.get_unverified_header.return_value = {"kid": "test_kid"}
 
@@ -165,7 +174,7 @@ class TestLogInOpenAuth(CustomTestCase):
 
         self.assertEqual(400, response.status_code)
         self.assertEqual(
-            response.json["error"], "Invalid token: Unknown key identifier (kid)"
+            response.json["error"], "Invalid token format, signature, or configuration"
         )
 
     @mock.patch("cornflow.shared.authentication.auth.jwt")
@@ -174,11 +183,20 @@ class TestLogInOpenAuth(CustomTestCase):
         Tests token validation failure when the token header is missing the kid
         """
         # Import the real exceptions to ensure they are preserved
-        from jwt.exceptions import InvalidTokenError, ExpiredSignatureError
+        from jwt.exceptions import (
+            InvalidTokenError,
+            ExpiredSignatureError,
+            InvalidIssuerError,
+            InvalidSignatureError,
+            DecodeError,
+        )
 
         # Keep the real exception classes in the mock
         mock_jwt.ExpiredSignatureError = ExpiredSignatureError
         mock_jwt.InvalidTokenError = InvalidTokenError
+        mock_jwt.InvalidIssuerError = InvalidIssuerError
+        mock_jwt.InvalidSignatureError = InvalidSignatureError
+        mock_jwt.DecodeError = DecodeError
 
         # Mock jwt.get_unverified_header to return a header without kid
         mock_jwt.get_unverified_header.return_value = {"alg": "RS256"}
@@ -194,8 +212,7 @@ class TestLogInOpenAuth(CustomTestCase):
 
         self.assertEqual(400, response.status_code)
         self.assertEqual(
-            response.json["error"],
-            "Invalid token: Missing key identifier (kid) in token header",
+            "Invalid token format, signature, or configuration", response.json["error"]
         )
 
     @mock.patch("cornflow.shared.authentication.auth.requests.get")
@@ -205,11 +222,20 @@ class TestLogInOpenAuth(CustomTestCase):
         Tests failure when trying to fetch public keys from the OIDC provider
         """
         # Import the real exceptions to ensure they are preserved
-        from jwt.exceptions import InvalidTokenError, ExpiredSignatureError
+        from jwt.exceptions import (
+            InvalidTokenError,
+            ExpiredSignatureError,
+            InvalidIssuerError,
+            InvalidSignatureError,
+            DecodeError,
+        )
 
         # Keep the real exception classes in the mock
         mock_jwt.ExpiredSignatureError = ExpiredSignatureError
         mock_jwt.InvalidTokenError = InvalidTokenError
+        mock_jwt.InvalidIssuerError = InvalidIssuerError
+        mock_jwt.InvalidSignatureError = InvalidSignatureError
+        mock_jwt.DecodeError = DecodeError
 
         # Clear the cache
         from cornflow.shared.authentication.auth import public_keys_cache
@@ -240,8 +266,8 @@ class TestLogInOpenAuth(CustomTestCase):
 
         self.assertEqual(400, response.status_code)
         self.assertEqual(
+            "Invalid token format, signature, or configuration",
             response.json["error"],
-            "Failed to fetch public keys from authentication provider",
         )
 
     @mock.patch("cornflow.shared.authentication.Auth.decode_token")
@@ -289,11 +315,20 @@ class TestLogInOpenAuth(CustomTestCase):
         the system fetches fresh keys from the provider.
         """
         # Import the real exceptions to ensure they are preserved
-        from jwt.exceptions import InvalidTokenError, ExpiredSignatureError
+        from jwt.exceptions import (
+            InvalidTokenError,
+            ExpiredSignatureError,
+            InvalidIssuerError,
+            InvalidSignatureError,
+            DecodeError,
+        )
 
         # Keep the real exception classes in the mock
         mock_jwt.ExpiredSignatureError = ExpiredSignatureError
         mock_jwt.InvalidTokenError = InvalidTokenError
+        mock_jwt.InvalidIssuerError = InvalidIssuerError
+        mock_jwt.InvalidSignatureError = InvalidSignatureError
+        mock_jwt.DecodeError = DecodeError
 
         # Mock jwt to return valid unverified header and payload
         mock_jwt.get_unverified_header.return_value = {"kid": "test_kid"}
@@ -531,4 +566,6 @@ class TestLogInOpenAuthService(CustomTestCase):
         )
 
         self.assertEqual(400, response.status_code)
-        self.assertEqual(response.json["error"], "Invalid token format or signature")
+        self.assertEqual(
+            response.json["error"], "Invalid token format, signature, or configuration"
+        )

cornflow/tests/unit/test_tables.py

@@ -11,7 +11,7 @@ from cornflow.app import create_app
 from cornflow.commands.access import access_init_command
 from cornflow.models import UserRoleModel
 from cornflow.shared import db
-from cornflow.shared.const import ADMIN_ROLE, SERVICE_ROLE
+from cornflow.shared.const import ADMIN_ROLE, SERVICE_ROLE, DATA_DOES_NOT_EXIST_MSG
 from cornflow.tests.const import LOGIN_URL, SIGNUP_URL, TABLES_URL
 
 
@@ -220,7 +220,7 @@ class TestTablesDetailEndpoint(TestCase):
             },
         )
         self.assertEqual(response.status_code, 404)
-        self.assertEqual(response.json["error"], "The object does not exist")
+        self.assertEqual(response.json["error"], DATA_DOES_NOT_EXIST_MSG)
 
 
 class TestTablesEndpointAdmin(TestCase):
@@ -291,4 +291,4 @@ class TestTablesEndpointAdmin(TestCase):
         self.assertEqual(response.status_code, 403)
         self.assertEqual(
             response.json["error"], "You do not have permission to access this endpoint"
-        )
+        )

cornflow/tests/unit/tools.py

@@ -3,7 +3,10 @@ from flask import Flask
 import json
 import os
 
-from cornflow.shared.const import DATABRICKS_TERMINATE_STATE, DATABRICKS_FINISH_TO_STATE_MAP
+from cornflow.shared.const import (
+    DATABRICKS_TERMINATE_STATE,
+    DATABRICKS_FINISH_TO_STATE_MAP,
+)
 
 
 def create_test_app():
@@ -13,9 +16,7 @@ def create_test_app():
 
 
 def patch_af_client(af_client_class):
-    with patch(
-        "cornflow.endpoints.execution.current_app.config"
-    ) as mock_config:
+    with patch("cornflow.endpoints.execution.current_app.config") as mock_config:
         mock_config.__getitem__.side_effect = lambda key: (
             1 if key == "CORNFLOW_BACKEND" else {}
         )
@@ -27,9 +28,24 @@ def patch_af_client(af_client_class):
             "state": "success",
         }
         af_client_mock.is_alive.return_value = True
-        af_client_mock.get_orch_info.return_value = responses_mock
-        af_client_mock.run_workflow.return_value = responses_mock
-        af_client_mock.get_run_status.return_value = responses_mock
+        af_client_mock.get_dag_info.return_value = responses_mock
+
+        # Configurar get_orch_info
+        schema_info_mock = Mock()
+        schema_info_mock.json.return_value = {"is_paused": False}
+        af_client_mock.get_orch_info.return_value = schema_info_mock
+
+        # Configurar run_workflow
+        run_response_mock = Mock()
+        run_response_mock.json.return_value = {"dag_run_id": "12345"}
+        af_client_mock.run_workflow.return_value = run_response_mock
+
+        # Configurar get_run_status
+        status_mock = Mock()
+        status_mock.json.return_value = {"state": "success"}
+        af_client_mock.get_run_status.return_value = status_mock
+
+        af_client_mock.run_dag.return_value = responses_mock
         af_client_mock.set_dag_run_to_fail.return_value = None
         af_client_class.from_config.return_value = af_client_mock
 
@@ -37,9 +53,7 @@ def patch_af_client(af_client_class):
 def patch_db_client(db_client_class):
     mock_config = {"CORNFLOW_BACKEND": 2}
 
-    with patch(
-        "cornflow.endpoints.execution.current_app.config", mock_config
-    ):
+    with patch("cornflow.endpoints.execution.current_app.config", mock_config):
         db_client_mock = Mock()
         responses_mock = Mock()
         responses_mock.json.return_value = {
@@ -65,10 +79,14 @@ def patch_db_client(db_client_class):
         state = response_get_run_status.json.return_value["status"]["state"]
         if state == DATABRICKS_TERMINATE_STATE:
             if (
-                    response_get_run_status.json.return_value["status"]["termination_details"]["code"]
-                    in DATABRICKS_FINISH_TO_STATE_MAP.keys()
+                response_get_run_status.json.return_value["status"][
+                    "termination_details"
+                ]["code"]
+                in DATABRICKS_FINISH_TO_STATE_MAP.keys()
             ):
-                response_get_run_status = response_get_run_status.json.return_value["status"]["termination_details"]["code"]
+                response_get_run_status = response_get_run_status.json.return_value[
+                    "status"
+                ]["termination_details"]["code"]
             else:
                 response_get_run_status = "OTHER_FINISH_ERROR"
         db_client_mock.is_alive.return_value = True