BasPH commented on code in PR #59224:
URL: https://github.com/apache/airflow/pull/59224#discussion_r2626708970
##########
providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py:
##########
@@ -199,14 +199,15 @@ def get_connection(self, conn_id: str) -> Connection |
None:
from airflow.models.connection import Connection
response = self.get_response(conn_id)
- if response is None:
+ if not response:
return None
+ try:
+ uri = response["conn_uri"]
Review Comment:
For readability sake I'd place the `return Connection(conn_id, uri=uri)`
inside the `try` clause, but the business logic looks ok.
##########
providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py:
##########
@@ -111,6 +111,56 @@ def
test_get_connection_without_predefined_mount_point(self, mock_hvac):
connection =
test_client.get_connection(conn_id="airflow/test_postgres")
assert connection.get_uri() ==
"postgresql://airflow:airflow@host:5432/airflow?foo=bar&baz=taz"
+
@mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
+ def test_get_connection_without_conn_uri_key(self, mock_hvac):
+ """
+ If a connection is found but the 'conn_uri' is missing,
+ a warning should be logged and the returned connection should be None
+ """
+ mock_client = mock.MagicMock()
+ mock_hvac.Client.return_value = mock_client
+ mock_client.secrets.kv.v2.read_secret_version.return_value = {
+ "request_id": "94011e25-f8dc-ec29-221b-1f9c1d9ad2ae",
+ "lease_id": "",
+ "renewable": False,
+ "lease_duration": 0,
+ "data": {
+ "data": {
+ "conn_type": "postgresql",
+ "login": "airflow",
+ "password": "airflow",
+ "host": "host",
+ "port": "5432",
+ "schema": "airflow",
+ "extra": '{"foo":"bar","baz":"taz"}',
+ },
+ "metadata": {
+ "created_time": "2020-03-16T21:01:43.331126Z",
+ "deletion_time": "",
+ "destroyed": False,
+ "version": 1,
+ },
+ },
+ "wrap_info": None,
+ "warnings": None,
+ "auth": None,
+ }
+
+ kwargs = {
+ "connections_path": "connections",
+ "mount_point": "airflow",
+ "auth_type": "token",
+ "url": "http://127.0.0.1:8200";,
+ "token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
+ }
+
+ test_client = VaultBackend(**kwargs)
+ response = {"test_key": "data"}
+ test_client.vault_client.get_response =
mock.MagicMock(return_value=response)
+ connection = test_client.get_connection(conn_id="test_postgres")
+ # How to test the KeyError branch of get_connection?
Review Comment:
I left a comment regarding the test for variables, the same approach can be
applied here.
##########
providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py:
##########
@@ -263,6 +313,52 @@ def test_get_variable_value_non_existent_key(self,
mock_hvac):
)
assert test_client.get_variable("hello") is None
+
@mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
+ def test_get_variable_does_not_contain_value_key(self, mock_hvac):
+ """
+ Test that if the 'value' key is not present in Vault,
_VaultClient.get_variable
+ should log a warning and return None
+ """
+ mock_client = mock.MagicMock()
+ mock_hvac.Client.return_value = mock_client
+
+ kwargs = {
+ "variables_path": "variables",
+ "mount_point": "airflow",
+ "auth_type": "token",
+ "url": "http://127.0.0.1:8200";,
+ "token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
+ }
+ test_client = VaultBackend(**kwargs)
+ response = {"test_key": "data"}
+ test_client.vault_client.get_secret =
mock.MagicMock(return_value=response)
+
+ result = test_client.get_variable("hello")
+ assert result is None
Review Comment:
Since we're changing the feedback to the user regarding incorrect keys in
Vault, IMO we should validate that _something_ returned to the user (in this
case, a log line). We can capture logs with
[caplog](https://docs.pytest.org/en/stable/how-to/logging.html#caplog-fixture).
It gets a bit lengthy, but what do you think of:
```python
@mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
def test_get_variable_does_not_contain_value_key(self, mock_hvac, caplog):
"""Test that if a variable isn't stored with key 'value', a warning is
logged."""
mock_client = mock.MagicMock()
mock_hvac.Client.return_value = mock_client
mock_client.secrets.kv.v2.read_secret_version.return_value = {
"request_id": "2d48a2ad-6bcb-e5b6-429d-da35fdf31f56",
"lease_id": "",
"renewable": False,
"lease_duration": 0,
"data": {
"data": {"not_value": "world"},
"metadata": {
"created_time": "2020-03-28T02:10:54.301784Z",
"deletion_time": "",
"destroyed": False,
"version": 1,
},
},
"wrap_info": None,
"warnings": None,
"auth": None,
}
kwargs = {
"variables_path": "variables",
"mount_point": "airflow",
"auth_type": "token",
"url": "http://127.0.0.1:8200";,
"token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
}
test_client = VaultBackend(**kwargs)
variable_name = "hello"
test_client.get_variable(variable_name)
expected_name = "airflow.providers.hashicorp.secrets.vault.VaultBackend"
expected_levelname = "WARNING"
expected_strings_in_msg = [variable_name, "value"]
log_matched = False
for record in caplog.records:
if (
record.name == expected_name and
record.levelname == expected_levelname and
all(substr in record.msg for substr in expected_strings_in_msg)
):
log_matched = True
break
assert log_matched, "No log found warning user that key 'value' is not
present in the Vault secret."
```
--
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]
