Filter fixes and test fixes

ministryofjustice · Dec 3, 2024 · 292a735 · 292a735
1 parent 4ba8eff
commit 292a735
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 16 deletions.
diff --git a/lib/datahub-client/data_platform_catalogue/client/search.py b/lib/datahub-client/data_platform_catalogue/client/search.py
@@ -299,12 +299,6 @@ def _map_filters(
             {
                 "and": [
                     {"field": filter.filter_name, "values": filter.included_values},
-                    {
-                        "field": "_entityType",
-                        "values": self.fmd_type_to_datahub_types_mapping[
-                            filter.included_values[0]
-                        ][0],
-                    },
                     {"field": "tags", "values": ["urn:li:tag:dc_display_in_catalogue"]},
                 ]
             }
@@ -386,7 +380,7 @@ def _parse_dataset(
             "total_parents": entity.get("relationships", {}).get("total", 0),
             "domain_name": domain.display_name,
             "domain_id": domain.urn,
-            "entity_types": self._parse_types_and_sub_types(entity, "Dataset"),
+            "entity_types": self._parse_types_and_sub_types(entity, result_type.value),
         }
         logger.debug(f"{metadata=}")
 
@@ -526,7 +520,7 @@ def _parse_container(
             "owner_email": owner.email,
             "domain_name": domain.display_name,
             "domain_id": domain.urn,
-            "entity_types": self._parse_types_and_sub_types(entity, "Container"),
+            "entity_types": self._parse_types_and_sub_types(entity, subtype.value),
         }
 
         metadata.update(custom_properties)

diff --git a/lib/datahub-client/tests/client/datahub/test_search.py b/lib/datahub-client/tests/client/datahub/test_search.py
@@ -139,7 +139,7 @@ def test_one_search_result(mock_graph, searcher):
                     "domain_name": "HMPPS",
                     "domain_id": "urn:li:domain:3dc18e48-c062-4407-84a9-73e23f768023",
                     "entity_types": {
-                        "entity_type": "DATASET",
+                        "entity_type": "Table",
                         "entity_sub_types": ["Table"],
                     },
                     "dpia_required": None,
@@ -234,7 +234,7 @@ def test_dataset_result(mock_graph, searcher):
                     "domain_name": "HMPPS",
                     "domain_id": "urn:li:domain:3dc18e48-c062-4407-84a9-73e23f768023",
                     "entity_types": {
-                        "entity_type": "DATASET",
+                        "entity_type": "Table",
                         "entity_sub_types": ["Table"],
                     },
                     "dpia_required": None,
@@ -610,7 +610,7 @@ def test_query_match(mock_graph, searcher):
                         "type": "DATASET",
                         "subTypes": {
                             "typeNames": [
-                                "Publication dataset"
+                                "Table"
                             ]
                         },
                         "urn": "urn:li:dataset:(urn:li:dataPlatform:bigquery,calm-pagoda-323403.jaffle_shop.customers,PROD)",  # noqa E501
@@ -649,7 +649,7 @@ def test_query_match(mock_graph, searcher):
                     "domain_name": "",
                     "domain_id": "",
                     "entity_types": {
-                        "entity_type": "Dataset",
+                        "entity_type": "Table",
                         "entity_sub_types": ["Table"],
                     },
                     "dpia_required": None,
@@ -683,7 +683,7 @@ def test_result_with_owner(mock_graph, searcher):
                         "type": "DATASET",
                         "subTypes": {
                             "typeNames": [
-                                "Publication dataset"
+                                "Table"
                             ]
                         },
                         "urn": "urn:li:dataset:(urn:li:dataPlatform:bigquery,calm-pagoda-323403.jaffle_shop.customers,PROD)",  # noqa E501
@@ -734,7 +734,7 @@ def test_result_with_owner(mock_graph, searcher):
                     "domain_name": "",
                     "domain_id": "",
                     "entity_types": {
-                        "entity_type": "Dataset",
+                        "entity_type": "Table",
                         "entity_sub_types": ["Table"],
                     },
                     "dpia_required": None,
@@ -1036,8 +1036,8 @@ def test_search_for_charts(mock_graph, searcher):
                     "domain_name": "",
                     "domain_id": "",
                     "entity_types": {
-                        "entity_type": "CHART",
-                        "entity_sub_types": ["CHART"],
+                        "entity_type": "Chart",
+                        "entity_sub_types": ["Chart"],
                     },
                     "dpia_required": None,
                     "dpia_location": "",