How to test a ‘while True’ in Python

This is a very short post to show the best way found to test a function or method that has a while true, a.k.a. infinite loop, in its implementation using Python.

Let’s go straight to the solution:

You will need a class attribute so that you can change while an instance of the class is executing. For example:

class Consumer:

    def __init__(self, service):
            self.__service = service
    
    RUNNING = True  # class attribute

    def start(self):
            
        while self.RUNNING:
            updates = self.__service.check_for_updates()
            print(updates)  # just and example

Given the class above, you could create a test similar to this:

class TestConsumer:

    @patch('path.to.Consumer')
    def test_should_check_for_updates_when_consumer_is_running(
        self,
        service_mock
    ):

        sentinel = PropertyMock(side_effect=[True, False])
        Consumer.RUNNING = sentinel

        my_consumer = Consumer(service_mock)
        my_consumer.start()

        service_mock.check_for_updates.assert_called()
        assert service_mock.check_for_updates.call_count == 1  # called only once because on the second lap RUNNING is False

You can run it:

 python -m pytest path/to/test-consumer.py

That’s it. There are other ways to do the same thing. However, if you need to get rid of this problem fast, this is the way to go.

Now, the problem that you might be facing before finding a solution is that your class probably looked something similar to this:

 class Consumer:

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

    def start(self):
            
        while True:
            updates = self.__service.check_for_updates()
            print(updates)

And when you run your test, the execution is hooked in while true loop forever, so you can’t run your test assertions. With the solution explained above, the class had to change a little by adding the RUNNING attribute, which was not your intention before. Nevertheless, it’s a tradeoff I think worth it.

If you have another approach to tackle this problem, please send me a message. I’ll be happy to learn new strategies.