KiteExtractor extends Extractor<LinkConfiguration,
    FromJobConfiguration, KiteDatasetPartition> { @Override public void extract(ExtractorContext context, LinkConfiguration linkConfig, FromJobConfiguration fromJobConfig, KiteDatasetPartition partition) { String uri = ConfigUtil.buildDatasetUri( linkConfig.linkConfig, partition.getUri()); LOG.info("Loading data from " + uri); KiteDatasetExecutor executor = getExecutor(uri); DataWriter writer = context.getDataWriter(); Object[] array; rowsRead = 0L; try { while ((array = executor.readRecord()) != null) { writer.writeArrayRecord(array); rowsRead++; } } finally { executor.closeReader(); } }  @Override void extract(ExtractorContext context, LinkConfiguration linkConfig,
      FromJobConfiguration fromJobConfig, KiteDatasetPartition partition); @Override long getRowsRead();  }
@Test public void testExtractor() throws Exception { Schema schema = new Schema("testExtractor"); schema.addColumn(new Text("TextCol")); ExtractorContext context = new ExtractorContext(null, writerMock, schema); LinkConfiguration linkConfig = new LinkConfiguration(); FromJobConfiguration jobConfig = new FromJobConfiguration(); KiteDatasetPartition partition = new KiteDatasetPartition(); partition.setUri("dataset:hdfs:/path/to/dataset"); OngoingStubbing<Object[]> readRecordMethodStub = when(executorMock.readRecord()); final int NUMBER_OF_ROWS = 1000; for (int i = 0; i < NUMBER_OF_ROWS; i++) { readRecordMethodStub = readRecordMethodStub.thenReturn(new Object[]{}); } readRecordMethodStub.thenReturn(null); extractor.extract(context, linkConfig, jobConfig, partition); verify(writerMock, times(NUMBER_OF_ROWS)).writeArrayRecord( any(Object[].class)); }