datahub/smoke-test/test_rapid.py
2025-10-16 10:36:38 +05:30

60 lines
2.0 KiB
Python

import pytest
import tenacity
from typing import Any, Dict
from tests.utils import (
execute_graphql,
ingest_file_via_rest,
with_test_retry,
)
bootstrap_small = "test_resources/bootstrap_single.json"
bootstrap_small_2 = "test_resources/bootstrap_single2.json"
@with_test_retry()
def _ensure_dataset_present_correctly(auth_session):
urn = "urn:li:dataset:(urn:li:dataPlatform:testPlatform,testDataset,PROD)"
query = """query getDataset($urn: String!) {
dataset(urn: $urn) {
urn
name
description
platform {
urn
}
schemaMetadata {
name
version
createdAt
}
outgoing: relationships(
input: { types: ["DownstreamOf", "Consumes", "Produces"], direction: OUTGOING, start: 0, count: 2000 }
) {
start
count
total
relationships {
type
direction
entity {
urn
type
}
}
}
}
}"""
variables: Dict[str, Any] = {"urn": urn}
res_data = execute_graphql(auth_session, query, variables)
assert res_data["data"]["dataset"]
assert res_data["data"]["dataset"]["urn"] == urn
assert len(res_data["data"]["dataset"]["outgoing"]["relationships"]) == 1
def test_ingestion_via_rest_rapid(auth_session):
ingest_file_via_rest(auth_session, bootstrap_small)
ingest_file_via_rest(auth_session, bootstrap_small_2)
_ensure_dataset_present_correctly(auth_session)