Use strings for Enum types in CustomEntityProperties

ministryofjustice · Nov 15, 2024 · 740e953 · 740e953
1 parent 8fccaca
commit 740e953
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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": "[email protected]",
                     "domain_name": "testdom",