[tests/fix] Fixed cm test and updated models validate_required_fields() method

Author: Aryamanz29

pyatlan/model/assets/badge.py

@@ -28,6 +28,7 @@ class Badge(Asset, type_name="Badge"):
     def creator(
         cls,
         *,
+        client: AtlanClient,
         name: StrictStr,
         cm_name: str,
         cm_attribute: str,
@@ -36,6 +37,7 @@ def creator(
         return cls(
             status=EntityStatus.ACTIVE,
             attributes=Badge.Attributes.create(
+                client=client,
                 name=name,
                 cm_name=cm_name,
                 cm_attribute=cm_attribute,
@@ -48,6 +50,7 @@ def creator(
     def create(
         cls,
         *,
+        client: AtlanClient,
         name: StrictStr,
         cm_name: str,
         cm_attribute: str,
@@ -62,6 +65,7 @@ def create(
             stacklevel=2,
         )
         return cls.creator(
+            client=client,
             name=name,
             cm_name=cm_name,
             cm_attribute=cm_attribute,
@@ -141,8 +145,8 @@ def create(
             badge_conditions: List[BadgeCondition],
         ) -> Badge.Attributes:
             validate_required_fields(
-                ["name", "cm_name", "cm_attribute", "badge_conditions"],
-                [name, cm_name, cm_attribute, badge_conditions],
+                ["client", "name", "cm_name", "cm_attribute", "badge_conditions"],
+                [client, name, cm_name, cm_attribute, badge_conditions],
             )
             cm_id = client.custom_metadata_cache.get_id_for_name(cm_name)
             cm_attr_id = client.custom_metadata_cache.get_attr_id_for_name(

pyatlan/model/assets/collection.py

@@ -103,7 +103,7 @@ def creator(
             client: AtlanClient,
             name: str,
         ) -> Collection.Attributes:
-            validate_required_fields(["name"], [name])
+            validate_required_fields(["client", "name"], [client, name])
             return Collection.Attributes(
                 name=name,
                 qualified_name=Collection._generate_qualified_name(client),

pyatlan/model/assets/connection.py

@@ -43,7 +43,9 @@ def creator(
         host: Optional[str] = None,
         port: Optional[int] = None,
     ) -> Connection:
-        validate_required_fields(["name", "connector_type"], [name, connector_type])
+        validate_required_fields(
+            ["client", "name", "connector_type"], [client, name, connector_type]
+        )
         if not admin_users and not admin_groups and not admin_roles:
             raise ValueError(
                 "One of admin_user, admin_groups or admin_roles is required"

pyatlan/model/assets/purpose.py

@@ -65,8 +65,8 @@ def create_metadata_policy(
         all_users: bool = False,
     ) -> AuthPolicy:
         validate_required_fields(
-            ["name", "purpose_id", "policy_type", "actions"],
-            [name, purpose_id, policy_type, actions],
+            ["client", "name", "purpose_id", "policy_type", "actions"],
+            [client, name, purpose_id, policy_type, actions],
         )
         target_found = False
         policy = AuthPolicy._AuthPolicy__create(name=name)  # type: ignore[attr-defined]
@@ -121,7 +121,8 @@ def create_data_policy(
         all_users: bool = False,
     ) -> AuthPolicy:
         validate_required_fields(
-            ["name", "purpose_id", "policy_type"], [name, purpose_id, policy_type]
+            ["client", "name", "purpose_id", "policy_type"],
+            [client, name, purpose_id, policy_type],
         )
         policy = AuthPolicy._AuthPolicy__create(name=name)  # type: ignore[attr-defined]
         policy.policy_actions = {DataAction.SELECT.value}

pyatlan/model/custom_metadata.py

@@ -118,13 +118,13 @@ def __init__(
         client: AtlanClient,
         business_attributes: Optional[Dict[str, Any]],
     ):
+        self._client = client
         self._metadata: Optional[Dict[str, CustomMetadataDict]] = None
         self._business_attributes = business_attributes
         self._modified = False
         if self._business_attributes is None:
             return
         self._metadata = {}
-        self._client = client
         for cm_id, cm_attributes in self._business_attributes.items():
             try:
                 cm_name = self._client.custom_metadata_cache.get_name_for_id(cm_id)

pyatlan/samples/custom_metadata/deploy_branded_cm.py

@@ -79,6 +79,7 @@ def create_badge():
     custom metadata structure.
     """
     badge = Badge.create(
+        client=client,
         name="Rating",
         cm_name=CUSTOM_METADATA_NAME,
         cm_attribute="Rating",

pyatlan/samples/events/lambda_scorer.py

@@ -71,6 +71,7 @@ def _create_cm_if_not_exists(client: AtlanClient) -> Optional[str]:
             client.typedef.create(cm_def)
             logger.info("Created DaaP custom metadata structure.")
             badge = Badge.create(
+                client=client,
                 name=CM_ATTR_DAAP_SCORE,
                 cm_name=CM_DAAP,
                 cm_attribute=CM_ATTR_DAAP_SCORE,

tests/integration/custom_metadata_test.py

@@ -1117,6 +1117,7 @@ def test_add_badge_cm_dq(
     cm_dq: CustomMetadataDef,
 ):
     badge = Badge.create(
+        client=client,
         name=CM_ATTR_QUALITY_COUNT,
         cm_name=CM_QUALITY,
         cm_attribute=CM_ATTR_QUALITY_COUNT,

tests/unit/model/badge_test.py

@@ -70,6 +70,7 @@ def test_create_when_required_parameters_are_missing_raises_value_error(
 ):
     with pytest.raises(ValueError, match=message):
         Badge.create(
+            client=client,
             name=name,
             cm_name=cm_name,
             cm_attribute=cm_attribute,