diff --git a/sdks/python/apache_beam/examples/matrix_power_test.py b/sdks/python/apache_beam/examples/matrix_power_test.py index ff17ca5ff123..b824ba7e24c2 100644 --- a/sdks/python/apache_beam/examples/matrix_power_test.py +++ b/sdks/python/apache_beam/examples/matrix_power_test.py @@ -17,6 +17,7 @@ """Test for the matrix power integration test.""" +import glob import logging import tempfile import unittest @@ -51,9 +52,15 @@ def test_basics(self): '--input_matrix=%s --input_vector=%s --exponent=%d --output=%s.result' % (matrix_path, vector_path, self.EXPONENT, vector_path)).split()) # Parse result file and compare. - with open(vector_path + '.result-00000-of-00001') as result_file: - results = result_file.read() - self.assertEqual(sorted(self.EXPECTED_OUTPUT), sorted(results)) + shard_paths = glob.glob(vector_path + '.result*') + self.assertTrue( + shard_paths, + 'No output shards found matching prefix: %s.result' % vector_path) + results = [] + for path in shard_paths: + with open(path) as result_file: + results.append(result_file.read()) + self.assertEqual(sorted(self.EXPECTED_OUTPUT), sorted(''.join(results))) if __name__ == '__main__':