<commit_before>import os
import shutil
import tempfile
import logging

logging.basicConfig(level=logging.DEBUG)

KEEP_FILES = os.environ.get("KUBEFACE_TEST_KEEP_FILES")


def with_bucket_storage(function):
    bucket = os.environ.get("KUBEFACE_BUCKET")
    if not bucket:
        logging.fatal("No bucket defined")

    def test_function():
        function("gs://" + bucket)
    return test_function


def with_local_storage(function):
    def test_function():
        tempdir = None
        try:
            tempdir = tempfile.mkdtemp(dir='/tmp')
            function(tempdir)
        finally:
            if tempdir and not KEEP_FILES:
                shutil.rmtree(tempdir)
    return test_function
<commit_msg>Wipe bucket after tests<commit_after>import os
import shutil
import tempfile
import logging

import kubeface
import kubeface.bucket_storage


logging.basicConfig(level=logging.DEBUG)

KEEP_FILES = os.environ.get("KUBEFACE_TEST_KEEP_FILES")


def wipe_bucket(bucket_url):
    objects = kubeface.bucket_storage.list_contents(bucket_url)
    for obj in objects:
        kubeface.bucket_storage.delete(bucket_url + "/" + obj)


def check_empty(bucket_url):
    assert not kubeface.bucket_storage.list_contents(bucket_url)


def with_bucket_storage(function):
    bucket = os.environ.get("KUBEFACE_BUCKET")
    if not bucket:
        logging.fatal("No bucket defined")

    def test_function():
        check_empty("gs://" + bucket)
        function("gs://" + bucket)
        wipe_bucket("gs://" + bucket)
    return test_function


def with_local_storage(function):
    def test_function():
        tempdir = None
        try:
            tempdir = tempfile.mkdtemp(dir='/tmp')
            function(tempdir)
        finally:
            if tempdir and not KEEP_FILES:
                shutil.rmtree(tempdir)
    return test_function

{'fraction_non_alphanumeric': 0.09332609875203472, 'fraction_numerical': 0.0, 'mean_word_length': 3.017429193899782, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6045845', 'n_tokens_mistral': 588, 'n_tokens_neox': 573, 'n_words': 118}