DataflowRunner extends PipelineRunner<DataflowPipelineJob> { @VisibleForTesting static String getContainerImageForJob(DataflowPipelineOptions options) { String workerHarnessContainerImage = options.getWorkerHarnessContainerImage(); Environments.JavaVersion javaVersion = Environments.getJavaVersion(); String javaVersionId = (javaVersion == Environments.JavaVersion.v8) ? "java" : javaVersion.toString(); if (!workerHarnessContainerImage.contains("IMAGE")) { return workerHarnessContainerImage; } else if (hasExperiment(options, "beam_fn_api")) { return workerHarnessContainerImage.replace("IMAGE", "java"); } else if (options.isStreaming()) { return workerHarnessContainerImage.replace( "IMAGE", String.format("beam-%s-streaming", javaVersionId)); } else { return workerHarnessContainerImage.replace( "IMAGE", String.format("beam-%s-batch", javaVersionId)); } } @VisibleForTesting protected  DataflowRunner(DataflowPipelineOptions options); static DataflowRunner fromOptions(PipelineOptions options); @Override DataflowPipelineJob run(Pipeline pipeline); static boolean hasExperiment(DataflowPipelineDebugOptions options, String experiment); DataflowPipelineTranslator getTranslator(); @Experimental void setHooks(DataflowRunnerHooks hooks); @Override String toString(); static final String PROJECT_ID_REGEXP; }
@Test public void testWorkerHarnessContainerImage() { DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); options.setWorkerHarnessContainerImage("some-container"); assertThat(getContainerImageForJob(options), equalTo("some-container")); options.setWorkerHarnessContainerImage("gcr.io/IMAGE/foo"); options.setExperiments(null); options.setStreaming(false); System.setProperty("java.specification.version", "1.8"); assertThat(getContainerImageForJob(options), equalTo("gcr.io/beam-java-batch/foo")); options.setStreaming(false); System.setProperty("java.specification.version", "11"); assertThat(getContainerImageForJob(options), equalTo("gcr.io/beam-java11-batch/foo")); System.setProperty("java.specification.version", "1.8"); options.setStreaming(true); assertThat(getContainerImageForJob(options), equalTo("gcr.io/beam-java-streaming/foo")); System.setProperty("java.specification.version", "11"); assertThat(getContainerImageForJob(options), equalTo("gcr.io/beam-java11-streaming/foo")); options.setExperiments(ImmutableList.of("experiment1", "beam_fn_api")); assertThat(getContainerImageForJob(options), equalTo("gcr.io/java/foo")); }