TaraReferenceExtractionIOUtils { public static void storeInOutput(SparkSession spark, Dataset<Row> resultsToOutputDF, String outputDocumentToProject) { storeInOutput(resultsToOutputDF, outputDocumentToProject, new AvroDataStoreWriter(spark)); } private  TaraReferenceExtractionIOUtils(); static void clearOutput(SparkSession spark,
                                   String outputDocumentToProject); static void clearOutput(String outputDocumentToProject,
                                   OutputCleaner cleaner); static void storeInOutput(SparkSession spark,
                                     Dataset<Row> resultsToOutputDF,
                                     String outputDocumentToProject); static void storeInOutput(Dataset<Row> resultsToOutputDF,
                                     String outputDocumentToProject,
                                     AvroDataStoreWriter writer);  }
@Test public void storeInOutputShouldRunProperly() { DocumentToProject documentToProject = DocumentToProject.newBuilder() .setDocumentId("docId-1") .setProjectId("projId-1") .setConfidenceLevel(1.0f) .build(); List<DocumentToProject> documentToProjectList = Collections.singletonList(documentToProject); Dataset<Row> documentToProjectDF = new AvroDataFrameSupport(spark).createDataFrame( documentToProjectList, DocumentToProject.SCHEMA$); AvroDataStoreWriter writer = mock(AvroDataStoreWriter.class); storeInOutput(documentToProjectDF, "path/to/output", writer); ArgumentCaptor<Dataset<Row>> dataFrameCaptor = new ArgumentCaptor<>(); verify(writer, atLeastOnce()).write(dataFrameCaptor.capture(), eq("path/to/output"), eq(DocumentToProject.SCHEMA$)); Dataset<Row> documentToProjectStoredDF = dataFrameCaptor.getValue(); assertEquals(SchemaConverters.toSqlType(DocumentToProject.SCHEMA$).dataType(), documentToProjectStoredDF.schema()); List<Row> documentToProjectRows = documentToProjectStoredDF.collectAsList(); assertEquals(1, documentToProjectRows.size()); Row documentToProjectRow = documentToProjectRows.get(0); assertEquals(documentToProject.getDocumentId(), documentToProjectRow.getAs("documentId")); assertEquals(documentToProject.getProjectId(), documentToProjectRow.getAs("projectId")); assertEquals(documentToProject.getConfidenceLevel(), documentToProjectRow.<Float>getAs("confidenceLevel"), 1e-3); assertNull(documentToProjectRow.getAs("textsnippet")); }