无法成功Mock Azure ContainerClient的list_blobs()函数，求排查解决

问题背景 #

我尝试在单元测试中Mock Azure ContainerClient的list_blobs()函数，期望让它返回['Blob1', 'Blob2']列表，但测试执行后输出的是MagicMock对象，函数并未按预期被Mock。

脚本文件 #

import sys
from datetime import datetime, timedelta
import pyspark
import pytz
import yaml

# from azure.storage.blob import BlobServiceClient, ContainerClient
from pyspark.dbutils import DBUtils as dbutils
import azure.storage.blob

# Open Config
def main():

    spark_context = pyspark.SparkContext.getOrCreate()
    spark_context.addFile(sys.argv[1])
    stream = None
    stream = open(sys.argv[1], "r")
    config = yaml.load(stream, Loader=yaml.FullLoader)
    stream.close()

    account_key = dbutils.secrets.get(scope=config["Secrets"]["Scope"], key=config["Secrets"]["Key Name"])

    target_container = config["Storage Configuration"]["Container"]
    target_account = config["Storage Configuration"]["Account"]
    days_history_to_keep = config["Storage Configuration"]["Days History To Keep"]

    connection_string = (
        "DefaultEndpointsProtocol=https;AccountName="
        + target_account
        + ";AccountKey="
        + account_key
        + ";EndpointSuffix=core.windows.net"
    )
    blob_service_client: azure.storage.blob.BlobServiceClient = (
        azure.storage.blob.BlobServiceClient.from_connection_string(connection_string)
    )
    container_client: azure.storage.blob.ContainerClient = (
        blob_service_client.get_container_client(target_container)
    )
    blobs = container_client.list_blobs()
    print(blobs)
    print(blobs)

    utc = pytz.UTC
    delete_before_date = utc.localize(
        datetime.today() - timedelta(days=days_history_to_keep)
    )

    for blob in blobs:
        if blob.creation_time < delete_before_date:
            print("Deleting Blob: " + blob.name)
            container_client.delete_blob(blob, delete_snapshots="include")


if __name__ == "__main__":
    main()

测试文件 #

import unittest
from unittest import mock
import DeleteOldBlobs


class DeleteBlobsTest(unittest.TestCase):
    def setUp(self):
        pass

    @mock.patch("DeleteOldBlobs.azure.storage.blob.ContainerClient")
    @mock.patch("DeleteOldBlobs.azure.storage.blob.BlobServiceClient")
    @mock.patch("DeleteOldBlobs.dbutils")
    @mock.patch("DeleteOldBlobs.sys")
    @mock.patch('DeleteOldBlobs.pyspark')
    def test_main(self, mock_pyspark, mock_sys, mock_dbutils, mock_blobserviceclient, mock_containerclient):

        # mock setup
        config_file = "Delete_Old_Blobs_UnitTest.yml"
        mock_sys.argv = ["unused_arg", config_file]
        mock_dbutils.secrets.get.return_value = "A Secret"
        mock_containerclient.list_blobs.return_value = ["ablob1", "ablob2"]

        # execute test
        DeleteOldBlobs.main()

        # TODO assert actions taken
        # mock_sys.argv.__get__.assert_called_with()
        # dbutils.secrets.get(scope=config['Secrets']['Scope'], key=config['Secrets']['Key Name'])


if __name__ == "__main__":

    unittest.main()

输出结果 #

<MagicMock name='BlobServiceClient.from_connection_string().get_container_client().list_blobs()' id='1143355577232'>

错误原因 #

你当前的Mock逻辑存在两个关键问题：

• Mock目标不匹配：你Mock的是ContainerClient类本身，但实际代码中container_client是通过blob_service_client.get_container_client()返回的实例，该实例并非直接由你Mock的ContainerClient类实例化而来，而是BlobServiceClient.from_connection_string()返回的实例调用方法得到的。
• 方法绑定错误：你给mock_containerclient.list_blobs.return_value赋值，是在设置类方法的返回值，但代码中调用的是具体实例的list_blobs()方法，两者没有关联。

修正方案 #

需要构建完整的Mock调用链，确保blob_service_client.get_container_client()返回的实例就是你配置了list_blobs返回值的Mock实例。以下是两种可行的修正方式：

方式一：直接构建BlobServiceClient调用链 #

移除对ContainerClient的单独patch，直接在BlobServiceClient的Mock链上关联ContainerClient的Mock实例：

import unittest
from unittest import mock
import DeleteOldBlobs


class DeleteBlobsTest(unittest.TestCase):
    def setUp(self):
        pass

    @mock.patch("DeleteOldBlobs.azure.storage.blob.BlobServiceClient")
    @mock.patch("DeleteOldBlobs.dbutils")
    @mock.patch("DeleteOldBlobs.sys")
    @mock.patch('DeleteOldBlobs.pyspark')
    def test_main(self, mock_pyspark, mock_sys, mock_dbutils, mock_blobserviceclient):
        # 创建ContainerClient的Mock实例并配置返回值
        mock_containerclient = mock.Mock()
        mock_containerclient.list_blobs.return_value = ["ablob1", "ablob2"]

        # 配置系统参数和密钥返回值
        config_file = "Delete_Old_Blobs_UnitTest.yml"
        mock_sys.argv = ["unused_arg", config_file]
        mock_dbutils.secrets.get.return_value = "A Secret"

        # 构建BlobServiceClient的调用链：from_connection_string返回的实例调用get_container_client返回mock_containerclient
        mock_blob_service_instance = mock_blobserviceclient.from_connection_string.return_value
        mock_blob_service_instance.get_container_client.return_value = mock_containerclient

        # 执行测试
        DeleteOldBlobs.main()

        # 验证list_blobs是否被正确调用
        mock_containerclient.list_blobs.assert_called_once()

方式二：保留ContainerClient的patch并关联调用链 #

如果需要保留对ContainerClient类的patch，需确保类的实例被正确关联到调用链中：

import unittest
from unittest import mock
import DeleteOldBlobs


class DeleteBlobsTest(unittest.TestCase):
    def setUp(self):
        pass

    @mock.patch("DeleteOldBlobs.azure.storage.blob.ContainerClient")
    @mock.patch("DeleteOldBlobs.azure.storage.blob.BlobServiceClient")
    @mock.patch("DeleteOldBlobs.dbutils")
    @mock.patch("DeleteOldBlobs.sys")
    @mock.patch('DeleteOldBlobs.pyspark')
    def test_main(self, mock_pyspark, mock_sys, mock_dbutils, mock_blobserviceclient, mock_containerclient_cls):
        # 获取ContainerClient类的Mock实例并配置返回值
        mock_containerclient = mock_containerclient_cls.return_value
        mock_containerclient.list_blobs.return_value = ["ablob1", "ablob2"]

        # 配置系统参数和密钥返回值
        config_file = "Delete_Old_Blobs_UnitTest.yml"
        mock_sys.argv = ["unused_arg", config_file]
        mock_dbutils.secrets.get.return_value = "A Secret"

        # 构建BlobServiceClient的调用链
        mock_blob_service_instance = mock_blobserviceclient.from_connection_string.return_value
        mock_blob_service_instance.get_container_client.return_value = mock_containerclient

        # 执行测试
        DeleteOldBlobs.main()

        # 验证调用
        mock_containerclient.list_blobs.assert_called_once()

关键说明 #

核心逻辑是确保代码中获取的container_client实例，与你配置了list_blobs返回值的Mock实例是同一个对象，这样调用list_blobs()时才会返回你指定的列表，而非默认的MagicMock对象。

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