diff --git a/app/services/dataset_manager/dataset_download.py b/app/services/dataset_manager/dataset_download.py index de36cb1d..5a60bdd1 100644 --- a/app/services/dataset_manager/dataset_download.py +++ b/app/services/dataset_manager/dataset_download.py @@ -5,7 +5,8 @@ import datetime import os import time -from urllib.parse import unquote +from typing import Any +from typing import Dict import requests from tqdm import tqdm @@ -18,6 +19,7 @@ from app.services.output_manager.error_handler import ECustomizedError from app.services.output_manager.error_handler import SrvErrorHandler from app.services.output_manager.message_handler import SrvOutPutHandler +from app.utils.aggregated import resilient_session from ..user_authentication.decorator import require_valid_token @@ -43,7 +45,7 @@ def pre_dataset_version_download(self): 'Session-ID': self.session_id, } payload = {'version': self.version} - response = requests.get(url, headers=headers, params=payload) + response = resilient_session().get(url, headers=headers, params=payload) res = response.json() code = response.status_code if code == 200: @@ -54,7 +56,7 @@ def pre_dataset_version_download(self): SrvErrorHandler.default_handle(response.content, True) @require_valid_token() - def pre_dataset_download(self): + def pre_dataset_download(self) -> Dict[str, Any]: url = AppConfig.Connections.url_dataset_v2download + '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/download/pre' headers = { 'Authorization': 'Bearer ' + self.user.access_token, @@ -63,36 +65,18 @@ def pre_dataset_download(self): } payload = {'dataset_code': self.dataset_code, 'session_id': self.session_id, 'operator': self.user.username} try: - response = requests.post(url, headers=headers, json=payload) + response = resilient_session().post(url, headers=headers, json=payload) res = response.json() return res - except Exception: - SrvErrorHandler.default_handle(response.content, True) - - def generate_download_url(self): - if self.version: - download_url = AppConfig.Connections.url_dataset_v2download + f'/download/{self.hash_code}' - else: - download_url = AppConfig.Connections.url_download_core + f'v1/download/{self.hash_code}' - headers = { - 'Authorization': 'Bearer ' + self.user.access_token, - } - res = requests.get(download_url, headers=headers) - res_json = res.json() - if self.version: - self.download_url = self.hash_code - default_filename = self.download_url.split('/')[-1].split('?')[0] - self.default_filename = unquote(default_filename) - else: - self.download_url = download_url - self.default_filename = res_json.get('error_msg').split('/')[-1].rstrip('.') + except Exception as e: + SrvErrorHandler.default_handle(f'error when pre dataset download:{e}', True) @require_valid_token() def download_status(self) -> EFileStatus: url = AppConfig.Connections.url_download_core + f'v1/download/status/{self.hash_code}' - res = requests.get(url) - res_json = res.json() - if res_json.get('code') == 200: + res = resilient_session().get(url) + if res.status_code == 200: + res_json = res.json() status = res_json.get('result').get('status') return EFileStatus(status) else: @@ -107,15 +91,17 @@ def check_download_preparing_status(self) -> EFileStatus: return status @require_valid_token() - def send_download_request(self): + def send_download_request(self) -> str: logger.info('start downloading...') + with requests.get(self.download_url, stream=True, allow_redirects=True) as r: r.raise_for_status() # Since version zip file was created by our system, thus no need to consider filename contain '?' if not self.default_filename: - filename = f'{self.dataset_code}_{self.version}_{str(datetime.datetime.now())}' + filename = f'{self.dataset_code}_{self.version}_{str(datetime.datetime.now())}.zip' else: filename = self.default_filename + output_path = self.avoid_duplicate_file_name(self.output.rstrip('/') + '/' + filename) self.total_size = int(r.headers.get('Content-length')) with open(output_path, 'wb') as file, tqdm( @@ -131,7 +117,7 @@ def send_download_request(self): bar.update(size) return output_path - def avoid_duplicate_file_name(self, filename): + def avoid_duplicate_file_name(self, filename) -> str: suffix = 1 original_filename = filename file, ext = os.path.splitext(original_filename) @@ -148,10 +134,15 @@ def avoid_duplicate_file_name(self, filename): return filename @require_valid_token() - def download_dataset(self): + def download_dataset(self) -> None: pre_result = self.pre_dataset_download() self.hash_code = pre_result.get('result').get('payload').get('hash_code') - self.generate_download_url() + self.download_url = AppConfig.Connections.url_download_core + f'v1/download/{self.hash_code}' + # format the naming for the default filename + self.default_filename = pre_result.get('result').get('target_names')[0] + self.default_filename = self.default_filename.split('/')[-1] + + # wait the download status to be ready status = self.check_download_preparing_status() SrvOutPutHandler.download_status(status) saved_filename = self.send_download_request() @@ -161,11 +152,11 @@ def download_dataset(self): SrvErrorHandler.customized_handle(ECustomizedError.DOWNLOAD_FAIL, True) @require_valid_token() - def download_dataset_version(self, version): + def download_dataset_version(self, version) -> None: self.version = version pre_result = self.pre_dataset_version_download() - self.hash_code = pre_result.get('result').get('source') - self.generate_download_url() + self.download_url = pre_result.get('result').get('source') + saved_filename = self.send_download_request() if os.path.isfile(saved_filename): SrvOutPutHandler.download_success(saved_filename) diff --git a/tests/app/commands/test_dataset.py b/tests/app/commands/test_dataset.py index bf043b9c..76e2ac5e 100644 --- a/tests/app/commands/test_dataset.py +++ b/tests/app/commands/test_dataset.py @@ -6,7 +6,7 @@ from app.configs.app_config import AppConfig -def test_download_not_exited_dataset_version(requests_mock, mocker, cli_runner, capsys): +def test_download_not_exited_dataset_version(requests_mock, httpx_mock, mocker, cli_runner, capsys): mocker.patch('app.services.user_authentication.token_manager.SrvTokenManager.check_valid', return_value=0) requests_mock.get( 'http://bff_cli' + '/v1/dataset/testdataset', @@ -39,8 +39,9 @@ def test_download_not_exited_dataset_version(requests_mock, mocker, cli_runner, }, ) - requests_mock.get( - AppConfig.Connections.url_dataset + '/fake-id/download/pre', + httpx_mock.add_response( + method='GET', + url=AppConfig.Connections.url_dataset + '/fake-id/download/pre?version=1.0', json={'error': 'version does not exist'}, status_code=404, ) diff --git a/tests/app/services/dataset_manager/test_dataset_download.py b/tests/app/services/dataset_manager/test_dataset_download.py index ff372835..0afea52e 100644 --- a/tests/app/services/dataset_manager/test_dataset_download.py +++ b/tests/app/services/dataset_manager/test_dataset_download.py @@ -1,3 +1,161 @@ # Copyright (C) 2022-2023 Indoc Research # # Contact Indoc Research for any questions regarding the use of this source code. + +from app.configs.app_config import AppConfig +from app.services.dataset_manager.dataset_download import SrvDatasetDownloadManager +from app.services.dataset_manager.model import EFileStatus + + +def test_dateset_pre_download_success(httpx_mock, mocker): + mocker.patch('app.services.user_authentication.token_manager.SrvTokenManager.check_valid', return_value=0) + + httpx_mock.add_response( + method='POST', + url=AppConfig.Connections.url_dataset_v2download + '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/download/pre', + json={'result': {'payload': {'hash_code': 'hash_code'}, 'target_names': ['test.txt']}}, + status_code=200, + ) + + dataset_download_cliet = SrvDatasetDownloadManager('output_path', 'dataset_code', 'dataset_geid') + + res = dataset_download_cliet.pre_dataset_download() + assert res.get('result').get('payload').get('hash_code') == 'hash_code' + + +def test_dateset_pre_download_status_waiting(httpx_mock, mocker): + dataset_download_cliet = SrvDatasetDownloadManager('output_path', 'dataset_code', 'dataset_geid') + + mocker.patch('app.services.user_authentication.token_manager.SrvTokenManager.check_valid', return_value=0) + + httpx_mock.add_response( + method='GET', + url=AppConfig.Connections.url_download_core + f'v1/download/status/{dataset_download_cliet.hash_code}', + json={'result': {'status': 'WAITING'}}, + status_code=200, + ) + + status = dataset_download_cliet.download_status() + assert status == EFileStatus.WAITING + + +def test_dateset_pre_download_status_success(httpx_mock, mocker): + dataset_download_cliet = SrvDatasetDownloadManager('output_path', 'dataset_code', 'dataset_geid') + + mocker.patch('app.services.user_authentication.token_manager.SrvTokenManager.check_valid', return_value=0) + + httpx_mock.add_response( + method='GET', + url=AppConfig.Connections.url_download_core + f'v1/download/status/{dataset_download_cliet.hash_code}', + json={'result': {'status': 'SUCCEED'}}, + status_code=200, + ) + + status = dataset_download_cliet.download_status() + assert status == EFileStatus.SUCCEED + + +def test_check_dateset_pre_download_status(httpx_mock, mocker): + dataset_download_cliet = SrvDatasetDownloadManager('output_path', 'dataset_code', 'dataset_geid') + + mocker.patch('app.services.user_authentication.token_manager.SrvTokenManager.check_valid', return_value=0) + + httpx_mock.add_response( + method='GET', + url=AppConfig.Connections.url_download_core + f'v1/download/status/{dataset_download_cliet.hash_code}', + json={'result': {'status': 'SUCCEED'}}, + status_code=200, + ) + + status = dataset_download_cliet.check_download_preparing_status() + assert status == EFileStatus.SUCCEED + + +def test_avoid_duplicate_file_name_no_duplicate(mocker): + dataset_download_cliet = SrvDatasetDownloadManager('output_path', 'dataset_code', 'dataset_geid') + + mocker.patch('os.path.isfile', side_effect=[False]) + file_name = dataset_download_cliet.avoid_duplicate_file_name('test.txt') + + assert file_name == 'test.txt' + + +def test_avoid_duplicate_file_name_duplicate_once(mocker): + dataset_download_cliet = SrvDatasetDownloadManager('output_path', 'dataset_code', 'dataset_geid') + + mocker.patch('os.path.isfile', side_effect=[True, False]) + file_name = dataset_download_cliet.avoid_duplicate_file_name('test.txt') + + assert file_name == 'test (1).txt' + + +def test_download_dataset(httpx_mock, mocker): + dataset_download_cliet = SrvDatasetDownloadManager('output_path', 'dataset_code', 'dataset_geid') + + mocker.patch('app.services.user_authentication.token_manager.SrvTokenManager.check_valid', return_value=0) + mocker.patch( + 'app.services.dataset_manager.dataset_download.SrvDatasetDownloadManager.send_download_request', + return_value='test.zip', + ) + mocker.patch('os.path.isfile', return_value=True) + success_msg = mocker.patch( + 'app.services.output_manager.message_handler.SrvOutPutHandler.download_success', return_value=None + ) + + httpx_mock.add_response( + method='POST', + url=AppConfig.Connections.url_dataset_v2download + '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/download/pre', + json={'result': {'payload': {'hash_code': 'hash_code'}, 'target_names': ['test.txt']}}, + status_code=200, + ) + + httpx_mock.add_response( + method='GET', + url=AppConfig.Connections.url_download_core + 'v1/download/status/hash_code', + json={'result': {'status': 'SUCCEED'}}, + status_code=200, + ) + + dataset_download_cliet.download_dataset() + success_msg.assert_called_once_with('test.zip') + + +def test_pre_dataset_version_download(httpx_mock, mocker): + dataset_download_cliet = SrvDatasetDownloadManager('output_path', 'dataset_code', 'dataset_geid') + + mocker.patch('app.services.user_authentication.token_manager.SrvTokenManager.check_valid', return_value=0) + + httpx_mock.add_response( + method='GET', + url=AppConfig.Connections.url_dataset + f'/{dataset_download_cliet.dataset_geid}/download/pre', + json={'result': {'source': 'test.txt'}}, + status_code=200, + ) + + res = dataset_download_cliet.pre_dataset_version_download() + assert res.get('result').get('source') == 'test.txt' + + +def test_download_dataset_version(httpx_mock, mocker): + dataset_download_cliet = SrvDatasetDownloadManager('output_path', 'dataset_code', 'dataset_geid') + + mocker.patch('app.services.user_authentication.token_manager.SrvTokenManager.check_valid', return_value=0) + mocker.patch( + 'app.services.dataset_manager.dataset_download.SrvDatasetDownloadManager.send_download_request', + return_value='test.zip', + ) + mocker.patch('os.path.isfile', return_value=True) + success_msg = mocker.patch( + 'app.services.output_manager.message_handler.SrvOutPutHandler.download_success', return_value=None + ) + + httpx_mock.add_response( + method='GET', + url=AppConfig.Connections.url_dataset + + f'/{dataset_download_cliet.dataset_geid}/download/pre?version=test_version', + json={'result': {'source': 'test.txt'}}, + status_code=200, + ) + + dataset_download_cliet.download_dataset_version('test_version') + success_msg.assert_called_once_with('test.zip') diff --git a/tests/conftest.py b/tests/conftest.py index aaee5cc6..e0f7bb59 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,6 +21,8 @@ def mock_settings(monkeypatch): monkeypatch.setattr(AppConfig.Connections, 'url_authn', 'http://service_auth') monkeypatch.setattr(AppConfig.Connections, 'url_bff', 'http://bff_cli') monkeypatch.setattr(AppConfig.Connections, 'url_dataset', 'http://url_dataset') + monkeypatch.setattr(AppConfig.Connections, 'url_dataset_v2download', 'http://url_dataset_download_v2') + monkeypatch.setattr(AppConfig.Connections, 'url_download_core', 'http://url_dataset_download_core') monkeypatch.setattr(AppConfig.Connections, 'url_upload_greenroom', 'http://upload_gr') monkeypatch.setattr(AppConfig.Connections, 'url_upload_core', 'http://upload_core') monkeypatch.setattr(UserConfig, 'username', 'test-user')