diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py index c533ef87b4bea..f87aa2f2bc93e 100644 --- a/sdks/python/apache_beam/io/fileio_test.py +++ b/sdks/python/apache_beam/io/fileio_test.py @@ -296,10 +296,11 @@ def test_write_to_single_file_batch(self): result = (p | fileio.MatchFiles(FileSystems.join(dir, '*')) | fileio.ReadMatches() - | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n'))) + | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n')) + | beam.Map(json.loads)) assert_that(result, - equal_to([json.dumps(row) for row in self.SIMPLE_COLLECTION])) + equal_to([row for row in self.SIMPLE_COLLECTION])) def test_write_to_different_file_types_some_spilling(self): @@ -321,7 +322,8 @@ def test_write_to_different_file_types_some_spilling(self): cncf_res = (p | fileio.MatchFiles(FileSystems.join(dir, 'cncf*')) | fileio.ReadMatches() - | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n'))) + | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n')) + | beam.Map(json.loads)) apache_res = (p | "MatchApache" >> fileio.MatchFiles( @@ -331,7 +333,7 @@ def test_write_to_different_file_types_some_spilling(self): lambda rf: csv.reader(_get_file_reader(rf)))) assert_that(cncf_res, - equal_to([json.dumps(row) + equal_to([row for row in self.SIMPLE_COLLECTION if row['foundation'] == 'cncf']), label='verifyCNCF') @@ -398,7 +400,8 @@ def test_write_to_different_file_types(self): cncf_res = (p | fileio.MatchFiles(FileSystems.join(dir, 'cncf*')) | fileio.ReadMatches() - | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n'))) + | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n')) + | beam.Map(json.loads)) apache_res = (p | "MatchApache" >> fileio.MatchFiles( @@ -408,7 +411,7 @@ def test_write_to_different_file_types(self): lambda rf: csv.reader(_get_file_reader(rf)))) assert_that(cncf_res, - equal_to([json.dumps(row) + equal_to([row for row in self.SIMPLE_COLLECTION if row['foundation'] == 'cncf']), label='verifyCNCF')