|
14 | 14 | import com.scalar.db.io.Column; |
15 | 15 | import com.scalar.db.io.IntColumn; |
16 | 16 | import com.scalar.db.io.Key; |
| 17 | +import java.io.BufferedReader; |
17 | 18 | import java.io.BufferedWriter; |
18 | 19 | import java.io.File; |
19 | 20 | import java.io.IOException; |
@@ -129,4 +130,72 @@ void startExport_givenPartitionKey_shouldGenerateOutputFile() |
129 | 130 | Assertions.assertTrue(file.exists()); |
130 | 131 | Assertions.assertTrue(file.delete()); |
131 | 132 | } |
| 133 | + |
| 134 | + @Test |
| 135 | + void startExport_givenNoHeaderRequired_shouldGenerateOutputFileWithoutHeader() throws Exception { |
| 136 | + String expectedFirstLine = |
| 137 | + "9007199254740992,2147483647,true,0.000000000000000000000000000000000000000000001401298464324817,0.0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000049,test value,YmxvYiB0ZXN0IHZhbHVl,2000-01-01,01:01:01,2000-01-01T01:01,1970-01-21T03:20:41.740Z"; |
| 138 | + |
| 139 | + runExportAndAssertFirstLine(true, expectedFirstLine); |
| 140 | + } |
| 141 | + |
| 142 | + @Test |
| 143 | + void startExport_givenHeaderRequired_shouldGenerateOutputFileWithHeader() throws Exception { |
| 144 | + String expectedFirstLine = "col1,col2,col3,col4,col5,col6,col7,col8,col9,col10,col11"; |
| 145 | + |
| 146 | + runExportAndAssertFirstLine(false, expectedFirstLine); |
| 147 | + } |
| 148 | + |
| 149 | + private void runExportAndAssertFirstLine(boolean excludeHeader, String expectedFirstLine) |
| 150 | + throws Exception { |
| 151 | + // Arrange |
| 152 | + producerTaskFactory = new ProducerTaskFactory(",", false, false); |
| 153 | + exportManager = new CsvExportManager(storage, dao, producerTaskFactory); |
| 154 | + Scanner scanner = Mockito.mock(Scanner.class); |
| 155 | + String filePath = Paths.get("").toAbsolutePath() + "/output.csv"; |
| 156 | + Map<String, Column<?>> values = UnitTestUtils.createTestValues(); |
| 157 | + Result result = new ResultImpl(values, mockData); |
| 158 | + List<Result> results = Collections.singletonList(result); |
| 159 | + |
| 160 | + ExportOptions exportOptions = |
| 161 | + ExportOptions.builder( |
| 162 | + "namespace", |
| 163 | + "table", |
| 164 | + Key.newBuilder().add(IntColumn.of("col1", 1)).build(), |
| 165 | + FileFormat.CSV) |
| 166 | + .sortOrders(Collections.emptyList()) |
| 167 | + .scanRange(new ScanRange(null, null, false, false)) |
| 168 | + .delimiter(",") |
| 169 | + .excludeHeaderRow(excludeHeader) |
| 170 | + .build(); |
| 171 | + |
| 172 | + Mockito.when( |
| 173 | + dao.createScanner( |
| 174 | + exportOptions.getNamespace(), |
| 175 | + exportOptions.getTableName(), |
| 176 | + exportOptions.getScanPartitionKey(), |
| 177 | + exportOptions.getScanRange(), |
| 178 | + exportOptions.getSortOrders(), |
| 179 | + exportOptions.getProjectionColumns(), |
| 180 | + exportOptions.getLimit(), |
| 181 | + storage)) |
| 182 | + .thenReturn(scanner); |
| 183 | + Mockito.when(scanner.iterator()).thenReturn(results.iterator()); |
| 184 | + try (BufferedWriter writer = |
| 185 | + Files.newBufferedWriter( |
| 186 | + Paths.get(filePath), |
| 187 | + Charset.defaultCharset(), |
| 188 | + StandardOpenOption.CREATE, |
| 189 | + StandardOpenOption.TRUNCATE_EXISTING)) { |
| 190 | + exportManager.startExport(exportOptions, mockData, writer); |
| 191 | + } |
| 192 | + File file = new File(filePath); |
| 193 | + Assertions.assertTrue(file.exists()); |
| 194 | + try (BufferedReader br = Files.newBufferedReader(file.toPath(), Charset.defaultCharset())) { |
| 195 | + String firstLine = br.readLine(); |
| 196 | + Assertions.assertEquals(expectedFirstLine, firstLine); |
| 197 | + } finally { |
| 198 | + Assertions.assertTrue(file.delete()); |
| 199 | + } |
| 200 | + } |
132 | 201 | } |
0 commit comments