common.py 文件源码

def clear_collections(*collections):
    """
    Clears collections listed after function has completed.
    Will throw an assertion if any collection is not empty when called.
    """

    def clear(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            db = api.common.get_conn()
            try:
                result = f(*args, **kwargs)
            finally:
                #Clear the collections.
                for collection in collections:
                    db[collection].remove()
                #Ensure they are then empty.
                for collection in collections:
                    collection_size = lambda collection: len(list(db[collection].find()))
                    assert collection_size(collection) == 0, "Collection: {} was not able to be cleared.".format(collection)
            return result
        return wrapper
    return clear