Cleanup tests for mypy

Signed-off-by: Christopher Grote <[email protected]>

Author: cmgrote

pyatlan/cache/custom_metadata_cache.py

@@ -65,12 +65,12 @@ def refresh_cache(cls) -> None:
                 applicable_types: set[str] = set()
                 if cm.attribute_defs:
                     for attr in cm.attribute_defs:
-                        if attr.options.custom_applicable_entity_types:
+                        if attr.options and attr.options.custom_applicable_entity_types:
                             applicable_types.update(
                                 json.loads(attr.options.custom_applicable_entity_types)
                             )
-                        attr_id = attr.name
-                        attr_name = attr.display_name
+                        attr_id = str(attr.name)
+                        attr_name = str(attr.display_name)
                         # Use a renamed attribute everywhere
                         attr_renamed = to_snake_case(attr_name.replace(" ", ""))
                         cls.map_attr_id_to_name[type_id][attr_id] = attr_renamed

pyatlan/model/typedef.py

@@ -427,7 +427,7 @@ def create(
             base_type = attribute_type.value
         if multi_valued:
             attr_def.type_name = f"array<{str(base_type)}>"
-            attr_def.options.multi_value_select = True
+            attr_def.options.multi_value_select = True  # type: ignore
         else:
             attr_def.type_name = base_type
         if add_enum_values:

tests/integration/conftest.py renamed to tests/conftest.py

@@ -1,14 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 Atlan Pte. Ltd.
-import time
 from typing import Generator, Callable
 
 import pytest
 
 from pyatlan.cache.classification_cache import ClassificationCache
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.error import AtlanError
-from pyatlan.model.enums import AtlanClassificationColor, AtlanTypeCategory
+from pyatlan.model.enums import AtlanClassificationColor
 from pyatlan.model.typedef import ClassificationDef
 
 from retry import retry

tests/integration/classification_test.py

@@ -163,16 +163,19 @@ def test_cm_ipr(cm_ipr: CustomMetadataDef):
     assert one.display_name == CM_ATTR_IPR_MANDATORY
     assert one.name != CM_ATTR_IPR_MANDATORY
     assert one.type_name == AtlanCustomAttributePrimitiveType.BOOLEAN.value
+    assert one.options
     assert not one.options.multi_value_select
     one = attributes[3]
     assert one.display_name == CM_ATTR_IPR_DATE
     assert one.name != CM_ATTR_IPR_DATE
     assert one.type_name == AtlanCustomAttributePrimitiveType.DATE.value
+    assert one.options
     assert not one.options.multi_value_select
     one = attributes[4]
     assert one.display_name == CM_ATTR_IPR_URL
     assert one.name != CM_ATTR_IPR_URL
     assert one.type_name == AtlanCustomAttributePrimitiveType.STRING.value
+    assert one.options
     assert not one.options.multi_value_select
 
 
@@ -237,16 +240,19 @@ def test_cm_raci(cm_raci: CustomMetadataDef):
     assert one.display_name == CM_ATTR_RACI_CONSULTED
     assert one.name != CM_ATTR_RACI_CONSULTED
     assert one.type_name == f"array<{AtlanCustomAttributePrimitiveType.STRING.value}>"
+    assert one.options
     assert one.options.multi_value_select
     one = attributes[3]
     assert one.display_name == CM_ATTR_RACI_INFORMED
     assert one.name != CM_ATTR_RACI_INFORMED
     assert one.type_name == f"array<{AtlanCustomAttributePrimitiveType.STRING.value}>"
+    assert one.options
     assert one.options.multi_value_select
     one = attributes[4]
     assert one.display_name == CM_ATTR_RACI_EXTRA
     assert one.name != CM_ATTR_RACI_EXTRA
     assert one.type_name == AtlanCustomAttributePrimitiveType.STRING.value
+    assert one.options
     assert not one.options.multi_value_select
 
 
@@ -324,6 +330,7 @@ def test_cm_dq(cm_dq: CustomMetadataDef):
     assert one.display_name == CM_ATTR_QUALITY_TYPE
     assert one.name != CM_ATTR_QUALITY_TYPE
     assert one.type_name == CM_ENUM_DQ_TYPE
+    assert one.options
     assert not one.options.multi_value_select
     assert one.options.primitive_type == AtlanCustomAttributePrimitiveType.OPTIONS.value
 
@@ -552,6 +559,7 @@ def test_remove_attribute(client: AtlanClient, cm_raci: CustomMetadataDef):
         to_keep = existing_attr
         if existing_attr.display_name == CM_ATTR_RACI_EXTRA:
             to_keep = existing_attr.archive(by="test-automation")
+            assert to_keep.options
             _removal_epoch = to_keep.options.archived_at
         updated_attrs.append(to_keep)
     existing.attribute_defs = updated_attrs
@@ -832,6 +840,7 @@ def _validate_raci_structure(
     assert one.display_name == CM_ATTR_RACI_RESPONSIBLE
     assert one.name != CM_ATTR_RACI_RESPONSIBLE
     assert one.type_name == f"array<{AtlanCustomAttributePrimitiveType.STRING.value}>"
+    assert one.options
     assert "Database" in str(one.options.custom_applicable_entity_types)
     assert not one.is_archived()
     assert one.options.multi_value_select
@@ -840,6 +849,7 @@ def _validate_raci_structure(
     assert one.display_name == CM_ATTR_RACI_ACCOUNTABLE
     assert one.name != CM_ATTR_RACI_ACCOUNTABLE
     assert one.type_name == AtlanCustomAttributePrimitiveType.STRING.value
+    assert one.options
     assert "Table" in str(one.options.custom_applicable_entity_types)
     assert not one.is_archived()
     assert not one.options.multi_value_select
@@ -848,6 +858,7 @@ def _validate_raci_structure(
     assert one.display_name == CM_ATTR_RACI_CONSULTED
     assert one.name != CM_ATTR_RACI_CONSULTED
     assert one.type_name == f"array<{AtlanCustomAttributePrimitiveType.STRING.value}>"
+    assert one.options
     assert "Column" in str(one.options.custom_applicable_entity_types)
     assert not one.is_archived()
     assert one.options.multi_value_select
@@ -856,6 +867,7 @@ def _validate_raci_structure(
     assert one.display_name == CM_ATTR_RACI_INFORMED
     assert not one.name == CM_ATTR_RACI_INFORMED
     assert one.type_name == f"array<{AtlanCustomAttributePrimitiveType.STRING.value}>"
+    assert one.options
     assert "MaterialisedView" in str(one.options.custom_applicable_entity_types)
     assert not one.is_archived()
     assert one.options.multi_value_select
@@ -868,6 +880,7 @@ def _validate_raci_structure(
         )
         assert one.name != CM_ATTR_RACI_EXTRA
         assert one.type_name == AtlanCustomAttributePrimitiveType.STRING.value
+        assert one.options
         assert "AtlasGlossaryTerm" in str(one.options.custom_applicable_entity_types)
         assert not one.options.multi_value_select
         assert one.is_archived()

tests/integration/custom_metadata_test.py

@@ -2,20 +2,24 @@
 # Copyright 2022 Atlan Pte. Ltd.
 from pyatlan.cache.role_cache import RoleCache
 
+import logging
+
+LOGGER = logging.getLogger(__name__)
+
 
 def test_retrieve_admin():
     admin_role_guid = RoleCache.get_id_for_name("$admin")
-    print(admin_role_guid)
     assert admin_role_guid
+    LOGGER.info(f"Admin role: {admin_role_guid}")
 
 
 def test_retrieve_guest():
     guest_role_guid = RoleCache.get_id_for_name("$guest")
-    print(guest_role_guid)
     assert guest_role_guid
+    LOGGER.info(f"Guest role: {guest_role_guid}")
 
 
 def test_retrieve_member():
     member_role_guid = RoleCache.get_id_for_name("$member")
-    print(member_role_guid)
     assert member_role_guid
+    LOGGER.info(f"Member role: {member_role_guid}")