diff options
| -rw-r--r-- | gn3/file_utils.py | 29 | ||||
| -rw-r--r-- | tests/unit/test_file_utils.py | 24 |
2 files changed, 24 insertions, 29 deletions
diff --git a/gn3/file_utils.py b/gn3/file_utils.py index cda4f2b..16b2f94 100644 --- a/gn3/file_utils.py +++ b/gn3/file_utils.py @@ -2,7 +2,6 @@ import hashlib import json import os -import shutil import random import string import tarfile @@ -59,27 +58,27 @@ def generate_random_n_string(n_length: int) -> str: for _ in range(n_length)) -def extract_uploaded_file(gzipped_file, target_dir: str) -> Dict: +def extract_uploaded_file(gzipped_file, target_dir: str, token="") -> Dict: """Get the (directory) hash of extracted contents of GZIPPED_FILE; and move contents to TARGET_DIR/<dir-hash>. """ - tar_target_loc = os.path.join(target_dir, - secure_filename(gzipped_file.filename)) - gzipped_file.save(tar_target_loc) + if not token: + token = (f"{generate_random_n_string(6)}-" + f"{generate_random_n_string(6)}") + tar_target_loc = os.path.join( + target_dir, + token, + secure_filename(gzipped_file.filename)) try: - # Extract to "tar_target_loc/tempdir" + if not os.path.exists(os.path.join(target_dir, token)): + os.mkdir(os.path.join(target_dir, token)) + gzipped_file.save(tar_target_loc) + # Extract to "tar_target_loc/token" tar = tarfile.open(tar_target_loc) - tar.extractall( - path=os.path.join(target_dir, "tempdir")) + tar.extractall(path=tar_target_loc) tar.close() # pylint: disable=W0703 except Exception: return {"status": 128, "error": "gzip failed to unpack file"} - dir_hash = get_dir_hash(tar_target_loc) - if os.path.exists(os.path.join(target_dir, dir_hash)): - shutil.rmtree(os.path.join(target_dir, 'tempdir')) - else: - os.rename(os.path.join(target_dir, "tempdir"), - os.path.join(target_dir, dir_hash)) - return {"status": 0, "token": dir_hash} + return {"status": 0, "token": token} diff --git a/tests/unit/test_file_utils.py b/tests/unit/test_file_utils.py index 166e576..7f4d83d 100644 --- a/tests/unit/test_file_utils.py +++ b/tests/unit/test_file_utils.py @@ -1,6 +1,5 @@ """Test cases for procedures defined in file_utils.py""" import os -import shutil import unittest from dataclasses import dataclass @@ -75,33 +74,30 @@ non-existent""" jsonfile_to_dict, "/non-existent-dir") + @mock.patch("gn3.file_utils.tarfile") @mock.patch("gn3.file_utils.secure_filename") - def test_extract_uploaded_file(self, mock_file): + def test_extract_uploaded_file(self, mock_file, mock_tarfile): """Test that the gzip file is extracted to the right location""" - file_loc = os.path.join( - os.path.dirname(__file__), - "upload-data.tar.gz") - mock_file.return_value = file_loc + mock_file.return_value = "upload-data.tar.gz" mock_fileobj = MockFile(save=mock.MagicMock(), filename="upload-data.tar.gz") - result = extract_uploaded_file(mock_fileobj, "/tmp") - mock_fileobj.save.assert_called_once_with(file_loc) + mock_tarfile.return_value = mock.Mock() + result = extract_uploaded_file(mock_fileobj, "/tmp", + token="abcdef-abcdef") + mock_fileobj.save.assert_called_once_with("/tmp/abcdef-abcdef/" + "upload-data.tar.gz") mock_file.assert_called_once_with("upload-data.tar.gz") - # Clean up! - shutil.rmtree(os.path.join("/tmp", - "d41d8cd98f00b204e9800998ecf8427e")) self.assertEqual(result, {"status": 0, - "token": "d41d8cd98f00b204e9800998ecf8427e"}) + "token": "abcdef-abcdef"}) @mock.patch("gn3.file_utils.secure_filename") def test_extract_uploaded_file_non_existent_gzip(self, mock_file): """Test that the right error message is returned when there is a problem extracting the file""" - file_loc = os.path.join( + mock_file.return_value = os.path.join( os.path.dirname(__file__), "CTtyodSTh5") # Does not exist! - mock_file.return_value = file_loc mock_fileobj = MockFile(save=mock.MagicMock(), filename="") result = extract_uploaded_file(mock_fileobj, "/tmp") |