Repository: apache/airflow
File: providers/hashicorp/tests/unit/hashicorp/_internal_client/test_vault_client.py (Python)

Code before:
        mock_hvac.Client.return_value = mock_client
        with pytest.raises(VaultError, match="requires 'kubernetes_jwt_path'"):
            _VaultClient(
                auth_type="kubernetes",
                kubernetes_role="kube_role",
                kubernetes_jwt_path=None,
                url="http://localhost:8180",
            )

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    def test_jwt_with_token(self, mock_hvac):
        mock_client = mock.MagicMock()
        mock_hvac.Client.return_value = mock_client
        vault_client = _VaultClient(
            auth_type="jwt",
            jwt_role="my-role",
            jwt_token="eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.test",
            url="http://localhost:8180",
            session=None,
        )
        client = vault_client.client
        mock_hvac.Client.assert_called_with(url="http://localhost:8180", session=None)
        client.auth.jwt.jwt_login.assert_called_with(
            role="my-role", jwt="eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.test"
        )
        client.is_authenticated.assert_called_with()
        assert vault_client.kv_engine_version == 2

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    def test_jwt_with_token_path(self, mock_hvac):
        mock_client = mock.MagicMock()
        mock_hvac.Client.return_value = mock_client
        vault_client = _VaultClient(
            auth_type="jwt",
            jwt_role="my-role",
            jwt_token_path="path/to/jwt",
            url="http://localhost:8180",
            session=None,
        )
        with patch("builtins.open", mock_open(read_data="eyJhbGciOiJSUzI1NiJ9.jwt-from-file")) as mock_file:
            client = vault_client.client
        mock_file.assert_called_with("path/to/jwt")
        mock_hvac.Client.assert_called_with(url="http://localhost:8180", session=None)
        client.auth.jwt.jwt_login.assert_called_with(role="my-role", jwt="eyJhbGciOiJSUzI1NiJ9.jwt-from-file")
        client.is_authenticated.assert_called_with()
        assert vault_client.kv_engine_version == 2

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    def test_jwt_with_default_token_path(self, mock_hvac):
        mock_client = mock.MagicMock()
        mock_hvac.Client.return_value = mock_client

Review comment (security):
```suggestion
        assert mock_hvac.Client.mock_calls == [
            call(
                url="http://localhost:8180",
                session=None
            )
        ]
        assert client.auth.jwt.jwt_login.mock_calls == [
            call(
                role="my-role",
                jwt="eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.test",
            )
        ]
        client.is_authenticated.mock_calls == [
            call()
        ]
```


Code after:
        mock_hvac.Client.return_value = mock_client
        with pytest.raises(VaultError, match="requires 'kubernetes_jwt_path'"):
            _VaultClient(
                auth_type="kubernetes",
                kubernetes_role="kube_role",
                kubernetes_jwt_path=None,
                url="http://localhost:8180",
            )

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    def test_jwt_with_token(self, mock_hvac):
        mock_client = mock.MagicMock()
        mock_hvac.Client.return_value = mock_client
        vault_client = _VaultClient(
            auth_type="jwt",
            jwt_role="my-role",
            jwt_token="eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.test",
            url="http://localhost:8180",
            session=None,
        )
        client = vault_client.client
        assert mock_hvac.Client.call_args_list == [call(url="http://localhost:8180", session=None)]
        assert client.auth.jwt.jwt_login.call_args_list == [
            call(role="my-role", jwt="eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.test")
        ]
        assert client.is_authenticated.call_args_list == [call(), call()]
        assert vault_client.kv_engine_version == 2

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    def test_jwt_with_token_path(self, mock_hvac):
        mock_client = mock.MagicMock()
        mock_hvac.Client.return_value = mock_client
        vault_client = _VaultClient(
            auth_type="jwt",
            jwt_role="my-role",
            jwt_token_path="path/to/jwt",
            url="http://localhost:8180",
            session=None,
        )
        with patch("builtins.open", mock_open(read_data="eyJhbGciOiJSUzI1NiJ9.jwt-from-file")) as mock_file:
            client = vault_client.client
        assert mock_file.call_args_list == [call("path/to/jwt")]
        assert mock_hvac.Client.call_args_list == [call(url="http://localhost:8180", session=None)]
        assert client.auth.jwt.jwt_login.call_args_list == [
            call(role="my-role", jwt="eyJhbGciOiJSUzI1NiJ9.jwt-from-file")
        ]
        assert client.is_authenticated.call_args_list == [call(), call()]
        assert vault_client.kv_engine_version == 2

    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
    def test_jwt_with_token_strips_whitespace(self, mock_hvac):