[tests/fix] Fixed failing cm test due to incorrect refactoring (save_replacing_cm)

Author: Aryamanz29

pyatlan/client/aio/asset.py

@@ -10,6 +10,7 @@
     Type,
     TypeVar,
     Union,
+    overload,
 )
 
 from pydantic.v1 import StrictStr, constr
@@ -460,8 +461,8 @@ async def save_replacing_cm(
         query_params, request = Save.prepare_request_replacing_cm(
             entity=entity,
             replace_atlan_tags=replace_atlan_tags,
+            client=self._async_client,
         )
-        Save.validate_and_flush_entities(request.entities, self._async_client)
         raw_json = await self._async_client._call_api(
             BULK_UPDATE, query_params, request
         )
@@ -929,6 +930,39 @@ async def _update_asset_by_attribute(
         # Process response using shared logic
         return UpdateAssetByAttribute.process_response(raw_json, asset_type)
 
+    @overload
+    async def update_certificate(
+        self,
+        asset_type: Type[AtlasGlossaryTerm],
+        qualified_name: str,
+        name: str,
+        certificate_status: CertificateStatus,
+        glossary_guid: str,
+        message: Optional[str] = None,
+    ) -> Optional[AtlasGlossaryTerm]: ...
+
+    @overload
+    async def update_certificate(
+        self,
+        asset_type: Type[AtlasGlossaryCategory],
+        qualified_name: str,
+        name: str,
+        certificate_status: CertificateStatus,
+        glossary_guid: str,
+        message: Optional[str] = None,
+    ) -> Optional[AtlasGlossaryCategory]: ...
+
+    @overload
+    async def update_certificate(
+        self,
+        asset_type: Type[A],
+        qualified_name: str,
+        name: str,
+        certificate_status: CertificateStatus,
+        glossary_guid: Optional[str] = None,
+        message: Optional[str] = None,
+    ) -> Optional[A]: ...
+
     async def update_certificate(
         self,
         asset_type: Type[A],
@@ -965,6 +999,33 @@ async def update_certificate(
         # Execute update using shared logic
         return await self._update_asset_by_attribute(asset, asset_type, qualified_name)
 
+    @overload
+    async def remove_certificate(
+        self,
+        asset_type: Type[AtlasGlossaryTerm],
+        qualified_name: str,
+        name: str,
+        glossary_guid: str,
+    ) -> Optional[AtlasGlossaryTerm]: ...
+
+    @overload
+    async def remove_certificate(
+        self,
+        asset_type: Type[AtlasGlossaryCategory],
+        qualified_name: str,
+        name: str,
+        glossary_guid: str,
+    ) -> Optional[AtlasGlossaryCategory]: ...
+
+    @overload
+    async def remove_certificate(
+        self,
+        asset_type: Type[A],
+        qualified_name: str,
+        name: str,
+        glossary_guid: Optional[str] = None,
+    ) -> Optional[A]: ...
+
     async def remove_certificate(
         self,
         asset_type: Type[A],
@@ -994,6 +1055,36 @@ async def remove_certificate(
         # Execute update using shared logic
         return await self._update_asset_by_attribute(asset, asset_type, qualified_name)
 
+    @overload
+    async def update_announcement(
+        self,
+        asset_type: Type[AtlasGlossaryTerm],
+        qualified_name: str,
+        name: str,
+        announcement: Announcement,
+        glossary_guid: str,
+    ) -> Optional[AtlasGlossaryTerm]: ...
+
+    @overload
+    async def update_announcement(
+        self,
+        asset_type: Type[AtlasGlossaryCategory],
+        qualified_name: str,
+        name: str,
+        announcement: Announcement,
+        glossary_guid: str,
+    ) -> Optional[AtlasGlossaryCategory]: ...
+
+    @overload
+    async def update_announcement(
+        self,
+        asset_type: Type[A],
+        qualified_name: str,
+        name: str,
+        announcement: Announcement,
+        glossary_guid: Optional[str] = None,
+    ) -> Optional[A]: ...
+
     async def update_announcement(
         self,
         asset_type: Type[A],
@@ -1026,6 +1117,33 @@ async def update_announcement(
         # Execute update using shared logic
         return await self._update_asset_by_attribute(asset, asset_type, qualified_name)
 
+    @overload
+    async def remove_announcement(
+        self,
+        asset_type: Type[AtlasGlossaryTerm],
+        qualified_name: str,
+        name: str,
+        glossary_guid: str,
+    ) -> Optional[AtlasGlossaryTerm]: ...
+
+    @overload
+    async def remove_announcement(
+        self,
+        asset_type: Type[AtlasGlossaryCategory],
+        qualified_name: str,
+        name: str,
+        glossary_guid: str,
+    ) -> Optional[AtlasGlossaryCategory]: ...
+
+    @overload
+    async def remove_announcement(
+        self,
+        asset_type: Type[A],
+        qualified_name: str,
+        name: str,
+        glossary_guid: Optional[str] = None,
+    ) -> Optional[A]: ...
+
     async def remove_announcement(
         self,
         asset_type: Type[A],

pyatlan/client/asset.py

@@ -550,8 +550,8 @@ def save_replacing_cm(
         query_params, request = Save.prepare_request_replacing_cm(
             entity=entity,
             replace_atlan_tags=replace_atlan_tags,
+            client=self._client,
         )
-        Save.validate_and_flush_entities(request.entities, self._client)
         raw_json = self._client._call_api(BULK_UPDATE, query_params, request)
         return Save.process_response_replacing_cm(raw_json)
 

pyatlan/client/common/asset.py

@@ -739,13 +739,15 @@ def log_connections_finished():
     def prepare_request_replacing_cm(
         entity: Union[Asset, List[Asset]],
         replace_atlan_tags: bool = False,
+        client=None,
     ) -> tuple[Dict[str, Any], BulkRequest[Asset]]:
         """
         Prepare the request for saving assets with replacing custom metadata.
         This uses different query parameter names than the regular save method.
 
         :param entity: one or more assets to save
         :param replace_atlan_tags: whether to replace AtlanTags during an update
+        :param client: the Atlan client instance for flushing custom metadata
         :returns: tuple of (query_params, bulk_request)
         """
         query_params = {
@@ -760,6 +762,12 @@ def prepare_request_replacing_cm(
         else:
             entities.append(entity)
 
+        # Validate and flush entities BEFORE creating the BulkRequest
+        if client:
+            for asset in entities:
+                asset.validate_required()
+                asset.flush_custom_metadata(client=client)
+
         return query_params, BulkRequest[Asset](entities=entities)
 
     @staticmethod