时间戳

首页 > 解决方案 > 在检查状态时对 While 循环进行单元测试

python - 在检查状态时对 While 循环进行单元测试

问题描述

我正在与 AWS Athena 合作以获得结果。我必须启动一个查询，然后检查它是否已完成。

我现在正在尝试为各种状态编写单元测试。这是一个示例代码。我从另一个函数生成 athena 连接并将其传递给该函数，以及执行 ID。

def check_athena_status(athena, execution):
    running = True
    print('Checking Athena Execution Running State')
    while running:
        running_state = athena.get_query_execution(QueryExecutionId=execution)['QueryExecution']['Status']['State']
        if running_state == 'SUCCEEDED':
            print('Run SUCCEEDED')
            running = False
        elif running_state == 'RUNNING':
            time.sleep(3)
            print('Athena Query Still Running')
        else:
            raise RuntimeError('Athena Query Failed')

    return True

我基本上是想弄清楚是否有一种方法可以将 running_state 的值从 RUNNING 更改为 SUCCEEDED。我目前将此用作成功运行的单元测试。

    athena_succeed = mock.Mock()
    execution_id = 'RandomID'
    athena_succeed.get_query_execution.return_value = test_data.athena_succeeded

    result = inventory_validator.check_athena_status(athena_succeed, execution_id)
    assert result == True

其中 test_data.athena_succeeded 基本上是一个字典

athena_succeed = {'QueryExecution': {
        'Status': {'State': 'SUCCEEDED',
                   'SubmissionDateTime': '2021-08-08'}
    }
}

我也有一个“正在运行”的。

athena_running = {'QueryExecution': {
        'Status': {'State': 'RUNNING',
                   'SubmissionDateTime': '2021-08-08'}
    }
}

我正在尝试测试分支，所以我想从跑步走向成功。我知道我可以更改 while 真实值，但我想在循环中间更改实际的“雅典娜响应”。我尝试使用 PropertyMock，但我不确定那是正确的用例。

标签： pythonunit-testingmocking

解决方案

用于side_effect更改连续调用的返回值。

class TestCheckAthenaStatus(unittest.TestCase):

    def test_check_athena_status_from_running_to_succeeded(self):
        athena_running_succeeded = mock.Mock()
        execution_id = 'RandomID'
        athena_running_succeeded.get_query_execution.side_effect = (test_data.athena_running, test_data.athena_succeeded)

        result = inventory_validator.check_athena_status(athena_running_succeeded, execution_id)
        assert result == True
        assert athena_running_succeeded.get_query_execution.call_count == 2

推荐阅读