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))