Use context manager to mock raw_input instead of patch

Bill Prin · Bill Prin · commit 045e6f879e9e · 2015-05-22T11:21:41.000-07:00
diff --git a/bigquery/requirements.txt b/bigquery/requirements.txt
@@ -12,4 +12,3 @@ tox==1.9.0
 uritemplate==0.6
 virtualenv==12.0.7
 wsgiref==0.1.2
-mock==1.0.1
diff --git a/bigquery/samples/async_query.py b/bigquery/samples/async_query.py
@@ -19,8 +19,6 @@
 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]
 def async_query(service, project_id, query, batch=False, num_retries=5):
@@ -71,14 +69,14 @@ def run(project_id, query_string, batch, num_retries, interval):
 
 # [START main]
 def main():
-    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 (
+    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 (
         'True', 'true', 'y', 'Y', 'yes', 'Yes')
 
-    num_retries = get_input(
+    num_retries = raw_input(
         "Enter number of times to retry in case of 500 error: ")
-    interval = get_input(
+    interval = raw_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)
diff --git a/bigquery/samples/utils.py b/bigquery/samples/utils.py
@@ -49,5 +49,3 @@ def paging(service, request_func, num_retries=5, **kwargs):
         yield response
 # [END paging]
 
-def get_input(text):
-    return input(text)
diff --git a/bigquery/tests/test_async_query.py b/bigquery/tests/test_async_query.py
@@ -16,9 +16,7 @@
 import unittest
 
 from bigquery.samples.async_query import run, main
-from tests import CloudBaseTest, BUCKET_NAME_ENV, PROJECT_ID_ENV
-from mock import patch
-
+from tests import CloudBaseTest, mock_raw_input, BUCKET_NAME_ENV, PROJECT_ID_ENV
 
 class TestAsyncQuery(CloudBaseTest):
 
@@ -31,23 +29,23 @@ def test_async_query(self):
             self.assertIsNotNone(json.loads(result))
 
 
-class TestAsyncRunner(CloudBaseTest):
+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
 
-    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
+class TestAsyncRunner(CloudBaseTest):
 
+    i = 0
 
-    @patch('bigquery.samples.async_query.get_input', new=mock_get_input)
     def test_async_query_runner(self):
-        main()
+        with mock_raw_input(mock_get_input):
+            main()
 
 
 if __name__ == '__main__':
diff --git a/tests/__init__.py b/tests/__init__.py
@@ -18,14 +18,23 @@
 import json
 import os
 import unittest
-
+from contextlib import contextmanager
+import __builtin__
 
 BUCKET_NAME_ENV = 'TEST_BUCKET_NAME'
 PROJECT_ID_ENV = 'TEST_PROJECT_ID'
 RESOURCE_PATH = os.path.join(
     os.path.abspath(os.path.dirname(__file__)), 'resources')
 
 
+@contextmanager
+def mock_raw_input(mock):
+    original_raw_input = __builtin__.raw_input
+    __builtin__.raw_input = mock
+    yield
+    __builtin__.raw_input = original_raw_input
+
+
 class CloudBaseTest(unittest.TestCase):
 
     def setUp(self):
