72 lines
2.8 KiB
Python
72 lines
2.8 KiB
Python
|
import unittest
|
||
|
|
from unittest.mock import patch, MagicMock
|
||
|
|
from botocore.exceptions import ClientError
|
||
|
|
import boto3
|
||
|
|
from pywxdump.file.S3Attachment import S3Attachment
|
||
|
|
|
||
|
|
|
||
|
|
class TestS3Attachment(unittest.TestCase):
|
||
|
|
|
||
|
|
@patch('boto3.client')
|
||
|
|
def setUp(self, mock_client):
|
||
|
|
self.mock_client = MagicMock()
|
||
|
|
mock_client.return_value = self.mock_client
|
||
|
|
s3_config = MagicMock()
|
||
|
|
self.attachment = S3Attachment(s3_config)
|
||
|
|
self.test_s3_url = "s3://test_bucket/test_file"
|
||
|
|
self.test_s3_dir = "s3://test_bucket/test_folder/"
|
||
|
|
|
||
|
|
@patch.object(S3Attachment, 'exists')
|
||
|
|
@patch.object(S3Attachment, 'isFolder')
|
||
|
|
@patch('boto3.client')
|
||
|
|
def test_removal_of_existing_file(self, mock_client, mock_isFolder, mock_exists):
|
||
|
|
mock_exists.return_value = True
|
||
|
|
mock_isFolder.return_value = False
|
||
|
|
mock_client.return_value = MagicMock()
|
||
|
|
self.assertTrue(self.attachment.remove(self.test_s3_url))
|
||
|
|
|
||
|
|
@patch.object(S3Attachment, 'exists')
|
||
|
|
def test_removal_of_non_existing_file(self, mock_exists):
|
||
|
|
mock_exists.return_value = False
|
||
|
|
with self.assertRaises(FileNotFoundError):
|
||
|
|
self.attachment.remove(self.test_s3_url)
|
||
|
|
|
||
|
|
@patch.object(S3Attachment, 'exists')
|
||
|
|
@patch.object(S3Attachment, 'isFolder')
|
||
|
|
def test_removal_of_folder_instead_of_file(self, mock_isFolder, mock_exists):
|
||
|
|
mock_exists.return_value = True
|
||
|
|
mock_isFolder.return_value = True
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
self.attachment.remove(self.test_s3_url)
|
||
|
|
|
||
|
|
@patch.object(S3Attachment, 'exists')
|
||
|
|
@patch.object(S3Attachment, 'isdir')
|
||
|
|
def test_removal_with_s3_error(self, mock_isdir, mock_exists):
|
||
|
|
mock_exists.return_value = True
|
||
|
|
mock_isdir.return_value = False
|
||
|
|
# 模拟 ClientError
|
||
|
|
error_response = {'Error': {'Code': 'InvalidRequest', 'Message': 'Some error message'}}
|
||
|
|
self.mock_client.delete_object.side_effect = ClientError(error_response, 'delete_object')
|
||
|
|
with self.assertRaises(ClientError):
|
||
|
|
self.attachment.remove(self.test_s3_url)
|
||
|
|
|
||
|
|
@patch.object(S3Attachment, 'exists')
|
||
|
|
def test_getsize_existing_file(self, mock_exists):
|
||
|
|
mock_exists.return_value = True
|
||
|
|
self.mock_client.head_object.return_value = {'ContentLength': 100}
|
||
|
|
self.assertEqual(self.attachment.getsize(self.test_s3_url), 100)
|
||
|
|
|
||
|
|
def test_getsize_non_existing_file(self):
|
||
|
|
with self.assertRaises(FileNotFoundError):
|
||
|
|
self.attachment.getsize('non_existing_file')
|
||
|
|
|
||
|
|
@patch.object(S3Attachment, 'isdir')
|
||
|
|
def test_getsize_existing_folder(self, mock_isdir):
|
||
|
|
mock_isdir.return_value = True
|
||
|
|
self.mock_client.list_objects_v2.return_value = {'Contents': [{'Size': 100}, {'Size': 300}]}
|
||
|
|
self.assertEqual(self.attachment.getsize(self.test_s3_dir), 400)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|