[BEAM-7389] Show code snippet outputs as stdout

sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_test.py

@@ -23,40 +23,34 @@
 
 import mock
 
+from apache_beam.examples.snippets.util import assert_matches_stdout
 from apache_beam.testing.test_pipeline import TestPipeline
-from apache_beam.testing.util import assert_that
-from apache_beam.testing.util import equal_to
 
 from . import filter
 
 
 def check_perennials(actual):
-  # [START perennials]
-  perennials = [
-      {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'},
-      {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'},
-      {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'},
-  ]
-  # [END perennials]
-  assert_that(actual, equal_to(perennials))
+  expected = '''[START perennials]
+{'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'}
+{'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'}
+{'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'}
+[END perennials]'''.splitlines()[1:-1]
+  assert_matches_stdout(actual, expected)
 
 
 def check_valid_plants(actual):
-  # [START valid_plants]
-  valid_plants = [
-      {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'},
-      {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'},
-      {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'},
-      {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'},
-  ]
-  # [END valid_plants]
-  assert_that(actual, equal_to(valid_plants))
+  expected = '''[START valid_plants]
+{'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'}
+{'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'}
+{'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'}
+{'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'}
+[END valid_plants]'''.splitlines()[1:-1]
+  assert_matches_stdout(actual, expected)
 
 
 @mock.patch('apache_beam.Pipeline', TestPipeline)
-# pylint: disable=line-too-long
-@mock.patch('apache_beam.examples.snippets.transforms.elementwise.filter.print', lambda elem: elem)
-# pylint: enable=line-too-long
+@mock.patch(
+    'apache_beam.examples.snippets.transforms.elementwise.filter.print', str)
 class FilterTest(unittest.TestCase):
   def test_filter_function(self):
     filter.filter_function(check_perennials)

sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py

@@ -23,42 +23,36 @@
 
 import mock
 
+from apache_beam.examples.snippets.util import assert_matches_stdout
 from apache_beam.testing.test_pipeline import TestPipeline
-from apache_beam.testing.util import assert_that
-from apache_beam.testing.util import equal_to
 
 from . import flatmap
 
 
 def check_plants(actual):
-  # [START plants]
-  plants = [
-      '🍓Strawberry',
-      '🥕Carrot',
-      '🍆Eggplant',
-      '🍅Tomato',
-      '🥔Potato',
-  ]
-  # [END plants]
-  assert_that(actual, equal_to(plants))
+  expected = '''[START plants]
+🍓Strawberry
+🥕Carrot
+🍆Eggplant
+🍅Tomato
+🥔Potato
+[END plants]'''.splitlines()[1:-1]
+  assert_matches_stdout(actual, expected)
 
 
 def check_valid_plants(actual):
-  # [START valid_plants]
-  valid_plants = [
-      {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'},
-      {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'},
-      {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'},
-      {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'},
-  ]
-  # [END valid_plants]
-  assert_that(actual, equal_to(valid_plants))
+  expected = '''[START valid_plants]
+{'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'}
+{'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'}
+{'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'}
+{'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'}
+[END valid_plants]'''.splitlines()[1:-1]
+  assert_matches_stdout(actual, expected)
 
 
 @mock.patch('apache_beam.Pipeline', TestPipeline)
-# pylint: disable=line-too-long
-@mock.patch('apache_beam.examples.snippets.transforms.elementwise.flatmap.print', lambda elem: elem)
-# pylint: enable=line-too-long
+@mock.patch(
+    'apache_beam.examples.snippets.transforms.elementwise.flatmap.print', str)
 class FlatMapTest(unittest.TestCase):
   def test_flatmap_simple(self):
     flatmap.flatmap_simple(check_plants)

sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys_test.py

@@ -23,30 +23,26 @@
 
 import mock
 
+from apache_beam.examples.snippets.util import assert_matches_stdout
 from apache_beam.testing.test_pipeline import TestPipeline
-from apache_beam.testing.util import assert_that
-from apache_beam.testing.util import equal_to
 
 from . import keys
 
 
 def check_icons(actual):
-  # [START icons]
-  icons = [
-      '🍓',
-      '🥕',
-      '🍆',
-      '🍅',
-      '🥔',
-  ]
-  # [END icons]
-  assert_that(actual, equal_to(icons))
+  expected = '''[START icons]
+🍓
+🥕
+🍆
+🍅
+🥔
+[END icons]'''.splitlines()[1:-1]
+  assert_matches_stdout(actual, expected)
 
 
 @mock.patch('apache_beam.Pipeline', TestPipeline)
-# pylint: disable=line-too-long
-@mock.patch('apache_beam.examples.snippets.transforms.elementwise.keys.print', lambda elem: elem)
-# pylint: enable=line-too-long
+@mock.patch(
+    'apache_beam.examples.snippets.transforms.elementwise.keys.print', str)
 class KeysTest(unittest.TestCase):
   def test_keys(self):
     keys.keys(check_icons)

sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py

@@ -23,30 +23,26 @@
 
 import mock
 
+from apache_beam.examples.snippets.util import assert_matches_stdout
 from apache_beam.testing.test_pipeline import TestPipeline
-from apache_beam.testing.util import assert_that
-from apache_beam.testing.util import equal_to
 
 from . import kvswap
 
 
 def check_plants(actual):
-  # [START plants]
-  plants = [
-      ('Strawberry', '🍓'),
-      ('Carrot', '🥕'),
-      ('Eggplant', '🍆'),
-      ('Tomato', '🍅'),
-      ('Potato', '🥔'),
-  ]
-  # [END plants]
-  assert_that(actual, equal_to(plants))
+  expected = '''[START plants]
+('Strawberry', '🍓')
+('Carrot', '🥕')
+('Eggplant', '🍆')
+('Tomato', '🍅')
+('Potato', '🥔')
+[END plants]'''.splitlines()[1:-1]
+  assert_matches_stdout(actual, expected)
 
 
 @mock.patch('apache_beam.Pipeline', TestPipeline)
-# pylint: disable=line-too-long
-@mock.patch('apache_beam.examples.snippets.transforms.elementwise.kvswap.print', lambda elem: elem)
-# pylint: enable=line-too-long
+@mock.patch(
+    'apache_beam.examples.snippets.transforms.elementwise.kvswap.print', str)
 class KvSwapTest(unittest.TestCase):
   def test_kvswap(self):
     kvswap.kvswap(check_plants)

sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_test.py

@@ -23,43 +23,37 @@
 
 import mock
 
+from apache_beam.examples.snippets.util import assert_matches_stdout
 from apache_beam.testing.test_pipeline import TestPipeline
-from apache_beam.testing.util import assert_that
-from apache_beam.testing.util import equal_to
 
 from . import map
 
 
 def check_plants(actual):
-  # [START plants]
-  plants = [
-      '🍓Strawberry',
-      '🥕Carrot',
-      '🍆Eggplant',
-      '🍅Tomato',
-      '🥔Potato',
-  ]
-  # [END plants]
-  assert_that(actual, equal_to(plants))
+  expected = '''[START plants]
+🍓Strawberry
+🥕Carrot
+🍆Eggplant
+🍅Tomato
+🥔Potato
+[END plants]'''.splitlines()[1:-1]
+  assert_matches_stdout(actual, expected)
 
 
 def check_plant_details(actual):
-  # [START plant_details]
-  plant_details = [
-      {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'},
-      {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'},
-      {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'},
-      {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'},
-      {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'},
-  ]
-  # [END plant_details]
-  assert_that(actual, equal_to(plant_details))
+  expected = '''[START plant_details]
+{'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'}
+{'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'}
+{'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'}
+{'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'}
+{'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'}
+[END plant_details]'''.splitlines()[1:-1]
+  assert_matches_stdout(actual, expected)
 
 
 @mock.patch('apache_beam.Pipeline', TestPipeline)
-# pylint: disable=line-too-long
-@mock.patch('apache_beam.examples.snippets.transforms.elementwise.map.print', lambda elem: elem)
-# pylint: enable=line-too-long
+@mock.patch(
+    'apache_beam.examples.snippets.transforms.elementwise.map.print', str)
 class MapTest(unittest.TestCase):
   def test_map_simple(self):
     map.map_simple(check_plants)

sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo.py

@@ -18,7 +18,6 @@
 
 from __future__ import absolute_import
 from __future__ import print_function
-from __future__ import unicode_literals
 
 
 def pardo_dofn(test=None):

sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py

@@ -18,38 +18,42 @@
 
 from __future__ import absolute_import
 from __future__ import print_function
-from __future__ import unicode_literals
 
-import io
 import platform
 import sys
 import unittest
 
 import mock
 
+from apache_beam.examples.snippets.util import assert_matches_stdout
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
 
 from . import pardo
 
+# TODO: Remove this after Python 2 deprecation.
+# https://issues.apache.org/jira/browse/BEAM-8124
+if sys.version_info[0] == 2:
+  from io import BytesIO as StringIO
+else:
+  from io import StringIO
+
 
 def check_plants(actual):
-  # [START plants]
-  plants = [
-      '🍓Strawberry',
-      '🥕Carrot',
-      '🍆Eggplant',
-      '🍅Tomato',
-      '🥔Potato',
-  ]
-  # [END plants]
-  assert_that(actual, equal_to(plants))
+  expected = '''[START plants]
+🍓Strawberry
+🥕Carrot
+🍆Eggplant
+🍅Tomato
+🥔Potato
+[END plants]'''.splitlines()[1:-1]
+  assert_matches_stdout(actual, expected)
 
 
 def check_dofn_params(actual):
   # pylint: disable=line-too-long
-  dofn_params = '\n'.join('''[START dofn_params]
+  expected = '\n'.join('''[START dofn_params]
 # timestamp
 type(timestamp) -> <class 'apache_beam.utils.timestamp.Timestamp'>
 timestamp.micros -> 1584675660000000
@@ -63,7 +67,7 @@ def check_dofn_params(actual):
 window.max_timestamp() -> Timestamp(1584675689.999999) (2020-03-20 03:41:29.999999)
 [END dofn_params]'''.splitlines()[1:-1])
   # pylint: enable=line-too-long
-  assert_that(actual, equal_to([dofn_params]))
+  assert_that(actual, equal_to([expected]))
 
 
 def check_dofn_methods(actual):
@@ -83,23 +87,22 @@ def check_dofn_methods(actual):
 
 
 @mock.patch('apache_beam.Pipeline', TestPipeline)
-# pylint: disable=line-too-long
-@mock.patch('apache_beam.examples.snippets.transforms.elementwise.pardo.print', lambda elem: elem)
-# pylint: enable=line-too-long
+@mock.patch(
+    'apache_beam.examples.snippets.transforms.elementwise.pardo.print', str)
 class ParDoTest(unittest.TestCase):
   def test_pardo_dofn(self):
     pardo.pardo_dofn(check_plants)
 
   # TODO: Remove this after Python 2 deprecation.
   # https://issues.apache.org/jira/browse/BEAM-8124
-  @unittest.skipIf(sys.version_info[0] < 3 and platform.system() == 'Windows',
+  @unittest.skipIf(sys.version_info[0] == 2 and platform.system() == 'Windows',
                    'Python 2 on Windows uses `long` rather than `int`')
   def test_pardo_dofn_params(self):
     pardo.pardo_dofn_params(check_dofn_params)
 
 
 @mock.patch('apache_beam.Pipeline', TestPipeline)
-@mock.patch('sys.stdout', new_callable=io.StringIO)
+@mock.patch('sys.stdout', new_callable=StringIO)
 class ParDoStdoutTest(unittest.TestCase):
   def test_pardo_dofn_methods(self, mock_stdout):
     expected = pardo.pardo_dofn_methods(check_dofn_methods)