How to create timeout wrapper for unittest class

Question

I am trying to create a timeout wrapper for a unittest class using this question as a template: PyUnit - How to unit test a method that runs into an infinite loop for some input?

The above solution works well if running from Main, but I'm trying to create a similar timeout by using a unittest class, but it fails.

PicklingError: Can't pickle : it's not found as TimeTest.test1

I'm open to any other solutions that are platform independent and work with the native 2.7 library

Here is the code

import unittest
import multiprocessing as mp
import functools as ft
import time

def timeout(seconds=5, error_message="Timeout"):
    def decorator(func):
        def wrapper(*args, **kwargs):
            process = mp.Process(None, func, None, args, kwargs)
            process.start()
            process.join(seconds)
            if process.is_alive():
                process.terminate()
                raise mp.TimeoutError(error_message)

        return ft.wraps(func)(wrapper)
    return decorator

class TestCases(unittest.TestCase):

    def setUp(self):
        self.a = 0

    @timeout(1)
    def test1(self):
        time.sleep(2)
        self.assertEquals(self.a, 0)

dranobob · Answer 1 · 2018-07-11T02:44:56.260

I ended up ditching the wrapper idea and creating a simple function and assert.

import unittest
import multiprocessing as mp
import time

TIME_LIMIT = 1

class TestCases(unittest.TestCase):

    def setUp(self):
        self.a = 0

    def my_func(self):
        time.sleep(2)
        self.q.put(self.a + 1)

    def run_case(self, func):
        self.q = mp.Queue()

        test_process = mp.Process(target=func)

        test_process.start()
        test_process.join(TIME_LIMIT)

        self.assertFalse(test_process.is_alive(), 'timeout exceeded')

    def test_case1(self):
        self.run_case(self.my_func)
        self.assertEquals(self.a + 1, self.q.get())

How to create timeout wrapper for unittest class

1 Answers1