Skip to content

Instantly share code, notes, and snippets.

@nhannguyen95
Last active August 2, 2018 10:25
Show Gist options
  • Save nhannguyen95/3ec1f98ec471f8cf9009d1eaff769aa4 to your computer and use it in GitHub Desktop.
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
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