import tempfile
import shutil
from django.conf import settings
from django.test.runner import DiscoverRunner

__author__ = 'zmbq'


class TempMediaDiscoverRunnerMixin(DiscoverRunner):
    """ A mixin class that makes sure all media created during the test is deleted """
    def setup_test_environment(self, **kwargs):
        self._old_media_root = settings.MEDIA_ROOT
        self._old_file_storage = settings.DEFAULT_FILE_STORAGE

        settings.DEFAULT_FILE_STORAGE = 'django.core.files.storage.FileSystemStorage'
        settings.MEDIA_ROOT = tempfile.mkdtemp(suffix='_test')
        print("Changing MEDIA_ROOT to %s" % settings.MEDIA_ROOT)

        super().setup_test_environment(**kwargs)

    def teardown_test_environment(self, **kwargs):
        to_be_deleted = settings.MEDIA_ROOT

        settings.MEDIA_ROOT = self._old_media_root
        settings.DEFAULT_FILE_STORAGE = self._old_file_storage

        print("Removing temporary media folder")
        # Delete the files created during the test
        try:
            shutil.rmtree(to_be_deleted)
        except:
            pass