test(search): search openAPI v3 test (#11049)

datahub-project · Aug 7, 2024 · 900c259 · 900c259
1 parent 40e61f9
commit 900c259
Showing 1 changed file with 57 additions and 18 deletions.
diff --git a/smoke-test/tests/read_only/test_search.py b/smoke-test/tests/read_only/test_search.py
@@ -1,10 +1,13 @@
 import pytest
+import requests
 
 from tests.test_result_msg import add_datahub_stats
-from tests.utils import get_frontend_session, get_frontend_url
+from tests.utils import get_frontend_session, get_frontend_url, get_gms_url
 
-restli_default_headers = {
-    "X-RestLi-Protocol-Version": "2.0.0",
+BASE_URL_V3 = f"{get_gms_url()}/openapi/v3"
+
+default_headers = {
+    "Content-Type": "application/json",
 }
 
 ENTITY_TO_MAP = {
@@ -59,16 +62,8 @@ def _get_search_result(frontend_session, entity: str):
         ("chart", "chart"),
         ("dataset", "dataset"),
         ("dashboard", "dashboard"),
-        (
-            # Task
-            "dataJob",
-            "dataJob",
-        ),
-        (
-            # Pipeline
-            "dataFlow",
-            "dataFlow",
-        ),
+        ("dataJob", "dataJob"),
+        ("dataFlow", "dataFlow"),
         ("container", "container"),
         ("tag", "tag"),
         ("corpUser", "corpUser"),
@@ -78,11 +73,7 @@ def _get_search_result(frontend_session, entity: str):
         ("mlPrimaryKey", "mlPrimaryKey"),
         ("corpGroup", "corpGroup"),
         ("mlFeatureTable", "mlFeatureTable"),
-        (
-            # Term group
-            "glossaryNode",
-            "glossaryNode",
-        ),
+        ("glossaryNode", "glossaryNode"),
         ("mlModel", "mlModel"),
     ],
 )
@@ -112,8 +103,56 @@ def test_search_works(entity_type, api_name):
         """,
         "variables": {"input": first_urn},
     }
+
     response = frontend_session.post(f"{get_frontend_url()}/api/v2/graphql", json=json)
     response.raise_for_status()
     res_data = response.json()
     assert res_data["data"], f"res_data was {res_data}"
     assert res_data["data"][api_name]["urn"] == first_urn, f"res_data was {res_data}"
+
+
+@pytest.mark.read_only
+@pytest.mark.parametrize(
+    "entity_type",
+    [
+        "chart",
+        "dataset",
+        "dashboard",
+        "dataJob",
+        "dataFlow",
+        "container",
+        "tag",
+        "corpUser",
+        "mlFeature",
+        "glossaryTerm",
+        "domain",
+        "mlPrimaryKey",
+        "corpGroup",
+        "mlFeatureTable",
+        "glossaryNode",
+        "mlModel",
+    ],
+)
+def test_openapi_v3_entity(entity_type):
+    frontend_session = get_frontend_session()
+    search_result = _get_search_result(frontend_session, entity_type)
+    num_entities = search_result["total"]
+    if num_entities == 0:
+        print(f"[WARN] No results for {entity_type}")
+        return
+    entities = search_result["searchResults"]
+
+    first_urn = entities[0]["entity"]["urn"]
+
+    session = requests.Session()
+    url = f"{BASE_URL_V3}/entity/{entity_type}/{first_urn}"
+    response = session.get(url, headers=default_headers)
+    response.raise_for_status()
+    actual_data = response.json()
+    print(f"Entity Data for URN {first_urn}: {actual_data}")
+
+    expected_data = {"urn": first_urn}
+
+    assert (
+        actual_data["urn"] == expected_data["urn"]
+    ), f"Mismatch: expected {expected_data}, got {actual_data}"