1717 */
1818package org .apache .beam .sdk .testing ;
1919
20- import static org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .base .Preconditions .checkState ;
20+ import static org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .base .Preconditions .checkNotNull ;
2121
2222import java .lang .annotation .Annotation ;
2323import java .lang .reflect .Method ;
24- import java .util .Arrays ;
25- import java .util .Optional ;
24+ import java .util .Collection ;
2625import org .apache .beam .sdk .options .ApplicationNameOptions ;
2726import org .apache .beam .sdk .options .PipelineOptions ;
28- import org .apache .beam .sdk .testing .TestPipeline .PipelineAbandonedNodeEnforcement ;
29- import org .apache .beam .sdk .testing .TestPipeline .PipelineRunEnforcement ;
30- import org .junit .experimental .categories .Category ;
27+ import org .apache .beam .vendor .grpc .v1p69p0 .com .google .common .collect .ImmutableList ;
28+ import org .checkerframework .checker .nullness .qual .Nullable ;
3129import org .junit .jupiter .api .extension .AfterEachCallback ;
3230import org .junit .jupiter .api .extension .BeforeEachCallback ;
3331import org .junit .jupiter .api .extension .ExtensionContext ;
@@ -86,16 +84,16 @@ public static TestPipelineExtension fromOptions(PipelineOptions options) {
8684 return new TestPipelineExtension (options );
8785 }
8886
89- private TestPipeline testPipeline ;
87+ private @ Nullable PipelineOptions options ;
9088
9189 /** Creates a TestPipelineExtension with default options. */
9290 public TestPipelineExtension () {
93- this .testPipeline = TestPipeline . create () ;
91+ this .options = null ;
9492 }
9593
9694 /** Creates a TestPipelineExtension with custom options. */
9795 public TestPipelineExtension (PipelineOptions options ) {
98- this .testPipeline = TestPipeline . fromOptions ( options ) ;
96+ this .options = options ;
9997 }
10098
10199 @ Override
@@ -107,52 +105,38 @@ public boolean supportsParameter(
107105 @ Override
108106 public Object resolveParameter (
109107 ParameterContext parameterContext , ExtensionContext extensionContext ) {
110- if (this .testPipeline == null ) {
111- return getOrCreateTestPipeline (extensionContext );
112- } else {
113- return this .testPipeline ;
114- }
108+ return getOrCreateTestPipeline (extensionContext );
115109 }
116110
117111 @ Override
118- public void beforeEach (ExtensionContext context ) throws Exception {
119- TestPipeline pipeline ;
120-
121- if (this .testPipeline != null ) {
122- pipeline = this .testPipeline ;
123- } else {
124- pipeline = getOrCreateTestPipeline (context );
125- }
112+ public void beforeEach (ExtensionContext context ) {
113+ TestPipeline pipeline = getOrCreateTestPipeline (context );
126114
127115 // Set application name based on test method
128116 String appName = getAppName (context );
129117 pipeline .getOptions ().as (ApplicationNameOptions .class ).setAppName (appName );
130118
131119 // Set up enforcement based on annotations
132- setDeducedEnforcementLevel (context , pipeline );
120+ pipeline . setDeducedEnforcementLevel (getAnnotations ( context ) );
133121 }
134122
135123 @ Override
136- public void afterEach (ExtensionContext context ) throws Exception {
137- Optional <PipelineRunEnforcement > enforcement = getEnforcement (context );
138- if (enforcement .isPresent ()) {
139- enforcement .get ().afterUserCodeFinished ();
140- }
124+ public void afterEach (ExtensionContext context ) {
125+ TestPipeline pipeline = getRequiredTestPipeline (context );
126+ pipeline .afterUserCodeFinished ();
141127 }
142128
143129 private TestPipeline getOrCreateTestPipeline (ExtensionContext context ) {
144130 return context
145131 .getStore (NAMESPACE )
146- .getOrComputeIfAbsent (PIPELINE_KEY , key -> TestPipeline .create (), TestPipeline .class );
147- }
148-
149- private Optional <PipelineRunEnforcement > getEnforcement (ExtensionContext context ) {
150- return Optional .ofNullable (
151- context .getStore (NAMESPACE ).get (ENFORCEMENT_KEY , PipelineRunEnforcement .class ));
132+ .getOrComputeIfAbsent (
133+ PIPELINE_KEY ,
134+ key -> options == null ? TestPipeline .create () : TestPipeline .fromOptions (options ),
135+ TestPipeline .class );
152136 }
153137
154- private void setEnforcement (ExtensionContext context , PipelineRunEnforcement enforcement ) {
155- context .getStore (NAMESPACE ).put ( ENFORCEMENT_KEY , enforcement );
138+ private TestPipeline getRequiredTestPipeline (ExtensionContext context ) {
139+ return checkNotNull ( context .getStore (NAMESPACE ).get ( PIPELINE_KEY , TestPipeline . class ) );
156140 }
157141
158142 private String getAppName (ExtensionContext context ) {
@@ -161,53 +145,10 @@ private String getAppName(ExtensionContext context) {
161145 return className + "-" + methodName ;
162146 }
163147
164- private void setDeducedEnforcementLevel (ExtensionContext context , TestPipeline pipeline ) {
165- // If enforcement level has not been set, do auto-inference
166- if (!getEnforcement (context ).isPresent ()) {
167- boolean annotatedWithNeedsRunner = hasNeedsRunnerAnnotation (context );
168-
169- PipelineOptions options = pipeline .getOptions ();
170- boolean crashingRunner = CrashingRunner .class .isAssignableFrom (options .getRunner ());
171-
172- checkState (
173- !(annotatedWithNeedsRunner && crashingRunner ),
174- "The test was annotated with a [@%s] / [@%s] while the runner "
175- + "was set to [%s]. Please re-check your configuration." ,
176- NeedsRunner .class .getSimpleName (),
177- ValidatesRunner .class .getSimpleName (),
178- CrashingRunner .class .getSimpleName ());
179-
180- if (annotatedWithNeedsRunner || !crashingRunner ) {
181- setEnforcement (context , new PipelineAbandonedNodeEnforcement (pipeline ));
182- }
183- }
184- }
185-
186- private boolean hasNeedsRunnerAnnotation (ExtensionContext context ) {
187- // Check method annotations
188- Method testMethod = context .getTestMethod ().orElse (null );
189- if (testMethod != null ) {
190- if (hasNeedsRunnerCategory (testMethod .getAnnotations ())) {
191- return true ;
192- }
193- }
194-
195- // Check class annotations
196- Class <?> testClass = context .getTestClass ().orElse (null );
197- if (testClass != null ) {
198- if (hasNeedsRunnerCategory (testClass .getAnnotations ())) {
199- return true ;
200- }
201- }
202-
203- return false ;
204- }
205-
206- private boolean hasNeedsRunnerCategory (Annotation [] annotations ) {
207- return Arrays .stream (annotations )
208- .filter (annotation -> annotation instanceof Category )
209- .map (annotation -> (Category ) annotation )
210- .flatMap (category -> Arrays .stream (category .value ()))
211- .anyMatch (categoryClass -> NeedsRunner .class .isAssignableFrom (categoryClass ));
148+ private static Collection <Annotation > getAnnotations (ExtensionContext context ) {
149+ ImmutableList .Builder <Annotation > builder = ImmutableList .builder ();
150+ context .getTestMethod ().ifPresent (testMethod -> builder .add (testMethod .getAnnotations ()));
151+ context .getTestClass ().ifPresent (testClass -> builder .add (testClass .getAnnotations ()));
152+ return builder .build ();
212153 }
213154}
0 commit comments