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

Question

我正在与 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，但我不确定那是正确的用例。

Answer 1

这将全面测试代码的所有分支：

import pytest

from your_module import check_athena_status


def test_check_athena_status(mocker, capsys):
    mock_sleep = mocker.patch("your_module.time.sleep")
    mock_athena = mocker.Mock()
    mock_athena.get_query_execution.side_effect = [
        {"QueryExecution": {"Status": {"State": "RUNNING"}}},
        {"QueryExecution": {"Status": {"State": "SUCCEEDED"}}},
    ]
    result = check_athena_status(mock_athena, execution="RandomID")
    assert result is True
    mock_sleep.assert_called_once_with(3)
    mock_athena.get_query_execution.assert_has_calls([
        mocker.call(QueryExecutionId="RandomID"),
        mocker.call(QueryExecutionId="RandomID"),
    ])
    out, err = capsys.readouterr()
    assert err == ""
    assert out.splitlines() == [
        "Checking Athena Execution Running State",
        "Athena Query Still Running",
        "Run SUCCEEDED",
    ]


def test_check_athena_status_error(mocker):
    mock_athena = mocker.Mock()
    other = {"QueryExecution": {"Status": {"State": "OTHER"}}}
    mock_athena.get_query_execution.return_value = other
    with pytest.raises(RuntimeError, match="^Athena Query Failed$"):
        check_athena_status(mock_athena, execution="RandomID")

需要注意的几点：

• time.sleep 被模拟出来，以便测试立即运行而不是花费 3 秒，但我们断言在连续调用查询之间存在延迟有意。
• 预期的 QueryExecutionId 参数已断言，我们检查它是否被调用了两次。
• 打印语句的状态输出文本被断言。

夹具mocker由插件pytest-mock提供。

Answer 2

使用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