Python修补类-方法return_value返回MagicMock

Question

I have a function that determines if a token can be pulled from environment variables, or from Google Cloud Secret Manager (via a custom class).

I am mocking the secret manager connection class so it doesn't actually connect to GCP, however the pull_secret() method I have from it is returning another MagicMock object.

Here's the function code ( 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

The test code I am using is as follows:

@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"

The Error I get is AssertionError: assert <MagicMock name='SecretManagerConnection().pull_secret()' id='2533589158400'> == '1234567890.abcdefghijklmnopqrstuvwxyz'

Is there something I am doing wrong in trying to assign the pull_secret() return value?

Answer 1

Cause of the problem

Patching SecretManagerConnection class will replace it with a magic mock instance. When this class is instantiated in the fetch_token function, the instantiated class will return the return_value of mock instance which in your case is not properly configured/set.

How to fix this?

Configure/set the return_value for pull_secret method on the return_value of mocked class rather than setting on the mocked class itself.

@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"