[beam] branch release-2.21.0 updated: [BEAM-9738] Update dataflow to setup correct docker environment options. (#11371)

[ib...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

ibzib pushed a commit to branch release-2.21.0
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/release-2.21.0 by this push:
     new fb58aa6  [BEAM-9738] Update dataflow to setup correct docker environment options. (#11371)
     new 3525205  Merge pull request #11390 from robertwb/cherry-pick-beam-9738
fb58aa6 is described below

commit fb58aa6c2b88297186e34b65445843fa433179ab
Author: ananvay <an...@google.com>
AuthorDate: Mon Apr 13 10:39:05 2020 -0700

    [BEAM-9738] Update dataflow to setup correct docker environment options. (#11371)
---
 sdks/python/apache_beam/runners/dataflow/dataflow_runner.py      | 4 ++--
 sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py | 4 +++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
index f92c7ca..5091d55 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
@@ -456,8 +456,8 @@ class DataflowRunner(PipelineRunner):
 
     use_fnapi = apiclient._use_fnapi(options)
     from apache_beam.transforms import environments
-    default_environment = environments.DockerEnvironment(
-        container_image=apiclient.get_container_image_from_options(options))
+    default_environment = environments.DockerEnvironment.from_container_image(
+        apiclient.get_container_image_from_options(options))
 
     # Snapshot the pipeline in a portable proto.
     self.proto_pipeline, self.proto_context = pipeline.to_runner_api(
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
index 6c3543c..c4f065c 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
@@ -51,6 +51,7 @@ from apache_beam.runners.dataflow.internal.clients import dataflow as dataflow_a
 from apache_beam.runners.runner import PipelineState
 from apache_beam.testing.extra_assertions import ExtraAssertionsMixin
 from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.transforms import environments
 from apache_beam.transforms import window
 from apache_beam.transforms.core import Windowing
 from apache_beam.transforms.core import _GroupByKeyOnly
@@ -219,7 +220,8 @@ class DataflowRunnerTest(unittest.TestCase, ExtraAssertionsMixin):
             beam_runner_api_pb2.Environment(
                 urn=common_urns.environments.DOCKER.urn,
                 payload=beam_runner_api_pb2.DockerPayload(
-                    container_image='FOO').SerializeToString())
+                    container_image='FOO').SerializeToString(),
+                capabilities=environments.python_sdk_capabilities())
         ])
 
   def test_remote_runner_translation(self):