Thursday, 14 October 2010

"with Directory" makes testing easier

One of the problems with writing tests for muddle was finding a clean way to write them in Python.
Testing muddle involves a lot of directory creation and moving between directories. Using os.chdir() and friends doesn't lead to easily readable code, and there is always the problem of forgetting to move back out of a particular directory.
My first attempt to make this easier was to write simple pushd() and popd() functions, which maintained a stack of directories. This was a little better, but still didn't solve the "forgetting" problem.
Of course, the solution is obvious -- use with.
Thus we have:
class Directory(object):
    """A class to facilitate pushd/popd behaviour

    It is intended for use with 'with', as in::

        with Directory('~'):
            print 'My home directory contains'
            print ' ',' '.join(os.listdir('.'))
    """
    def __init__(self, where, verbose=True):
        self.start = normalise(os.getcwd())
        self.where = normalise(where)
        self.verbose = verbose
        os.chdir(self.where)
        if verbose:
            print '++ pushd to %s'%self.where

    def close(self):
        os.chdir(self.start)
        if self.verbose:
            print '++ popd to %s'%self.start

    def __enter__(self):
        return self

    def __exit__(self, etype, value, tb):
        if tb is None:
            # No exception, so just finish normally
            self.close()
        else:
            # An exception occurred, so do any tidying up necessary
            if self.verbose:
                print '** Oops, an exception occurred - %s tidying up'%self.__class__.__name__
            # well, there isn't anything special to do, really
            self.close()
            if self.verbose:
                print '** ----------------------------------------------------------------------'
            # And allow the exception to be re-raised
            return False
and its obvious friends:
class NewDirectory(Directory):
    """A pushd/popd directory that gets created first.

    It is an Error if the directory already exists.
    """
    def __init__(self, where, verbose=True):
        where = normalise(where)
        if os.path.exists(where):
            raise Error('Directory %s already exists'%where)
        if verbose:
            print '++ mkdir %s'%where
        os.makedirs(where)
        super(NewDirectory, self).__init__(where, verbose)

class TransientDirectory(NewDirectory):
    """A pushd/popd directory that gets created first and deleted afterwards

    If 'keep_on_error' is True, then the directory will not be deleted
    if an exception occurs in its 'with' clause.

    It is an Error if the directory already exists.
    """
    def __init__(self, where, keep_on_error=False, verbose=True):
        self.rmtree_on_error = not keep_on_error
        super(TransientDirectory, self).__init__(where, verbose)

    def close(self, delete_tree):
        super(NewDirectory, self).close()
        if delete_tree:
            if self.verbose:
                # The extra space after 'rmtree' is so the directory name
                # left aligns with a previous 'popd to' message
                print '++ rmtree  %s'%self.where
            shutil.rmtree(self.where)

    def __exit__(self, etype, value, tb):
        if tb is None:
            # No exception, so just finish normally
            self.close(True)
        else:
            # An exception occurred, so do any tidying up necessary
            if self.verbose:
                print '** Oops, an exception occurred - %s tidying up'%self.__class__.__name__
            # but don't delete the tree if we've been asked not to
            self.close(self.rmtree_on_error)
            if self.verbose:
                print '** ----------------------------------------------------------------------'
            # And allow the exception to be re-raised
            return False
and I can now write code like:
root_repo = 'file://' + os.path.join(root_dir, 'repo')
with NewDirectory('test_build1'):
    banner('Bootstrapping checkout build')
    muddle(['bootstrap', 'git+%s'%root_repo, 'test_build'])
    cat('src/builds/01.py')

    banner('Setting up src/')
    with Directory('src'):
        with Directory('builds'):
            touch('01.py', CHECKOUT_BUILD)
            git('add 01.py')
            git('commit -m "New build"')
            git('push %s/builds HEAD'%root_repo)

        with NewDirectory('checkout1'):
            touch('Makefile.muddle', MUDDLE_MAKEFILE)
            git('init')
            git('add Makefile.muddle')
            git('commit -m "Add muddle makefile"')
            git('push %s/checkout1 HEAD'%root_repo)
            muddle(['assert', 'checkout:checkout1/checked_out'])
which seems to fit with how I want to think about the tests.

No comments:

Post a Comment