Reference JOBS_API_VERSION instead of hardcoding value in asserts

tests/databricks/test_notebook.py

@@ -3,6 +3,7 @@
 
 import pytest
 from airflow.exceptions import AirflowException
+from astro_databricks.constants import JOBS_API_VERSION
 from astro_databricks.operators.notebook import DatabricksNotebookOperator
 from astro_databricks.operators.workflow import (
     DatabricksWorkflowTaskGroup,
@@ -152,7 +153,7 @@ def test_databricks_notebook_operator_without_taskgroup_new_cluster(
             "timeout_seconds": 0,
             "email_notifications": {},
         },
-        version="2.1",
+        version=JOBS_API_VERSION,
     )
     mock_monitor.assert_called_once()
 
@@ -199,7 +200,7 @@ def test_databricks_notebook_operator_without_taskgroup_existing_cluster(
             "timeout_seconds": 0,
             "email_notifications": {},
         },
-        version="2.1",
+        version=JOBS_API_VERSION,
     )
     mock_monitor.assert_called_once()
 
@@ -296,7 +297,7 @@ def test_wait_for_pending_task(mock_sleep, mock_runs_api, databricks_notebook_op
         {"state": {"life_cycle_state": "RUNNING"}},
     ]
     databricks_notebook_operator._wait_for_pending_task(current_task, mock_runs_api)
-    mock_runs_api.get_run.assert_called_with("123", version="2.1")
+    mock_runs_api.get_run.assert_called_with("123", version=JOBS_API_VERSION)
     assert mock_runs_api.get_run.call_count == 2
     mock_runs_api.reset_mock()
 
@@ -313,7 +314,7 @@ def test_wait_for_terminating_task(
         {"state": {"life_cycle_state": "TERMINATED"}},
     ]
     databricks_notebook_operator._wait_for_terminating_task(current_task, mock_runs_api)
-    mock_runs_api.get_run.assert_called_with("123", version="2.1")
+    mock_runs_api.get_run.assert_called_with("123", version=JOBS_API_VERSION)
     assert mock_runs_api.get_run.call_count == 3
     mock_runs_api.reset_mock()
 
@@ -328,7 +329,7 @@ def test_wait_for_running_task(mock_sleep, mock_runs_api, databricks_notebook_op
         {"state": {"life_cycle_state": "TERMINATED"}},
     ]
     databricks_notebook_operator._wait_for_running_task(current_task, mock_runs_api)
-    mock_runs_api.get_run.assert_called_with("123", version="2.1")
+    mock_runs_api.get_run.assert_called_with("123", version=JOBS_API_VERSION)
     assert mock_runs_api.get_run.call_count == 3
     mock_runs_api.reset_mock()
 
@@ -382,7 +383,7 @@ def test_monitor_databricks_job_success(
     databricks_notebook_operator.databricks_run_id = "1"
     databricks_notebook_operator.monitor_databricks_job()
     mock_runs_api.return_value.get_run.assert_called_with(
-        databricks_notebook_operator.databricks_run_id, version="2.1"
+        databricks_notebook_operator.databricks_run_id, version=JOBS_API_VERSION
     )
     assert (
         "Check the job run in Databricks: https://databricks-instance-xyz.cloud.databricks.com/#job/1234/run/1"

tests/databricks/test_workflow.py

@@ -6,6 +6,7 @@
 import pytest
 from airflow.exceptions import AirflowException
 from airflow.utils.task_group import TaskGroup
+from astro_databricks.constants import JOBS_API_VERSION
 from astro_databricks.operators.notebook import DatabricksNotebookOperator
 from astro_databricks.operators.workflow import DatabricksWorkflowTaskGroup
 
@@ -142,15 +143,15 @@ def test_create_workflow_from_notebooks_with_create(
     task_group.children["test_workflow.launch"].execute(context={})
     mock_jobs_api.return_value.create_job.assert_called_once_with(
         json=expected_workflow_json,
-        version="2.1",
+        version=JOBS_API_VERSION,
     )
     mock_jobs_api.return_value.run_now.assert_called_once_with(
         job_id=1,
         jar_params=[],
         notebook_params={"notebook_path": "/foo/bar"},
         python_params=[],
         spark_submit_params=[],
-        version="2.1",
+        version=JOBS_API_VERSION,
     )