Fix flaky MatrixPowerTest.test_basics by reading all generated shards (apache#38585)

* Fix flaky MatrixPowerTest.test_basics by reading all generated shards

* Address comments.

Author: shunping

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__':