From 29d138d16d0c3b27f931a535c829e85618069013 Mon Sep 17 00:00:00 2001
From: Tim Heap <tim@timheap.me>
Date: Thu, 13 Oct 2016 12:11:32 +1100
Subject: [PATCH] Move get_image_model, add get_image_model_string

Having `get_image_model` at `wagtail.wagtailimages` is consistent with
django.contrib.auth.get_user_model, and means developers can import
`get_image_model` in an environment where models are not yet ready.

`get_image_model_string` has been added for uses where the model itself
is not required, and might not be available, but a reference to the
possibly swapped out Image model is required.
---
 wagtail/wagtailimages/__init__.py             | 25 ++++++++++
 wagtail/wagtailimages/models.py               | 31 ++++--------
 wagtail/wagtailimages/tests/test_rich_text.py | 13 +++--
 wagtail/wagtailimages/tests/tests.py          | 50 +++++++++++++++++++
 wagtail/wagtailimages/tests/utils.py          |  2 +-
 5 files changed, 91 insertions(+), 30 deletions(-)

diff --git a/wagtail/wagtailimages/__init__.py b/wagtail/wagtailimages/__init__.py
index 61397a911..01092a060 100644
--- a/wagtail/wagtailimages/__init__.py
+++ b/wagtail/wagtailimages/__init__.py
@@ -1 +1,26 @@
+from __future__ import absolute_import, unicode_literals
+
+from django.conf import settings
+from django.core.exceptions import ImproperlyConfigured
+
+
 default_app_config = 'wagtail.wagtailimages.apps.WagtailImagesAppConfig'
+
+
+def get_image_model_string():
+    """Get the dotted app.Model name for the image model"""
+    return getattr(settings, 'WAGTAILIMAGES_IMAGE_MODEL', 'wagtailimages.Image')
+
+
+def get_image_model():
+    """Get the image model from WAGTAILIMAGES_IMAGE_MODEL."""
+    from django.apps import apps
+    model_string = get_image_model_string()
+    try:
+        return apps.get_model(model_string)
+    except ValueError:
+        raise ImproperlyConfigured("WAGTAILIMAGES_IMAGE_MODEL must be of the form 'app_label.model_name'")
+    except LookupError:
+        raise ImproperlyConfigured(
+            "WAGTAILIMAGES_IMAGE_MODEL refers to model '%s' that has not been installed" % model_string
+        )
diff --git a/wagtail/wagtailimages/models.py b/wagtail/wagtailimages/models.py
index b4511b220..36e7700f6 100644
--- a/wagtail/wagtailimages/models.py
+++ b/wagtail/wagtailimages/models.py
@@ -10,7 +10,6 @@ from contextlib import contextmanager
 import django
 from django.conf import settings
 from django.core import checks
-from django.core.exceptions import ImproperlyConfigured
 from django.core.files import File
 from django.core.urlresolvers import reverse
 from django.db import models
@@ -27,7 +26,7 @@ from taggit.managers import TaggableManager
 from unidecode import unidecode
 from willow.image import Image as WillowImage
 
-from wagtail.utils.deprecation import RemovedInWagtail19Warning
+from wagtail.utils.deprecation import RemovedInWagtail19Warning, RemovedInWagtail110Warning
 from wagtail.wagtailadmin.utils import get_object_usage
 from wagtail.wagtailcore import hooks
 from wagtail.wagtailcore.models import CollectionMember
@@ -48,6 +47,14 @@ class ImageQuerySet(SearchableQuerySetMixin, models.QuerySet):
     pass
 
 
+def get_image_model():
+    warnings.warn("wagtail.wagtailimages.models.get_image_model "
+                  "has been moved to wagtail.wagtailimages.get_image_model",
+                  RemovedInWagtail110Warning)
+    from wagtail.wagtailimages import get_image_model
+    return get_image_model()
+
+
 def get_upload_to(instance, filename):
     """
     Obtain a valid upload path for an image file.
@@ -359,26 +366,6 @@ def image_delete(sender, instance, **kwargs):
     instance.file.delete(False)
 
 
-def get_image_model():
-    from django.conf import settings
-    from django.apps import apps
-
-    try:
-        app_label, model_name = settings.WAGTAILIMAGES_IMAGE_MODEL.split('.')
-    except AttributeError:
-        return Image
-    except ValueError:
-        raise ImproperlyConfigured("WAGTAILIMAGES_IMAGE_MODEL must be of the form 'app_label.model_name'")
-
-    image_model = apps.get_model(app_label, model_name)
-    if image_model is None:
-        raise ImproperlyConfigured(
-            "WAGTAILIMAGES_IMAGE_MODEL refers to model '%s' that has not been installed" %
-            settings.WAGTAILIMAGES_IMAGE_MODEL
-        )
-    return image_model
-
-
 class Filter(models.Model):
     """
     Represents one or more operations that can be applied to an Image to produce a rendition
diff --git a/wagtail/wagtailimages/tests/test_rich_text.py b/wagtail/wagtailimages/tests/test_rich_text.py
index f446a544a..318f24a18 100644
--- a/wagtail/wagtailimages/tests/test_rich_text.py
+++ b/wagtail/wagtailimages/tests/test_rich_text.py
@@ -2,10 +2,11 @@ from __future__ import absolute_import, unicode_literals
 
 from bs4 import BeautifulSoup
 from django.test import TestCase
-from mock import patch
 
 from wagtail.wagtailimages.rich_text import ImageEmbedHandler
 
+from .utils import Image, get_test_image_file
+
 
 class TestImageEmbedHandler(TestCase):
     def test_get_db_attributes(self):
@@ -27,9 +28,8 @@ class TestImageEmbedHandler(TestCase):
         )
         self.assertEqual(result, '<img>')
 
-    @patch('wagtail.wagtailimages.models.Image')
-    @patch('django.core.files.File')
-    def test_expand_db_attributes_not_for_editor(self, mock_file, mock_image):
+    def test_expand_db_attributes_not_for_editor(self):
+        Image.objects.create(id=1, title='Test', file=get_test_image_file())
         result = ImageEmbedHandler.expand_db_attributes(
             {'id': 1,
              'alt': 'test-alt',
@@ -38,9 +38,8 @@ class TestImageEmbedHandler(TestCase):
         )
         self.assertIn('<img class="richtext-image left"', result)
 
-    @patch('wagtail.wagtailimages.models.Image')
-    @patch('django.core.files.File')
-    def test_expand_db_attributes_for_editor(self, mock_file, mock_image):
+    def test_expand_db_attributes_for_editor(self):
+        Image.objects.create(id=1, title='Test', file=get_test_image_file())
         result = ImageEmbedHandler.expand_db_attributes(
             {'id': 1,
              'alt': 'test-alt',
diff --git a/wagtail/wagtailimages/tests/tests.py b/wagtail/wagtailimages/tests/tests.py
index 2255b81b4..00e623625 100644
--- a/wagtail/wagtailimages/tests/tests.py
+++ b/wagtail/wagtailimages/tests/tests.py
@@ -2,8 +2,10 @@ from __future__ import absolute_import, unicode_literals
 
 import os
 import unittest
+import warnings
 
 from django import forms, template
+from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 from django.core.urlresolvers import reverse
 from django.test import TestCase, override_settings
@@ -13,6 +15,8 @@ from taggit.forms import TagField, TagWidget
 
 from wagtail.tests.testapp.models import CustomImage, CustomImageFilePath
 from wagtail.tests.utils import WagtailTestUtils
+from wagtail.utils.deprecation import RemovedInWagtail110Warning
+from wagtail.wagtailimages import get_image_model, get_image_model_string
 from wagtail.wagtailimages.fields import WagtailImageField
 from wagtail.wagtailimages.formats import Format, get_image_format, register_image_format
 from wagtail.wagtailimages.forms import get_image_form
@@ -580,3 +584,49 @@ class TestDifferentUpload(TestCase):
         # The files should be uploaded based on it's content, not just
         # it's filename
         self.assertFalse(image.file.url == second_image.file.url)
+
+
+class TestGetImageModel(WagtailTestUtils, TestCase):
+    @override_settings(WAGTAILIMAGES_IMAGE_MODEL='tests.CustomImage')
+    def test_custom_get_image_model(self):
+        """Test get_image_model with a custom image model"""
+        self.assertIs(get_image_model(), CustomImage)
+
+    @override_settings(WAGTAILIMAGES_IMAGE_MODEL='tests.CustomImage')
+    def test_custom_get_image_model_string(self):
+        """Test get_image_model_string with a custom image model"""
+        self.assertEqual(get_image_model_string(), 'tests.CustomImage')
+
+    @override_settings()
+    def test_standard_get_image_model(self):
+        """Test get_image_model with no WAGTAILIMAGES_IMAGE_MODEL"""
+        del settings.WAGTAILIMAGES_IMAGE_MODEL
+        from wagtail.wagtailimages.models import Image
+        self.assertIs(get_image_model(), Image)
+
+    @override_settings()
+    def test_standard_get_image_model_string(self):
+        """Test get_image_model_STRING with no WAGTAILIMAGES_IMAGE_MODEL"""
+        del settings.WAGTAILIMAGES_IMAGE_MODEL
+        self.assertEqual(get_image_model_string(), 'wagtailimages.Image')
+
+    @override_settings(WAGTAILIMAGES_IMAGE_MODEL='tests.UnknownModel')
+    def test_unknown_get_image_model(self):
+        """Test get_image_model with an unknown model"""
+        with self.assertRaises(ImproperlyConfigured):
+            get_image_model()
+
+    @override_settings(WAGTAILIMAGES_IMAGE_MODEL='invalid-string')
+    def test_invalid_get_image_model(self):
+        """Test get_image_model with an invalid model string"""
+        with self.assertRaises(ImproperlyConfigured):
+            get_image_model()
+
+    def test_deprecated_get_image_model(self):
+        from wagtail.wagtailimages.models import get_image_model
+        with warnings.catch_warnings(record=True) as ws:
+            warnings.simplefilter('always')
+
+            self.assertIs(Image, get_image_model())
+            self.assertEqual(len(ws), 1)
+            self.assertIs(ws[0].category, RemovedInWagtail110Warning)
diff --git a/wagtail/wagtailimages/tests/utils.py b/wagtail/wagtailimages/tests/utils.py
index a8566f4a4..32b058873 100644
--- a/wagtail/wagtailimages/tests/utils.py
+++ b/wagtail/wagtailimages/tests/utils.py
@@ -4,7 +4,7 @@ import PIL.Image
 from django.core.files.images import ImageFile
 from django.utils.six import BytesIO
 
-from wagtail.wagtailimages.models import get_image_model
+from wagtail.wagtailimages import get_image_model
 
 Image = get_image_model()