Merge pull request #132 from treasure-data/copilot/fix-fe93b6aa-303e-4886-bf56-e2cd32aad076

Add timeout parameter to BulkImport.perform() method

Author: chezou

tdclient/bulk_import_model.py

@@ -101,14 +101,15 @@ def unfreeze(self):
         self.update()
         return response
 
-    def perform(self, wait=False, wait_interval=5, wait_callback=None):
+    def perform(self, wait=False, wait_interval=5, wait_callback=None, timeout=None):
         """Perform bulk import
 
         Args:
             wait (bool, optional): Flag for wait bulk import job. Default `False`
             wait_interval (int, optional): wait interval in second. Default `5`.
             wait_callback (callable, optional): A callable to be called on every tick of
                 wait interval.
+            timeout (int, optional): Timeout in seconds. No timeout by default.
         """
         self.update()
         if not self.upload_frozen:
@@ -117,7 +118,7 @@ def perform(self, wait=False, wait_interval=5, wait_callback=None):
             )
         job = self._client.perform_bulk_import(self.name)
         if wait:
-            job.wait(wait_interval=wait_interval, wait_callback=wait_callback)
+            job.wait(timeout=timeout, wait_interval=wait_interval, wait_callback=wait_callback)
         self.update()
         return job
 

tdclient/test/bulk_import_model_test.py

@@ -120,6 +120,32 @@ def test_bulk_import_perfom():
     assert bulk_import.update.called
 
 
+def test_bulk_import_perform_with_timeout():
+    client = mock.MagicMock()
+    job_mock = mock.MagicMock()
+    client.perform_bulk_import.return_value = job_mock
+    
+    bulk_import = models.BulkImport(
+        client,
+        name="name",
+        database="database",
+        table="table",
+        status="status",
+        upload_frozen=True,
+        job_id="job_id",
+        valid_records="valid_records",
+        error_records="error_records",
+        valid_parts="valid_parts",
+        error_parts="error_parts",
+    )
+    bulk_import.update = mock.MagicMock()
+    bulk_import.perform(wait=True, timeout=300, wait_interval=10)
+    
+    client.perform_bulk_import.assert_called_with("name")
+    job_mock.wait.assert_called_with(timeout=300, wait_interval=10, wait_callback=None)
+    assert bulk_import.update.called
+
+
 def test_bulk_import_commit():
     client = mock.MagicMock()
     bulk_import = models.BulkImport(