Pilot 2999: fixup the dataset download issue

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 + '/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)

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,
     )

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 + '/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 + '/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')

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')