def test_download_table_as_df(self, mocker):
self._stop_time(mocker)
mocked_gcs_to_df = mocker.patch("feast.sdk.utils.bq_util.gcs_to_df",
return_value=None)
staging_path = "gs://temp/"
staging_file_name = "temp_0"
table_id = "project_id.dataset_id.table_id"
table_dldr = TableDownloader()
exp_staging_path = os.path.join(staging_path, staging_file_name)
table_dldr._bq = _Mock_BQ_Client()
mocker.patch.object(table_dldr._bq,
"extract_table",
return_value=_Job())
table_dldr.download_table_as_df(table_id,
staging_location=staging_path)
assert len(table_dldr._bq.extract_table.call_args_list) == 1
args, kwargs = table_dldr._bq.extract_table.call_args_list[0]
assert args[0].full_table_id == Table.from_string(
table_id).full_table_id
assert args[1] == exp_staging_path
assert kwargs["job_config"].destination_format == "CSV"
mocked_gcs_to_df.assert_called_once_with(exp_staging_path)
def _test_download_file(self, mocker, type):
staging_path = "gs://temp/"
staging_file_name = "temp_0"
dst_path = "/tmp/myfile.csv"
table_id = "project_id.dataset_id.table_id"
table_dldr = TableDownloader()
mock_blob = _Blob()
mocker.patch.object(mock_blob, "download_to_filename")
table_dldr._bq = _Mock_BQ_Client()
mocker.patch.object(table_dldr._bq,
"extract_table",
return_value=_Job())
table_dldr._gcs = _Mock_GCS_Client()
mocker.patch.object(table_dldr._gcs,
"get_bucket",
return_value=_Bucket(mock_blob))
table_dldr.download_table_as_file(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._bq.extract_table.call_args_list) == 1
args, kwargs = table_dldr._bq.extract_table.call_args_list[0]
assert args[0].full_table_id == Table.from_string(
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)