Using a generator with Python mock to replicate server responses

Question

I would like to use a list (converted into a generator) to serve as a mock for my API calls (using unittest.mock). My function is:

def monitor_order(order_id)
    order_info = client.get_order_status(order_id)
    order_status = order_info['status']

    while order_status != 'filled':
        print('order_status: ', order_status)
        time.sleep(5)
        order_info = client.get_order_status(order_id)
        order_status = order_info['status']

    return order_info

My test function is:

@patch('my_package.client.get_order_status')
def test_monitor_order(mocked_get_order_status):
    order_states = [
        dict(status='open'),
        dict(status='open'),
        dict(status='filled'),
    ]

    # Make into a generator
    status_changes = (status for status in order_states)
    mocked_get_order_status.return_value = next(order_states)

    # Execute function to test
    monitor_order("dummy_order")

However, I can see that the status is always 'open' when executing the test:

order_status:  open
order_status:  open
order_status:  open

I think I understand why it's wrong, but how could I implement it correctly?

Answer 1

To achieve what you want, you can rewrite your test as follows:

@patch('my_package.client.get_order_status')
def test_monitor_order(mocked_get_order_status):
    order_states = [
        dict(status='open'),
        dict(status='open'),
        dict(status='filled'),
    ]

    mocked_get_order_status.side_effect = order_states

    # Execute function to test
    monitor_order("dummy_order")