def test_default_capabilities(self):
environment = DockerEnvironment.from_options(
PortableOptions(sdk_location='container'))
context = pipeline_context.PipelineContext()
proto = environment.to_runner_api(context)
self.assertEqual(set(proto.capabilities),
set(environments.python_sdk_docker_capabilities()))
def test_environment_override_translation_sdk_container_image(self):
self.default_properties.append('--experiments=beam_fn_api')
self.default_properties.append('--sdk_container_image=FOO')
remote_runner = DataflowRunner()
with Pipeline(remote_runner,
options=PipelineOptions(self.default_properties)) as p:
( # pylint: disable=expression-not-assigned
p | ptransform.Create([1, 2, 3])
| 'Do' >> ptransform.FlatMap(lambda x: [(x, x)])
| ptransform.GroupByKey())
self.assertEqual(
list(
remote_runner.proto_pipeline.components.environments.values()),
[
beam_runner_api_pb2.Environment(
urn=common_urns.environments.DOCKER.urn,
payload=beam_runner_api_pb2.DockerPayload(
container_image='FOO').SerializeToString(),
capabilities=environments.python_sdk_docker_capabilities())
])