JUnit5 support

CHANGES.md

@@ -74,6 +74,8 @@
 
 * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 * BigtableRead Connector for BeamYaml added with new Config Param ([#35696](https://github.com/apache/beam/pull/35696))
+* Introduced a dedicated module for JUnit-based testing support: `sdks/java/testing/junit`, which provides `TestPipelineExtension` for JUnit 5 while maintaining backward compatibility with existing JUnit 4 `TestRule`-based tests (Java) ([#18733](https://github.com/apache/beam/issues/18733), [#35688](https://github.com/apache/beam/pull/35688)).
+  - To use JUnit 5 with Beam tests, add a test-scoped dependency on `org.apache.beam:beam-sdks-java-testing-junit`.
 
 ## Breaking Changes
 

sdks/java/core/build.gradle

@@ -130,3 +130,11 @@ project.tasks.compileTestJava {
   // TODO: fix other places with warnings in tests and delete this option
   options.compilerArgs += ['-Xlint:-rawtypes']
 }
+
+// Configure test task to use JUnit 4. JUnit 5 support is provided in module
+// sdks/java/testing/junit, which configures useJUnitPlatform(). Submodules that
+// need to run both JUnit 4 and 5 via the JUnit Platform must also add the
+// Vintage engine explicitly.
+test {
+  useJUnit()
+}

sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestPipeline.java

@@ -82,7 +82,11 @@
  * </ul>
  *
  * <p>Use {@link PAssert} for tests, as it integrates with this test harness in both direct and
- * remote execution modes. For example:
+ * remote execution modes.
+ *
+ * <h3>JUnit 4 Usage</h3>
+ *
+ * For JUnit 4 tests, use this class as a TestRule:
  *
  * <pre><code>
  * {@literal @Rule}
@@ -97,6 +101,25 @@
  *  }
  * </code></pre>
  *
+ * <h3>JUnit5 Usage</h3>
+ *
+ * For JUnit5 tests, use {@link TestPipelineExtension} from the module <code>
+ * sdks/java/testing/junit</code> (artifact <code>org.apache.beam:beam-sdks-java-testing-junit
+ * </code>):
+ *
+ * <pre><code>
+ * {@literal @ExtendWith}(TestPipelineExtension.class)
+ * class MyPipelineTest {
+ *   {@literal @Test}
+ *   {@literal @Category}(NeedsRunner.class)
+ *   void myPipelineTest(TestPipeline pipeline) {
+ *     final PCollection&lt;String&gt; pCollection = pipeline.apply(...)
+ *     PAssert.that(pCollection).containsInAnyOrder(...);
+ *     pipeline.run();
+ *   }
+ * }
+ * </code></pre>
+ *
  * <p>For pipeline runners, it is required that they must throw an {@link AssertionError} containing
  * the message from the {@link PAssert} that failed.
  *
@@ -108,7 +131,7 @@ public class TestPipeline extends Pipeline implements TestRule {
 
   private final PipelineOptions options;
 
-  private static class PipelineRunEnforcement {
+  static class PipelineRunEnforcement {
 
     @SuppressWarnings("WeakerAccess")
     protected boolean enableAutoRunIfMissing;
@@ -117,7 +140,7 @@ private static class PipelineRunEnforcement {
 
     protected boolean runAttempted;
 
-    private PipelineRunEnforcement(final Pipeline pipeline) {
+    PipelineRunEnforcement(final Pipeline pipeline) {
       this.pipeline = pipeline;
     }
 
@@ -138,7 +161,7 @@ protected void afterUserCodeFinished() {
     }
   }
 
-  private static class PipelineAbandonedNodeEnforcement extends PipelineRunEnforcement {
+  static class PipelineAbandonedNodeEnforcement extends PipelineRunEnforcement {
 
     // Null until the pipeline has been run
     private @MonotonicNonNull List<TransformHierarchy.Node> runVisitedNodes;
@@ -164,7 +187,7 @@ public void visitPrimitiveTransform(final TransformHierarchy.Node node) {
       }
     }
 
-    private PipelineAbandonedNodeEnforcement(final TestPipeline pipeline) {
+    PipelineAbandonedNodeEnforcement(final TestPipeline pipeline) {
       super(pipeline);
       runVisitedNodes = null;
     }
@@ -574,7 +597,7 @@ public static void verifyPAssertsSucceeded(Pipeline pipeline, PipelineResult pip
     }
   }
 
-  private static class IsEmptyVisitor extends PipelineVisitor.Defaults {
+  static class IsEmptyVisitor extends PipelineVisitor.Defaults {
     private boolean empty = true;
 
     public boolean isEmpty() {

sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOWriteTest.java

@@ -731,11 +731,20 @@ public void testWriteUnboundedWithCustomBatchParameters() throws Exception {
     input.apply(write);
     p.run();
 
+    // On some environments/runners, the exact shard filenames may not be materialized
+    // deterministically by the time we assert. Verify shard count via a glob, then
+    // validate contents using pattern matching.
+    String pattern = baseFilename.toString() + "*";
+    List<MatchResult> matches = FileSystems.match(Collections.singletonList(pattern));
+    List<Metadata> found = new ArrayList<>(Iterables.getOnlyElement(matches).metadata());
+    assertEquals(3, found.size());
+
+    // Now assert file contents irrespective of exact shard indices.
     assertOutputFiles(
         LINES2_ARRAY,
         null,
         null,
-        3,
+        0, // match all files by prefix
         baseFilename,
         DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE,
         false);

sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ExecutionStateSamplerTest.java

@@ -809,8 +809,12 @@ public Long answer(InvocationOnMock invocation) throws Throwable {
                   // and unblock the state transition once a certain number of samples
                   // have been taken.
                   waitTillActive.await();
-                  waitForSamples.countDown();
-                  currentTime += Duration.standardMinutes(1).getMillis();
+                  // Freeze time after the desired number of samples to avoid races where
+                  // the sampling loop spins and exceeds the timeout before we deactivate.
+                  if (waitForSamples.getCount() > 0) {
+                    waitForSamples.countDown();
+                    currentTime += Duration.standardMinutes(1).getMillis();
+                  }
                   return currentTime;
                 }
               }

sdks/java/testing/junit/build.gradle

@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+plugins { id 'org.apache.beam.module' }
+
+applyJavaNature(
+    exportJavadoc: false,
+    automaticModuleName: 'org.apache.beam.sdk.testing.junit',
+    archivesBaseName: 'beam-sdks-java-testing-junit'
+)
+
+description = "Apache Beam :: SDKs :: Java :: Testing :: JUnit"
+
+dependencies {
+  implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom)
+  implementation project(path: ":sdks:java:core", configuration: "shadow")
+  implementation library.java.vendored_guava_32_1_2_jre
+  // Needed to resolve TestPipeline's JUnit 4 TestRule type and @Category at compile time,
+  // but should not leak to consumers at runtime.
+  provided library.java.junit
+
+  // JUnit 5 API needed to compile the extension; not packaged for consumers of core.
+  provided library.java.jupiter_api
+
+  testImplementation project(path: ":sdks:java:core", configuration: "shadow")
+  testImplementation library.java.jupiter_api
+  testImplementation library.java.junit
+  testRuntimeOnly library.java.jupiter_engine
+  testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")
+}
+
+// This module runs JUnit 5 tests using the JUnit Platform.
+test {
+  useJUnitPlatform()
+}