Examples with PlanUtils io.openlineage.spark.agent.util.PlanUtils used on opensource projects

Example 1 with PlanUtils

use of io.openlineage.spark.agent.util.PlanUtils in project OpenLineage by OpenLineage.

the class PlanUtils3Test method testFromDataSourceV2Relation.

@Test
public void testFromDataSourceV2Relation() {
    try (MockedStatic<CatalogUtils3> mocked = mockStatic(CatalogUtils3.class)) {
        try (MockedStatic<PlanUtils> mockedPlanUtils = mockStatic(PlanUtils.class)) {
            DatasetIdentifier di = mock(DatasetIdentifier.class);
            when(di.getNamespace()).thenReturn("file://tmp");
            when(di.getName()).thenReturn("name");
            OpenLineage.DatasetFacets datasetFacets = mock(OpenLineage.DatasetFacets.class);
            OpenLineage.Dataset dataset = mock(OpenLineage.Dataset.class);
            OpenLineage.SchemaDatasetFacet schemaDatasetFacet = mock(OpenLineage.SchemaDatasetFacet.class);
            OpenLineage.DatasourceDatasetFacet datasourceDatasetFacet = mock(OpenLineage.DatasourceDatasetFacet.class);
            when(PlanUtils.schemaFacet(openLineage, schema)).thenReturn(schemaDatasetFacet);
            when(PlanUtils.datasourceFacet(openLineage, di.getNamespace())).thenReturn(datasourceDatasetFacet);
            when(datasetFacetsBuilder.schema(schemaDatasetFacet)).thenReturn(datasetFacetsBuilder);
            when(datasetFacetsBuilder.dataSource(datasourceDatasetFacet)).thenReturn(datasetFacetsBuilder);
            when(datasetFacetsBuilder.build()).thenReturn(datasetFacets);
            when(CatalogUtils3.getDatasetIdentifier(sparkSession, tableCatalog, identifier, tableProperties)).thenReturn(di);
            when(datasetFactory.getDataset(di.getName(), di.getNamespace(), datasetFacets)).thenReturn(dataset);
            assertEquals(Collections.singletonList(dataset), PlanUtils3.fromDataSourceV2Relation(datasetFactory, openLineageContext, dataSourceV2Relation));
        }
    }
}