Examples with DataflowRunner.getContainerImageForJob org.apache.beam.runners.dataflow.DataflowRunner.getContainerImageForJob used on opensource projects

Example 1 with DataflowRunner.getContainerImageForJob

use of org.apache.beam.runners.dataflow.DataflowRunner.getContainerImageForJob in project beam by apache.

the class DataflowRunnerTest method testSdkHarnessConfiguration.

@Test
public void testSdkHarnessConfiguration() throws IOException {
    DataflowPipelineOptions options = buildPipelineOptions();
    ExperimentalOptions.addExperiment(options, "use_runner_v2");
    Pipeline p = Pipeline.create(options);
    p.apply(Create.of(Arrays.asList(1, 2, 3)));
    String defaultSdkContainerImage = DataflowRunner.getContainerImageForJob(options);
    SdkComponents sdkComponents = SdkComponents.create();
    RunnerApi.Environment defaultEnvironmentForDataflow = Environments.createDockerEnvironment(defaultSdkContainerImage);
    sdkComponents.registerEnvironment(defaultEnvironmentForDataflow.toBuilder().build());
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);
    Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, pipelineProto, sdkComponents, DataflowRunner.fromOptions(options), Collections.emptyList()).getJob();
    DataflowRunner.configureSdkHarnessContainerImages(options, pipelineProto, job);
    List<SdkHarnessContainerImage> sdks = job.getEnvironment().getWorkerPools().get(0).getSdkHarnessContainerImages();
    Map<String, String> expectedEnvIdsAndContainerImages = pipelineProto.getComponents().getEnvironmentsMap().entrySet().stream().filter(x -> BeamUrns.getUrn(RunnerApi.StandardEnvironments.Environments.DOCKER).equals(x.getValue().getUrn())).collect(Collectors.toMap(x -> x.getKey(), x -> {
        RunnerApi.DockerPayload payload;
        try {
            payload = RunnerApi.DockerPayload.parseFrom(x.getValue().getPayload());
        } catch (InvalidProtocolBufferException e) {
            throw new RuntimeException(e);
        }
        return payload.getContainerImage();
    }));
    assertEquals(1, expectedEnvIdsAndContainerImages.size());
    assertEquals(1, sdks.size());
    assertEquals(expectedEnvIdsAndContainerImages, sdks.stream().collect(Collectors.toMap(SdkHarnessContainerImage::getEnvironmentId, SdkHarnessContainerImage::getContainerImage)));
}