Repository: beam Updated Branches: refs/heads/master 87a12af6c -> 4d6f6a10f
[BEAM-1316] Remove the usage of mock from ptransform tests Project: http://git-wip-us.apache.org/repos/asf/beam/repo Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/ab55ef3b Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/ab55ef3b Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/ab55ef3b Branch: refs/heads/master Commit: ab55ef3b4248789aa4c9eaf4a6dab7262d673819 Parents: 87a12af Author: Sourabh Bajaj <sourabhba...@google.com> Authored: Tue May 2 11:48:19 2017 -0700 Committer: Ahmet Altay <al...@google.com> Committed: Tue May 2 13:59:09 2017 -0700 ---------------------------------------------------------------------- .../apache_beam/transforms/ptransform_test.py | 60 +++++++++++++------- 1 file changed, 40 insertions(+), 20 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/beam/blob/ab55ef3b/sdks/python/apache_beam/transforms/ptransform_test.py ---------------------------------------------------------------------- diff --git a/sdks/python/apache_beam/transforms/ptransform_test.py b/sdks/python/apache_beam/transforms/ptransform_test.py index 80c9768..46c340c 100644 --- a/sdks/python/apache_beam/transforms/ptransform_test.py +++ b/sdks/python/apache_beam/transforms/ptransform_test.py @@ -22,7 +22,6 @@ from __future__ import absolute_import import operator import re import unittest -import mock import hamcrest as hc from nose.plugins.attrib import attr @@ -47,16 +46,6 @@ from apache_beam.utils.pipeline_options import TypeOptions # Disable frequent lint warning due to pipe operator for chaining transforms. # pylint: disable=expression-not-assigned -class MyDoFn(beam.DoFn): - def start_bundle(self): - pass - - def process(self, element): - pass - - def finish_bundle(self): - yield 'finish' - class PTransformTest(unittest.TestCase): # Enable nose tests running in parallel @@ -286,6 +275,13 @@ class PTransformTest(unittest.TestCase): self.assertStartswith(cm.exception.message, expected_error_prefix) def test_do_fn_with_finish(self): + class MyDoFn(beam.DoFn): + def process(self, element): + pass + + def finish_bundle(self): + yield 'finish' + pipeline = TestPipeline() pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) result = pcoll | 'Do' >> beam.ParDo(MyDoFn()) @@ -300,22 +296,46 @@ class PTransformTest(unittest.TestCase): assert_that(result, matcher()) pipeline.run() - @mock.patch.object(MyDoFn, 'start_bundle') - def test_do_fn_with_start(self, mock_method): - mock_method.return_value = None + def test_do_fn_with_start(self): + class MyDoFn(beam.DoFn): + def __init__(self): + self.state = 'init' + + def start_bundle(self): + self.state = 'started' + return None + + def process(self, element): + if self.state == 'started': + yield 'started' + self.state = 'process' + pipeline = TestPipeline() - pipeline | 'Start' >> beam.Create([1, 2, 3]) | 'Do' >> beam.ParDo(MyDoFn()) + pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) + result = pcoll | 'Do' >> beam.ParDo(MyDoFn()) + + # May have many bundles, but each has a start and finish. + def matcher(): + def match(actual): + equal_to(['started'])(list(set(actual))) + equal_to([1])([actual.count('started')]) + return match + + assert_that(result, matcher()) pipeline.run() - self.assertTrue(mock_method.called) - @mock.patch.object(MyDoFn, 'start_bundle') - def test_do_fn_with_start_error(self, mock_method): - mock_method.return_value = [1] + def test_do_fn_with_start_error(self): + class MyDoFn(beam.DoFn): + def start_bundle(self): + return [1] + + def process(self, element): + pass + pipeline = TestPipeline() pipeline | 'Start' >> beam.Create([1, 2, 3]) | 'Do' >> beam.ParDo(MyDoFn()) with self.assertRaises(RuntimeError): pipeline.run() - self.assertTrue(mock_method.called) def test_filter(self): pipeline = TestPipeline()