diff --git a/src/org/labkey/test/tests/assay/UploadLargeExcelAssayTest.java b/src/org/labkey/test/tests/assay/UploadLargeExcelAssayTest.java index 167abd2662..1070df1e83 100644 --- a/src/org/labkey/test/tests/assay/UploadLargeExcelAssayTest.java +++ b/src/org/labkey/test/tests/assay/UploadLargeExcelAssayTest.java @@ -39,7 +39,7 @@ protected void doCleanup(boolean afterTest) @BeforeClass public static void setupProject() throws Exception { - UploadLargeExcelAssayTest init = (UploadLargeExcelAssayTest) getCurrentTest(); + UploadLargeExcelAssayTest init = getCurrentTest(); init.doSetup(); } @@ -85,9 +85,8 @@ public void testUpload200kRows() throws Exception String fileName = "200kXlsxFile.xlsx"; var dgen = new TestDataGenerator("samples", "chaos_sample", getProjectName()) .withColumns(ASSAY_FIELDS); - dgen.generateRows(200_000); log("writing large .xlsx file"); - var largeExcelFile = dgen.writeData(fileName); + var largeExcelFile = dgen.writeData(fileName, 200_000); log("finished writing large .xlsx file"); // import large generated excel to assay1 diff --git a/src/org/labkey/test/util/TestDataGenerator.java b/src/org/labkey/test/util/TestDataGenerator.java index 4a0fc70ad1..f7edcb775a 100644 --- a/src/org/labkey/test/util/TestDataGenerator.java +++ b/src/org/labkey/test/util/TestDataGenerator.java @@ -16,11 +16,8 @@ package org.labkey.test.util; import org.apache.commons.csv.CSVFormat; -import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.time.DateUtils; -import org.apache.poi.xssf.streaming.SXSSFRow; -import org.apache.poi.xssf.streaming.SXSSFWorkbook; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import org.junit.Assert; @@ -37,7 +34,6 @@ import org.labkey.remoteapi.query.SelectRowsResponse; import org.labkey.remoteapi.query.Sort; import org.labkey.serverapi.reader.TabLoader; -import org.labkey.test.TestFileUtils; import org.labkey.test.WebTestHelper; import org.labkey.test.params.FieldDefinition; import org.labkey.test.util.data.ColumnNameMapper; @@ -45,7 +41,6 @@ import org.labkey.test.util.query.QueryApiHelper; import java.io.File; -import java.io.FileOutputStream; import java.io.IOException; import java.text.SimpleDateFormat; import java.util.ArrayList; @@ -54,8 +49,11 @@ import java.util.Collections; import java.util.Date; import java.util.HashMap; +import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; import java.util.Set; import java.util.concurrent.ThreadLocalRandom; import java.util.function.Supplier; @@ -313,7 +311,7 @@ public TestDataGenerator withGeneratedRows(int desiredRowCount) public TestDataGenerator addDataSupplier(String columnName, Supplier supplier) { - _dataSuppliers.put(columnName, ()-> supplier.get()); + _dataSuppliers.put(columnName, supplier); return this; } @@ -400,7 +398,7 @@ public List getColumns() public void generateRows(int numberOfRowsToGenerate) { - if (_columns.keySet().size() == 0) + if (_columns.isEmpty()) throw new IllegalStateException("can't generate row data without column definitions"); for (int i= 0; i < numberOfRowsToGenerate; i++) @@ -444,7 +442,7 @@ private Supplier getDefaultDataSupplier(String columnType) case "double": return ()-> randomDouble(0, 20); case "boolean": - return ()-> randomBoolean(); + return this::randomBoolean; case "date": case "datetime": return ()-> randomDateString(DateUtils.addWeeks(new Date(), -39), new Date()); @@ -670,40 +668,6 @@ public String getDataAsTsv() return TestDataUtils.stringFromRowMaps(_rows, getFieldsForFile(), true, CSVFormat.TDF); } - public File writeGeneratedDataToExcel(String sheetName, String fileName) throws IOException - { - File file = new File(TestFileUtils.getTestTempDir(), fileName); - FileUtils.forceMkdirParent(file); - - try (SXSSFWorkbook workbook = new SXSSFWorkbook(1000); // only holds 1000 rows in memory - FileOutputStream out = new FileOutputStream(file)) - { - var sheet = workbook.createSheet(sheetName); - - // write headers as row 0 - List columnNames = getFieldsForFile(); - var headerRow = sheet.createRow(0); - for (int i = 0; i < columnNames.size(); i++) - { - headerRow.createCell(i).setCellValue(columnNames.get(i)); - } - - // write content - for (int i = 0; i < _rows.size(); i++) - { - Map row = _rows.get(i); - SXSSFRow currentRow = sheet.createRow(i + 1); - for (int j = 0; j < columnNames.size(); j++) - { - currentRow.createCell(j).setCellValue(row.getOrDefault(columnNames.get(j), "").toString()); - } - } - workbook.write(out); - } - - return file; - } - /** * Creates a file containing the contents of the current rows, formatted in TSV, CSV, or xlsx. * The file is written to the test temp dir @@ -712,33 +676,43 @@ public File writeGeneratedDataToExcel(String sheetName, String fileName) throws */ public File writeData(String fileName) { - String fileExtension = fileName.toLowerCase().substring(fileName.lastIndexOf('.') + 1); - switch (fileExtension) - { - case "xlsx": - case "xls": - try - { - return writeGeneratedDataToExcel("sheet1", fileName); - } - catch (IOException e) - { - throw new RuntimeException(e); - } - case "csv": - return writeData(fileName, CSVFormat.DEFAULT); - case "tsv": - return writeData(fileName, CSVFormat.TDF); - default: - throw new IllegalArgumentException("Unsupported file extension: " + fileExtension); - } + return writeData(fileName, new FileRowIterator(getFieldsForFile(), _rows)); } - public File writeData(String fileName, CSVFormat format) + /** + * Generate rows and write to file without saving in data generator + * @param fileName the name of the file, e.g. 'testDataFileForMyTest.tsv' + * @param numRows the number of rows to generate + * @return object pointing at created file + */ + public File writeData(String fileName, int numRows) + { + return writeData(fileName, new FileRowIterator(getFieldsForFile(), this::generateRow, numRows)); + } + + /** + * Creates a file containing the contents of the current rows, formatted in TSV, CSV, or xlsx. + * The file is written to the test temp dir + * @param fileName the name of the file, e.g. 'testDataFileForMyTest.tsv' + * @return File object pointing at created file + */ + public File writeData(String fileName, Iterator> rowIterator) { try { - return TestDataUtils.writeRowsToFile(fileName, TestDataUtils.rowListsFromMaps(_rows, getFieldsForFile()), format); + String fileExtension = fileName.toLowerCase().substring(fileName.lastIndexOf('.') + 1); + switch (fileExtension) + { + case "xlsx": + case "xls": + return TestDataUtils.writeRowsToExcel(fileName, rowIterator); + case "csv": + return TestDataUtils.writeRowsToFile(fileName, rowIterator, CSVFormat.DEFAULT); + case "tsv": + return TestDataUtils.writeRowsToFile(fileName, rowIterator, CSVFormat.TDF); + default: + throw new IllegalArgumentException("Unsupported file extension: " + fileExtension); + } } catch (IOException e) { @@ -915,3 +889,71 @@ public static boolean doesDomainExists(final String containerPath, final String } } + +class FileRowIterator implements Iterator> +{ + private final List headers; + private final Iterator> rows; + + private boolean firstRow = true; + + public FileRowIterator(@NotNull List headers, @NotNull Iterator> rows) + { + this.headers = Objects.requireNonNull(headers); + this.rows = Objects.requireNonNull(rows); + } + + public FileRowIterator(@NotNull List headers, @NotNull Supplier> rowSupplier, final int rowCount) + { + this(headers, new Iterator<>() + { + int count = 0; + + @Override + public boolean hasNext() + { + return count < rowCount; + } + + @Override + public Map next() + { + count++; + return rowSupplier.get(); + } + }); + } + + public FileRowIterator(@NotNull List headers, @NotNull List> rows) + { + this(headers, rows.iterator()); + } + + @Override + public boolean hasNext() + { + return firstRow || rows.hasNext(); + } + + @Override + public List next() + { + if (!hasNext()) + throw new NoSuchElementException(); + + if (firstRow) + { + firstRow = false; + return Collections.unmodifiableList(headers); + } + else + { + return rowMapToList(rows.next()); + } + } + + private List rowMapToList(Map row) + { + return headers.stream().map(h -> row.getOrDefault(h, "")).toList(); + } +} \ No newline at end of file diff --git a/src/org/labkey/test/util/data/TestDataUtils.java b/src/org/labkey/test/util/data/TestDataUtils.java index e813c1f9cf..30739cca0d 100644 --- a/src/org/labkey/test/util/data/TestDataUtils.java +++ b/src/org/labkey/test/util/data/TestDataUtils.java @@ -7,14 +7,18 @@ import org.apache.commons.io.FileUtils; import org.apache.commons.io.IOUtils; import org.apache.commons.lang3.StringUtils; +import org.apache.poi.xssf.streaming.SXSSFRow; +import org.apache.poi.xssf.streaming.SXSSFWorkbook; import org.jetbrains.annotations.NotNull; import org.labkey.serverapi.reader.DataLoader; import org.labkey.serverapi.reader.TabLoader; import org.labkey.test.TestFileUtils; import org.labkey.test.params.FieldDefinition; import org.labkey.test.util.TestDataGenerator; +import org.labkey.test.util.TestLogger; import java.io.File; +import java.io.FileOutputStream; import java.io.FileReader; import java.io.FileWriter; import java.io.IOException; @@ -24,6 +28,7 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Collection; +import java.util.Iterator; import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; @@ -314,15 +319,70 @@ public static File writeRowsToCsv(String fileName, List> rows) throw } public static @NotNull File writeRowsToFile(String fileName, List> rows, CSVFormat format) throws IOException + { + return writeRowsToFile(fileName, rows.iterator(), format); + } + + public static @NotNull File writeRowsToFile(String fileName, Iterator> rowIterator, CSVFormat format) throws IOException { File file = new File(TestFileUtils.getTestTempDir(), fileName); FileUtils.forceMkdirParent(file); + TestLogger.log("Writing data file " + file.getAbsolutePath()); + try (CSVPrinter printer = new CSVPrinter(new FileWriter(file, StandardCharsets.UTF_8), format)) { - for (List row : rows) + while (rowIterator.hasNext()) { - printer.printRecord(row); + printer.printRecord(rowIterator.next()); + } + } + + return file; + } + + public static @NotNull File writeRowsToExcel(String fileName, List> rows) throws IOException + { + return writeRowsToExcel(fileName, rows.iterator()); + } + + public static @NotNull File writeRowsToExcel(String fileName, Iterator> rowIterator) throws IOException + { + File file = new File(TestFileUtils.getTestTempDir(), fileName); + FileUtils.forceMkdirParent(file); + + TestLogger.log("Writing data file " + file.getAbsolutePath()); + + try (SXSSFWorkbook workbook = new SXSSFWorkbook(1000); // only holds 1000 rows in memory + FileOutputStream out = new FileOutputStream(file)) + { + var sheet = workbook.createSheet("sheet1"); + + // write headers as row 0 + List columnNames = rowIterator.next(); + var headerRow = sheet.createRow(0); + for (int col = 0; col < columnNames.size(); col++) + { + if (columnNames.get(col) instanceof String s) + { + headerRow.createCell(col).setCellValue(s); + } + else + { + throw new IllegalArgumentException("Expected column headers to be Strings but got " + columnNames.get(col).getClass().getSimpleName()); + } + } + + // write content + for (int rowNum = 1; rowIterator.hasNext(); rowNum++) + { + List row = rowIterator.next(); + SXSSFRow currentRow = sheet.createRow(rowNum + 1); + for (int col = 0; col < columnNames.size(); col++) + { + currentRow.createCell(col).setCellValue(row.get(col).toString()); + } } + workbook.write(out); } return file;