Last active
August 2, 2018 10:25
-
-
Save nhannguyen95/3ec1f98ec471f8cf9009d1eaff769aa4 to your computer and use it in GitHub Desktop.
Use file storage on the local filesystem in Django unit tests, worked with Django >= 1.9
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import shutil | |
import tempfile | |
from django.conf import settings | |
from django.core.files.storage import FileSystemStorage | |
from django.db.models import FileField | |
from django.apps.apps import get_model, get_models | |
from django.test.runner import DiscoverRunner | |
class LocalStorageDiscoverRunner(DiscoverRunner): | |
""" | |
Use file storage on the local filesystem in unit tests, worked with Django >= 1.9. | |
Inspiration: https://www.caktusgroup.com/blog/2013/06/26/media-root-and-django-tests/ | |
""" | |
def setup_test_environment(self): | |
super().setup_test_environment() | |
# Keep track of original storages. | |
settings._original_media_root = settings.MEDIA_ROOT | |
settings._original_file_storage = settings.DEFAULT_FILE_STORAGE | |
settings._original_fields_storages = {} | |
# Creates a temporary directory. | |
settings._temp_media_dir = tempfile.mkdtemp(dir=settings.BASE_DIR) | |
# Use the FileSystemStorage for tests. | |
settings.MEDIA_ROOT = settings._temp_media_dir | |
settings.DEFAULT_FILE_STORAGE = 'django.core.files.storage.FileSystemStorage' | |
# Use the FileSystemStorage for all model fields. | |
for model in get_models(): | |
fields = [f for f in model._meta.fields if isinstance(f, FileField)] | |
for field in fields: | |
model_path = '%s.%s' % (model._meta.app_label, model._meta.model_name) | |
original_storage = (field.name, field.storage) | |
original_storages = settings._original_fields_storages.setdefault(model_path, []) | |
original_storages.append(original_storage) | |
field.storage = FileSystemStorage(location=settings.MEDIA_ROOT) | |
def teardown_test_environment(self): | |
super().teardown_test_environment() | |
# Delete the temporary directory. | |
shutil.rmtree(settings.MEDIA_ROOT, ignore_errors=True) | |
# Restore original storage. | |
settings.MEDIA_ROOT = settings._original_media_root | |
settings.DEFAULT_FILE_STORAGE = settings._original_file_storage | |
# Restore original storages for all model fields. | |
for model_path, original_storages in settings._original_fields_storages.items(): | |
model = get_model(model_path) | |
for field_name, original_storage in original_storages: | |
field = model._meta.get_field(field_name) | |
field.storage = original_storage | |
del settings._original_media_root | |
del settings._original_file_storage | |
del settings._original_fields_storages |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment