diff --git a/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java b/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java index bdc76824f5be..773b4485a8cd 100644 --- a/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java +++ b/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java @@ -34,7 +34,11 @@ import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Count; import org.apache.beam.sdk.transforms.Create; -import org.apache.beam.sdk.transforms.DoFnTester; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Values; +import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.BaseEncoding; import org.apache.solr.SolrTestCaseJ4; @@ -205,46 +209,43 @@ public void testWrite() throws Exception { assertEquals(NUM_DOCS / NUM_SCIENTISTS, response.getResults().getNumFound()); } + private static class EmitDocumentsFn + extends DoFn, SolrInputDocument> { + @ProcessElement + public void processElement(ProcessContext c) { + for (SolrInputDocument doc : c.element()) { + c.output(doc); + } + } + } + @Test public void testWriteWithMaxBatchSize() throws Exception { + int maxBatchSize = 1; + List data = SolrIOTestUtils.createDocuments(NUM_DOCS); SolrIO.Write write = SolrIO.write() .withConnectionConfiguration(connectionConfiguration) .to(SOLR_COLLECTION) - .withMaxBatchSize(BATCH_SIZE); - // write bundles size is the runner decision, we cannot force a bundle size, - // so we test the Writer as a DoFn outside of a runner. - try (DoFnTester fnTester = - DoFnTester.of(new SolrIO.Write.WriteFn(write))) { - List input = SolrIOTestUtils.createDocuments(NUM_DOCS); - long numDocsProcessed = 0; - long numDocsInserted = 0; - for (SolrInputDocument document : input) { - fnTester.processElement(document); - numDocsProcessed++; - // test every 100 docs to avoid overloading Solr - if ((numDocsProcessed % 100) == 0) { - // force the index to upgrade after inserting for the inserted docs - // to be searchable immediately - long currentNumDocs = - SolrIOTestUtils.commitAndGetCurrentNumDocs(SOLR_COLLECTION, solrClient); - if ((numDocsProcessed % BATCH_SIZE) == 0) { - /* bundle end */ - assertEquals( - "we are at the end of a bundle, we should have inserted all processed documents", - numDocsProcessed, - currentNumDocs); - numDocsInserted = currentNumDocs; - } else { - /* not bundle end */ - assertEquals( - "we are not at the end of a bundle, we should have inserted no more documents", - numDocsInserted, - currentNumDocs); - } - } - } - } + .withMaxBatchSize(maxBatchSize); + + // Force all elements into a single bundle via GroupByKey and use a maxBatchSize of 1 + // to guarantee that batch.size() >= spec.getMaxBatchSize() is triggered and flushes + // mid-bundle during @ProcessElement, rather than only flushing at @FinishBundle. + pipeline + .apply(Create.of(data)) + .apply(WithKeys.of(0)) + .apply(GroupByKey.create()) + .apply(Values.create()) + .apply(ParDo.of(new EmitDocumentsFn())) + .apply(write); + pipeline.run(); + + long currentNumDocs = SolrIOTestUtils.commitAndGetCurrentNumDocs(SOLR_COLLECTION, solrClient); + assertEquals(NUM_DOCS, currentNumDocs); + + QueryResponse response = solrClient.query(SOLR_COLLECTION, new SolrQuery("scientist:Lovelace")); + assertEquals(NUM_DOCS / NUM_SCIENTISTS, response.getResults().getNumFound()); } /**