无法模拟python unittest中私有方法的返回值

Question

我的核心逻辑如下:

class DBTProjectUploader:
def __init__(self, file_manager: FileManager, s3_client):
self.file_manager = file_manager
self.s3_client = s3_client
def copy(self, dbt_attr: DbtAttr):
old_bucket, old_bucket_name, old_prefix, new_bucket, new_prefix = self.__get_bucket_name_and_prefixes(dbt_attr)
for f in self.file_manager.list_objects(old_bucket, old_prefix):
old_source = {'Bucket': old_bucket_name, 'Key': f.key}
new_key = f.key.replace(old_prefix, new_prefix, 1)
self.file_manager.copy_object(new_bucket, new_key, old_source)
return "Success"
def __get_bucket_name_and_prefixes(self, dbt_attr: DbtAttr):
old_bucket_name, old_prefix = self.__get_bucket_and_key(dbt_attr.dag_file_location)
new_bucket_name, new_prefix = self.__get_bucket_and_key("s3://" + old_bucket_name + "/dags/dbt/")
old_bucket = self.s3_client.Bucket(old_bucket_name)
new_bucket = self.s3_client.Bucket(new_bucket_name)
return old_bucket, old_bucket_name, old_prefix, new_bucket, new_prefix
def __get_bucket_and_key(self, path: str):
return path[5:].split('/', 1)
class FileManager:
logger = get_provisioner_logger()
def __init__(self, s3_client):
self.s3_client = s3_client

def list_objects(self, old_bucket, old_prefix):
return old_bucket.objects.filter(Prefix=old_prefix)
def copy_object(self, new_bucket, new_key, old_source):
new_obj = new_bucket.Object(new_key)
new_obj.copy(old_source)
return "Success"  

其中DBTProjectUploader.copy负责将S3嵌套对象从一个路径复制到另一个路径。现在，我正试图使用Python unit test为copy方法编写测试用例。但是，我无法模拟从私有方法返回的值，因为我在返回值中获得模拟id。

def mock_s3_object(self, key, body):
obj = MagicMock()
obj.key = key
obj.get.return_value = {'Body': MagicMock(read=MagicMock(return_value=body))}
return obj

@patch('boto3.resource')
@patch('provisioner.src.services.dbt_project_uploader.DBTProjectUploader._DBTProjectUploader__get_bucket_name_and_prefixes',
return_value="bucket-name")
def test_upload_dbt_project(self, mock_resource, mock_get_bucket_name_and_prefixes):
source_bucket = MagicMock()
dest_bucket = MagicMock()
mock_s3_client = MagicMock()
file_manager = MagicMock()
dbt_attr = DbtAttr("path", "s3://bucket-name/old-prefix/", "dbt_mwaa01")
mock_resource.Bucket.side_effect = [source_bucket, dest_bucket]
print(source_bucket)
source_bucket.objects.filter.return_value = [
self.mock_s3_object('new_prefix/dags/dbt/dbt_mwaa01/', b'Sample data')
]
uploader = DBTProjectUploader(file_manager, mock_s3_client)
result = uploader.copy(dbt_attr)
file_manager.list_objects.assert_called_once_with('bucket-name', "old-prefix/")
# source_bucket.objects.filter.assert_called_once_with("old-prefix/")
dest_bucket.Object.assert_called_once_with('new_prefix/dags/dbt/dbt_mwaa01/')
dest_bucket.Object.return_value.copy.assert_called_once_with(
{'Bucket': 'bucket-name', 'Key': 'new_prefix/dags/dbt/dbt_mwaa01/'})
self.assertEqual(result, 'Success')

下面给出错误:

Expected: list_objects('bucket-name', 'old-prefix/')
Actual: list_objects(<MagicMock name='mock.Bucket()' id='4572034384'>, 'old-prefix/')  

我已经尝试模拟私有方法，但是当我试图打印原始方法中的变量时，它只给出mock-id。它不返回值。我还遗漏了什么?TIA

Answer 1

注意__get_bucket_name_and_prefixes方法是私有的，它是由copy方法调用的，它在你的测试中没有被正确地模拟。因此，它返回的是一个模拟对象，而不是预期的值，这导致了测试失败。

试

@patch('boto3.resource')
def test_upload_dbt_project(self, mock_resource):
source_bucket = MagicMock()
dest_bucket = MagicMock()
mock_s3_client = MagicMock()
file_manager = FileManager(mock_s3_client)
dbt_attr = DbtAttr("path", "s3://bucket-name/old-prefix/", "dbt_mwaa01")
mock_resource.return_value.Bucket.side_effect = [source_bucket, dest_bucket]
mock_resource.return_value.meta.client.head_bucket.side_effect = botocore.exceptions.ClientError(
{'Error': {'Code': '404', 'Message': 'Not Found'}},
'HeadBucket'
)
source_bucket.objects.filter.return_value = [
self.mock_s3_object('new_prefix/dags/dbt/dbt_mwaa01/', b'Sample data')
]
with patch.object(DBTProjectUploader, '_DBTProjectUploader__get_bucket_name_and_prefixes') as mock_get_bucket_name_and_prefixes:
mock_get_bucket_name_and_prefixes.return_value = (source_bucket, 'bucket-name', 'old-prefix', dest_bucket, 'new_prefix/')
uploader = DBTProjectUploader(file_manager, mock_s3_client)
result = uploader.copy(dbt_attr)
file_manager.list_objects.assert_called_once_with(source_bucket, "old-prefix")
source_bucket.objects.filter.assert_called_once_with(Prefix="old-prefix")
dest_bucket.Object.assert_called_once_with('new_prefix/dags/dbt/dbt_mwaa01/')
dest_bucket.Object.return_value.copy.assert_called_once_with(
{'Bucket': 'bucket-name', 'Key': 'new_prefix/dags/dbt/dbt_mwaa01/'})
self.assertEqual(result, 'Success')