Skip to content

Commit

Permalink
Use wrapper for initialise calls in tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@patchwork01
patchwork01 committed Feb 21, 2025
1 parent 94b879a commit bb3f6d2
Show file tree
Hide file tree
Showing 39 changed files with 152 additions and 136 deletions.
3 changes: 2 additions & 1 deletion ...ort/bulk-import-runner/src/test/java/sleeper/bulkimport/runner/BulkImportJobDriverIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,7 @@
import static sleeper.core.properties.table.TableProperty.TABLE_NAME;
import static sleeper.core.properties.testutils.InstancePropertiesTestHelper.createTestInstanceProperties;
import static sleeper.core.properties.testutils.TablePropertiesTestHelper.createTestTableProperties;
import static sleeper.core.statestore.testutils.StateStoreUpdatesWrapper.update;
import static sleeper.core.tracker.ingest.job.IngestJobStatusTestData.ingestFinishedStatus;
import static sleeper.core.tracker.job.run.JobRunSummaryTestHelper.summary;
import static sleeper.core.tracker.job.run.JobRunTestData.jobRunOnTask;
Expand Down Expand Up @@ -500,7 +501,7 @@ private static void writeRecordsToFile(List<Record> records, String file) throws
private StateStore createTable(InstanceProperties instanceProperties, TableProperties tableProperties, List<Object> splitPoints) {
tablePropertiesStore(instanceProperties).save(tableProperties);
StateStore stateStore = new StateStoreFactory(instanceProperties, s3Client, dynamoClient, hadoopConf).getStateStore(tableProperties);
stateStore.initialise(new PartitionsFromSplitPoints(getSchema(), splitPoints).construct());
update(stateStore).initialise(new PartitionsFromSplitPoints(tableProperties.getSchema(), splitPoints).construct());
return stateStore;
}

Expand Down
2 changes: 1 addition & 1 deletion java/clients/src/test/java/sleeper/clients/admin/FilesStatusReportScreenTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ class FilesStatusReportScreenTest extends AdminClientMockStoreBase {
@BeforeEach
void setUp() {
setInstanceProperties(instanceProperties, tableProperties);
stateStore.initialise(PartitionsBuilderSplitsFirst.leavesWithSplits(
update(stateStore).initialise(PartitionsBuilderSplitsFirst.leavesWithSplits(
schema, List.of("A", "B"), List.of("aaa"))
.parentJoining("parent", "A", "B").buildList());
update(stateStore).addFiles(FileReferenceFactory.from(stateStore).singleFileInEachLeafPartitionWithRecords(5).toList());
Expand Down
2 changes: 1 addition & 1 deletion java/clients/src/test/java/sleeper/clients/admin/PartitionsStatusReportScreenTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ void setUp() {
@Test
void shouldRunPartitionStatusReport() throws Exception {
// Given
stateStore.initialise(PartitionsBuilderSplitsFirst.leavesWithSplits(
update(stateStore).initialise(PartitionsBuilderSplitsFirst.leavesWithSplits(
schema, List.of("A", "B"), List.of("aaa"))
.parentJoining("parent", "A", "B").buildList());
update(stateStore).addFiles(FileReferenceFactory.from(stateStore).singleFileInEachLeafPartitionWithRecords(5).toList());
Expand Down
9 changes: 5 additions & 4 deletions java/clients/src/test/java/sleeper/clients/status/partitions/ExportSplitPointsTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
import static org.assertj.core.api.Assertions.assertThat;
import static sleeper.core.properties.testutils.InstancePropertiesTestHelper.createTestInstanceProperties;
import static sleeper.core.properties.testutils.TablePropertiesTestHelper.createTestTableProperties;
import static sleeper.core.statestore.testutils.StateStoreUpdatesWrapper.update;

public class ExportSplitPointsTest {
InstanceProperties instanceProperties = createTestInstanceProperties();
Expand All @@ -62,7 +63,7 @@ public void shouldExportCorrectSplitPointsIntType() {
List<Object> splitPoints = new ArrayList<>();
splitPoints.add(-10);
splitPoints.add(1000);
stateStore.initialise(new PartitionsFromSplitPoints(schema, splitPoints).construct());
update(stateStore).initialise(new PartitionsFromSplitPoints(schema, splitPoints).construct());
ExportSplitPoints exportSplitPoints = new ExportSplitPoints(stateStore, schema);

// When
Expand All @@ -80,7 +81,7 @@ public void shouldExportCorrectSplitPointsLongType() {
List<Object> splitPoints = new ArrayList<>();
splitPoints.add(-10L);
splitPoints.add(1000L);
stateStore.initialise(new PartitionsFromSplitPoints(schema, splitPoints).construct());
update(stateStore).initialise(new PartitionsFromSplitPoints(schema, splitPoints).construct());
ExportSplitPoints exportSplitPoints = new ExportSplitPoints(stateStore, schema);

// When
Expand All @@ -98,7 +99,7 @@ public void shouldExportCorrectSplitPointsStringType() {
List<Object> splitPoints = new ArrayList<>();
splitPoints.add("A");
splitPoints.add("T");
stateStore.initialise(new PartitionsFromSplitPoints(schema, splitPoints).construct());
update(stateStore).initialise(new PartitionsFromSplitPoints(schema, splitPoints).construct());
ExportSplitPoints exportSplitPoints = new ExportSplitPoints(stateStore, schema);

// When
Expand All @@ -116,7 +117,7 @@ public void shouldExportCorrectSplitPointsByteArrayType() {
List<Object> splitPoints = new ArrayList<>();
splitPoints.add(new byte[]{10});
splitPoints.add(new byte[]{100});
stateStore.initialise(new PartitionsFromSplitPoints(schema, splitPoints).construct());
update(stateStore).initialise(new PartitionsFromSplitPoints(schema, splitPoints).construct());
ExportSplitPoints exportSplitPoints = new ExportSplitPoints(stateStore, schema);

// When
Expand Down
14 changes: 7 additions & 7 deletions ...clients/src/test/java/sleeper/clients/status/report/filestatus/FilesStatusReportTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ public void shouldReportFilesStatusGivenOneActiveFilePerLeafPartition() throws E
.splitToNewChildren("5", "6", "C", "bbb")
.splitToNewChildren("6", "A", "B", "aaa")
.buildTree();
stateStore.initialise(partitions.getAllPartitions());
update(stateStore).initialise(partitions.getAllPartitions());
FileReferenceFactory fileReferenceFactory = FileReferenceFactory.from(partitions);
update(stateStore).addFiles(List.of(
fileReferenceFactory.partitionFile("A", 50000001),
Expand All @@ -73,7 +73,7 @@ public void shouldReportFilesStatusGivenActiveFileInLeafAndMiddlePartition() thr
.splitToNewChildren("A", "B", "C", "mmm")
.splitToNewChildren("B", "D", "E", "ggg")
.buildTree();
stateStore.initialise(partitions.getAllPartitions());
update(stateStore).initialise(partitions.getAllPartitions());
FileReferenceFactory fileReferenceFactory = FileReferenceFactory.from(partitions);
update(stateStore).addFiles(List.of(
fileReferenceFactory.partitionFile("D", 50000001),
Expand All @@ -93,7 +93,7 @@ public void shouldReportFilesStatusGivenFilesWithNoReferencesBelowMaxCount() thr
.rootFirst("A")
.splitToNewChildren("A", "B", "C", "mmm")
.buildTree();
stateStore.initialise(partitions.getAllPartitions());
update(stateStore).initialise(partitions.getAllPartitions());
FileReferenceFactory fileReferenceFactory = FileReferenceFactory.from(partitions);
update(stateStore).addFiles(List.of(
fileReferenceFactory.partitionFile("B", "file1.parquet", 100),
Expand All @@ -117,7 +117,7 @@ public void shouldReportFilesStatusGivenFilesWithNoReferencesAboveMaxCount() thr
.rootFirst("A")
.splitToNewChildren("A", "B", "C", "mmm")
.buildTree();
stateStore.initialise(partitions.getAllPartitions());
update(stateStore).initialise(partitions.getAllPartitions());
FileReferenceFactory fileReferenceFactory = FileReferenceFactory.from(partitions);
update(stateStore).addFiles(List.of(
fileReferenceFactory.partitionFile("B", "file1.parquet", 100),
Expand Down Expand Up @@ -146,7 +146,7 @@ public void shouldReportFilesStatusWhenSomeFilesHaveBeenSplit() throws Exception
.rootFirst("A")
.splitToNewChildren("A", "B", "C", "mmm")
.buildTree();
stateStore.initialise(partitions.getAllPartitions());
update(stateStore).initialise(partitions.getAllPartitions());
FileReferenceFactory fileReferenceFactory = FileReferenceFactory.fromUpdatedAt(partitions, lastStateStoreUpdate);
FileReference rootFile = fileReferenceFactory.partitionFile("A", "not-split.parquet", 1000);
FileReference pendingSplit = fileReferenceFactory.partitionFile("B", "pending-split.parquet", 2000);
Expand All @@ -165,11 +165,11 @@ public void shouldReportFilesStatusWhenSomeFilesHaveBeenSplit() throws Exception
public void shouldReportFilesStatusWhenPartitionsNoLongerExist() throws Exception {
// Given
PartitionTree partitions = new PartitionsBuilder(schema).rootFirst("A").buildTree();
stateStore.initialise(partitions.getAllPartitions());
update(stateStore).initialise(partitions.getAllPartitions());
FileReferenceFactory fileReferenceFactory = FileReferenceFactory.fromUpdatedAt(partitions, lastStateStoreUpdate);
FileReference file1 = fileReferenceFactory.rootFile("file1.parquet", 1000L);
FileReference file2 = fileReferenceFactory.rootFile("file2.parquet", 2000L);
stateStore.initialise(new PartitionsBuilder(schema).rootFirst("B").buildList());
update(stateStore).initialise(new PartitionsBuilder(schema).rootFirst("B").buildList());
update(stateStore).addFiles(List.of(file1, file2));

// When / Then
Expand Down
2 changes: 1 addition & 1 deletion ...ents/src/test/java/sleeper/clients/status/report/partitions/PartitionsStatusReportIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ void shouldGetReportWhenTwoLeafPartitionsBothNeedSplitting() throws Exception {
.splitToNewChildren("parent", "A", "B", "aaa")
.buildTree();
StateStore stateStore = stateStore();
stateStore.initialise(tree.getAllPartitions());
update(stateStore).initialise(tree.getAllPartitions());
update(stateStore).addFiles(FileReferenceFactory.from(tree).singleFileInEachLeafPartitionWithRecords(100).toList());

// When / Then
Expand Down
4 changes: 2 additions & 2 deletions java/clients/src/test/java/sleeper/clients/status/update/DeleteTableIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ void shouldDeleteOnlyTable() throws Exception {
// Given
TableProperties table = createTable(uniqueIdAndName("test-table-1", "table-1"));
StateStore stateStoreBefore = createStateStore(table);
stateStoreBefore.initialise(new PartitionsBuilder(schema)
update(stateStoreBefore).initialise(new PartitionsBuilder(schema)
.rootFirst("root")
.splitToNewChildren("root", "L", "R", 50L)
.buildList());
Expand Down Expand Up @@ -154,7 +154,7 @@ void shouldDeleteTableWhenSnapshotIsPresent() throws Exception {
// Given
TableProperties table = createTable(uniqueIdAndName("test-table-1", "table-1"));
StateStore stateStore = createStateStore(table);
stateStore.initialise(new PartitionsBuilder(schema)
update(stateStore).initialise(new PartitionsBuilder(schema)
.rootFirst("root")
.splitToNewChildren("root", "L", "R", 50L)
.buildList());
Expand Down
2 changes: 1 addition & 1 deletion ...a/src/test/java/sleeper/compaction/job/creation/lambda/CompactionJobDispatchLambdaIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ private TableProperties addTable(InstanceProperties instanceProperties, Schema s
TableProperties tableProperties = createTestTableProperties(instanceProperties, schema);
S3TableProperties.createStore(instanceProperties, s3Client, dynamoClient)
.createTable(tableProperties);
stateStoreProvider.getStateStore(tableProperties)
update(stateStoreProvider.getStateStore(tableProperties))
.initialise(partitions.getAllPartitions());
return tableProperties;
}
Expand Down
7 changes: 4 additions & 3 deletions ...ion/src/test/java/sleeper/compaction/job/execution/JavaCompactionRunnerEmptyOutputIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import static sleeper.compaction.job.execution.testutils.CompactionRunnerTestData.writeRootFile;
import static sleeper.compaction.job.execution.testutils.CompactionRunnerTestUtils.assignJobIdToInputFiles;
import static sleeper.compaction.job.execution.testutils.CompactionRunnerTestUtils.createSchemaWithTypesForKeyAndTwoValues;
import static sleeper.core.statestore.testutils.StateStoreUpdatesWrapper.update;

class JavaCompactionRunnerEmptyOutputIT extends CompactionRunnerTestBase {

Expand All @@ -43,7 +44,7 @@ void shouldMergeFilesWhenSomeAreEmpty() throws Exception {
// Given
Schema schema = createSchemaWithTypesForKeyAndTwoValues(new LongType(), new LongType(), new LongType());
tableProperties.setSchema(schema);
stateStore.initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());
update(stateStore).initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());

List<Record> data = keyAndTwoValuesSortedEvenLongs();
FileReference file1 = ingestRecordsGetFile(data);
Expand All @@ -67,7 +68,7 @@ void shouldMergeFilesWhenAllAreEmpty() throws Exception {
// Given
Schema schema = createSchemaWithTypesForKeyAndTwoValues(new LongType(), new LongType(), new LongType());
tableProperties.setSchema(schema);
stateStore.initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());
update(stateStore).initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());

FileReference file1 = writeRootFile(schema, stateStore, dataFolderName + "/file1.parquet", List.of());
FileReference file2 = writeRootFile(schema, stateStore, dataFolderName + "/file2.parquet", List.of());
Expand All @@ -90,7 +91,7 @@ void shouldWriteSketchWhenWritingEmptyFile() throws Exception {
// Given
Schema schema = createSchemaWithTypesForKeyAndTwoValues(new LongType(), new LongType(), new LongType());
tableProperties.setSchema(schema);
stateStore.initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());
update(stateStore).initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());

FileReference file1 = writeRootFile(schema, stateStore, dataFolderName + "/file1.parquet", List.of());
FileReference file2 = writeRootFile(schema, stateStore, dataFolderName + "/file2.parquet", List.of());
Expand Down
9 changes: 5 additions & 4 deletions ...-job-execution/src/test/java/sleeper/compaction/job/execution/JavaCompactionRunnerIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
import static org.assertj.core.api.Assertions.assertThat;
import static sleeper.compaction.job.execution.testutils.CompactionRunnerTestUtils.assignJobIdToInputFiles;
import static sleeper.compaction.job.execution.testutils.CompactionRunnerTestUtils.createSchemaWithTypesForKeyAndTwoValues;
import static sleeper.core.statestore.testutils.StateStoreUpdatesWrapper.update;

class JavaCompactionRunnerIT extends CompactionRunnerTestBase {

Expand All @@ -45,7 +46,7 @@ void shouldMergeFilesWithLongKey() throws Exception {
// Given
Schema schema = createSchemaWithTypesForKeyAndTwoValues(new LongType(), new LongType(), new LongType());
tableProperties.setSchema(schema);
stateStore.initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());
update(stateStore).initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());

List<Record> data1 = CompactionRunnerTestData.keyAndTwoValuesSortedEvenLongs();
List<Record> data2 = CompactionRunnerTestData.keyAndTwoValuesSortedOddLongs();
Expand All @@ -71,7 +72,7 @@ void shouldWriteSketchWhenMergingFiles() throws Exception {
// Given
Schema schema = createSchemaWithTypesForKeyAndTwoValues(new LongType(), new LongType(), new LongType());
tableProperties.setSchema(schema);
stateStore.initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());
update(stateStore).initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());

List<Record> data1 = CompactionRunnerTestData.keyAndTwoValuesSortedEvenLongs();
List<Record> data2 = CompactionRunnerTestData.keyAndTwoValuesSortedOddLongs();
Expand Down Expand Up @@ -120,7 +121,7 @@ void shouldMergeFilesWithStringKey() throws Exception {
// Given
Schema schema = createSchemaWithTypesForKeyAndTwoValues(new StringType(), new StringType(), new LongType());
tableProperties.setSchema(schema);
stateStore.initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());
update(stateStore).initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());

List<Record> data1 = CompactionRunnerTestData.keyAndTwoValuesSortedEvenStrings();
List<Record> data2 = CompactionRunnerTestData.keyAndTwoValuesSortedOddStrings();
Expand Down Expand Up @@ -173,7 +174,7 @@ void shouldMergeFilesWithByteArrayKey() throws Exception {
// Given
Schema schema = createSchemaWithTypesForKeyAndTwoValues(new ByteArrayType(), new ByteArrayType(), new LongType());
tableProperties.setSchema(schema);
stateStore.initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());
update(stateStore).initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());

List<Record> data1 = CompactionRunnerTestData.keyAndTwoValuesSortedEvenByteArrays();
List<Record> data2 = CompactionRunnerTestData.keyAndTwoValuesSortedOddByteArrays();
Expand Down
3 changes: 2 additions & 1 deletion ...cution/src/test/java/sleeper/compaction/job/execution/JavaCompactionRunnerIteratorIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import static sleeper.compaction.job.execution.testutils.CompactionRunnerTestUtils.assignJobIdToInputFiles;
import static sleeper.core.properties.table.TableProperty.ITERATOR_CLASS_NAME;
import static sleeper.core.properties.table.TableProperty.ITERATOR_CONFIG;
import static sleeper.core.statestore.testutils.StateStoreUpdatesWrapper.update;

class JavaCompactionRunnerIteratorIT extends CompactionRunnerTestBase {

Expand All @@ -43,7 +44,7 @@ void shouldApplyIteratorDuringCompaction() throws Exception {
// Given
Schema schema = CompactionRunnerTestUtils.createSchemaWithKeyTimestampValue();
tableProperties.setSchema(schema);
stateStore.initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());
update(stateStore).initialise(new PartitionsBuilder(schema).singlePartition("root").buildList());

List<Record> data1 = CompactionRunnerTestData.specifiedFromEvens((even, record) -> {
record.put("key", (long) even);
Expand Down
3 changes: 2 additions & 1 deletion ...tion/src/test/java/sleeper/compaction/job/execution/JavaCompactionRunnerLocalStackIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
import static sleeper.core.properties.instance.CommonProperty.ID;
import static sleeper.core.properties.instance.TableDefaultProperty.DEFAULT_INGEST_PARTITION_FILE_WRITER_TYPE;
import static sleeper.core.properties.table.TableProperty.GARBAGE_COLLECTOR_DELAY_BEFORE_DELETION;
import static sleeper.core.statestore.testutils.StateStoreUpdatesWrapper.update;

public class JavaCompactionRunnerLocalStackIT extends CompactionRunnerTestBase {

Expand All @@ -74,7 +75,7 @@ public void shouldRunCompactionJob() throws Exception {
tableProperties.setSchema(schema);
createStateStore();
PartitionTree tree = new PartitionsBuilder(schema).singlePartition("root").buildTree();
stateStore.initialise(tree.getAllPartitions());
update(stateStore).initialise(tree.getAllPartitions());

List<Record> data1 = CompactionRunnerTestData.keyAndTwoValuesSortedEvenLongs();
List<Record> data2 = CompactionRunnerTestData.keyAndTwoValuesSortedOddLongs();
Expand Down
Loading

0 comments on commit bb3f6d2

Please sign in to comment.