def test_pack_combiners_with_missing_environment_capability(self):
class MultipleCombines(beam.PTransform):
def expand(self, pcoll):
_ = pcoll | 'mean-perkey' >> combiners.Mean.PerKey()
_ = pcoll | 'count-perkey' >> combiners.Count.PerKey()
_ = pcoll | 'largest-perkey' >> core.CombinePerKey(combiners.Largest(1))
pipeline = beam.Pipeline()
vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]
_ = pipeline | Create([('a', x) for x in vals]) | MultipleCombines()
environment = environments.DockerEnvironment(capabilities=())
pipeline_proto = pipeline.to_runner_api(default_environment=environment)
_, stages = translations.create_and_optimize_stages(
pipeline_proto, [translations.pack_combiners],
known_runner_urns=frozenset())
combine_per_key_stages = []
for stage in stages:
for transform in stage.transforms:
if transform.spec.urn == common_urns.composites.COMBINE_PER_KEY.urn:
combine_per_key_stages.append(stage)
# Combiner packing should be skipped because the environment is missing
# the beam:combinefn:packed_python:v1 capability.
self.assertEqual(len(combine_per_key_stages), 3)
for combine_per_key_stage in combine_per_key_stages:
self.assertNotIn('Packed', combine_per_key_stage.name)
self.assertNotIn(
'Packed', combine_per_key_stage.transforms[0].unique_name)
def test__create_default_environment(self):
docker_image = environments.DockerEnvironment.default_docker_image()
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({'sdk_location':
'container'})),
environments.DockerEnvironment(container_image=docker_image))
def test__create_docker_environment(self):
docker_image = 'py-docker'
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type': 'DOCKER',
'environment_config': docker_image,
})),
environments.DockerEnvironment(container_image=docker_image))
