diff --git a/lib/datahub-client/data_platform_catalogue/entities.py b/lib/datahub-client/data_platform_catalogue/entities.py index d2c5f6f9..69cee134 100644 --- a/lib/datahub-client/data_platform_catalogue/entities.py +++ b/lib/datahub-client/data_platform_catalogue/entities.py @@ -355,9 +355,12 @@ class CustomEntityProperties(BaseModel): ) audience: Audience = Field( description="If the data is published or not", - default=Audience.INTERNAL, + default="Internal", ) + class Config: + use_enum_values = True + class Entity(BaseModel): """ diff --git a/lib/datahub-client/tests/client/datahub/test_datahub_client.py b/lib/datahub-client/tests/client/datahub/test_datahub_client.py index 2faadbc6..b1a0d051 100644 --- a/lib/datahub-client/tests/client/datahub/test_datahub_client.py +++ b/lib/datahub-client/tests/client/datahub/test_datahub_client.py @@ -13,6 +13,7 @@ ReferencedEntityMissing, ) from data_platform_catalogue.entities import ( + Audience, AccessInformation, Chart, Column, @@ -331,7 +332,6 @@ def test_get_dataset( }, "lastIngested": 1709619407814, "domain": None, - "audience": "Internal", "provider": "LAA", "schemaMetadata": { "fields": [ @@ -376,6 +376,7 @@ def test_get_dataset( fully_qualified_name="Foo.Dataset", description="Dataset", relationships={ + RelationshipType.DATA_LINEAGE: [], RelationshipType.PARENT: [ EntitySummary( entity_ref=EntityRef( @@ -391,7 +392,6 @@ def test_get_dataset( entity_type="Database", ) ], - RelationshipType.DATA_LINEAGE: [], }, domain=DomainRef(display_name="", urn=""), governance=Governance( @@ -400,7 +400,6 @@ def test_get_dataset( ), tags=[TagRef(display_name="some-tag", urn="urn:li:tag:Entity")], last_modified=1709619407814, - audience="Internal", provider="LAA", created=None, platform=EntityRef(urn="datahub", display_name="datahub"), @@ -483,6 +482,7 @@ def test_get_dataset_minimal_properties( ), data_summary=DataSummary(), further_information=FurtherInformation(), + audience=Audience.INTERNAL, ), column_details=[], ) @@ -542,6 +542,7 @@ def test_get_chart_details(self, datahub_client, base_mock_graph): ), data_summary=DataSummary(), further_information=FurtherInformation(), + audience=Audience.INTERNAL, ), external_url="https://data.justice.gov.uk/prisons/public-protection/absconds", ) diff --git a/lib/datahub-client/tests/client/datahub/test_graphql_helpers.py b/lib/datahub-client/tests/client/datahub/test_graphql_helpers.py index d279611b..96c09293 100644 --- a/lib/datahub-client/tests/client/datahub/test_graphql_helpers.py +++ b/lib/datahub-client/tests/client/datahub/test_graphql_helpers.py @@ -19,6 +19,7 @@ parse_updated, ) from data_platform_catalogue.entities import ( + Audience, AccessInformation, Column, ColumnRef, @@ -285,6 +286,7 @@ def test_parse_properties(): {"key": "s3_location", "value": "s3://databucket/"}, {"key": "row_count", "value": 100}, {"key": "Not_IN", "value": "dddd"}, + {"key": "audience", "value": "Internal"}, ], "name": "test", "description": "test description", @@ -313,6 +315,7 @@ def test_parse_properties(): further_information=FurtherInformation( dc_slack_channel_name="test-channel", dc_slack_channel_url="test-url" ), + audience=Audience.INTERNAL, ) @@ -328,6 +331,7 @@ def test_parse_properties_with_none_values(): {"key": "s3_location", "value": "s3://databucket/"}, {"key": "row_count", "value": 100}, {"key": "Not_IN", "value": "dddd"}, + {"key": "audience", "value": "Internal"}, ], "name": "test", "description": None, @@ -356,6 +360,7 @@ def test_parse_properties_with_none_values(): ), data_summary=DataSummary(row_count=100), further_information=FurtherInformation(), + audience=Audience.INTERNAL, ) diff --git a/lib/datahub-client/tests/client/datahub/test_search.py b/lib/datahub-client/tests/client/datahub/test_search.py index 36106962..93b31dbe 100644 --- a/lib/datahub-client/tests/client/datahub/test_search.py +++ b/lib/datahub-client/tests/client/datahub/test_search.py @@ -1179,7 +1179,7 @@ def test_search_for_container(mock_graph, searcher): "name": "test_db", }, metadata={ - "audience": Audience.INTERNAL, + "audience": "Internal", "owner": "Shannon Lovett", "owner_email": "shannon@longtail.com", "domain_name": "testdom",