def test_reuse_cloned_custom_transform_instance(self):
pipeline = Pipeline(DirectPipelineRunner())
pcoll1 = pipeline | Create('pcoll1', [1, 2, 3])
pcoll2 = pipeline | Create('pcoll2', [4, 5, 6])
transform = PipelineTest.CustomTransform()
result1 = pcoll1 | transform
result2 = pcoll2 | transform.clone('new label')
assert_that(result1, equal_to([2, 3, 4]), label='r1')
assert_that(result2, equal_to([5, 6, 7]), label='r2')
pipeline.run()
def test_reuse_custom_transform_instance(self):
pipeline = Pipeline(DirectPipelineRunner())
pcoll1 = pipeline | Create('pcoll1', [1, 2, 3])
pcoll2 = pipeline | Create('pcoll2', [4, 5, 6])
transform = PipelineTest.CustomTransform()
pcoll1 | transform
with self.assertRaises(RuntimeError) as cm:
pipeline.apply(transform, pcoll2)
self.assertEqual(
cm.exception.message,
'Transform "CustomTransform" does not have a stable unique label. '
'This will prevent updating of pipelines. '
'To clone a transform with a new label use: '
'transform.clone("NEW LABEL").')
def test_visit_entire_graph(self):
pipeline = Pipeline(DirectPipelineRunner())
pcoll1 = pipeline | Create('pcoll', [1, 2, 3])
pcoll2 = pcoll1 | FlatMap('do1', lambda x: [x + 1])
pcoll3 = pcoll2 | FlatMap('do2', lambda x: [x + 1])
pcoll4 = pcoll2 | FlatMap('do3', lambda x: [x + 1])
transform = PipelineTest.CustomTransform()
pcoll5 = pcoll4 | transform
visitor = PipelineTest.Visitor(visited=[])
pipeline.visit(visitor)
self.assertEqual(set([pcoll1, pcoll2, pcoll3, pcoll4, pcoll5]),
set(visitor.visited))
self.assertEqual(set(visitor.enter_composite),
set(visitor.leave_composite))
self.assertEqual(2, len(visitor.enter_composite))
self.assertEqual(visitor.enter_composite[1].transform, transform)
self.assertEqual(visitor.leave_composite[0].transform, transform)
def test_create_singleton_pcollection(self):
pipeline = Pipeline(DirectPipelineRunner())
pcoll = pipeline | Create('label', [[1, 2, 3]])
assert_that(pcoll, equal_to([[1, 2, 3]]))
pipeline.run()
def test_apply_custom_transform(self):
pipeline = Pipeline(DirectPipelineRunner())
pcoll = pipeline | Create('pcoll', [1, 2, 3])
result = pcoll | PipelineTest.CustomTransform()
assert_that(result, equal_to([2, 3, 4]))
pipeline.run()