def test_download_table_as_df(self, mocker):
self._stop_time(mocker)
mocked_gcs_folder_to_df = mocker.patch(
"feast.sdk.utils.bq_util.gcs_folder_to_df", return_value=None)
staging_path = "gs://temp"
temp_folder = "temp_0"
full_table_id = "project_id.dataset_id.table_id"
table_dldr = TableDownloader()
exp_staging_folder = os.path.join(staging_path, temp_folder)
exp_staging_path = os.path.join(exp_staging_folder, "shard_*")
table_dldr._bqclient = _Mock_BQ_Client()
mocker.patch.object(table_dldr._bqclient,
"extract_table",
return_value=_Job())
table_dldr.download_table_as_df(full_table_id,
staging_location=staging_path)
assert len(table_dldr._bqclient.extract_table.call_args_list) == 1
args, kwargs = table_dldr._bqclient.extract_table.call_args_list[0]
assert args[0].full_table_id == Table.from_string(
full_table_id).full_table_id
assert args[1] == exp_staging_path
assert kwargs["job_config"].destination_format == "CSV"
mocked_gcs_folder_to_df.assert_called_once_with(exp_staging_folder)
def _test_download_file(self, mocker, type):
staging_path = "gs://temp/"
staging_file_name = "temp_0"
dst_path = "/tmp/myfile.csv"
full_table_id = "project_id.dataset_id.table_id"
table_dldr = TableDownloader()
mock_blob = _Blob()
mocker.patch.object(mock_blob, "download_to_filename")
table_dldr._bqclient = _Mock_BQ_Client()
mocker.patch.object(table_dldr._bqclient, "extract_table", return_value=_Job())
table_dldr._storageclient = _Mock_GCS_Client()
mocker.patch.object(
table_dldr._storageclient, "get_bucket", return_value=_Bucket(mock_blob)
)
table_dldr.download_table_as_file(
full_table_id, dst_path, staging_location=staging_path, file_type=type
)
exp_staging_path = os.path.join(staging_path, staging_file_name)
assert len(table_dldr._bqclient.extract_table.call_args_list) == 1
args, kwargs = table_dldr._bqclient.extract_table.call_args_list[0]
assert args[0].full_table_id == Table.from_string(full_table_id).full_table_id
assert args[1] == exp_staging_path
assert kwargs["job_config"].destination_format == str(type)
mock_blob.download_to_filename.assert_called_once_with(dst_path)
