Python unittest中无法Mock私有方法返回值的问题排查

问题场景 #

我需要为DBTProjectUploader.copy方法编写单元测试，该方法负责将S3对象从一个路径复制到另一个路径。尝试Mock私有方法__get_bucket_name_and_prefixes后，测试断言失败，实际调用参数显示为Mock ID而非预期的字符串值。

核心逻辑代码 #

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"  

测试代码 #

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

问题原因及解决方法 #

1. 私有方法Mock返回值不匹配 #

原私有方法__get_bucket_name_and_prefixes返回5个值的元组：(old_bucket, old_bucket_name, old_prefix, new_bucket, new_prefix)，但测试代码中只返回了字符串"bucket-name"，导致copy方法解构赋值时出错，最终传给list_objects的第一个参数是Mock对象（而非预期字符串）。

2. 修正Mock返回值 #

调整patch的return_value，使其匹配原方法的返回结构：

@patch('provisioner.src.services.dbt_project_uploader.DBTProjectUploader._DBTProjectUploader__get_bucket_name_and_prefixes',
    return_value=(source_bucket, 'bucket-name', 'old-prefix/', dest_bucket, 'new_prefix/dags/dbt/'))

3. 修正断言逻辑 #

file_manager.list_objects的第一个参数是old_bucket（Mock的Bucket对象），而非字符串'bucket-name'，所以断言需要修改为：

file_manager.list_objects.assert_called_once_with(source_bucket, "old-prefix/")

4. 可选：直接Mock__get_bucket_and_key简化测试 #

如果不需要测试私有方法__get_bucket_name_and_prefixes的逻辑，可以同时Mock__get_bucket_and_key，避免依赖内部实现：

@patch('provisioner.src.services.dbt_project_uploader.DBTProjectUploader._DBTProjectUploader__get_bucket_and_key')
def test_upload_dbt_project(self, mock_get_bucket_and_key):
    # 配置Mock返回值
    mock_get_bucket_and_key.side_effect = [
        ('bucket-name', 'old-prefix/'),  # 第一次调用：处理dag_file_location
        ('bucket-name', 'new_prefix/dags/dbt/')  # 第二次调用：处理新路径
    ]
    # 后续测试逻辑不变

内容的提问来源于stack exchange，提问作者RushHour