TableInputSource extends InputSource { @Override public Object[] getNextRow() throws IOException { if (totalCount == 0 || index >= totalCount) { return null; } return dataList.get(index++); } TableInputSource(TableInfo tableInfo);  TableInputSource(String project, String table, String[] partitions, String[] columns); @Override Object[] getNextRow();  }
@Test public void testGetNextRow() throws IOException, LocalRunException { String project = "project_name"; String table = "wc_in2"; String[] partitions = null; String[] columns = null; tableInputSource = new TableInputSource(project, table, partitions, columns); Object[] data; int count = 0; while ((data = tableInputSource.getNextRow()) != null) { count++; } Assert.assertEquals(10, count); project = "project_name"; table = "wc_in2"; partitions = new String[]{"p2=2"}; columns = null; tableInputSource = new TableInputSource(project, table, partitions, columns); count = 0; while ((data = tableInputSource.getNextRow()) != null) { count++; } Assert.assertEquals(6, count); project = "project_name"; table = "wc_in2"; partitions = new String[]{"p2=1", "p1=2"}; columns = null; tableInputSource = new TableInputSource(project, table, partitions, columns); count = 0; while ((data = tableInputSource.getNextRow()) != null) { count++; } Assert.assertEquals(3, count); project = "project_name"; table = "wc_in2"; partitions = new String[]{"p2=2"}; columns = new String[]{"colc", "cola"}; tableInputSource = new TableInputSource(project, table, partitions, columns); count = 0; while ((data = tableInputSource.getNextRow()) != null) { count++; } Assert.assertEquals(6, count); project = "project_name"; table = "wc_in1"; partitions = null; columns = new String[]{"col2", "col3"}; tableInputSource = new TableInputSource(project, table, partitions, columns); count = 0; while ((data = tableInputSource.getNextRow()) != null) { count++; } Assert.assertEquals(4, count); }