Is it possible to test a while True loop with pytest (I try with a timeout)?

Question

I have a python function foo with a while True loop inside. For background: It is expected do stream info from the web, do some writing and run indefinitely. The asserts test if the writing was done correctly.

Clearly I need it to stop sometime, in order to test.

What I did was to run via multirpocessing and introduce a timeout there, however when I see the test coverage, the function which ran through the multiprocessing, are not marked as covered.

Question 1: Why does pytest now work this way?

Question 2: How can I make this work?

I was thinking it's probably because I technically exit the loop, so maybe pytest does not mark this as tested....

import time
import multiprocessing

def test_a_while_loop():
    # Start through multiprocessing in order to have a timeout.
    p = multiprocessing.Process(
        target=foo
        name="Foo",
    )
    try:
        p.start()
        # my timeout
        time.sleep(10)
        p.terminate()
    finally:
        # Cleanup.
        p.join()

    # Asserts below
    ...

More info

1. I looked into adding a decorator such as @pytest.mark.timeout(5), but that did not work and it stops the whole function, so I never get to the asserts. (as suggested here).
2. If I don't find a way, I will just test the parts, but ideally I would like to find a way to test by breaking the loop.
3. I know I can re-write my code in order to make it have a timeout, but that would mean changing the code to make it testable, which I don't think is a good design.
   4. Mocks I have not tried (as suggested here), because I don't believe I can mock what I do, since it writes info from the web. I need to actually see the "original" working.

Vile

"that would mean changing the code to make it testable, which I don't think is a good design" - I don't think so. Writing code with testability in mind is the best way to go, but refactoring afterwards to make it testable is a good thing, too, IMO - mostly it also means to make the code more modular. – Allurement

asyncio and websockets - what framework are you using, aiohttp? You don't need to run an endless loop for that. Can you show a bit of code that is more related to your use case? – Yclept

@Yclept the code looks very similar to this one: speakerdeck.com/pydataamsterdam/… Let me know if this helps. – Vile

Use an async for msg in conn: ... loop instead of an endless loop. The loop stops iterating when the connection is closed, which is pretty easy to test with a custom server mock. – Yclept

@Yclept that's a good idea. I did that before I changed to while True. What is the benefit of using this vs while True? – Vile

Because it's a lot easier to test? Replacing the websockets connection with an async iterator is easier than coding a custom object that memorizes the receive calls and raises to simulate a connection closing. Also an async for is simply more pythonic than a while True if you ask me. – Yclept

Okay will check that. Thanks. My progress on the issue, I found out that it has to do with coverage not accounting for it, but the unit test actually works. I describe it here: #61144358 – Vile

Answer 1

Break out the functionality you want to test into a helper method. Test the helper method.

def scrape_web_info(url):
  data = get_it(url)
  return data

# In production:

while True:
  scrape_web_info(...)

# During test:

def test_web_info():
  assert scrape_web_info(...) == ...

Bioscopy

Hm... Mine is a bit more complicated because in short it: connects to a websocket and then fetches the data. This is done asynchronously with async/await. I will see if it can be broken like that... : ) Thank for the input – Vile

Answer 2

Actually, I had the same problem with an endless task to test and make coverage. However, In my code, there is a .run_forever() method that runs a .run_once() method inside an infinite loop. So, I can write a unit test for the .run_once() method to test its functionality. Nevertheless, if you want to test your forever function despite the Halting Problem for getting more extent code coverage, I propose the following approach using a timeout regardless of tools you've mentioned including multiprocessing or @pytest.mark.timeout(5) which didn't work for me either:

• First, install the interruptingcow PyPI package to have a nice timeout for raising an optional exception: pip install interruptingcow
• Then:

import pytest
import asyncio
from interruptingcow import timeout
from <path-to-the-loop-module> import EventLoop

class TestCase:
    @pytest.mark.parametrize("test_case", ['none'])
    def test_events(self, test_case: list):
        assert EventLoop().run_once()  # It's usual

    @pytest.mark.parametrize("test_case", ['none'])
    def test_events2(self, test_case: list):
        try:
            with timeout(10, exception=asyncio.CancelledError):
                EventLoop().run_forever()
                assert False
        except asyncio.CancelledError:
            assert True

Lease

Answer 3

If you can modify the code under test then you can use a class variable for the while loop condition. Then your test can mock that variable to cause the loop to exit.

from unittest import mock


class Consumer:
    RUN = True

    def __init__(self, service):
        self._service = service

    def poll_forever(self):
        i = 0
        while self.RUN:
            # do the work
            self._service.update(i)
            i += 1


@mock.patch.object(Consumer, "RUN", new_callable=mock.PropertyMock)
def test_consumer(mocked):
    service_mock = mock.Mock()
    service_mock.update = mock.MagicMock()
    mocked.side_effect = [True, False] # will cause the loop to exit on the second iteration
    consumer = Consumer(service_mock)
    consumer.poll_forever()
    service_mock.update.assert_called_with(0)

Shahjahanpur

RUN seems to become a mock method, so when you end up actually checking self.RUN it is not a bool value that you patch in the side_effect – Meggy

When inspecting the type of self.RUN in the poll_forever method, I see that the type of self.RUN is bool. I'm using a Python 3.10 environment. – Shahjahanpur

Answer 4

Yes, it is possible and the code above shows one way to do it (run through a multiprocessing with a timeout).

Since the asserts were running fine, I found out that the issue was not the pytest, but the coverage report not accounting for the multiprocessing properly.

I describe how I fix this (now separate) issue question here.

Vile

Answer 5

Another option if you can modify the code under test is to pass an optional keep_looping function to be called as the while loop condition. The default value can always return True, and your test can pass in a function that returns False at some point to cause the loop to exit.

For example:

from unittest import mock


def poll_forever(keep_looping=lambda: True):
    while keep_looping():
        # do the work
        pass


def test_poll_forever():
    # this will cause the loop to exit on the second iteration
    run_once = mock.Mock(side_effect=[True, False])
    poll_forever(keep_looping=run_once)

Paternoster