diff --git a/pyatlan/model/assets/core/data_contract.py b/pyatlan/model/assets/core/data_contract.py index f8817a812..ae60e7b4a 100644 --- a/pyatlan/model/assets/core/data_contract.py +++ b/pyatlan/model/assets/core/data_contract.py @@ -5,7 +5,7 @@ from __future__ import annotations from json import JSONDecodeError, loads -from typing import ClassVar, List, Optional, Union, overload +from typing import TYPE_CHECKING, ClassVar, List, Optional, Tuple, Type, Union, overload from pydantic.v1 import Field, validator @@ -21,18 +21,32 @@ from .catalog import Catalog +if TYPE_CHECKING: + from pyatlan.client.atlan import AtlanClient + from pyatlan.model.response import AssetMutationResponse + class DataContract(Catalog): """Description""" @overload @classmethod - def creator(cls, asset_qualified_name: str, contract_json: str) -> DataContract: ... + def creator( + cls, + *, + asset_qualified_name: str, + asset_type: Type[Asset], + contract_json: str, + ) -> DataContract: ... @overload @classmethod def creator( - cls, asset_qualified_name: str, contract_spec: Union[DataContractSpec, str] + cls, + *, + asset_qualified_name: str, + asset_type: Type[Asset], + contract_spec: Union[DataContractSpec, str], ) -> DataContract: ... @classmethod @@ -41,6 +55,7 @@ def creator( cls, *, asset_qualified_name: str, + asset_type: Type[Asset], contract_json: Optional[str] = None, contract_spec: Optional[Union[DataContractSpec, str]] = None, ) -> DataContract: @@ -50,11 +65,71 @@ def creator( ) attributes = DataContract.Attributes.creator( asset_qualified_name=asset_qualified_name, + asset_type=asset_type, contract_json=contract_json, contract_spec=contract_spec, ) return cls(attributes=attributes) + @staticmethod + def save( + client: "AtlanClient", + contract: "DataContract", + ) -> "AssetMutationResponse": + """Save a DataContract. + + The contract's ``data_contract_asset_latest`` relationship (set by + ``creator()``) links the contract to the governed asset automatically. + + :param client: connectivity to an Atlan tenant + :param contract: DataContract to save (from ``DataContract.creator()``) + :returns: the result of the save + """ + return client.asset.save(contract) + + @staticmethod + def delete( + client: "AtlanClient", + contract_guid: str, + ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": + """Delete (purge) a DataContract and clean up the linked asset. + + Retrieves the contract to find the linked asset, clears + ``hasContract``, ``dataContractLatest``, and + ``dataContractLatestCertified`` on it, then hard-deletes the contract. + + :param client: connectivity to an Atlan tenant + :param contract_guid: GUID of the DataContract to delete + :returns: tuple of (contract delete response, asset update response) + """ + from pyatlan.model.assets.core.asset import Asset + from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset + + contract = client.asset.get_by_guid( + contract_guid, + asset_type=DataContract, + attributes=[DataContract.DATA_CONTRACT_ASSET_LATEST], + related_attributes=[Asset.NAME, Asset.QUALIFIED_NAME, Asset.TYPE_NAME], + ) + linked_asset = contract.data_contract_asset_latest + if not linked_asset or not linked_asset.guid: + raise ValueError( + "Cannot determine the linked asset for this contract. " + "Ensure the contract has a valid data_contract_asset_latest relationship." + ) + + delete_response = client.asset.purge_by_guid(contract_guid) + asset_update = IndistinctAsset() + asset_update.type_name = linked_asset.type_name + asset_update.guid = linked_asset.guid + asset_update.qualified_name = linked_asset.qualified_name + asset_update.name = linked_asset.name + asset_update.has_contract = False + asset_update.data_contract_latest = None # type: ignore[assignment] + asset_update.data_contract_latest_certified = None # type: ignore[assignment] + asset_response = client.asset.save(asset_update) + return delete_response, asset_response + type_name: str = Field(default="DataContract", allow_mutation=False) @validator("type_name") @@ -261,6 +336,7 @@ def creator( cls, *, asset_qualified_name: str, + asset_type: Type[Asset], contract_json: Optional[str] = None, contract_spec: Optional[Union[DataContractSpec, str]] = None, ) -> DataContract.Attributes: @@ -309,6 +385,9 @@ def creator( qualified_name=f"{asset_qualified_name}/contract", data_contract_json=contract_json, data_contract_spec=contract_spec, # type: ignore[arg-type] + data_contract_asset_latest=asset_type.ref_by_qualified_name( + asset_qualified_name + ), ) attributes: DataContract.Attributes = Field( diff --git a/pyatlan_v9/model/assets/asset.py b/pyatlan_v9/model/assets/asset.py index aeb1fbe1b..5ce7a7690 100644 --- a/pyatlan_v9/model/assets/asset.py +++ b/pyatlan_v9/model/assets/asset.py @@ -31,6 +31,7 @@ from .anomalo_related import RelatedAnomaloCheck from .app_related import RelatedApplication, RelatedApplicationField from .asset_related import RelatedAsset +from .catalog_related import RelatedCatalog from .data_mesh_related import RelatedDataProduct from .data_quality_related import RelatedDataQualityRule, RelatedMetric from .gtc_related import RelatedAtlasGlossaryTerm @@ -990,6 +991,12 @@ class Asset(Referenceable): soda_checks: Union[List[RelatedSodaCheck], None, UnsetType] = UNSET """""" + data_contract_latest: Union[RelatedCatalog, None, UnsetType] = UNSET + """Latest data contract for this asset.""" + + data_contract_latest_certified: Union[RelatedCatalog, None, UnsetType] = UNSET + """Latest certified data contract for this asset.""" + def __post_init__(self) -> None: if self.type_name is UNSET: self.type_name = "Asset" @@ -2043,6 +2050,12 @@ class AssetRelationshipAttributes(ReferenceableRelationshipAttributes): soda_checks: Union[List[RelatedSodaCheck], None, UnsetType] = UNSET """""" + data_contract_latest: Union[RelatedCatalog, None, UnsetType] = UNSET + """Latest data contract for this asset.""" + + data_contract_latest_certified: Union[RelatedCatalog, None, UnsetType] = UNSET + """Latest certified data contract for this asset.""" + class AssetNested(ReferenceableNested): """Asset in nested API format for high-performance serialization.""" @@ -2081,6 +2094,8 @@ class AssetNested(ReferenceableNested): "readme", "schema_registry_subjects", "soda_checks", + "data_contract_latest", + "data_contract_latest_certified", ] diff --git a/pyatlan_v9/model/assets/data_contract.py b/pyatlan_v9/model/assets/data_contract.py index 7230be37e..2a465d9b5 100644 --- a/pyatlan_v9/model/assets/data_contract.py +++ b/pyatlan_v9/model/assets/data_contract.py @@ -7,7 +7,7 @@ import re from json import JSONDecodeError, loads -from typing import Union +from typing import TYPE_CHECKING, Tuple, Type, Union from msgspec import UNSET, UnsetType @@ -32,6 +32,10 @@ ) from .catalog_related import RelatedCatalog +if TYPE_CHECKING: + from pyatlan_v9.client.atlan import AtlanClient + from pyatlan_v9.model.response import AssetMutationResponse + @register_asset class DataContract(Catalog): @@ -63,28 +67,118 @@ class DataContract(Catalog): data_contract_previous_version: Union[RelatedCatalog, None, UnsetType] = UNSET """Previous version in this contract chain.""" + def __post_init__(self) -> None: + self.type_name = "DataContract" + @classmethod @init_guid def creator( cls, *, asset_qualified_name: str, + asset_type: Type["Asset"], contract_json: Union[str, None] = None, contract_spec: Union[DataContractSpec, str, None] = None, ) -> "DataContract": - """Create a new DataContract asset.""" + """Create a new DataContract asset. + + :param asset_qualified_name: qualified name of the asset this contract governs + :param asset_type: the type of the governed asset (e.g. ``Table``, ``View``) + :param contract_json: deprecated JSON representation of the contract + :param contract_spec: YAML contract spec (string or ``DataContractSpec``) + """ attrs = DataContract.Attributes.creator( asset_qualified_name=asset_qualified_name, contract_json=contract_json, contract_spec=contract_spec, ) + asset_ref = RelatedAsset(qualified_name=asset_qualified_name) + asset_ref.type_name = asset_type.__name__ return cls( name=attrs.name, qualified_name=attrs.qualified_name, data_contract_json=attrs.data_contract_json, data_contract_spec=attrs.data_contract_spec, + data_contract_asset_latest=asset_ref, ) + @staticmethod + def save( + client: "AtlanClient", + contract: "DataContract", + ) -> "AssetMutationResponse": + """Save a DataContract. + + The contract's ``data_contract_asset_latest`` relationship (set by + ``creator()``) links the contract to the governed asset automatically. + + :param client: connectivity to an Atlan tenant + :param contract: DataContract to save (from ``DataContract.creator()``) + :returns: the result of the save + """ + return client.asset.save(contract) + + @staticmethod + def delete( + client: "AtlanClient", + contract_guid: str, + ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": + """Delete (purge) a DataContract and clean up the linked asset. + + Retrieves the contract to find the linked asset, clears + ``hasContract``, ``dataContractLatest``, and + ``dataContractLatestCertified`` on it, then hard-deletes the contract. + + :param client: connectivity to an Atlan tenant + :param contract_guid: GUID of the DataContract to delete + :returns: tuple of (contract delete response, asset update response) + """ + from pyatlan.client.constants import BULK_UPDATE + from pyatlan_v9.client.asset import _parse_mutation_response + from pyatlan_v9.model.assets.asset import Asset + from pyatlan_v9.model.assets.referenceable import Referenceable + + contract = client.asset.get_by_guid( + contract_guid, + asset_type=DataContract, + attributes=["dataContractAssetLatest"], + related_attributes=[ + Asset.NAME, + Referenceable.QUALIFIED_NAME, + Referenceable.TYPE_NAME, + ], + ) + linked_asset = contract.data_contract_asset_latest + if not linked_asset or not linked_asset.guid: + raise ValueError( + "Cannot determine the linked asset for this contract. " + "Ensure the contract has a valid data_contract_asset_latest relationship." + ) + + delete_response = client.asset.purge_by_guid(contract_guid) + + # Build raw payload because v9 Asset doesn't model these fields + asset_payload = { + "entities": [ + { + "guid": linked_asset.guid, + "typeName": linked_asset.type_name, + "attributes": { + "qualifiedName": linked_asset.qualified_name, + "name": linked_asset.name, + "hasContract": False, + }, + "relationshipAttributes": { + "dataContractLatest": None, + "dataContractLatestCertified": None, + }, + } + ] + } + raw_json = client._call_api(BULK_UPDATE, {}, asset_payload) + asset_response = _parse_mutation_response(raw_json) + return delete_response, asset_response + @classmethod def updater(cls, *, qualified_name: str, name: str) -> "DataContract": """Create a DataContract instance for update operations.""" diff --git a/pyatlan_v9/model/assets/entity.py b/pyatlan_v9/model/assets/entity.py index d7c5ba58f..201a44f1b 100644 --- a/pyatlan_v9/model/assets/entity.py +++ b/pyatlan_v9/model/assets/entity.py @@ -123,6 +123,9 @@ class Entity(msgspec.Struct, kw_only=True, omit_defaults=True, rename="camel"): status: Union[str, UnsetType] = UNSET """Entity status (ACTIVE, DELETED, PURGED).""" + delete_handler: Union[str, UnsetType] = UNSET + """Handler used for deletion of this entity (e.g. HARD, SOFT, PURGE).""" + version: Union[int, UnsetType] = UNSET """Version number of this entity.""" diff --git a/pyatlan_v9/model/assets/referenceable.py b/pyatlan_v9/model/assets/referenceable.py index 9949a9ed6..2413f5ded 100644 --- a/pyatlan_v9/model/assets/referenceable.py +++ b/pyatlan_v9/model/assets/referenceable.py @@ -290,6 +290,7 @@ class ReferenceableNested( guid: Union[Any, UnsetType] = UNSET type_name: Union[Any, UnsetType] = UNSET status: Union[Any, UnsetType] = UNSET + delete_handler: Union[Any, UnsetType] = UNSET version: Union[Any, UnsetType] = UNSET create_time: Union[Any, UnsetType] = UNSET update_time: Union[Any, UnsetType] = UNSET @@ -368,6 +369,7 @@ def _referenceable_to_nested(referenceable: Referenceable) -> ReferenceableNeste guid=referenceable.guid, type_name=referenceable.type_name, status=referenceable.status, + delete_handler=referenceable.delete_handler, version=referenceable.version, create_time=referenceable.create_time, update_time=referenceable.update_time, @@ -413,6 +415,7 @@ def _referenceable_from_nested(nested: ReferenceableNested) -> Referenceable: guid=nested.guid, type_name=nested.type_name, status=nested.status, + delete_handler=nested.delete_handler, version=nested.version, create_time=nested.create_time, update_time=nested.update_time, diff --git a/tests/integration/data_mesh_test.py b/tests/integration/data_mesh_test.py index d79f1c08b..d615503f9 100644 --- a/tests/integration/data_mesh_test.py +++ b/tests/integration/data_mesh_test.py @@ -305,12 +305,13 @@ def contract( } contract = DataContract.creator( asset_qualified_name=table.qualified_name, + asset_type=Table, contract_json=dumps(contract_json), ) - response = client.asset.save(contract) - result = response.assets_created(asset_type=DataContract)[0] + contract_response = DataContract.save(client=client, contract=contract) + result = contract_response.assets_created(asset_type=DataContract)[0] yield result - delete_asset(client, guid=result.guid, asset_type=DataContract) + # Cleanup is handled by updated_contract fixture since it updates the same entity @pytest.fixture(scope="module") @@ -330,17 +331,29 @@ def updated_contract( } contract = DataContract.creator( asset_qualified_name=table.qualified_name, + asset_type=Table, contract_json=dumps(contract_json), ) response = client.asset.save(contract) result = response.assets_updated(asset_type=DataContract)[0] yield result - delete_asset(client, guid=result.guid, asset_type=DataContract) + delete_response, asset_response = DataContract.delete( + client=client, contract_guid=result.guid + ) + assert delete_response + assert asset_response + deleted = delete_response.assets_deleted(asset_type=DataContract) + assert deleted and len(deleted) == 1 + assert deleted[0].guid == result.guid def test_contract( client: AtlanClient, table: Table, product: DataProduct, contract: DataContract ): + """Verify DataContract.save() sets dataContractLatest and hasContract on the + linked asset — the behaviour that was broken when using client.asset.save() + alone (Bug 1 / BLDX-721). + """ assert product and product.guid product = client.asset.get_by_guid( guid=product.guid, asset_type=DataProduct, ignore_relationships=False @@ -504,21 +517,6 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): assert isinstance(asset_list, IndexSearchResults) -@pytest.mark.order(after="test_retrieve_contract") -def test_delete_contract(client: AtlanClient, contract: DataContract): - response = client.asset.purge_by_guid(contract.guid) - assert response - assert not response.assets_created(asset_type=DataContract) - assert not response.assets_updated(asset_type=DataContract) - deleted = response.assets_deleted(asset_type=DataContract) - assert deleted - assert len(deleted) == 1 - assert deleted[0].guid == contract.guid - assert deleted[0].qualified_name == contract.qualified_name - assert deleted[0].delete_handler == "PURGE" - assert deleted[0].status == EntityStatus.DELETED - - @pytest.mark.order(after="test_retrieve_product") def test_delete_product(client: AtlanClient, product: DataProduct): response = client.asset.purge_by_guid(product.guid) diff --git a/tests/unit/model/data_contract_test.py b/tests/unit/model/data_contract_test.py index d9888eba2..6d56ac412 100644 --- a/tests/unit/model/data_contract_test.py +++ b/tests/unit/model/data_contract_test.py @@ -1,11 +1,13 @@ from json import dumps from typing import Union +from unittest.mock import MagicMock import pytest from pyatlan.errors import InvalidRequestError -from pyatlan.model.assets import DataContract +from pyatlan.model.assets import DataContract, Table from pyatlan.model.contract import DataContractSpec +from pyatlan.model.response import AssetMutationResponse from tests.unit.model.constants import ( ASSET_QUALIFIED_NAME, DATA_CONTRACT_JSON, @@ -42,6 +44,7 @@ def test_creator_with_missing_parameters_raise_value_error( with pytest.raises(ValueError, match=message): DataContract.creator( # type: ignore asset_qualified_name=asset_qualified_name, + asset_type=Table, contract_json=contract_json, contract_spec=contract_spec, ) @@ -68,6 +71,7 @@ def test_creator_with_invalid_contract_json_raises_error( with pytest.raises(InvalidRequestError, match=error_msg): DataContract.creator( asset_qualified_name=asset_qualified_name, + asset_type=Table, contract_json=contract_json, ) @@ -75,6 +79,7 @@ def test_creator_with_invalid_contract_json_raises_error( def test_creator_atttributes_with_required_parameters(): attributes = DataContract.Attributes.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_json=dumps(DATA_CONTRACT_JSON), ) _assert_contract(attributes, is_json=True) @@ -83,6 +88,7 @@ def test_creator_atttributes_with_required_parameters(): def test_creator_with_required_parameters_json(): test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_json=dumps(DATA_CONTRACT_JSON), ) _assert_contract(test_contract) @@ -91,6 +97,7 @@ def test_creator_with_required_parameters_json(): def test_creator_with_required_parameters_spec_str(): test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=DATA_CONTRACT_SPEC_STR, ) _assert_contract(test_contract) @@ -99,6 +106,7 @@ def test_creator_with_required_parameters_spec_str(): def test_creator_with_required_parameters_spec_str_without_dataset(): test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=DATA_CONTRACT_SPEC_STR_WITHOUT_DATASET, ) # Ensure the default contract name is extracted from the table's qualified name (QN). @@ -109,11 +117,24 @@ def test_creator_with_required_parameters_spec_model(): spec = DataContractSpec.from_yaml(DATA_CONTRACT_SPEC_STR) test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=spec, ) _assert_contract(test_contract) +def test_creator_sets_data_contract_asset_latest(): + """Creator should set data_contract_asset_latest with the correct asset ref.""" + test_contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + ref = test_contract.data_contract_asset_latest + assert ref is not None + assert ref.unique_attributes == {"qualifiedName": ASSET_QUALIFIED_NAME} + + def test_create_for_modification(): test_contract = DataContract.create_for_modification( name=DATA_CONTRACT_NAME, @@ -128,3 +149,76 @@ def test_trim_to_required(): qualified_name=DATA_CONTRACT_QUALIFIED_NAME, ).trim_to_required() _assert_contract(test_contract, False) + + +class TestSaveContract: + def test_save_delegates_to_asset_client(self): + """save() should delegate to client.asset.save().""" + mock_client = MagicMock() + contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + expected_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.save.return_value = expected_response + + result = DataContract.save(client=mock_client, contract=contract) + + assert result == expected_response + mock_client.asset.save.assert_called_once_with(contract) + + +class TestDeleteContract: + def test_delete_retrieves_contract_and_clears_linked_asset(self): + """delete() should retrieve the contract, find the linked asset, clear state, and purge.""" + mock_client = MagicMock() + + # Mock the contract retrieved by get_by_guid + mock_linked_asset = MagicMock() + mock_linked_asset.type_name = "Table" + mock_linked_asset.guid = "asset-guid-456" + mock_linked_asset.qualified_name = "default/test/table-qn" + mock_linked_asset.name = "test-table" + mock_contract = MagicMock(spec=DataContract) + mock_contract.data_contract_asset_latest = mock_linked_asset + mock_client.asset.get_by_guid.return_value = mock_contract + + delete_response = MagicMock(spec=AssetMutationResponse) + asset_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.purge_by_guid.return_value = delete_response + mock_client.asset.save.return_value = asset_response + + result = DataContract.delete( + client=mock_client, + contract_guid="contract-guid-123", + ) + + assert result == (delete_response, asset_response) + from pyatlan.model.assets import Asset + + mock_client.asset.get_by_guid.assert_called_once_with( + "contract-guid-123", + asset_type=DataContract, + attributes=[DataContract.DATA_CONTRACT_ASSET_LATEST], + related_attributes=[Asset.NAME, Asset.QUALIFIED_NAME, Asset.TYPE_NAME], + ) + mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") + asset_update = mock_client.asset.save.call_args[0][0] + assert asset_update.guid == "asset-guid-456" + assert asset_update.has_contract is False + assert asset_update.data_contract_latest is None + assert asset_update.data_contract_latest_certified is None + + def test_delete_raises_when_no_linked_asset(self): + """delete() should raise ValueError when contract has no linked asset.""" + mock_client = MagicMock() + mock_contract = MagicMock(spec=DataContract) + mock_contract.data_contract_asset_latest = None + mock_client.asset.get_by_guid.return_value = mock_contract + + with pytest.raises(ValueError, match="Cannot determine the linked asset"): + DataContract.delete( + client=mock_client, + contract_guid="contract-guid-123", + ) diff --git a/tests_v9/integration/data_mesh_test.py b/tests_v9/integration/data_mesh_test.py index 8f0035881..dd9970517 100644 --- a/tests_v9/integration/data_mesh_test.py +++ b/tests_v9/integration/data_mesh_test.py @@ -308,12 +308,13 @@ def contract( } contract = DataContract.creator( asset_qualified_name=table.qualified_name, + asset_type=Table, contract_json=dumps(contract_json), ) - response = client.asset.save(contract) - result = response.assets_created(asset_type=DataContract)[0] + contract_response = DataContract.save(client=client, contract=contract) + result = contract_response.assets_created(asset_type=DataContract)[0] yield result - delete_asset(client, guid=result.guid, asset_type=DataContract) + # Cleanup is handled by updated_contract fixture since it updates the same entity @pytest.fixture(scope="module") @@ -333,12 +334,20 @@ def updated_contract( } contract = DataContract.creator( asset_qualified_name=table.qualified_name, + asset_type=Table, contract_json=dumps(contract_json), ) response = client.asset.save(contract) result = response.assets_updated(asset_type=DataContract)[0] yield result - delete_asset(client, guid=result.guid, asset_type=DataContract) + delete_response, asset_response = DataContract.delete( + client=client, contract_guid=result.guid + ) + assert delete_response + assert asset_response + deleted = delete_response.assets_deleted(asset_type=DataContract) + assert deleted and len(deleted) == 1 + assert deleted[0].guid == result.guid def test_contract( @@ -446,7 +455,7 @@ def test_update_product( assert len(products) == 1 assert products[ 0 - ].data_product_assets_d_s_l == DataProductsAssetsDSL.get_asset_selection(assets) + ].data_product_assets_dsl == DataProductsAssetsDSL.get_asset_selection(assets) # Test the product.updater() method without assets # (ensure asset selection remains unchanged) @@ -505,7 +514,7 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): product.guid, asset_type=DataProduct, ignore_relationships=False ) assert test_product - assert test_product.data_product_assets_d_s_l + assert test_product.data_product_assets_dsl asset_list = test_product.get_assets(client=client) assert asset_list.count and asset_list.count > 0 TOTAL_ASSETS = asset_list.count @@ -517,21 +526,6 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): assert isinstance(asset_list, IndexSearchResults) -@pytest.mark.order(after="test_retrieve_contract") -def test_delete_contract(client: AtlanClient, contract: DataContract): - response = client.asset.purge_by_guid(contract.guid) - assert response - assert not response.assets_created(asset_type=DataContract) - assert not response.assets_updated(asset_type=DataContract) - deleted = response.assets_deleted(asset_type=DataContract) - assert deleted - assert len(deleted) == 1 - assert deleted[0].guid == contract.guid - assert deleted[0].qualified_name == contract.qualified_name - assert deleted[0].delete_handler == "PURGE" - assert deleted[0].status == EntityStatus.DELETED - - @pytest.mark.order(after="test_retrieve_product") def test_delete_product(client: AtlanClient, product: DataProduct): response = client.asset.purge_by_guid(product.guid) diff --git a/tests_v9/unit/model/data_contract_test.py b/tests_v9/unit/model/data_contract_test.py index 47929177a..ed9024c43 100644 --- a/tests_v9/unit/model/data_contract_test.py +++ b/tests_v9/unit/model/data_contract_test.py @@ -5,12 +5,14 @@ from json import dumps from typing import Union +from unittest.mock import MagicMock import pytest from pyatlan_v9.errors import InvalidRequestError -from pyatlan_v9.model import DataContract +from pyatlan_v9.model import DataContract, Table from pyatlan_v9.model.contract import DataContractSpec +from pyatlan_v9.model.response import AssetMutationResponse from tests_v9.unit.model.constants import ( ASSET_QUALIFIED_NAME, DATA_CONTRACT_JSON, @@ -48,6 +50,7 @@ def test_creator_with_missing_parameters_raise_value_error( with pytest.raises(ValueError, match=message): DataContract.creator( # type: ignore[arg-type] asset_qualified_name=asset_qualified_name, + asset_type=Table, contract_json=contract_json, contract_spec=contract_spec, ) @@ -75,6 +78,7 @@ def test_creator_with_invalid_contract_json_raises_error( with pytest.raises(InvalidRequestError, match=error_msg): DataContract.creator( asset_qualified_name=asset_qualified_name, + asset_type=Table, contract_json=contract_json, ) @@ -92,6 +96,7 @@ def test_creator_with_required_parameters_json(): """Test DataContract.creator for JSON payload.""" test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_json=dumps(DATA_CONTRACT_JSON), ) _assert_contract(test_contract) @@ -101,6 +106,7 @@ def test_creator_with_required_parameters_spec_str(): """Test DataContract.creator for YAML string payload.""" test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=DATA_CONTRACT_SPEC_STR, ) _assert_contract(test_contract) @@ -110,6 +116,7 @@ def test_creator_with_required_parameters_spec_str_without_dataset(): """Test creator defaults contract name from asset QN when dataset is absent.""" test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=DATA_CONTRACT_SPEC_STR_WITHOUT_DATASET, ) _assert_contract(test_contract, contract_name=DATA_CONTRACT_NAME_DEFAULT) @@ -120,11 +127,25 @@ def test_creator_with_required_parameters_spec_model(): spec = DataContractSpec.from_yaml(DATA_CONTRACT_SPEC_STR) test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=spec, ) _assert_contract(test_contract) +def test_creator_sets_data_contract_asset_latest(): + """Creator should set data_contract_asset_latest with the correct asset ref.""" + test_contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + ref = test_contract.data_contract_asset_latest + assert ref is not None + assert ref.type_name == "Table" + assert ref.unique_attributes == {"qualifiedName": ASSET_QUALIFIED_NAME} + + @pytest.mark.parametrize( "qualified_name, name, message", [ @@ -156,3 +177,86 @@ def test_trim_to_required(): qualified_name=DATA_CONTRACT_QUALIFIED_NAME, ).trim_to_required() _assert_contract(test_contract, False) + + +class TestSaveContract: + def test_save_delegates_to_asset_client(self): + """save() should delegate to client.asset.save().""" + mock_client = MagicMock() + contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + expected_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.save.return_value = expected_response + + result = DataContract.save(client=mock_client, contract=contract) + + assert result == expected_response + mock_client.asset.save.assert_called_once_with(contract) + + +class TestDeleteContract: + def test_delete_retrieves_contract_and_clears_linked_asset(self): + """delete() should retrieve the contract, find the linked asset, clear state, and purge.""" + mock_client = MagicMock() + + # Mock the contract retrieved by get_by_guid + mock_linked_asset = MagicMock() + mock_linked_asset.type_name = "Table" + mock_linked_asset.guid = "asset-guid-456" + mock_linked_asset.qualified_name = "default/test/table-qn" + mock_linked_asset.name = "test-table" + mock_contract = MagicMock(spec=DataContract) + mock_contract.data_contract_asset_latest = mock_linked_asset + mock_client.asset.get_by_guid.return_value = mock_contract + + delete_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.purge_by_guid.return_value = delete_response + mock_client._call_api.return_value = {"mutatedEntities": {}} + + result = DataContract.delete( + client=mock_client, + contract_guid="contract-guid-123", + ) + + from pyatlan_v9.model.assets.asset import Asset + from pyatlan_v9.model.assets.referenceable import Referenceable + + assert result[0] == delete_response + mock_client.asset.get_by_guid.assert_called_once_with( + "contract-guid-123", + asset_type=DataContract, + attributes=["dataContractAssetLatest"], + related_attributes=[ + Asset.NAME, + Referenceable.QUALIFIED_NAME, + Referenceable.TYPE_NAME, + ], + ) + mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") + # Verify the raw API call clears contract state + mock_client._call_api.assert_called_once() + payload = mock_client._call_api.call_args[0][2] + entity = payload["entities"][0] + assert entity["guid"] == "asset-guid-456" + assert entity["typeName"] == "Table" + assert entity["attributes"]["qualifiedName"] == "default/test/table-qn" + assert entity["attributes"]["name"] == "test-table" + assert entity["attributes"]["hasContract"] is False + assert entity["relationshipAttributes"]["dataContractLatest"] is None + assert entity["relationshipAttributes"]["dataContractLatestCertified"] is None + + def test_delete_raises_when_no_linked_asset(self): + """delete() should raise ValueError when contract has no linked asset.""" + mock_client = MagicMock() + mock_contract = MagicMock(spec=DataContract) + mock_contract.data_contract_asset_latest = None + mock_client.asset.get_by_guid.return_value = mock_contract + + with pytest.raises(ValueError, match="Cannot determine the linked asset"): + DataContract.delete( + client=mock_client, + contract_guid="contract-guid-123", + )