From 8b6f84d899c8239b8b511757cffb7637eb1923e2 Mon Sep 17 00:00:00 2001
From: Tim Heap <tim@timheap.me>
Date: Tue, 29 Mar 2016 06:34:25 +0200
Subject: [PATCH] Add tests for snippet edit handler form classes

---
 wagtail/tests/snippets/forms.py               |  7 +++++
 .../0003_fancysnippet_standardsnippet.py      | 28 +++++++++++++++++++
 wagtail/tests/snippets/models.py              | 12 ++++++++
 wagtail/wagtailsnippets/tests.py              | 19 ++++++++++++-
 4 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 wagtail/tests/snippets/forms.py
 create mode 100644 wagtail/tests/snippets/migrations/0003_fancysnippet_standardsnippet.py

diff --git a/wagtail/tests/snippets/forms.py b/wagtail/tests/snippets/forms.py
new file mode 100644
index 0000000000..1ce9403937
--- /dev/null
+++ b/wagtail/tests/snippets/forms.py
@@ -0,0 +1,7 @@
+from wagtail.wagtailadmin.forms import WagtailAdminModelForm
+
+
+class FancySnippetForm(WagtailAdminModelForm):
+    """
+    A custom form class for FancySnippets in the admin
+    """
diff --git a/wagtail/tests/snippets/migrations/0003_fancysnippet_standardsnippet.py b/wagtail/tests/snippets/migrations/0003_fancysnippet_standardsnippet.py
new file mode 100644
index 0000000000..e528434f03
--- /dev/null
+++ b/wagtail/tests/snippets/migrations/0003_fancysnippet_standardsnippet.py
@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.9.4 on 2016-03-29 04:28
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('snippetstests', '0002_searchablesnippet'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='FancySnippet',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+            ],
+        ),
+        migrations.CreateModel(
+            name='StandardSnippet',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('text', models.CharField(max_length=255)),
+            ],
+        ),
+    ]
diff --git a/wagtail/tests/snippets/models.py b/wagtail/tests/snippets/models.py
index 93382ca69e..d2c9af0b68 100644
--- a/wagtail/tests/snippets/models.py
+++ b/wagtail/tests/snippets/models.py
@@ -4,6 +4,8 @@ from django.utils.encoding import python_2_unicode_compatible
 from wagtail.wagtailsearch import index
 from wagtail.wagtailsnippets.models import register_snippet
 
+from .forms import FancySnippetForm
+
 
 # AlphaSnippet and ZuluSnippet are for testing ordering of
 # snippets when registering.  They are named as such to ensure
@@ -52,3 +54,13 @@ class SearchableSnippet(models.Model, index.Indexed):
 
     def __str__(self):
         return self.text
+
+
+@register_snippet
+class StandardSnippet(models.Model):
+    text = models.CharField(max_length=255)
+
+
+@register_snippet
+class FancySnippet(models.Model):
+    base_form_class = FancySnippetForm
diff --git a/wagtail/wagtailsnippets/tests.py b/wagtail/wagtailsnippets/tests.py
index c918ce5d75..8afa8f4a90 100644
--- a/wagtail/wagtailsnippets/tests.py
+++ b/wagtail/wagtailsnippets/tests.py
@@ -6,10 +6,13 @@ from django.test import TestCase
 from django.test.utils import override_settings
 from taggit.models import Tag
 
+from wagtail.tests.snippets.forms import FancySnippetForm
 from wagtail.tests.snippets.models import (
-    AlphaSnippet, RegisterDecorator, RegisterFunction, SearchableSnippet, ZuluSnippet)
+    AlphaSnippet, FancySnippet, RegisterDecorator, RegisterFunction, SearchableSnippet,
+    StandardSnippet, ZuluSnippet)
 from wagtail.tests.testapp.models import Advert, AdvertWithTabbedInterface, SnippetChooserModel
 from wagtail.tests.utils import WagtailTestUtils
+from wagtail.wagtailadmin.forms import WagtailAdminModelForm
 from wagtail.wagtailcore.models import Page
 from wagtail.wagtailsnippets.edit_handlers import SnippetChooserPanel
 from wagtail.wagtailsnippets.models import SNIPPET_MODELS, register_snippet
@@ -628,3 +631,17 @@ class TestDeleteOnlyPermissions(TestCase, WagtailTestUtils):
         response = self.client.get(reverse('wagtailsnippets:delete', args=('tests', 'advert', self.test_snippet.id, )))
         self.assertEqual(response.status_code, 200)
         self.assertTemplateUsed(response, 'wagtailsnippets/snippets/confirm_delete.html')
+
+
+class TestSnippetEditHandlers(TestCase, WagtailTestUtils):
+    def test_standard_edit_handler(self):
+        edit_handler_class = get_snippet_edit_handler(StandardSnippet)
+        form_class = edit_handler_class.get_form_class(StandardSnippet)
+        self.assertTrue(issubclass(form_class, WagtailAdminModelForm))
+        self.assertFalse(issubclass(form_class, FancySnippetForm))
+
+    def test_fancy_edit_handler(self):
+        edit_handler_class = get_snippet_edit_handler(FancySnippet)
+        form_class = edit_handler_class.get_form_class(FancySnippet)
+        self.assertTrue(issubclass(form_class, WagtailAdminModelForm))
+        self.assertTrue(issubclass(form_class, FancySnippetForm))