fmd-574 - Add extra test cases

lib/datahub-client/tests/client/datahub/test_datahub_client.py

@@ -8,7 +8,7 @@
     DataHubCatalogueClient,
     InvalidDomain,
 )
-from data_platform_catalogue.client.exceptions import ReferencedEntityMissing
+from data_platform_catalogue.client.exceptions import ReferencedEntityMissing, EntityDoesNotExist
 from data_platform_catalogue.entities import (
     AccessInformation,
     Chart,
@@ -697,3 +697,33 @@ def test_database_not_exist_given_error(self, datahub_client, table, database):
 
     def test_get_custom_property_key_value_pairs(self, datahub_client, database):
         datahub_client._get_custom_property_key_value_pairs(database.custom_properties)
+
+    def test_get_dataset_invalid_urn(self, datahub_client, base_mock_graph):
+        urn = "invalid_urn"
+        datahub_response = {"dataset": None}
+        base_mock_graph.execute_graphql = MagicMock(return_value=datahub_response)
+
+        with patch(
+                "data_platform_catalogue.client.datahub_client.DataHubCatalogueClient.check_entity_exists_by_urn") as mock_exists:
+            mock_exists.return_value = False  # Entity does not exist
+            with pytest.raises(EntityDoesNotExist):
+                datahub_client.get_table_details(urn)
+
+    def test_get_dataset_missing_properties(self, datahub_client, base_mock_graph):
+        urn = "urn:li:dataset:missing_props"
+        datahub_response = {
+            "dataset": {
+                "platform": {"name": "datahub"},
+                "name": "Dataset",
+                "properties": {}
+            }
+        }
+        base_mock_graph.execute_graphql = MagicMock(return_value=datahub_response)
+
+        with patch(
+                "data_platform_catalogue.client.datahub_client.DataHubCatalogueClient.check_entity_exists_by_urn") as mock_exists:
+            mock_exists.return_value = True
+            dataset = datahub_client.get_table_details(urn)
+
+        assert dataset.name == "Dataset"
+        assert dataset.description == ""

lib/datahub-client/tests/client/datahub/test_graphql_helpers.py

@@ -351,6 +351,114 @@ def test_parse_properties_with_none_values():
     )
 
 
+def test_parse_columns_with_empty_fields():
+    entity = {
+        "schemaMetadata": {
+            "fields": [],
+            "primaryKeys": [],
+            "foreignKeys": [],
+        }
+    }
+
+    assert parse_columns(entity) == []
+
+
+def test_parse_columns_with_unrecognized_type():
+    entity = {
+        "schemaMetadata": {
+            "fields": [
+                {
+                    "fieldPath": "unknownField",
+                    "label": None,
+                    "nullable": True,
+                    "description": "An unknown field type",
+                    "type": "UNKNOWN_TYPE",
+                    "nativeDataType": "unknown",
+                }
+            ],
+            "primaryKeys": [],
+            "foreignKeys": [],
+        }
+    }
+
+    assert parse_columns(entity) == [
+        Column(
+            name="unknownField",
+            display_name="unknownField",
+            type="unknown",
+            description="An unknown field type",
+            nullable=True,
+            is_primary_key=False,
+            foreign_keys=[],
+        )
+    ]
+
+
+def test_parse_relations_multiple_relationships():
+    relations = {
+        "relationships": {
+            "total": 2,
+            "relationships": [
+                {
+                    "entity": {
+                        "urn": "urn:li:dataProduct:test1",
+                        "type": "DATA_PRODUCT",
+                        "properties": {"name": "test1", "description": "first test entity"},
+                        "tags": {
+                            "tags": [
+                                {"tag": {"urn": "urn:li:tag:dc_display_in_catalogue"}}
+                            ]
+                        },
+                    }
+                },
+                {
+                    "entity": {
+                        "urn": "urn:li:dataProduct:test2",
+                        "type": "DATA_PRODUCT",
+                        "properties": {"name": "test2", "description": "second test entity"},
+                        "tags": {
+                            "tags": [
+                                {"tag": {"urn": "urn:li:tag:dc_display_in_catalogue2"}}
+                            ]
+                        },
+                    }
+                }
+            ],
+        }
+    }
+    result = parse_relations(RelationshipType.PARENT, [relations["relationships"]])
+    assert result == {
+        RelationshipType.PARENT: [
+            EntitySummary(
+                entity_ref=EntityRef(
+                    urn="urn:li:dataProduct:test1", display_name="test1"
+                ),
+                description="first test entity",
+                entity_type="DATA_PRODUCT",
+                tags=[
+                    TagRef(
+                        urn="urn:li:tag:dc_display_in_catalogue",
+                        display_name="dc_display_in_catalogue",
+                    )
+                ],
+            ),
+            EntitySummary(
+                entity_ref=EntityRef(
+                    urn="urn:li:dataProduct:test2", display_name="test2"
+                ),
+                description="second test entity",
+                entity_type="DATA_PRODUCT",
+                tags=[
+                    TagRef(
+                        urn="urn:li:tag:dc_display_in_catalogue2",
+                        display_name="dc_display_in_catalogue2",
+                    )
+                ],
+            ),
+        ]
+    }
+
+
 def test_parse_tags():
     tag = TagRef(display_name="abc", urn="urn:tag:abc")
     result = parse_tags(