[beam] branch master updated: [BEAM-7437] Raise RuntimeError for PY2 in BigqueryFullResultStreamingMatcher (#9044)

[udim]

This is an automated email from the ASF dual-hosted git repository.

udim pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new f7cbf88  [BEAM-7437] Raise RuntimeError for PY2 in 
BigqueryFullResultStreamingMatcher (#9044)
f7cbf88 is described below

commit f7cbf88f550c8918b99a13af4182d6efa07cd2b5
Author: Tanay Tummalapalli <ttanay...@gmail.com>
AuthorDate: Sat Jul 13 02:32:03 2019 +0530

    [BEAM-7437] Raise RuntimeError for PY2 in 
BigqueryFullResultStreamingMatcher (#9044)
    
    TimeoutError is not a built-in exception in Python 2.
    Raise RuntimeError in Python 2.
---
 .../apache_beam/io/gcp/tests/bigquery_matcher.py   |  6 +++++-
 .../io/gcp/tests/bigquery_matcher_test.py          | 22 ++++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py 
b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
index 216963a..6709d6a 100644
--- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
+++ b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
@@ -20,6 +20,7 @@
 from __future__ import absolute_import
 
 import logging
+import sys
 import time
 
 from hamcrest.core.base_matcher import BaseMatcher
@@ -192,7 +193,10 @@ class 
BigqueryFullResultStreamingMatcher(BigqueryFullResultMatcher):
       if len(response) >= len(self.expected_data):
         return response
       time.sleep(1)
-    raise TimeoutError('Timeout exceeded for matcher.')
+    if sys.version_info >= (3,):
+      raise TimeoutError('Timeout exceeded for matcher.') # noqa: F821
+    else:
+      raise RuntimeError('Timeout exceeded for matcher.')
 
 
 class BigQueryTableMatcher(BaseMatcher):
diff --git a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py 
b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
index c8315db..2a0db08 100644
--- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
+++ b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
@@ -20,6 +20,7 @@
 from __future__ import absolute_import
 
 import logging
+import sys
 import unittest
 
 import mock
@@ -113,6 +114,27 @@ class BigqueryTableMatcherTest(unittest.TestCase):
     self.assertEqual(bq_verifier.MAX_RETRIES + 1, mock_query.call_count)
 
 
+@unittest.skipIf(bigquery is None, 'Bigquery dependencies are not installed.')
+@mock.patch.object(
+    bq_verifier.BigqueryFullResultStreamingMatcher,
+    '_query_with_retry')
+class BigqueryFullResultStreamingMatcher(unittest.TestCase):
+
+  def setUp(self):
+    self.timeout = 5
+
+  def test__get_query_result_timeout(self, mock__query_with_retry):
+    mock__query_with_retry.side_effect = lambda _: []
+    matcher = bq_verifier.BigqueryFullResultStreamingMatcher(
+        'some-project', 'some-query', [1, 2, 3], timeout=self.timeout)
+    if sys.version_info >= (3,):
+      with self.assertRaises(TimeoutError): # noqa: F821
+        matcher._get_query_result(None)
+    else:
+      with self.assertRaises(RuntimeError):
+        matcher._get_query_result(None)
+
+
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)
   unittest.main()