Repository: apache/airflow
File: providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py (Python)

Code before:
        )
        assert test_client.get_variable("hello") is None

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    @mock.patch("airflow.utils.log.logging_mixin.structlog.get_logger")
    def test_get_variable_does_not_contain_value_key(self, mock_hvac, mock_get_logger):
        """
        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
        mock_logger = mock.MagicMock()
        mock_get_logger.return_value = mock_logger

        kwargs = {
            "variables_path": "variables",
            "mount_point": "airflow",
            "auth_type": "token",
            "url": "http://127.0.0.1:8200",
            "token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
        }
        test_client = VaultBackend(**kwargs)
        test_client._log = mock_logger
        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
        mock_logger.warning.assert_called_with(
            'Vault secret %s fetched but does not have required key "value"', "hello"
        )

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    @mock.patch("airflow.utils.log.logging_mixin.structlog.get_logger")
    def test_get_config_does_not_contain_value_key(self, mock_hvac, mock_get_logger):
        """
        Test that if the 'value' key is not present in Vault, _VaultClient.get_config
        should log a warning and return None
        """
        mock_client = mock.MagicMock()
        mock_hvac.Client.return_value = mock_client
        mock_logger = mock.MagicMock()
        mock_get_logger.return_value = mock_logger

        kwargs = {
            "variables_path": "variables",
            "mount_point": "airflow",
            "auth_type": "token",
            "url": "http://127.0.0.1:8200",
            "token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
        }

Review comment (security):
```suggestion
        test_client = VaultBackend(
            **{
                "variables_path": "variables",
                "mount_point": "airflow",
                "auth_type": "token",
                "url": "http://127.0.0.1:8200",
                "token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
            }
        )
        test_client._log = mock_logger
        test_client.vault_client.get_secret = mock.MagicMock(return_value={"test_key": "data"})
```

Code after:
        )
        assert test_client.get_variable("hello") is None

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    @mock.patch("airflow.utils.log.logging_mixin.logging")
    def test_get_variable_does_not_contain_value_key(self, mock_hvac, mock_get_logger):
        """
        Test that if the 'value' key is not present in Vault, _VaultClient.get_variable
        should log a warning and return None
        """
        mock_hvac.Client.return_value = mock.MagicMock()
        mock_logger = mock.MagicMock()
        mock_get_logger.getLogger.return_value = mock_logger

        test_client = VaultBackend(
            **{
                "variables_path": "variables",
                "mount_point": "airflow",
                "auth_type": "token",
                "url": "http://127.0.0.1:8200",
                "token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
            }
        )
        test_client._log = mock_logger
        test_client.vault_client.get_secret = mock.MagicMock(return_value={"test_key": "data"})
        result = test_client.get_variable("hello")
        assert result is None
        mock_logger.warning.assert_called_with(
            'Vault secret %s fetched but does not have required key "value"', "hello"
        )

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    @mock.patch("airflow.utils.log.logging_mixin.logging")
    def test_get_config_does_not_contain_value_key(self, mock_hvac, mock_get_logger):
        """
        Test that if the 'value' key is not present in Vault, _VaultClient.get_config
        should log a warning and return None
        """
        mock_client = mock.MagicMock()
        mock_hvac.Client.return_value = mock_client
        mock_logger = mock.MagicMock()
        mock_get_logger.getLogger.return_value = mock_logger

        kwargs = {
            "variables_path": "variables",
            "mount_point": "airflow",
            "auth_type": "token",
            "url": "http://127.0.0.1:8200",
            "token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
        }
        test_client = VaultBackend(**kwargs)