This is an automated email from the ASF dual-hosted git repository.
kennknowles pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new b7f6a6b1d2e [#18690] Port SolrIOTest off deprecated DoFnTester (#40819)
b7f6a6b1d2e is described below
commit b7f6a6b1d2e926ac0fbf049ccde4f987f41382c5
Author: Anavay Sharma <[email protected]>
AuthorDate: Thu Oct 29 00:19:20 2026 +0530
[#18690] Port SolrIOTest off deprecated DoFnTester (#40819)
---
.../org/apache/beam/sdk/io/solr/SolrIOTest.java | 71 +++++++++++-----------
1 file changed, 36 insertions(+), 35 deletions(-)
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 bdc76824f5b..773b4485a8c 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.PAssert;
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 class SolrIOTest extends SolrCloudTestCase {
assertEquals(NUM_DOCS / NUM_SCIENTISTS,
response.getResults().getNumFound());
}
+ private static class EmitDocumentsFn
+ extends DoFn<Iterable<SolrInputDocument>, 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<SolrInputDocument> 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<SolrInputDocument, Void> fnTester =
- DoFnTester.of(new SolrIO.Write.WriteFn(write))) {
- List<SolrInputDocument> 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());
}
/**
