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)