diff --git a/smoke-test/tests/data_process_instance/test_data_process_instance.py b/smoke-test/tests/data_process_instance/test_data_process_instance.py index ad11dab0bb1f29..d45158e3f62600 100644 --- a/smoke-test/tests/data_process_instance/test_data_process_instance.py +++ b/smoke-test/tests/data_process_instance/test_data_process_instance.py @@ -72,6 +72,9 @@ def create_test_data(filename: str): output_model_urn = ( "urn:li:mlModel:(urn:li:dataPlatform:mlflow,my_output_model,PROD)" ) + data_platform_instance_urn = ( + "urn:li:dataPlatformInstance:(urn:li:dataPlatform:airflow,1234567890)" + ) container_urn = "urn:li:container:testGroup1" mcps = [ create_status_mcp(urn) @@ -80,6 +83,7 @@ def create_test_data(filename: str): input_model_urn, output_dataset_urn, output_model_urn, + data_platform_instance_urn, ] ] mcps += [ @@ -88,6 +92,12 @@ def create_test_data(filename: str): aspect=ContainerPropertiesClass(name="testGroup1"), ) ] + mcps += [ + MetadataChangeProposalWrapper( + entityUrn=data_platform_instance_urn, + aspect=DataPlatformInstancePropertiesClass(name="my process instance"), + ) + ] mcps += [ e for e in MetadataChangeProposalWrapper.construct_many( @@ -101,7 +111,7 @@ def create_test_data(filename: str): time=1640692800000, actor="urn:li:corpuser:datahub" ), ), - # Run Event aspect + # # Run Event aspect DataProcessInstanceRunEventClass( timestampMillis=1704067200000, eventGranularity=TimeWindowSizeClass(unit="WEEK", multiple=1), @@ -112,9 +122,6 @@ def create_test_data(filename: str): platform="urn:li:dataPlatform:airflow", instance="urn:li:dataPlatformInstance:(urn:li:dataPlatform:airflow,1234567890)", ), - DataPlatformInstancePropertiesClass( - name="my process instance", - ), # SubTypes aspect SubTypesClass(typeNames=["TEST", "BATCH_JOB"]), ContainerClass(container="urn:li:container:testGroup1"), @@ -183,6 +190,15 @@ def ingest_cleanup_data(auth_session, graph_client, request): def test_search_dpi(auth_session, ingest_cleanup_data): """Test DPI search and validation of returned fields using GraphQL.""" + # TODO: Commented out section of the query that is not yet supported in the backend + # platform { + # urn + # name + # properties { + # type + # } + # } + json = { "query": """query scrollAcrossEntities($input: ScrollAcrossEntitiesInput!) { scrollAcrossEntities(input: $input) { @@ -209,13 +225,6 @@ def test_search_dpi(auth_session, ingest_cleanup_data): container { urn } - platform { - urn - name - properties { - type - } - } mlTrainingRunProperties { id trainingMetrics {