Adds Regression Test By Mocking Input

bigquery/requirements.txt

@@ -12,3 +12,4 @@ tox==1.9.0
 uritemplate==0.6
 virtualenv==12.0.7
 wsgiref==0.1.2
+mock==1.0.1

bigquery/samples/async_query.py

@@ -19,6 +19,7 @@
 from bigquery.samples.utils import get_service
 from bigquery.samples.utils import paging
 from bigquery.samples.utils import poll_job
+from bigquery.samples.utils import get_input
 
 
 # [START async_query]
@@ -70,16 +71,15 @@ def run(project_id, query_string, batch, num_retries, interval):
 
 # [START main]
 def main():
-    project_id = raw_input("Enter the project ID: ")
-    query_string = raw_input("Enter the Bigquery SQL Query: ")
-    batch = raw_input("Run query as batch (y/n)?: ") in (
+    project_id = get_input("Enter the project ID: ")
+    query_string = get_input("Enter the Bigquery SQL Query: ")
+    batch = get_input("Run query as batch (y/n)?: ") in (
         'True', 'true', 'y', 'Y', 'yes', 'Yes')
 
-    num_retries = raw_input(
+    num_retries = get_input(
         "Enter number of times to retry in case of 500 error: ")
-    interval = raw_input(
+    interval = get_input(
         "Enter how often to poll the query for completion (seconds): ")
-
     for result in run(project_id, query_string, batch, num_retries, interval):
         print(result)
 # [END main]

bigquery/samples/utils.py

@@ -48,3 +48,6 @@ def paging(service, request_func, num_retries=5, **kwargs):
             has_next = False
         yield response
 # [END paging]
+
+def get_input(text):
+    return input(text)

bigquery/tests/test_async_query.py

@@ -12,10 +12,12 @@
 # limitations under the License.
 #
 import json
+import os
 import unittest
 
-from bigquery.samples.async_query import run
-from tests import CloudBaseTest
+from bigquery.samples.async_query import run, main
+from tests import CloudBaseTest, BUCKET_NAME_ENV, PROJECT_ID_ENV
+from mock import patch
 
 
 class TestAsyncQuery(CloudBaseTest):
@@ -29,5 +31,24 @@ def test_async_query(self):
             self.assertIsNotNone(json.loads(result))
 
 
+class TestAsyncRunner(CloudBaseTest):
+
+    i = 0
+
+    def mock_get_input(input):
+        test_bucket_name = os.environ.get(BUCKET_NAME_ENV)
+        test_project_id = os.environ.get(PROJECT_ID_ENV)
+        answers = [test_bucket_name, test_project_id, 'n',
+                   '1', '1']
+        ret = answers[TestAsyncRunner.i]
+        TestAsyncRunner.i += 1
+        return ret
+
+
+    @patch('bigquery.samples.async_query.get_input', new=mock_get_input)
+    def test_async_query_runner(self):
+        main()
+
+
 if __name__ == '__main__':
     unittest.main()