CODEWITHSUNDEEP
pythonunit-testingmocking
Mallachar
Mallachar

Reputation: 241

Unit Test a While Loop while checking state

I am working with AWS Athena to get results. I have to initiate a query and then check to see if its completed.

I am now trying to write a unit test for the various states. Here is a sample code. I generate the athena connection from another function and hand it off to this function, as well as the execution 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

I am basically trying to figure out is there a way where I can change the value of running_state from RUNNING to SUCCEEDED. I currently use this as the unit test for a successful run.

    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

where test_data.athena_succeeded is basically a dict

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

I also have a "RUNNING" one.

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

I am trying to test branches so I want to go from running to succeed. I know I can change the while true value, but I want to change the actual "athena response" in the middle of the loop. I tried with PropertyMock but I am not sure thats the right use case.

Upvotes: 2

Views: 974

Answers (2)

wim
wim

Reputation: 362756

This will fully test all branches of the code:

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

A couple of points to note:

  • time.sleep is mocked out so that the test runs immediately rather than taking 3 seconds, yet we assert that there was a delay intended between consecutive calls to the query.
  • the expected QueryExecutionId argument was asserted, and we check that it was called twice.
  • the status output text from the print statements is asserted.

The fixture mocker is provided by the plugin pytest-mock.

Upvotes: 1

aaron
aaron

Reputation: 43083

Use side_effect to change the return value on successive calls.

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

Upvotes: 1

Related Questions