tests: add unit tests for helpers (#40)

* tests: add unit tests for helpers

* Update tests/test_mcp_server_helpers.py

Co-authored-by: Harshal Sheth <hsheth2@gmail.com>

* fix chart urn in test

---------

Co-authored-by: Harshal Sheth <hsheth2@gmail.com>

Author: mayurinehate

src/mcp_server_datahub/mcp_server.py

@@ -96,7 +96,7 @@ def _execute_graphql(
     )
 
 
-def _inject_urls_for_urns(
+def inject_urls_for_urns(
     graph: DataHubGraph, response: Any, json_paths: List[str]
 ) -> None:
     if not _is_datahub_cloud(graph):
@@ -112,7 +112,7 @@ def _inject_urls_for_urns(
                 item.update(new_item)
 
 
-def _maybe_convert_to_schema_field_urn(urn: str, column: Optional[str]) -> str:
+def maybe_convert_to_schema_field_urn(urn: str, column: Optional[str]) -> str:
     if column is not None:
         maybe_dataset_urn = Urn.from_string(urn)
         if not isinstance(maybe_dataset_urn, DatasetUrn):
@@ -131,7 +131,7 @@ def _maybe_convert_to_schema_field_urn(urn: str, column: Optional[str]) -> str:
 queries_gql = (pathlib.Path(__file__).parent / "gql/queries.gql").read_text()
 
 
-def _clean_gql_response(response: Any) -> Any:
+def clean_gql_response(response: Any) -> Any:
     if isinstance(response, dict):
         banned_keys = {
             "__typename",
@@ -141,19 +141,19 @@ def _clean_gql_response(response: Any) -> Any:
         for k, v in response.items():
             if k in banned_keys or v is None or v == []:
                 continue
-            cleaned_v = _clean_gql_response(v)
+            cleaned_v = clean_gql_response(v)
             if cleaned_v is not None and cleaned_v != {}:
                 cleaned_response[k] = cleaned_v
 
         return cleaned_response
     elif isinstance(response, list):
-        return [_clean_gql_response(item) for item in response]
+        return [clean_gql_response(item) for item in response]
     else:
         return response
 
 
-def _clean_get_entity_response(raw_response: dict) -> dict:
-    response = _clean_gql_response(raw_response)
+def clean_get_entity_response(raw_response: dict) -> dict:
+    response = clean_gql_response(raw_response)
 
     if response and (schema_metadata := response.get("schemaMetadata")):
         # Remove empty platformSchema to reduce response clutter
@@ -191,9 +191,9 @@ def get_entity(urn: str) -> dict:
         operation_name="GetEntity",
     )["entity"]
 
-    _inject_urls_for_urns(client._graph, result, [""])
+    inject_urls_for_urns(client._graph, result, [""])
 
-    return _clean_get_entity_response(result)
+    return clean_get_entity_response(result)
 
 
 @mcp.tool(
@@ -274,7 +274,7 @@ def search(
         response.pop("searchResults", None)
         response.pop("count", None)
 
-    return _clean_gql_response(response)
+    return clean_gql_response(response)
 
 
 @mcp.tool(
@@ -286,7 +286,7 @@ def get_dataset_queries(
 ) -> dict:
     client = get_datahub_client()
 
-    urn = _maybe_convert_to_schema_field_urn(urn, column)
+    urn = maybe_convert_to_schema_field_urn(urn, column)
 
     entities_filter = FilterDsl.custom_filter(
         field="entities", condition="EQUAL", values=[urn]
@@ -310,7 +310,7 @@ def get_dataset_queries(
         if query.get("subjects"):
             query["subjects"] = _deduplicate_subjects(query["subjects"])
 
-    return _clean_gql_response(result)
+    return clean_gql_response(result)
 
 
 def _deduplicate_subjects(subjects: list[dict]) -> list[str]:
@@ -374,7 +374,7 @@ def get_lineage(
             "searchFlags": {"skipHighlighting": True, "maxAggValues": 3},
         }
         if asset_lineage_directive.upstream:
-            result["upstreams"] = _clean_gql_response(
+            result["upstreams"] = clean_gql_response(
                 _execute_graphql(
                     self.graph,
                     query=entity_details_fragment_gql,
@@ -388,7 +388,7 @@ def get_lineage(
                 )["searchAcrossLineage"]
             )
         if asset_lineage_directive.downstream:
-            result["downstreams"] = _clean_gql_response(
+            result["downstreams"] = clean_gql_response(
                 _execute_graphql(
                     self.graph,
                     query=entity_details_fragment_gql,
@@ -430,7 +430,7 @@ def get_lineage(
 
     lineage_api = AssetLineageAPI(client._graph)
 
-    urn = _maybe_convert_to_schema_field_urn(urn, column)
+    urn = maybe_convert_to_schema_field_urn(urn, column)
     asset_lineage_directive = AssetLineageDirective(
         urn=urn,
         upstream=upstream,
@@ -439,5 +439,5 @@ def get_lineage(
         extra_filters=filters,
     )
     lineage = lineage_api.get_lineage(asset_lineage_directive)
-    _inject_urls_for_urns(client._graph, lineage, ["*.searchResults[].entity"])
+    inject_urls_for_urns(client._graph, lineage, ["*.searchResults[].entity"])
     return lineage

tests/test_mcp_server_helpers.py

@@ -0,0 +1,197 @@
+import pytest
+from unittest.mock import Mock, patch
+from mcp_server_datahub.mcp_server import (
+    inject_urls_for_urns,
+    maybe_convert_to_schema_field_urn,
+    clean_gql_response,
+    clean_get_entity_response,
+)
+from datahub.ingestion.graph.links import make_url_for_urn
+
+
+def test_inject_urls_for_urns():
+    mock_graph = Mock()
+    mock_graph.url_for.side_effect = lambda urn: make_url_for_urn(
+        "https://xyz.com", urn
+    )
+
+    with patch("mcp_server_datahub.mcp_server._is_datahub_cloud", return_value=True):
+        response = {
+            "searchResults": [
+                {
+                    "entity": {
+                        "urn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)",
+                        "name": "users",
+                    }
+                },
+                {
+                    "entity": {
+                        "urn": "urn:li:chart:(looker,baz)",
+                        "name": "baz",
+                    }
+                },
+            ]
+        }
+
+        inject_urls_for_urns(mock_graph, response, ["searchResults[].entity"])
+
+        expected_response = {
+            "searchResults": [
+                {
+                    "entity": {
+                        "urn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)",
+                        "url": "https://xyz.com/dataset/urn%3Ali%3Adataset%3A%28urn%3Ali%3AdataPlatform%3Asnowflake%2Canalytics_db.raw_schema.users%2CPROD%29/",
+                        "name": "users",
+                    }
+                },
+                {
+                    "entity": {
+                        "urn": "urn:li:chart:(looker,baz)",
+                        "url": "https://xyz.com/chart/urn%3Ali%3Achart%3A%28looker%2Cbaz%29/",
+                        "name": "baz",
+                    }
+                },
+            ]
+        }
+
+        assert response == expected_response
+        assert mock_graph.url_for.call_count == 2
+
+
+def test_maybe_convert_to_schema_field_urn_with_column():
+    dataset_urn = "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)"
+    column = "user_id"
+
+    result = maybe_convert_to_schema_field_urn(dataset_urn, column)
+
+    assert (
+        result
+        == "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD),user_id)"
+    )
+
+
+def test_maybe_convert_to_schema_field_urn_without_column():
+    original_urn = "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)"
+
+    result = maybe_convert_to_schema_field_urn(original_urn, None)
+
+    assert result == original_urn
+
+
+def test_maybe_convert_to_schema_field_urn_with_incorrect_entity():
+    chart_urn = "urn:li:chart:(looker,baz)"
+
+    # Ok if no column is provided
+    result = maybe_convert_to_schema_field_urn(chart_urn, None)
+    assert result == chart_urn
+
+    # Fail if column is provided
+    column = "user_id"
+    with pytest.raises(ValueError):
+        maybe_convert_to_schema_field_urn(chart_urn, column)
+
+
+def test_clean_gql_response_with_dict():
+    response = {
+        "__typename": "Dataset",
+        "urn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)",
+        "name": "users",
+        "description": None,
+        "tags": [],
+        "properties": {"__typename": "Properties", "key1": "value1", "key2": None},
+    }
+
+    result = clean_gql_response(response)
+
+    expected_result = {
+        "urn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)",
+        "name": "users",
+        "properties": {"key1": "value1"},
+    }
+
+    assert result == expected_result
+
+
+def test_clean_gql_response_with_nested_empty_objects():
+    response = {
+        "urn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)",
+        "name": "users",
+        "empty_object": {},
+        "empty_array": [],
+        "nested_object": {
+            "empty_object": {},
+            "empty_array": [],
+            "valid_data": "value",
+            "null_value": None,
+        },
+        "array_of_objects": [
+            {"valid": "data", "empty_object": {}, "empty_array": [], "null_value": None}
+        ],
+    }
+
+    result = clean_gql_response(response)
+
+    expected_result = {
+        "urn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)",
+        "name": "users",
+        "nested_object": {"valid_data": "value"},
+        "array_of_objects": [{"valid": "data"}],
+    }
+
+    assert result == expected_result
+
+
+def test_clean_get_entity_response_with_schema_metadata():
+    raw_response = {
+        "urn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)",
+        "name": "users",
+        "schemaMetadata": {
+            "platformSchema": {
+                "schema": ""  # Empty schema should be removed
+            },
+            "fields": [
+                {
+                    "fieldPath": "user_id",
+                    "recursive": False,  # Should be removed
+                    "isPartOfKey": False,  # Should be removed
+                    "type": "STRING",
+                },
+                {
+                    "fieldPath": "email",
+                    "recursive": True,  # Should be kept
+                    "isPartOfKey": True,  # Should be kept
+                    "type": "STRING",
+                },
+                {
+                    "fieldPath": "created_at",
+                    "recursive": None,  # Should be removed
+                    "isPartOfKey": None,  # Should be removed
+                    "type": "TIMESTAMP",
+                },
+            ],
+        },
+    }
+
+    result = clean_get_entity_response(raw_response)
+
+    expected_result = {
+        "urn": "urn:li:dataset:(urn:li:dataPlatform:snowflake,analytics_db.raw_schema.users,PROD)",
+        "name": "users",
+        "schemaMetadata": {
+            "fields": [
+                {"fieldPath": "user_id", "type": "STRING"},
+                {
+                    "fieldPath": "email",
+                    "recursive": True,
+                    "isPartOfKey": True,
+                    "type": "STRING",
+                },
+                {
+                    "fieldPath": "created_at",
+                    "type": "TIMESTAMP",
+                },
+            ]
+        },
+    }
+
+    assert result == expected_result