datahub/metadata-ingestion/tests/unit/sdk_v2/test_client_v2.py

from unittest.mock import Mock

import pytest

from datahub.errors import ItemNotFoundError, MultipleItemsFoundError, SdkUsageError
from datahub.ingestion.graph.client import DataHubGraph
from datahub.ingestion.graph.config import DatahubClientConfig
from datahub.metadata.urns import CorpUserUrn
from datahub.sdk.main_client import DataHubClient


@pytest.fixture
def mock_graph() -> Mock:
    graph = Mock(spec=DataHubGraph)
    graph.exists.return_value = False
    return graph


def test_client_creation(mock_graph: Mock) -> None:
    assert DataHubClient(graph=mock_graph)
    assert DataHubClient(server="https://example.com", token="token")


def test_client_init_errors(mock_graph: Mock) -> None:
    config = DatahubClientConfig(server="https://example.com", token="token")

    with pytest.raises(SdkUsageError):
        DataHubClient(server="https://example.com", graph=mock_graph)  # type: ignore
    with pytest.raises(SdkUsageError):
        DataHubClient(server="https://example.com", config=config)  # type: ignore
    with pytest.raises(SdkUsageError):
        DataHubClient(config=config, graph=mock_graph)  # type: ignore
    with pytest.raises(SdkUsageError):
        DataHubClient()  # type: ignore


def test_resolve_user(mock_graph: Mock) -> None:
    client = DataHubClient(graph=mock_graph)

    # This test doesn't really validate the graphql query or vars.
    # It probably makes more sense to test via smoke-tests.

    mock_graph.get_urns_by_filter.return_value = []
    with pytest.raises(ItemNotFoundError):
        client.resolve.user(name="User")

    mock_graph.get_urns_by_filter.return_value = ["urn:li:corpuser:user"]
    assert client.resolve.user(name="User") == CorpUserUrn("urn:li:corpuser:user")

    mock_graph.get_urns_by_filter.return_value = [
        "urn:li:corpuser:user",
        "urn:li:corpuser:user2",
    ]
    with pytest.raises(MultipleItemsFoundError):
        client.resolve.user(name="User")