uranusjr commented on a change in pull request #20361:
URL: https://github.com/apache/airflow/pull/20361#discussion_r778592427
##########
File path: tests/utils/test_operator_helpers.py
##########
@@ -71,6 +71,29 @@ def test_context_to_airflow_vars_all_context(self):
'AIRFLOW_CTX_DAG_EMAIL': '[email protected]',
}
+ def test_context_to_airflow_vars_with_default_context_vars(self):
+ with mock.patch('airflow.settings.get_airflow_context_vars') as
mock_method:
+ airflow_cluster = 'cluster-a'
+ mock_method.return_value = {'airflow_cluster': airflow_cluster}
+
+ context_vars =
operator_helpers.context_to_airflow_vars(self.context)
+ assert context_vars['airflow.ctx.airflow_cluster'] ==
airflow_cluster
+
+ context_vars =
operator_helpers.context_to_airflow_vars(self.context, in_env_var_format=True)
+ assert context_vars['AIRFLOW_CTX_AIRFLOW_CLUSTER'] ==
airflow_cluster
+
+ with mock.patch('airflow.settings.get_airflow_context_vars') as
mock_method:
+ mock_method.return_value = {'airflow_cluster': [1, 2]}
+ with pytest.raises(TypeError) as error:
+ assert "value of key <airflow_cluster> must be string, not
<class 'list'>" == error.value
+ operator_helpers.context_to_airflow_vars(self.context)
Review comment:
```suggestion
with pytest.raises(TypeError) as error:
operator_helpers.context_to_airflow_vars(self.context)
assert "value of key <airflow_cluster> must be string, not
<class 'list'>" == error.value
```
##########
File path: tests/utils/test_operator_helpers.py
##########
@@ -71,6 +71,29 @@ def test_context_to_airflow_vars_all_context(self):
'AIRFLOW_CTX_DAG_EMAIL': '[email protected]',
}
+ def test_context_to_airflow_vars_with_default_context_vars(self):
+ with mock.patch('airflow.settings.get_airflow_context_vars') as
mock_method:
+ airflow_cluster = 'cluster-a'
+ mock_method.return_value = {'airflow_cluster': airflow_cluster}
+
+ context_vars =
operator_helpers.context_to_airflow_vars(self.context)
+ assert context_vars['airflow.ctx.airflow_cluster'] ==
airflow_cluster
+
+ context_vars =
operator_helpers.context_to_airflow_vars(self.context, in_env_var_format=True)
+ assert context_vars['AIRFLOW_CTX_AIRFLOW_CLUSTER'] ==
airflow_cluster
+
+ with mock.patch('airflow.settings.get_airflow_context_vars') as
mock_method:
+ mock_method.return_value = {'airflow_cluster': [1, 2]}
+ with pytest.raises(TypeError) as error:
+ assert "value of key <airflow_cluster> must be string, not
<class 'list'>" == error.value
+ operator_helpers.context_to_airflow_vars(self.context)
+
+ with mock.patch('airflow.settings.get_airflow_context_vars') as
mock_method:
+ mock_method.return_value = {1: "value"}
+ with pytest.raises(TypeError) as error:
+ assert 'key <1> must be string' == error.value
+ operator_helpers.context_to_airflow_vars(self.context)
Review comment:
```suggestion
with pytest.raises(TypeError) as error:
operator_helpers.context_to_airflow_vars(self.context)
assert 'key <1> must be string' == error.value
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
