模块级缓存AWS STS凭证的单元测试隔离问题及优化问询

这个问题太常见了——全局单例在测试环境里就是容易踩状态污染的坑！咱们不用粗暴地把CREDENTIALS_设为None，有几个更优雅的方式来解决这个测试隔离问题，一个个说：

方案1：给Credentials类添加重置方法（最小改动） #

最直接的思路是把缓存状态的重置逻辑封装到Credentials类内部，既符合OOP封装原则，又能方便地重置状态，不用在测试里直接修改它的私有属性。

修改你的Credentials类：

from datetime import datetime, timedelta

class Credentials:
    def __init__(self):
        self.sts_credentials = None
        self.token_expiry_time = None

    def is_token_expired(self):
        current_time_with_buffer = datetime.now() + timedelta(minutes=2)
        return not self.token_expiry_time or self.token_expiry_time < current_time_with_buffer

    # 新增重置方法
    def reset(self):
        """重置缓存的凭证和过期时间，用于测试隔离"""
        self.sts_credentials = None
        self.token_expiry_time = None

然后在你的测试用例前后调用这个方法，确保每次测试都从干净的状态开始：

def test_get_credentials_refreshes_token_if_about_to_expire(sts_response, credentials):
    # 测试前重置缓存状态
    CREDENTIALS_.reset()
    
    with mock.patch("boto3.client") as mock_boto_client:
        mock_assume_role = mock_boto_client.return_value.assume_role
        mock_assume_role.return_value = sts_response
        
        get_credentials()
        actual_credentials = get_credentials()
        
        calls = [
            call('sts'), 
            call().assume_role(RoleArn='arn:aws:iam::000000000000:role/dummyarn', RoleSessionName=ANY), 
            call('sts'), 
            call().assume_role(RoleArn='arn:aws:iam::000000000000:role/dummyarn', RoleSessionName=ANY)
        ]
        assert credentials == actual_credentials
        mock_boto_client.assert_has_calls(calls)

方案2：用测试框架的自动Fixture实现全局重置（适合pytest用户） #

如果你的项目用pytest，可以写一个自动生效的fixture，让它在每个测试用例前后自动重置缓存状态，不用手动在每个测试里写reset()，更省心：

import pytest
from your_module import CREDENTIALS_

@pytest.fixture(autouse=True)
def isolate_credentials_cache():
    """自动重置凭证缓存，确保测试间状态隔离"""
    # 测试前重置
    CREDENTIALS_.reset()
    yield
    # 测试后也可以再重置一次，双重保险
    CREDENTIALS_.reset()

这个fixture加上autouse=True后，会自动作用于所有测试用例，彻底解决跨测试的状态污染问题。

方案3：重构为可注入的缓存实例（最解耦，可扩展性强） #

如果想从根本上避免全局单例带来的问题，可以把get_credentials函数改成支持依赖注入的形式，让它可以接受自定义的Credentials实例，而不是硬编码使用模块级的CREDENTIALS_。

修改get_credentials函数：

CREDENTIALS_ = Credentials()

def get_credentials(credentials=CREDENTIALS_):
    if credentials.is_token_expired():
        sts_client = boto3.client('sts')
        LOGGER.info("The credentials are either empty or expiring, refreshing")
        try:
            sts_token = sts_client.assume_role(
                RoleArn=os.environ["KINESIS_ASSUME_ROLE"],
                RoleSessionName=str(uuid.uuid4()))
        except Exception as e:
            LOGGER.error(f"Error occurred while trying to assume role with {os.environ['KINESIS_ASSUME_ROLE']}", e)
            raise e
        credentials.sts_credentials = {
            "aws_access_key_id": sts_token['Credentials']['AccessKeyId'],
            "aws_secret_access_key": sts_token['Credentials']['SecretAccessKey'],
            "aws_session_token": sts_token['Credentials']['SessionToken']
        }
        credentials.token_expiry_time = sts_token["Credentials"]["Expiration"]
    return credentials.sts_credentials

然后在测试用例里创建独立的Credentials实例，完全隔离测试状态：

def test_get_credentials_refreshes_token_if_about_to_expire(sts_response, credentials):
    # 创建测试专用的凭证实例，和全局实例完全隔离
    test_credentials = Credentials()
    
    with mock.patch("boto3.client") as mock_boto_client:
        mock_assume_role = mock_boto_client.return_value.assume_role
        mock_assume_role.return_value = sts_response
        
        get_credentials(test_credentials)
        actual_credentials = get_credentials(test_credentials)
        
        calls = [
            call('sts'), 
            call().assume_role(RoleArn='arn:aws:iam::000000000000:role/dummyarn', RoleSessionName=ANY), 
            call('sts'), 
            call().assume_role(RoleArn='arn:aws:iam::000000000000:role/dummyarn', RoleSessionName=ANY)
        ]
        assert credentials == actual_credentials
        mock_boto_client.assert_has_calls(calls)

这种方式不仅解决了测试隔离问题，还让代码更灵活——以后如果需要多套独立的凭证缓存，直接创建新的Credentials实例即可，不用修改核心逻辑。

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