def test_copy_datasets(self, mock_client, mock_list_tables, mock_copy_table): full_table_ids = [ f'{self.project_id}.{self.dataset_id}.{table_id}' for table_id in CDM_TABLES ] list_tables_results = [ list_item_from_table_id(table_id) for table_id in full_table_ids ] mock_list_tables.return_value = list_tables_results bq.copy_datasets(mock_client, self.dataset_id, f'{self.dataset_id}_snapshot') mock_list_tables.assert_called_once_with(self.dataset_id) self.assertEqual(mock_copy_table.call_count, len(list_tables_results))
def _mock_client_with(self, table_ids): """ Get a mock client :param table_ids: :return: """ full_table_ids = [ f'{self.project_id}.{self.dataset_id}.{table_id}' for table_id in table_ids ] list_tables_results = [ list_item_from_table_id(table_id) for table_id in full_table_ids ] table_count_query_results = [ dict(table_count=len(list_tables_results)) ] mock_client = MagicMock() mock_client.list_tables.return_value = list_tables_results mock_client.query.return_value = mock_query_result( table_count_query_results) return mock_client
def _table_id_to_list_item(self, table_id): full_table_id = self._full_table_id(table_id) return list_item_from_table_id(full_table_id)