[tests] Added unit and integration tests for custom connector type

Author: Aryamanz29

tests/integration/connection_test.py

@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 Atlan Pte. Ltd.
+from typing import Generator
+
 import pytest
 
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.assets import Connection
-from pyatlan.model.enums import AtlanConnectorType
-from tests.integration.client import TestId
+from pyatlan.model.enums import AtlanConnectionCategory, AtlanConnectorType
+from tests.integration.client import TestId, delete_asset
 
 MODULE_NAME = TestId.make_unique("CONN")
 
@@ -24,6 +26,29 @@ def create_connection(
     )
 
 
+@pytest.fixture(scope="module")
+def custom_connection(client: AtlanClient) -> Generator[Connection, None, None]:
+    CUSTOM_CONNECTOR_TYPE = AtlanConnectorType.CREATE_CUSTOM(
+        name=f"{MODULE_NAME}_NAME",
+        value=f"{MODULE_NAME}_type",
+        category=AtlanConnectionCategory.API,
+    )
+    result = create_connection(
+        client=client, name=MODULE_NAME, connector_type=CUSTOM_CONNECTOR_TYPE
+    )
+    yield result
+    # TODO: proper connection delete workflow
+    delete_asset(client, guid=result.guid, asset_type=Connection)
+
+
+def test_custom_connection(custom_connection: Connection):
+    assert custom_connection.name == MODULE_NAME
+    assert custom_connection.connector_name == f"{MODULE_NAME}_type"
+    assert custom_connection.qualified_name
+    assert f"default/{MODULE_NAME}_type" in custom_connection.qualified_name
+    assert AtlanConnectorType[f"{MODULE_NAME}_NAME"].value == f"{MODULE_NAME}_type"
+
+
 def test_invalid_connection(client: AtlanClient):
     with pytest.raises(
         ValueError, match="One of admin_user, admin_groups or admin_roles is required"

tests/unit/model/a_d_l_s_object_test.py

@@ -134,7 +134,7 @@ def test_create():
         ),
         (
             ADLS_OBJECT_NAME,
-            "default/adls-invalid/production",
+            "default/adls/invalid/production",
             "abc",
             ADLS_CONTAINER_NAME,
             ADLS_CONTAINER_QUALIFIED_NAME,

tests/unit/model/connection_test.py

@@ -6,7 +6,7 @@
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.client.token import TokenClient
 from pyatlan.model.assets import Connection
-from pyatlan.model.enums import AtlanConnectorType
+from pyatlan.model.enums import AtlanConnectionCategory, AtlanConnectorType
 from tests.unit.model.constants import CONNECTION_NAME, CONNECTION_QUALIFIED_NAME
 
 
@@ -174,6 +174,69 @@ def test_create(
     assert sut.admin_roles == set(admin_roles)
 
 
+@pytest.mark.parametrize(
+    "name, connector_type, admin_users, admin_groups, admin_roles",
+    [
+        (
+            CONNECTION_NAME,
+            AtlanConnectorType.CREATE_CUSTOM(
+                name="FOO", value="foo", category=AtlanConnectionCategory.BI
+            ),
+            ["ernest"],
+            [],
+            [],
+        ),
+        (
+            CONNECTION_NAME,
+            AtlanConnectorType.CREATE_CUSTOM(
+                name="BAR", value="bar", category=AtlanConnectionCategory.API
+            ),
+            [],
+            ["ernest"],
+            [],
+        ),
+        (
+            CONNECTION_NAME,
+            AtlanConnectorType.CREATE_CUSTOM(
+                name="BAZ", value="baz", category=AtlanConnectionCategory.WAREHOUSE
+            ),
+            [],
+            [],
+            ["ernest"],
+        ),
+    ],
+)
+def test_creator_with_custom_type(
+    name: str,
+    connector_type: AtlanConnectorType,
+    admin_users: List[str],
+    admin_groups: List[str],
+    admin_roles: List[str],
+    mock_role_cache,
+    mock_user_cache,
+    mock_group_cache,
+):
+    mock_role_cache.validate_idstrs
+    mock_user_cache.validate_names
+    mock_group_cache.validate_aliases
+
+    sut = Connection.create(
+        name=name,
+        connector_type=connector_type,
+        admin_users=admin_users,
+        admin_groups=admin_groups,
+        admin_roles=admin_roles,
+    )
+
+    assert sut.name == name
+    assert sut.qualified_name
+    assert sut.qualified_name[:20] == connector_type.to_qualified_name()[:20]
+    assert sut.connector_name == connector_type.value
+    assert sut.admin_users == set(admin_users)
+    assert sut.admin_groups == set(admin_groups)
+    assert sut.admin_roles == set(admin_roles)
+
+
 @pytest.mark.parametrize(
     "qualified_name, name, message",
     [

tests/unit/test_utils.py

@@ -5,6 +5,7 @@
 import pytest
 
 from pyatlan.errors import InvalidRequestError
+from pyatlan.model.enums import AtlanConnectionCategory, AtlanConnectorType
 from pyatlan.model.utils import construct_object_key
 from pyatlan.utils import (
     ComparisonCategory,
@@ -292,3 +293,120 @@ def test_validate_type_with_valid_values(name, the_type, value):
 def test_contruct_object_key(prefix, name, expected_key):
     key = construct_object_key(prefix, name)
     assert key == expected_key
+
+
+@pytest.mark.parametrize(
+    "custom_connectors",
+    [
+        [
+            AtlanConnectorType.CREATE_CUSTOM(
+                name="FOO", value="foo", category=AtlanConnectionCategory.BI
+            ),
+            AtlanConnectorType.CREATE_CUSTOM(
+                name="BAR", value="bar", category=AtlanConnectionCategory.API
+            ),
+            AtlanConnectorType.CREATE_CUSTOM(
+                name="BAZ", value="baz", category=AtlanConnectionCategory.WAREHOUSE
+            ),
+        ]
+    ],
+)
+def test_atlan_connector_type_create_custom(custom_connectors):
+    for custom_connector in custom_connectors:
+        assert custom_connector and custom_connector.category
+        assert custom_connector.value in AtlanConnectorType.get_values()
+        assert custom_connector.strip() in AtlanConnectorType.get_values()
+        assert custom_connector.name in AtlanConnectorType.get_names()
+        assert (
+            custom_connector.name,
+            custom_connector.value,
+        ) in AtlanConnectorType.get_items()
+
+
+@pytest.mark.parametrize(
+    "custom_asset_qns",
+    [
+        [
+            "default/c1/1234567890/asset/name",
+            "default/c2/1234567890/asset/name",
+            "default/c3/1234567890/asset/name",
+            # Duplicate custom connector names
+            "default/c1/1234567890/asset/name",
+            "default/c2/1234567890/asset/name",
+            "default/c3/1234567890/asset/name",
+            # Duplicate custom connector names
+            "default/c1/1234567890/asset/name",
+            "default/c2/1234567890/asset/name",
+            "default/c3/1234567890/asset/name",
+        ]
+    ],
+)
+def test_atlan_connector_type_custom_asset_qn(custom_asset_qns):
+    # Calculate the unique custom QNs
+    unique_custom_qns = set(custom_asset_qns)
+
+    # Calculate the initial lengths of predefined values, names, and items
+    len_get_values = len(AtlanConnectorType.get_values())
+    len_get_names = len(AtlanConnectorType.get_names())
+    len_get_items = len(AtlanConnectorType.get_items())
+
+    # Check each custom QN
+    for custom_qn in custom_asset_qns:
+        custom_connector = AtlanConnectorType._get_connector_type_from_qualified_name(
+            custom_qn
+        )
+        custom_connector.category = AtlanConnectionCategory.CUSTOM
+        assert custom_connector.value in AtlanConnectorType.get_values()
+        assert custom_connector.strip() in AtlanConnectorType.get_values()
+        assert (
+            custom_connector.name,
+            custom_connector.value,
+        ) in AtlanConnectorType.get_items()
+
+    # Ensure the value is only added once (i.e no duplicates)
+    assert len(AtlanConnectorType.get_values()) == len_get_values + len(
+        unique_custom_qns
+    )
+    assert len(AtlanConnectorType.get_names()) == len_get_names + len(unique_custom_qns)
+    assert len(AtlanConnectorType.get_items()) == len_get_items + len(unique_custom_qns)
+
+
+@pytest.mark.parametrize(
+    "custom_connection_qns",
+    [
+        [
+            "default/cm1/1234567890",
+            "default/cm2/1234567890",
+            "default/cm3/1234567890",
+            # Duplicate custom connector names
+            "default/cm1/1234567890",
+            "default/cm2/1234567890",
+            "default/cm3/1234567890",
+            # Duplicate custom connector names
+            "default/cm1/1234567890",
+            "default/cm2/1234567890",
+            "default/cm3/1234567890",
+        ]
+    ],
+)
+def test_atlan_connector_type_custom_connection_qn(custom_connection_qns):
+    # Calculate the unique custom QNs
+    unique_custom_qns = set(custom_connection_qns)
+
+    # Calculate the initial lengths of predefined values, names, and items
+    len_get_values = len(AtlanConnectorType.get_values())
+    len_get_names = len(AtlanConnectorType.get_names())
+    len_get_items = len(AtlanConnectorType.get_items())
+
+    # Check each custom QN
+    for custom_qn in custom_connection_qns:
+        custom_connector_name = AtlanConnectorType.get_connector_name(custom_qn)
+        assert custom_connector_name in AtlanConnectorType.get_values()
+        assert custom_connector_name.strip() in AtlanConnectorType.get_values()
+
+    # Ensure the value is only added once (i.e no duplicates)
+    assert len(AtlanConnectorType.get_values()) == len_get_values + len(
+        unique_custom_qns
+    )
+    assert len(AtlanConnectorType.get_names()) == len_get_names + len(unique_custom_qns)
+    assert len(AtlanConnectorType.get_items()) == len_get_items + len(unique_custom_qns)