From 9e6a56765f0229184859a547fee76a7516628a50 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              | 18 +++++++++++-
 4 files changed, 64 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 000000000..1ce940393
--- /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 000000000..e528434f0
--- /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 dbb8a274b..11a6f42eb 100644
--- a/wagtail/tests/snippets/models.py
+++ b/wagtail/tests/snippets/models.py
@@ -5,6 +5,8 @@ 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
@@ -53,3 +55,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 353001dfa..42aedb1de 100644
--- a/wagtail/wagtailsnippets/tests.py
+++ b/wagtail/wagtailsnippets/tests.py
@@ -9,10 +9,12 @@ from taggit.models import Tag
 
 from wagtail.tests.utils import WagtailTestUtils
 from wagtail.tests.testapp.models import Advert, SnippetChooserModel, AdvertWithTabbedInterface
+from wagtail.tests.snippets.forms import FancySnippetForm
 from wagtail.tests.snippets.models import (
-    AlphaSnippet, ZuluSnippet, RegisterDecorator, RegisterFunction, SearchableSnippet
+    AlphaSnippet, ZuluSnippet, RegisterDecorator, RegisterFunction, SearchableSnippet, StandardSnippet, FancySnippet
 )
 from wagtail.wagtailsnippets.models import register_snippet, SNIPPET_MODELS
+from wagtail.wagtailadmin.forms import WagtailAdminModelForm
 from wagtail.wagtailsnippets.edit_handlers import SnippetChooserPanel
 
 from wagtail.wagtailsnippets.views.snippets import (
@@ -633,3 +635,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))