def test_backend_stage_name_must_be_unique(self):
backend_1 = BeamBackend("SAME_backend_SUFFIX")
backend_2 = BeamBackend("SAME_backend_SUFFIX")
with test_pipeline.TestPipeline() as p:
col = p | f"UNIQUE_BEAM_CREATE_NAME" >> beam.Create([(6, 1),
(6, 2)])
backend_1.map(col, lambda x: x, "SAME_MAP_NAME")
with self.assertRaisesRegex(
RuntimeError,
expected_regex="A transform with label "
"\"SAME_MAP_NAME_SAME_backend_SUFFIX\" "
"already exists in the pipeline"):
backend_2.map(col, lambda x: x, "SAME_MAP_NAME")
def test_one_suffix_multiple_same_stage_name(self):
backend = BeamBackend("UNIQUE_BACKEND_SUFFIX")
with test_pipeline.TestPipeline() as p:
col = p | f"UNIQUE_BEAM_CREATE_NAME" >> beam.Create([(6, 1),
(6, 2)])
backend.map(col, lambda x: x, "SAME_MAP_NAME")
backend.map(col, lambda x: x, "SAME_MAP_NAME")
backend.map(col, lambda x: x, "SAME_MAP_NAME")
self.assertEqual("UNIQUE_BACKEND_SUFFIX", backend._ulg._suffix)
self.assertEqual(3, len(backend._ulg._labels))
self.assertIn("SAME_MAP_NAME_UNIQUE_BACKEND_SUFFIX",
backend._ulg._labels)
self.assertIn("SAME_MAP_NAME_1_UNIQUE_BACKEND_SUFFIX",
backend._ulg._labels)
self.assertIn("SAME_MAP_NAME_2_UNIQUE_BACKEND_SUFFIX",
backend._ulg._labels)
def setUpClass(cls):
cls.backend = BeamBackend()
cls.data_extractors = DataExtractors(
partition_extractor=lambda x: x[1],
privacy_id_extractor=lambda x: x[0],
value_extractor=lambda x: x[2])
def _test_helper():
mock_pcollection = BeamBackendStageNameTest._create_mock_pcollection()
backend = BeamBackend()
backend._ulg = BeamBackendStageNameTest.MockUniqueLabelGenerators()
return mock_pcollection, backend