datahub/metadata-ingestion/tests/unit/secret/test_datahub_secret_store.py

from unittest.mock import Mock, patch

import pytest

from datahub.ingestion.graph.client import DataHubGraph
from datahub.ingestion.graph.config import DatahubClientConfig
from datahub.secret.datahub_secret_store import (
    DataHubSecretStore,
    DataHubSecretStoreConfig,
)
from datahub.secret.datahub_secrets_client import DataHubSecretsClient


class TestDataHubSecretStore:
    def test_init_with_graph_client(self):
        mock_graph = Mock(spec=DataHubGraph)
        mock_graph.test_connection.return_value = True

        config = DataHubSecretStoreConfig(graph_client=mock_graph)
        store = DataHubSecretStore(config)

        assert store.client is not None
        assert isinstance(store.client, DataHubSecretsClient)
        mock_graph.test_connection.assert_called_once()

    def test_init_with_graph_client_config(self):
        mock_client_config = Mock(spec=DatahubClientConfig)

        with patch(
            "datahub.secret.datahub_secret_store.DataHubGraph"
        ) as mock_graph_class:
            mock_graph = Mock(spec=DataHubGraph)
            mock_graph_class.return_value = mock_graph

            config = DataHubSecretStoreConfig(graph_client_config=mock_client_config)
            store = DataHubSecretStore(config)

            assert store.client is not None
            mock_graph_class.assert_called_once_with(mock_client_config)

    def test_init_with_no_config_raises_exception(self):
        config = DataHubSecretStoreConfig()

        with pytest.raises(Exception, match="Invalid configuration provided"):
            DataHubSecretStore(config)

    def test_get_secret_values_success(self):
        mock_graph = Mock(spec=DataHubGraph)
        mock_graph.test_connection.return_value = True

        expected_secrets = {"secret1": "value1", "secret2": "value2"}

        with patch(
            "datahub.secret.datahub_secret_store.DataHubSecretsClient"
        ) as mock_client_class:
            mock_client = Mock(spec=DataHubSecretsClient)
            mock_client.get_secret_values.return_value = expected_secrets
            mock_client_class.return_value = mock_client

            config = DataHubSecretStoreConfig(graph_client=mock_graph)
            store = DataHubSecretStore(config)

            result = store.get_secret_values(["secret1", "secret2"])

            assert result == expected_secrets
            mock_client.get_secret_values.assert_called_once_with(
                ["secret1", "secret2"]
            )

    def test_get_secret_values_exception_handling(self):
        mock_graph = Mock(spec=DataHubGraph)
        mock_graph.test_connection.return_value = True

        with patch(
            "datahub.secret.datahub_secret_store.DataHubSecretsClient"
        ) as mock_client_class:
            mock_client = Mock(spec=DataHubSecretsClient)
            mock_client.get_secret_values.side_effect = Exception("Connection failed")
            mock_client_class.return_value = mock_client

            config = DataHubSecretStoreConfig(graph_client=mock_graph)
            store = DataHubSecretStore(config)

            with patch("datahub.secret.datahub_secret_store.logger") as mock_logger:
                result = store.get_secret_values(["secret1"])

                assert result == {}
                mock_logger.exception.assert_called_once()

    def test_get_secret_value(self):
        mock_graph = Mock(spec=DataHubGraph)
        mock_graph.test_connection.return_value = True

        with patch(
            "datahub.secret.datahub_secret_store.DataHubSecretsClient"
        ) as mock_client_class:
            mock_client = Mock(spec=DataHubSecretsClient)
            mock_client.get_secret_values.return_value = {"secret1": "value1"}
            mock_client_class.return_value = mock_client

            config = DataHubSecretStoreConfig(graph_client=mock_graph)
            store = DataHubSecretStore(config)

            result = store.get_secret_value("secret1")

            assert result == "value1"
            mock_client.get_secret_values.assert_called_once_with(["secret1"])

    def test_get_secret_value_not_found(self):
        mock_graph = Mock(spec=DataHubGraph)
        mock_graph.test_connection.return_value = True

        with patch(
            "datahub.secret.datahub_secret_store.DataHubSecretsClient"
        ) as mock_client_class:
            mock_client = Mock(spec=DataHubSecretsClient)
            mock_client.get_secret_values.return_value = {}
            mock_client_class.return_value = mock_client

            config = DataHubSecretStoreConfig(graph_client=mock_graph)
            store = DataHubSecretStore(config)

            result = store.get_secret_value("nonexistent")

            assert result is None

    def test_get_id(self):
        mock_graph = Mock(spec=DataHubGraph)
        mock_graph.test_connection.return_value = True

        config = DataHubSecretStoreConfig(graph_client=mock_graph)
        store = DataHubSecretStore(config)

        assert store.get_id() == "datahub"

    def test_create_classmethod(self):
        mock_graph = Mock(spec=DataHubGraph)
        mock_graph.test_connection.return_value = True

        config_dict = {"graph_client": mock_graph}

        store = DataHubSecretStore.create(config_dict)

        assert isinstance(store, DataHubSecretStore)
        assert store.client is not None

    def test_close(self):
        mock_graph = Mock(spec=DataHubGraph)
        mock_graph.test_connection.return_value = True

        with patch(
            "datahub.secret.datahub_secret_store.DataHubSecretsClient"
        ) as mock_client_class:
            mock_client = Mock(spec=DataHubSecretsClient)
            mock_client.graph = mock_graph
            mock_client_class.return_value = mock_client

            config = DataHubSecretStoreConfig(graph_client=mock_graph)
            store = DataHubSecretStore(config)

            store.close()

            mock_graph.close.assert_called_once()

    def test_config_validator_with_working_connection(self):
        mock_graph = Mock(spec=DataHubGraph)
        mock_graph.test_connection.return_value = True

        config = DataHubSecretStoreConfig(graph_client=mock_graph)

        assert config.graph_client == mock_graph
        mock_graph.test_connection.assert_called_once()

    def test_config_validator_with_none_graph_client(self):
        config = DataHubSecretStoreConfig(graph_client=None)

        assert config.graph_client is None