时间戳

首页 > 解决方案 > 在模块级别缓存值和单元测试

python-3.x - 在模块级别缓存值和单元测试

问题描述

下面是一个用于查询和缓存 AWS STS 令牌的模块，目的是避免在存在有效令牌时查询 STS。

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


CREDENTIALS_ = Credentials()


def get_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_可变的，由其他测试修改，我可以将此值设置为None，但我想知道什么是更清洁的清除缓存值的方法

def test_get_credentials_refreshes_token_if_about_to_expire(sts_response, 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()
            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)

标签： python-3.xpytest

解决方案

更简洁的方法是确保您的unit测试正在执行单元测试。这意味着对于每个单元，都不应与其他单元交互。由于您使用的是全局变量CREDENTIALS_，因此这几乎是不可能的。

1）容易修复

CREDENTIALS_一个简单的解决方法是作为输入参数传递。CREDENTIALS_然后，您可以在每次测试期间创建一个根据您的测试条件量身定制的假对象。

2）更好的修复

除了使用credential输入参数之外，更好的解决方案是分解get_credentials. 通过将其拆分为更小的函数，您可以将服务器逻辑和凭证更新分开。使其更容易模拟和测试。整个功能的可能划分是：

  • get_sts_token
  • update_credentials
  • 获取凭据

现在get_sts_token有与服务器的连接，但不必直接与它交互update_credentials。get_credentials

代码

示例 1)

def update_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


# Where you need the credentials
CREDENTIALS_ = update_credentials(CREDENTIALS_)
CREDENTIALS_.sts_credentials

现在您可以CREDENTIALS_在测试中插入您自己的对象。

例 2)

def get_sts_token():
    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
    return sts_token

def update_credentials(credentials, sts_token):
    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']
        }
    return credentials

def get_credentials(credentials: Credentials):
    if credentials.is_token_expired():
        sts_token = get_sts_token()
        credentials = update_credentials(credentials, sts_token)
    return credentials.sts_credentials

推荐阅读