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)