Python 修补类 - 方法 return_value 返回 MagicMock

Question

我有一个函数可以确定是否可以从环境变量或 Google Cloud Secret Manager(通过自定义类)中提取 token 。

我正在模拟 secret 管理器连接类，因此它实际上并没有连接到 GCP，但是我从中获得的 pull_secret() 方法正在返回另一个 MagicMock 对象。

这是函数代码(utility_functions.py):

def fetch_token(secret_name: str):
try:
    token = os.environ['TOKEN']
except KeyError:
    print("TOKEN environment variable not set. Pulling Secret from Google Cloud Secret Manager...")
    token = None

if token is None or token == '':
    gcp_sm = SecretManagerConnection()
    token = gcp_sm.pull_secret(project_id=os.environ['GCP_PROJECT_ID'], secret_name=secret_name)

return token

我使用的测试代码如下:

@patch("app.utilities.utility_functions.SecretManagerConnection")
def test_fetch_token_google_secret(mock_sm):
    os.environ["GCP_PROJECT_ID"] = "test"

    mock_sm.pull_secret().return_value = "1234567890.abcdefghijklmnopqrstuvwxyz"

    del os.environ["TOKEN"]

    assert fetch_token("test") == "1234567890.abcdefghijklmnopqrstuvwxyz"

我得到的错误是 AssertionError: assert <MagicMock name='SecretManagerConnection().pull_secret()' id='2533589158400'> == '1234567890.abcdefghijklmnopqrstuvwxyz'

我在尝试分配 pull_secret() 返回值时是否做错了什么？

Answer 1

问题原因

修补 SecretManagerConnection 类将用一个神奇的模拟实例替换它。当此类在 fetch_token 函数中实例化时，实例化的类将返回模拟实例的 return_value，在您的情况下未正确配置/设置。

如何解决这个问题？

在模拟类的 return_value 上为 pull_secret 方法配置/设置 return_value 而不是在模拟类本身上设置。

@patch("app.utilities.utility_functions.SecretManagerConnection")
def test_fetch_token_google_secret(mock_sm):
    os.environ["GCP_PROJECT_ID"] = "test"

    instance = mock_sm.return_value
    instance.pull_secret.return_value = "1234567890.abcdefghijklmnopqrstuvwxyz"

    del os.environ["TOKEN"]
    assert fetch_token("test") == "1234567890.abcdefghijklmnopqrstuvwxyz"