Context manager for regression testing makedirs race condition

easy_race.py

import os
import errno
import time
import shutil
from threading import Timer
from contextlib import contextmanager

@contextmanager
def makedirs_with_easy_race(pth, wait=0.5):
    """ Context manager for testing code that might suffer from a
        race condition when using os.makedirs. 
        It slows down os.makedirs and creates the directory in 
        another thread. 
        Make the wait long enough that it will predictably fail on 
        unpatched code.

        Use this like so:

        >>> with makedirs_with_easy_race('/tmp/123'):
        ...    os.makedirs('/tmp/123')
        ...    print "All clear"
        RaceCondition
        >>> with makedirs_with_easy_race('/tmp/234'):
        ...    safe_makedirs('/tmp/234')
        ...    print "All clear"
        All clear

        This is just a test, you have to write safe_makedirs yourself:-)

    """
    original_makedirs = os.makedirs

    def slow_makedirs(path, mode=0777, exists_ok=False):
        if not exists_ok and pth.startswith(path): 
            if not os.path.exists(path):
                time.sleep(wait)
                if os.path.exists(path):
                    err = OSError()
                    err.errno = errno.EEXIST
                    raise err
        original_makedirs(path, mode)

    # Setup another thread that creates the same directory
    def _in_another_thread():
        assert not os.path.exists(pth), "Precondition: path should not exist %s" % pth
        os.makedirs = original_makedirs
        os.makedirs(pth)
        os.makedirs = slow_makedirs
    Timer(0.5*wait, _in_another_thread).start()

    # Monkeypath and run relevant tests that call os.makedirs
    os.makedirs = slow_makedirs
    try:
        yield
    except OSError, e:
        if e.errno == errno.EEXIST:
            raise makedirs_with_easy_race.RaceCondition()
        else:
            raise
    os.makedirs = original_makedirs
makedirs_with_easy_race.RaceCondition = type('RaceCondition', (Exception,), {})