You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2024/03/27 08:03:42 UTC
(airflow) branch main updated: Resolve `PT012` in `databricks` provider tests (#38517)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bf9bb72bb6 Resolve `PT012` in `databricks` provider tests (#38517)
bf9bb72bb6 is described below
commit bf9bb72bb631c098bfee34a372aefb332c3ffd64
Author: Shahar Epstein <60...@users.noreply.github.com>
AuthorDate: Wed Mar 27 10:03:35 2024 +0200
Resolve `PT012` in `databricks` provider tests (#38517)
* Resolve `PT012` in `databricks` provider tests
* Fix format error
---------
Co-authored-by: Tzu-ping Chung <ur...@gmail.com>
---
pyproject.toml | 4 --
.../providers/databricks/hooks/test_databricks.py | 15 +++----
.../databricks/operators/test_databricks.py | 4 +-
.../databricks/operators/test_databricks_repos.py | 2 +-
.../sensors/test_databricks_partition.py | 51 ++++++++++------------
5 files changed, 31 insertions(+), 45 deletions(-)
diff --git a/pyproject.toml b/pyproject.toml
index cc397bf075..6c5bb5bd9d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -381,10 +381,6 @@ combine-as-imports = true
"tests/providers/cncf/kubernetes/operators/test_pod.py" = ["PT012"]
"tests/providers/cncf/kubernetes/utils/test_k8s_resource_iterator.py" = ["PT012"]
"tests/providers/cncf/kubernetes/utils/test_pod_manager.py" = ["PT012"]
-"tests/providers/databricks/hooks/test_databricks.py" = ["PT012"]
-"tests/providers/databricks/operators/test_databricks.py" = ["PT012"]
-"tests/providers/databricks/operators/test_databricks_repos.py" = ["PT012"]
-"tests/providers/databricks/sensors/test_databricks_partition.py" = ["PT012"]
"tests/providers/sftp/hooks/test_sftp.py" = ["PT012"]
"tests/providers/sftp/operators/test_sftp.py" = ["PT012"]
"tests/providers/sftp/sensors/test_sftp.py" = ["PT012"]
diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py
index c9004e7175..64d4de1d37 100644
--- a/tests/providers/databricks/hooks/test_databricks.py
+++ b/tests/providers/databricks/hooks/test_databricks.py
@@ -888,13 +888,10 @@ class TestDatabricksHook:
with pytest.raises(AirflowException):
self.hook._is_oauth_token_valid({})
- def test_is_oauth_token_valid_raises_invalid_type(self):
- token_missing_type = {"access_token": "my_token"}
- token_wrong_type = {"access_token": "my_token", "token_type": "not bearer"}
-
+ @pytest.mark.parametrize("access_token, token_type", [("my_token", None), ("my_token", "not bearer")])
+ def test_is_oauth_token_valid_raises_invalid_type(self, access_token, token_type):
with pytest.raises(AirflowException):
- self.hook._is_oauth_token_valid(token_missing_type)
- self.hook._is_oauth_token_valid(token_wrong_type)
+ self.hook._is_oauth_token_valid({"access_token": access_token, "token_type": token_type})
def test_is_oauth_token_valid_raises_wrong_time_key(self):
token = {
@@ -1220,8 +1217,7 @@ class TestRunState:
def test_is_terminal_with_nonexistent_life_cycle_state(self):
with pytest.raises(AirflowException):
- run_state = RunState("blah", "", "")
- assert run_state.is_terminal
+ RunState("blah", "", "")
def test_is_successful(self):
run_state = RunState("TERMINATED", "SUCCESS", "")
@@ -1255,8 +1251,7 @@ class TestClusterState:
def test_is_terminal_with_nonexistent_life_cycle_state(self):
with pytest.raises(AirflowException):
- cluster_state = ClusterState("blah", "")
- assert cluster_state.is_terminal
+ ClusterState("blah", "")
def test_is_running(self):
running_states = ["RUNNING", "RESIZING"]
diff --git a/tests/providers/databricks/operators/test_databricks.py b/tests/providers/databricks/operators/test_databricks.py
index d0b8199701..a50610ce61 100644
--- a/tests/providers/databricks/operators/test_databricks.py
+++ b/tests/providers/databricks/operators/test_databricks.py
@@ -1340,12 +1340,12 @@ class TestDatabricksRunNowOperator:
with pytest.raises(AirflowException, match=exception_message):
DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID, job_name=JOB_NAME)
+ run = {"job_id": JOB_ID, "job_name": JOB_NAME}
with pytest.raises(AirflowException, match=exception_message):
- run = {"job_id": JOB_ID, "job_name": JOB_NAME}
DatabricksRunNowOperator(task_id=TASK_ID, json=run)
+ run = {"job_id": JOB_ID}
with pytest.raises(AirflowException, match=exception_message):
- run = {"job_id": JOB_ID}
DatabricksRunNowOperator(task_id=TASK_ID, json=run, job_name=JOB_NAME)
@mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook")
diff --git a/tests/providers/databricks/operators/test_databricks_repos.py b/tests/providers/databricks/operators/test_databricks_repos.py
index a1b5ae3c51..078519531e 100644
--- a/tests/providers/databricks/operators/test_databricks_repos.py
+++ b/tests/providers/databricks/operators/test_databricks_repos.py
@@ -220,8 +220,8 @@ class TestDatabricksReposCreateOperator:
f"repo_path should have form of /Repos/{{folder}}/{{repo-name}}, got '{repo_path}'"
)
+ op = DatabricksReposCreateOperator(task_id=TASK_ID, git_url=git_url, repo_path=repo_path)
with pytest.raises(AirflowException, match=exception_message):
- op = DatabricksReposCreateOperator(task_id=TASK_ID, git_url=git_url, repo_path=repo_path)
op.execute(None)
with pytest.raises(
diff --git a/tests/providers/databricks/sensors/test_databricks_partition.py b/tests/providers/databricks/sensors/test_databricks_partition.py
index 779b3c4093..96789c41db 100644
--- a/tests/providers/databricks/sensors/test_databricks_partition.py
+++ b/tests/providers/databricks/sensors/test_databricks_partition.py
@@ -104,14 +104,12 @@ class TestDatabricksPartitionSensor:
@pytest.mark.parametrize(
"soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
)
- def test_fail__generate_partition_query(self, soft_fail, expected_exception):
+ @patch("airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor._sql_sensor")
+ def test_fail__generate_partition_query(self, _sql_sensor, soft_fail, expected_exception):
self.partition_sensor.soft_fail = soft_fail
table_name = "test"
- with pytest.raises(expected_exception, match=f"Table {table_name} does not have partitions"), patch(
- "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
- "._sql_sensor"
- ) as _sql_sensor:
- _sql_sensor.return_value = [[[], [], [], [], [], [], [], []]]
+ _sql_sensor.return_value = [[[], [], [], [], [], [], [], []]]
+ with pytest.raises(expected_exception, match=f"Table {table_name} does not have partitions"):
self.partition_sensor._generate_partition_query(
prefix="", suffix="", joiner_val="", table_name=table_name
)
@@ -119,17 +117,15 @@ class TestDatabricksPartitionSensor:
@pytest.mark.parametrize(
"soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
)
- def test_fail__generate_partition_query_with_partition_col_mismatch(self, soft_fail, expected_exception):
+ @patch("airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor._sql_sensor")
+ def test_fail__generate_partition_query_with_partition_col_mismatch(
+ self, _sql_sensor, soft_fail, expected_exception
+ ):
self.partition_sensor.soft_fail = soft_fail
partition_col = "non_existent_col"
partition_columns = ["col1", "col2"]
- with pytest.raises(
- expected_exception, match=f"Column {partition_col} not part of table partitions"
- ), patch(
- "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
- "._sql_sensor"
- ) as _sql_sensor:
- _sql_sensor.return_value = [[[], [], [], [], [], [], [], partition_columns]]
+ _sql_sensor.return_value = [[[], [], [], [], [], [], [], partition_columns]]
+ with pytest.raises(expected_exception, match=f"Column {partition_col} not part of table partitions"):
self.partition_sensor._generate_partition_query(
prefix="", suffix="", joiner_val="", table_name="", opts={partition_col: "1"}
)
@@ -137,15 +133,13 @@ class TestDatabricksPartitionSensor:
@pytest.mark.parametrize(
"soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
)
- def test_fail__generate_partition_query_with_missing_opts(self, soft_fail, expected_exception):
+ @patch("airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor._sql_sensor")
+ def test_fail__generate_partition_query_with_missing_opts(
+ self, _sql_sensor, soft_fail, expected_exception
+ ):
self.partition_sensor.soft_fail = soft_fail
- with pytest.raises(
- expected_exception, match="No partitions specified to check with the sensor."
- ), patch(
- "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
- "._sql_sensor"
- ) as _sql_sensor:
- _sql_sensor.return_value = [[[], [], [], [], [], [], [], ["col1", "col2"]]]
+ _sql_sensor.return_value = [[[], [], [], [], [], [], [], ["col1", "col2"]]]
+ with pytest.raises(expected_exception, match="No partitions specified to check with the sensor."):
self.partition_sensor._generate_partition_query(
prefix="", suffix="", joiner_val="", table_name=""
)
@@ -153,15 +147,16 @@ class TestDatabricksPartitionSensor:
@pytest.mark.parametrize(
"soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException))
)
- def test_fail_poke(self, soft_fail, expected_exception):
+ @patch(
+ "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
+ "._check_table_partitions"
+ )
+ def test_fail_poke(self, _check_table_partitions, soft_fail, expected_exception):
self.partition_sensor.soft_fail = soft_fail
partitions = "test"
self.partition_sensor.partitions = partitions
+ _check_table_partitions.return_value = False
with pytest.raises(
expected_exception, match=f"Specified partition\(s\): {partitions} were not found."
- ), patch(
- "airflow.providers.databricks.sensors.databricks_partition.DatabricksPartitionSensor"
- "._check_table_partitions"
- ) as _check_table_partitions:
- _check_table_partitions.return_value = False
+ ):
self.partition_sensor.poke(context={})