From 2086c407b8c550de3d287b564079db2c623990f7 Mon Sep 17 00:00:00 2001
From: Ahmet Altay <altay@google.com>
Date: Fri, 27 Jan 2017 17:35:24 -0800
Subject: [PATCH] Add mock time to slow bigquery unit tests.

Unit tests, testing retries does not need to use real time. This change
reduces the total tox time for unit tests from 235 seconds to 73 seconds
locally.
---
 sdks/python/apache_beam/io/bigquery_test.py | 26 ++++++++++++++-------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/sdks/python/apache_beam/io/bigquery_test.py b/sdks/python/apache_beam/io/bigquery_test.py
index b8682d112397..14eb0353bdfe 100644
--- a/sdks/python/apache_beam/io/bigquery_test.py
+++ b/sdks/python/apache_beam/io/bigquery_test.py
@@ -539,7 +539,8 @@ def test_table_schema_without_project(self):
 
 class TestBigQueryWriter(unittest.TestCase):
 
-  def test_no_table_and_create_never(self):
+  @mock.patch('time.sleep', return_value=None)
+  def test_no_table_and_create_never(self, patched_time_sleep):
     client = mock.Mock()
     client.tables.Get.side_effect = HttpError(
         response={'status': '404'}, url='', content='')
@@ -572,7 +573,9 @@ def test_no_table_and_create_if_needed(self):
     self.assertTrue(client.tables.Get.called)
     self.assertTrue(client.tables.Insert.called)
 
-  def test_no_table_and_create_if_needed_and_no_schema(self):
+  @mock.patch('time.sleep', return_value=None)
+  def test_no_table_and_create_if_needed_and_no_schema(
+      self, patched_time_sleep):
     client = mock.Mock()
     client.tables.Get.side_effect = HttpError(
         response={'status': '404'}, url='', content='')
@@ -587,7 +590,9 @@ def test_no_table_and_create_if_needed_and_no_schema(self):
         'Table project:dataset.table requires a schema. None can be inferred '
         'because the table does not exist.')
 
-  def test_table_not_empty_and_write_disposition_empty(self):
+  @mock.patch('time.sleep', return_value=None)
+  def test_table_not_empty_and_write_disposition_empty(
+      self, patched_time_sleep):
     client = mock.Mock()
     client.tables.Get.return_value = bigquery.Table(
         tableReference=bigquery.TableReference(
@@ -712,7 +717,8 @@ def test_delete_non_existing_dataset(self):
     wrapper._delete_dataset('', '')
     self.assertTrue(client.datasets.Delete.called)
 
-  def test_delete_dataset_retries_fail(self):
+  @mock.patch('time.sleep', return_value=None)
+  def test_delete_dataset_retries_fail(self, patched_time_sleep):
     client = mock.Mock()
     client.datasets.Delete.side_effect = ValueError("Cannot delete")
     wrapper = beam.io.bigquery.BigQueryWrapper(client)
@@ -730,7 +736,8 @@ def test_delete_non_existing_table(self):
     wrapper._delete_table('', '', '')
     self.assertTrue(client.tables.Delete.called)
 
-  def test_delete_table_retries_fail(self):
+  @mock.patch('time.sleep', return_value=None)
+  def test_delete_table_retries_fail(self, patched_time_sleep):
     client = mock.Mock()
     client.tables.Delete.side_effect = ValueError("Cannot delete")
     wrapper = beam.io.bigquery.BigQueryWrapper(client)
@@ -738,7 +745,8 @@ def test_delete_table_retries_fail(self):
       wrapper._delete_table('', '', '')
     self.assertTrue(client.tables.Delete.called)
 
-  def test_delete_dataset_retries_for_timeouts(self):
+  @mock.patch('time.sleep', return_value=None)
+  def test_delete_dataset_retries_for_timeouts(self, patched_time_sleep):
     client = mock.Mock()
     client.datasets.Delete.side_effect = [
         HttpError(
@@ -749,7 +757,8 @@ def test_delete_dataset_retries_for_timeouts(self):
     wrapper._delete_dataset('', '')
     self.assertTrue(client.datasets.Delete.called)
 
-  def test_delete_table_retries_for_timeouts(self):
+  @mock.patch('time.sleep', return_value=None)
+  def test_delete_table_retries_for_timeouts(self, patched_time_sleep):
     client = mock.Mock()
     client.tables.Delete.side_effect = [
         HttpError(
@@ -760,7 +769,8 @@ def test_delete_table_retries_for_timeouts(self):
     wrapper._delete_table('', '', '')
     self.assertTrue(client.tables.Delete.called)
 
-  def test_temporary_dataset_is_unique(self):
+  @mock.patch('time.sleep', return_value=None)
+  def test_temporary_dataset_is_unique(self, patched_time_sleep):
     client = mock.Mock()
     client.datasets.Get.return_value = bigquery.Dataset(
         datasetReference=bigquery.DatasetReference(