def test_load_cloudasset_data_cai_timeout(self): """Validate load_cloud_asset handles a timeout error.""" self.mock_export_assets.side_effect = ( api_errors.OperationTimeoutError('organizations/987654321', {})) results = cloudasset.load_cloudasset_data(self.engine, self.inventory_config) self.assertIsNone(results) self.assertFalse(self.mock_download.called) self.validate_no_data_in_table()
def wait_for_completion(self, parent, operation, timeout=0, initial_delay=None, retry_delay=None): """Wait for the operation to complete. Args: parent (str): The name of the parent resource to export assests under. operation (dict): The operation response from an API call. timeout (float): The maximum time to wait for the operation to complete. initial_delay (float): The time to wait before first checking if the API has completed. If None then the default value, configured as CloudAssetClient.OPERATION_DELAY_IN_SEC, is used. retry_delay (float): The time to wait between checking if the API has completed. If None then the default value, configured as CloudAssetClient.OPERATION_DELAY_IN_SEC, is used. Returns: dict: Operation status and info. Raises: OperationTimeoutError: Raised if the operation times out. """ if operation.get('done', False): return operation if initial_delay is None: initial_delay = self.OPERATION_DELAY_IN_SEC if retry_delay is None: retry_delay = self.OPERATION_DELAY_IN_SEC started_timestamp = time.time() time.sleep(initial_delay) while True: operation_name = operation['name'] operation = self.get_operation(operation_name) if operation.get('done', False): return operation if timeout and time.time() - started_timestamp > timeout: raise api_errors.OperationTimeoutError(parent, operation) time.sleep(retry_delay)