From b2d4c1e313dc21d7d62d7dcf979694cf34181225 Mon Sep 17 00:00:00 2001
From: Ernest Hill <ernest@atlan.com>
Date: Mon, 5 Jun 2023 14:52:43 +0300
Subject: [PATCH] Added validators to Asset Request and Bulk request to flush
 the custom metadat of the assets

---
 pyatlan/model/assets.py                   |  3 +++
 pyatlan/model/core.py                     | 18 +++++++++++++++++-
 tests/integration/custom_metadata_test.py | 14 +++++++-------
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/pyatlan/model/assets.py b/pyatlan/model/assets.py
index a3972cf6f..bcb348515 100644
--- a/pyatlan/model/assets.py
+++ b/pyatlan/model/assets.py
@@ -268,6 +268,9 @@ def get_custom_metadata(self, name: str) -> CustomMetadataDict:
     def set_custom_metadata(self, custom_metadata: CustomMetadataDict):
         return self._metadata_proxy.set_custom_metadata(custom_metadata=custom_metadata)
 
+    def flush_custom_metadata(self):
+        self.business_attributes = self._metadata_proxy.business_attributes
+
 
 class Asset(Referenceable):
     """Description"""
diff --git a/pyatlan/model/core.py b/pyatlan/model/core.py
index 2528fe6ba..763961335 100644
--- a/pyatlan/model/core.py
+++ b/pyatlan/model/core.py
@@ -2,7 +2,7 @@
 # Copyright 2022 Atlan Pte. Ltd.
 from typing import TYPE_CHECKING
 
-from pydantic import BaseModel, Extra, Field
+from pydantic import BaseModel, Extra, Field, validator
 
 if TYPE_CHECKING:
     from dataclasses import dataclass
@@ -184,6 +184,22 @@ class AssetResponse(AtlanObject, GenericModel, Generic[T]):
 class AssetRequest(AtlanObject, GenericModel, Generic[T]):
     entity: T
 
+    @validator("entity")
+    def flush_custom_metadata(cls, v):
+        from pyatlan.model.assets import Asset
+
+        if isinstance(v, Asset):
+            v.flush_custom_metadata()
+        return v
+
 
 class BulkRequest(AtlanObject, GenericModel, Generic[T]):
     entities: list[T]
+
+    @validator("entities", each_item=True)
+    def flush_custom_metadata(cls, v):
+        from pyatlan.model.assets import Asset
+
+        if isinstance(v, Asset):
+            v.flush_custom_metadata()
+        return v
diff --git a/tests/integration/custom_metadata_test.py b/tests/integration/custom_metadata_test.py
index 2049bc807..47ab7504a 100644
--- a/tests/integration/custom_metadata_test.py
+++ b/tests/integration/custom_metadata_test.py
@@ -404,7 +404,7 @@ def test_add_term_cm_raci(
     groups: List[AtlanGroup],
 ):
     cm_name = make_unique("RACI")
-    raci_attrs = term.get_custom_metadata(cm_name)
+    raci_attrs = CustomMetadataDict(cm_name)
     _validate_raci_empty(raci_attrs)
     group1, group2 = _get_groups(client, make_unique)
     raci_attrs[CM_ATTR_RACI_RESPONSIBLE] = [FIXED_USER]
@@ -424,7 +424,7 @@ def test_add_term_cm_ipr(
     term: AtlasGlossaryTerm,
 ):
     cm_name = make_unique("IPR")
-    ipr_attrs = term.get_custom_metadata(cm_name)
+    ipr_attrs = CustomMetadataDict(cm_name)
     _validate_ipr_empty(ipr_attrs)
     ipr_attrs[CM_ATTR_IPR_LICENSE] = "CC BY"
     ipr_attrs[CM_ATTR_IPR_VERSION] = 2.0
@@ -445,7 +445,7 @@ def test_add_term_cm_dq(
     term: AtlasGlossaryTerm,
 ):
     cm_name = make_unique("DQ")
-    dq_attrs = term.get_custom_metadata(cm_name)
+    dq_attrs = CustomMetadataDict(cm_name)
     _validate_dq_empty(dq_attrs)
     dq_attrs[CM_ATTR_QUALITY_COUNT] = 42
     dq_attrs[CM_ATTR_QUALITY_SQL] = "SELECT * from SOMEWHERE;"
@@ -464,7 +464,7 @@ def test_update_term_cm_ipr(
     term: AtlasGlossaryTerm,
 ):
     cm_name = make_unique("IPR")
-    ipr = term.get_custom_metadata(cm_name)
+    ipr = CustomMetadataDict(cm_name)
     # Note: MUST access the getter / setter, not the underlying store
     ipr[CM_ATTR_IPR_MANDATORY] = False
     client.update_custom_metadata_attributes(term.guid, ipr)
@@ -487,7 +487,7 @@ def test_replace_term_cm_raci(
     CM_RACI = make_unique("RACI")
     CM_IPR = make_unique("IPR")
     CM_QUALITY = make_unique("DQ")
-    raci = term.get_custom_metadata(CM_RACI)
+    raci = CustomMetadataDict(CM_RACI)
     group1, group2 = _get_groups(client, make_unique)
     raci[CM_ATTR_RACI_RESPONSIBLE] = [FIXED_USER]
     raci[CM_ATTR_RACI_ACCOUNTABLE] = FIXED_USER
@@ -513,7 +513,7 @@ def test_replace_term_cm_ipr(
     CM_RACI = make_unique("RACI")
     CM_IPR = make_unique("IPR")
     CM_QUALITY = make_unique("DQ")
-    term_cm_ipr = term.get_custom_metadata(CM_IPR)
+    term_cm_ipr = CustomMetadataDict(CM_IPR)
     client.replace_custom_metadata(term.guid, term_cm_ipr)
     t = client.retrieve_minimal(guid=term.guid, asset_type=AtlasGlossaryTerm)
     assert t
@@ -800,7 +800,7 @@ def test_update_replacing_cm(
     CM_RACI = make_unique("RACI")
     CM_IPR = make_unique("IPR")
     CM_QUALITY = make_unique("DQ")
-    raci = term.get_custom_metadata(CM_RACI)
+    raci = CustomMetadataDict(CM_RACI)
     group1, group2 = _get_groups(client, make_unique)
     raci[CM_ATTR_RACI_RESPONSIBLE] = [FIXED_USER]
     raci[CM_ATTR_RACI_ACCOUNTABLE] = FIXED_USER