content: /*
   * 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.
   */
  package org.apache.beam.runners.flink;
  
  import static org.hamcrest.CoreMatchers.allOf;
  import static org.junit.Assert.assertThrows;
  
  import org.apache.beam.sdk.Pipeline;
  import org.apache.beam.sdk.PipelineResult;
  import org.apache.beam.sdk.io.GenerateSequence;
  import org.apache.beam.sdk.options.PipelineOptions;
  import org.apache.flink.client.program.OptimizerPlanEnvironment;
  import org.apache.flink.client.program.PackagedProgram;
  import org.apache.flink.client.program.PackagedProgramUtils;
  import org.apache.flink.client.program.ProgramInvocationException;
  import org.apache.flink.configuration.Configuration;
  import org.hamcrest.MatcherAssert;
  import org.hamcrest.core.StringContains;
  import org.junit.Test;
  
  /** Test for {@link FlinkRunner}. */
  public class FlinkRunnerTest {
  
    @Test
    public void testEnsureStdoutStdErrIsRestored() throws Exception {
      PackagedProgram packagedProgram =
          PackagedProgram.newBuilder().setEntryPointClassName(getClass().getName()).build();
      int parallelism = Runtime.getRuntime().availableProcessors();
      OptimizerPlanEnvironment env =
          new OptimizerPlanEnvironment(new Configuration(), getClass().getClassLoader(), parallelism);
      Exception e =
          assertThrows(
              ProgramInvocationException.class,
              () -> {
                // Flink will throw an error because no job graph will be generated by the main method
                PackagedProgramUtils.getPipelineFromProgram(
                    packagedProgram, env.getConfiguration(), parallelism, true);
              });
      // Test that Flink wasn't able to intercept the stdout/stderr and we printed to the regular
      // output instead
      MatcherAssert.assertThat(
          e.getMessage(),
          allOf(
              StringContains.containsString("System.out: "),
              StringContains.containsString("System.err: ")));
    }
  
    /** Main method for {@code testEnsureStdoutStdErrIsRestored()}. */
    public static void main(String[] args) {
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setRunner(NotExecutingFlinkRunner.class);
      Pipeline p = Pipeline.create(options);
      p.apply(GenerateSequence.from(0));
  
      // This will call Workarounds.restoreOriginalStdOutAndStdErr() through the constructor of
      // FlinkRunner
      p.run();
    }
  
    private static class NotExecutingFlinkRunner extends FlinkRunner {
  
      protected NotExecutingFlinkRunner(FlinkPipelineOptions options) {
        // Stdout/Stderr is restored here
        super(options);
      }
  
      @SuppressWarnings("unused")
      public static NotExecutingFlinkRunner fromOptions(PipelineOptions options) {
        return new NotExecutingFlinkRunner(options.as(FlinkPipelineOptions.class));
      }
  
      @Override
      public PipelineResult run(Pipeline pipeline) {
        // Do not execute to test the stdout printing
        return null;
      }
    }
  }
  