diff --git a/smoke-test/tests/read_only/test_search.py b/smoke-test/tests/read_only/test_search.py index 90385c5228bc1a..3b9635f3da2cd5 100644 --- 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}"