Add tests for snippets with DraftStateMixin enabled

wagtail/snippets/tests/test_snippets.py

@@ -23,7 +23,7 @@ from wagtail.admin.admin_url_finder import AdminURLFinder
 from wagtail.admin.forms import WagtailAdminModelForm
 from wagtail.admin.panels import FieldPanel, ObjectList, Panel, get_edit_handler
 from wagtail.blocks.field_block import FieldBlockAdapter
-from wagtail.models import Locale, ModelLogEntry, Page
+from wagtail.models import Locale, ModelLogEntry, Page, Revision
 from wagtail.snippets.action_menu import (
     ActionMenuItem,
     get_base_snippet_action_menu_items,
@@ -54,6 +54,7 @@ from wagtail.test.testapp.models import (
     AdvertWithCustomPrimaryKey,
     AdvertWithCustomUUIDPrimaryKey,
     AdvertWithTabbedInterface,
+    DraftStateModel,
     RevisableChildModel,
     RevisableModel,
     SnippetChooserModel,
@@ -181,6 +182,28 @@ class TestSnippetListView(TestCase, WagtailTestUtils):
         self.assertContains(response, "Dummy Button")
         self.assertContains(response, "/dummy-button")
 
+    def test_use_latest_draft_as_title(self):
+        snippet = DraftStateModel.objects.create(text="Draft-enabled Foo, Published")
+        snippet.save_revision().publish()
+        snippet.text = "Draft-enabled Bar, In Draft"
+        snippet.save_revision()
+
+        response = self.client.get(
+            reverse("wagtailsnippets_tests_draftstatemodel:list"),
+        )
+
+        edit_url = reverse(
+            "wagtailsnippets_tests_draftstatemodel:edit",
+            args=[quote(snippet.pk)],
+        )
+
+        # Should use the latest draft title in the listing
+        self.assertContains(
+            response,
+            f'<a href="{edit_url}">Draft-enabled Bar, In Draft</a>',
+            html=True,
+        )
+
 
 @override_settings(WAGTAIL_I18N_ENABLED=True)
 class TestLocaleSelectorOnList(TestCase, WagtailTestUtils):
@@ -616,6 +639,99 @@ class TestLocaleSelectorOnCreate(TestCase, WagtailTestUtils):
         self.assertNotContains(response, "Switch locales")
 
 
+class TestCreateDraftStateSnippet(TestCase, WagtailTestUtils):
+    def setUp(self):
+        self.user = self.login()
+
+    def get(self):
+        return self.client.get(reverse("wagtailsnippets_tests_draftstatemodel:add"))
+
+    def post(self, post_data={}):
+        return self.client.post(
+            reverse("wagtailsnippets_tests_draftstatemodel:add"),
+            post_data,
+        )
+
+    def test_get(self):
+        response = self.get()
+
+        self.assertEqual(response.status_code, 200)
+        self.assertTemplateUsed(response, "wagtailsnippets/snippets/create.html")
+
+        # The save button should be labelled "Save draft"
+        self.assertContains(response, "Save draft")
+        # The publish button should exist
+        self.assertContains(response, "Publish")
+        # The publish button should have name="action-publish"
+        self.assertContains(
+            response,
+            '<button type="submit" name="action-publish" value="action-publish" class="button action-save button-longrunning" data-clicked-text="Publishing…">',
+        )
+        # The status side panel should not be shown
+        self.assertNotContains(
+            response,
+            '<div class="form-side__panel" data-side-panel="status">',
+        )
+
+    def test_save_draft(self):
+        response = self.post(post_data={"text": "Draft-enabled Foo"})
+        snippet = DraftStateModel.objects.get(text="Draft-enabled Foo")
+
+        self.assertRedirects(
+            response, reverse("wagtailsnippets_tests_draftstatemodel:list")
+        )
+
+        # The instance should be created
+        self.assertEqual(snippet.text, "Draft-enabled Foo")
+
+        # The instance should be a draft
+        self.assertFalse(snippet.live)
+        self.assertTrue(snippet.has_unpublished_changes)
+        self.assertIsNone(snippet.first_published_at)
+        self.assertIsNone(snippet.last_published_at)
+        self.assertIsNone(snippet.live_revision)
+
+        # A revision should be created and set as latest_revision
+        self.assertIsNotNone(snippet.latest_revision)
+
+        # The revision content should contain the data
+        self.assertEqual(snippet.latest_revision.content["text"], "Draft-enabled Foo")
+
+    def test_publish(self):
+        timestamp = now()
+        with freeze_time(timestamp):
+            response = self.post(
+                post_data={
+                    "text": "Draft-enabled Foo, Published",
+                    "action-publish": "action-publish",
+                }
+            )
+        snippet = DraftStateModel.objects.get(text="Draft-enabled Foo, Published")
+
+        self.assertRedirects(
+            response, reverse("wagtailsnippets_tests_draftstatemodel:list")
+        )
+
+        # The instance should be created
+        self.assertEqual(snippet.text, "Draft-enabled Foo, Published")
+
+        # The instance should be live
+        self.assertTrue(snippet.live)
+        self.assertFalse(snippet.has_unpublished_changes)
+        self.assertEqual(snippet.first_published_at, timestamp)
+        self.assertEqual(snippet.last_published_at, timestamp)
+
+        # A revision should be created and set as both latest_revision and live_revision
+        self.assertIsNotNone(snippet.live_revision)
+        self.assertEqual(snippet.live_revision, snippet.latest_revision)
+
+        # The revision content should contain the new data
+        self.assertEqual(
+            snippet.live_revision.content["text"],
+            "Draft-enabled Foo, Published",
+        )
+
+
 class BaseTestSnippetEditView(TestCase, WagtailTestUtils):
     def get(self, params={}):
         snippet = self.test_snippet
@@ -970,6 +1086,327 @@ class TestEditRevisionSnippet(BaseTestSnippetEditView):
         self.assertEqual(log_entries.first().revision, revision)
 
 
+class TestEditDraftStateSnippet(BaseTestSnippetEditView):
+    def setUp(self):
+        super().setUp()
+        self.test_snippet = DraftStateModel.objects.create(
+            text="Draft-enabled Foo", live=False
+        )
+
+    def test_get(self):
+        response = self.get()
+
+        self.assertEqual(response.status_code, 200)
+        self.assertTemplateUsed(response, "wagtailsnippets/snippets/edit.html")
+
+        # The save button should be labelled "Save draft"
+        self.assertContains(response, "Save draft")
+        # The publish button should exist
+        self.assertContains(response, "Publish")
+        # The publish button should have name="action-publish"
+        self.assertContains(
+            response,
+            '<button type="submit" name="action-publish" value="action-publish" class="button action-save button-longrunning" data-clicked-text="Publishing…">',
+        )
+
+    def test_save_draft(self):
+        response = self.post(post_data={"text": "Draft-enabled Bar"})
+        self.test_snippet.refresh_from_db()
+        revisions = Revision.objects.for_instance(self.test_snippet)
+        latest_revision = self.test_snippet.latest_revision
+
+        self.assertRedirects(
+            response, reverse("wagtailsnippets_tests_draftstatemodel:list")
+        )
+
+        # The instance should not be updated
+        self.assertEqual(self.test_snippet.text, "Draft-enabled Foo")
+
+        # The instance should be a draft
+        self.assertFalse(self.test_snippet.live)
+        self.assertTrue(self.test_snippet.has_unpublished_changes)
+        self.assertIsNone(self.test_snippet.first_published_at)
+        self.assertIsNone(self.test_snippet.last_published_at)
+        self.assertIsNone(self.test_snippet.live_revision)
+
+        # The revision should be created and set as latest_revision
+        self.assertEqual(revisions.count(), 1)
+        self.assertEqual(latest_revision, revisions.first())
+
+        # The revision content should contain the new data
+        self.assertEqual(latest_revision.content["text"], "Draft-enabled Bar")
+
+    def test_publish(self):
+        timestamp = now()
+        with freeze_time(timestamp):
+            response = self.post(
+                post_data={
+                    "text": "Draft-enabled Bar, Published",
+                    "action-publish": "action-publish",
+                }
+            )
+
+        self.test_snippet.refresh_from_db()
+        revisions = Revision.objects.for_instance(self.test_snippet)
+        latest_revision = self.test_snippet.latest_revision
+
+        log_entries = ModelLogEntry.objects.filter(
+            content_type=ContentType.objects.get_for_model(DraftStateModel),
+            action="wagtail.publish",
+            object_id=self.test_snippet.pk,
+        )
+        log_entry = log_entries.first()
+
+        self.assertRedirects(
+            response, reverse("wagtailsnippets_tests_draftstatemodel:list")
+        )
+
+        # The instance should be updated
+        self.assertEqual(self.test_snippet.text, "Draft-enabled Bar, Published")
+
+        # The instance should be live
+        self.assertTrue(self.test_snippet.live)
+        self.assertFalse(self.test_snippet.has_unpublished_changes)
+        self.assertEqual(self.test_snippet.first_published_at, timestamp)
+        self.assertEqual(self.test_snippet.last_published_at, timestamp)
+        self.assertEqual(self.test_snippet.live_revision, latest_revision)
+
+        # The revision should be created and set as latest_revision
+        self.assertEqual(revisions.count(), 1)
+        self.assertEqual(latest_revision, revisions.first())
+
+        # The revision content should contain the new data
+        self.assertEqual(
+            latest_revision.content["text"],
+            "Draft-enabled Bar, Published",
+        )
+
+        # A log entry with wagtail.publish action should be created
+        self.assertEqual(log_entries.count(), 1)
+        self.assertEqual(log_entry.timestamp, timestamp)
+
+    def test_save_draft_then_publish(self):
+        save_timestamp = now()
+        with freeze_time(save_timestamp):
+            self.test_snippet.text = "Draft-enabled Bar, In Draft"
+            self.test_snippet.save_revision()
+
+        publish_timestamp = now()
+        with freeze_time(publish_timestamp):
+            response = self.post(
+                post_data={
+                    "text": "Draft-enabled Bar, Now Published",
+                    "action-publish": "action-publish",
+                }
+            )
+
+        self.test_snippet.refresh_from_db()
+        revisions = Revision.objects.for_instance(self.test_snippet).order_by("pk")
+        latest_revision = self.test_snippet.latest_revision
+
+        self.assertRedirects(
+            response, reverse("wagtailsnippets_tests_draftstatemodel:list")
+        )
+
+        # The instance should be updated
+        self.assertEqual(self.test_snippet.text, "Draft-enabled Bar, Now Published")
+
+        # The instance should be live
+        self.assertTrue(self.test_snippet.live)
+        self.assertFalse(self.test_snippet.has_unpublished_changes)
+        self.assertEqual(self.test_snippet.first_published_at, publish_timestamp)
+        self.assertEqual(self.test_snippet.last_published_at, publish_timestamp)
+        self.assertEqual(self.test_snippet.live_revision, latest_revision)
+
+        # The revision should be created and set as latest_revision
+        self.assertEqual(revisions.count(), 2)
+        self.assertEqual(latest_revision, revisions.last())
+
+        # The revision content should contain the new data
+        self.assertEqual(
+            latest_revision.content["text"],
+            "Draft-enabled Bar, Now Published",
+        )
+
+    def test_publish_then_save_draft(self):
+        publish_timestamp = now()
+        with freeze_time(publish_timestamp):
+            self.test_snippet.text = "Draft-enabled Bar, Published"
+            self.test_snippet.save_revision().publish()
+
+        save_timestamp = now()
+        with freeze_time(save_timestamp):
+            response = self.post(
+                post_data={"text": "Draft-enabled Bar, Published and In Draft"}
+            )
+
+        self.test_snippet.refresh_from_db()
+        revisions = Revision.objects.for_instance(self.test_snippet).order_by("pk")
+        latest_revision = self.test_snippet.latest_revision
+
+        self.assertRedirects(
+            response, reverse("wagtailsnippets_tests_draftstatemodel:list")
+        )
+
+        # The instance should be updated with the last published changes
+        self.assertEqual(self.test_snippet.text, "Draft-enabled Bar, Published")
+
+        # The instance should be live
+        self.assertTrue(self.test_snippet.live)
+        # The instance should have unpublished changes
+        self.assertTrue(self.test_snippet.has_unpublished_changes)
+
+        self.assertEqual(self.test_snippet.first_published_at, publish_timestamp)
+        self.assertEqual(self.test_snippet.last_published_at, publish_timestamp)
+
+        # The live revision should be the first revision
+        self.assertEqual(self.test_snippet.live_revision, revisions.first())
+
+        # The second revision should be created and set as latest_revision
+        self.assertEqual(revisions.count(), 2)
+        self.assertEqual(latest_revision, revisions.last())
+
+        # The revision content should contain the new data
+        self.assertEqual(
+            latest_revision.content["text"],
+            "Draft-enabled Bar, Published and In Draft",
+        )
+
+    def test_publish_twice(self):
+        first_timestamp = now()
+        with freeze_time(first_timestamp):
+            self.test_snippet.text = "Draft-enabled Bar, Published Once"
+            self.test_snippet.save_revision().publish()
+
+        second_timestamp = now() + datetime.timedelta(days=1)
+        with freeze_time(second_timestamp):
+            response = self.post(
+                post_data={
+                    "text": "Draft-enabled Bar, Published Twice",
+                    "action-publish": "action-publish",
+                }
+            )
+
+        self.test_snippet.refresh_from_db()
+
+        revisions = Revision.objects.for_instance(self.test_snippet).order_by("pk")
+        latest_revision = self.test_snippet.latest_revision
+
+        self.assertRedirects(
+            response, reverse("wagtailsnippets_tests_draftstatemodel:list")
+        )
+
+        # The instance should be updated with the last published changes
+        self.assertEqual(self.test_snippet.text, "Draft-enabled Bar, Published Twice")
+
+        # The instance should be live
+        self.assertTrue(self.test_snippet.live)
+        self.assertFalse(self.test_snippet.has_unpublished_changes)
+
+        # The first_published_at and last_published_at should be set correctly
+        self.assertEqual(self.test_snippet.first_published_at, first_timestamp)
+        self.assertEqual(self.test_snippet.last_published_at, second_timestamp)
+
+        # The live revision should be the second revision
+        self.assertEqual(self.test_snippet.live_revision, revisions.last())
+
+        # The second revision should be created and set as latest_revision
+        self.assertEqual(revisions.count(), 2)
+        self.assertEqual(latest_revision, revisions.last())
+
+        # The revision content should contain the new data
+        self.assertEqual(
+            latest_revision.content["text"],
+            "Draft-enabled Bar, Published Twice",
+        )
+
+    def test_get_after_save_draft(self):
+        self.post(post_data={"text": "Draft-enabled Bar"})
+        response = self.get()
+
+        self.assertEqual(response.status_code, 200)
+        self.assertTemplateUsed(response, "wagtailsnippets/snippets/edit.html")
+
+        # Should not show the Live status
+        self.assertNotContains(
+            response,
+            '<h3 id="status-sidebar-live" class="w-label-1 !w-mt-0 w-mb-1"><span class="w-sr-only">Status: </span>Live</h3>',
+            html=True,
+        )
+        # Should show the Draft status
+        self.assertContains(
+            response,
+            '<h3 id="status-sidebar-draft" class="w-label-1 !w-mt-0 w-mb-1"><span class="w-sr-only">Status: </span>Draft</h3>',
+            html=True,
+        )
+
+    def test_get_after_publish(self):
+        self.post(
+            post_data={
+                "text": "Draft-enabled Bar, Published",
+                "action-publish": "action-publish",
+            }
+        )
+        response = self.get()
+
+        self.assertEqual(response.status_code, 200)
+        self.assertTemplateUsed(response, "wagtailsnippets/snippets/edit.html")
+
+        # Should show the Live status
+        self.assertContains(
+            response,
+            '<h3 id="status-sidebar-live" class="w-label-1 !w-mt-0 w-mb-1"><span class="w-sr-only">Status: </span>Live</h3>',
+            html=True,
+        )
+        # Should not show the Draft status
+        self.assertNotContains(
+            response,
+            '<h3 id="status-sidebar-draft" class="w-label-1 !w-mt-0 w-mb-1"><span class="w-sr-only">Status: </span>Draft</h3>',
+            html=True,
+        )
+
+    def test_get_after_publish_and_save_draft(self):
+        self.post(
+            post_data={
+                "text": "Draft-enabled Bar, Published",
+                "action-publish": "action-publish",
+            }
+        )
+        self.post(post_data={"text": "Draft-enabled Bar, In Draft"})
+        response = self.get()
+        html = response.content.decode()
+
+        self.assertEqual(response.status_code, 200)
+        self.assertTemplateUsed(response, "wagtailsnippets/snippets/edit.html")
+
+        # Should show the Live status
+        self.assertContains(
+            response,
+            '<h3 id="status-sidebar-live" class="w-label-1 !w-mt-0 w-mb-1"><span class="w-sr-only">Status: </span>Live</h3>',
+            html=True,
+        )
+        # Should show the Draft status
+        self.assertContains(
+            response,
+            '<h3 id="status-sidebar-draft" class="w-label-1 !w-mt-0 w-mb-1"><span class="w-sr-only">Status: </span>Draft</h3>',
+            html=True,
+        )
+
+        # Should use the latest draft content for the title
+        self.assertContains(
+            response,
+            '<h1 class="w-header__title"><svg class="icon icon-snippet w-header__glpyh" aria-hidden="true"><use href="#icon-snippet"></use></svg>Draft-enabled Bar, In Draft</h1>',
+            html=True,
+        )
+
+        # Should use the latest draft content for the form
+        self.assertTagInHTML(
+            '<textarea name="text">Draft-enabled Bar, In Draft</textarea>',
+            html,
+            allow_extra_attrs=True,
+        )
+
+
 class TestSnippetDelete(TestCase, WagtailTestUtils):
     fixtures = ["test.json"]
 
@@ -1300,6 +1737,31 @@ class TestUsedBy(TestCase):
         self.assertEqual(type(advert.get_usage()[0]), Page)
 
 
+@override_settings(WAGTAIL_USAGE_COUNT_ENABLED=True)
+class TestSnippetUsageView(TestCase, WagtailTestUtils):
+    def setUp(self):
+        self.user = self.login()
+
+    def test_use_latest_draft_as_title(self):
+        snippet = DraftStateModel.objects.create(text="Draft-enabled Foo, Published")
+        snippet.save_revision().publish()
+        snippet.text = "Draft-enabled Bar, In Draft"
+        snippet.save_revision()
+
+        response = self.client.get(
+            reverse(
+                "wagtailsnippets_tests_draftstatemodel:usage",
+                args=[quote(snippet.pk)],
+            )
+        )
+
+        # Should use the latest draft title in the header subtitle
+        self.assertContains(
+            response,
+            '<span class="w-header__subtitle">Draft-enabled Bar, In Draft</span>',
+        )
+
+
 class TestSnippetHistory(TestCase, WagtailTestUtils):
     fixtures = ["test.json"]
 
@@ -1397,6 +1859,20 @@ class TestSnippetHistory(TestCase, WagtailTestUtils):
             count=1,
         )
 
+    def test_use_latest_draft_as_title(self):
+        snippet = DraftStateModel.objects.create(text="Draft-enabled Foo, Published")
+        snippet.save_revision().publish()
+        snippet.text = "Draft-enabled Bar, In Draft"
+        snippet.save_revision()
+
+        response = self.get(snippet)
+
+        # Should use the latest draft title in the header subtitle
+        self.assertContains(
+            response,
+            '<span class="w-header__subtitle">Draft-enabled Bar, In Draft</span>',
+        )
+
     @override_settings(WAGTAIL_I18N_ENABLED=True)
     def test_get_with_i18n_enabled(self):
         response = self.get(self.non_revisable_snippet)
@@ -1485,6 +1961,24 @@ class TestSnippetRevisions(TestCase, WagtailTestUtils):
         response = self.get()
         self.assertEqual(response.status_code, 302)
 
+    def test_get_with_draft_state_snippet(self):
+        self.snippet = DraftStateModel.objects.create(text="Draft-enabled Foo")
+        self.initial_revision = self.snippet.save_revision()
+        response = self.get()
+
+        self.assertEqual(response.status_code, 200)
+        self.assertTemplateUsed(response, "wagtailsnippets/snippets/edit.html")
+
+        # The save button should be labelled "Replace current draft"
+        self.assertContains(response, "Replace current draft")
+        # The publish button should exist
+        self.assertContains(response, "Publish this version")
+        # The publish button should have name="action-publish"
+        self.assertContains(
+            response,
+            '<button type="submit" name="action-publish" value="action-publish" class="button action-save button-longrunning warning" data-clicked-text="Publishing…">',
+        )
+
     def test_replace_revision(self):
         get_response = self.get()
         text_from_revision = get_response.context["form"].initial["text"]
@@ -1535,6 +2029,104 @@ class TestSnippetRevisions(TestCase, WagtailTestUtils):
         # Only the initial revision and edited revision, no revert revision
         self.assertEqual(self.snippet.revisions.count(), 2)
 
+    def test_replace_draft(self):
+        self.snippet = DraftStateModel.objects.create(
+            text="Draft-enabled Foo", live=False
+        )
+        self.initial_revision = self.snippet.save_revision()
+        self.snippet.text = "Draft-enabled Foo edited"
+        self.edit_revision = self.snippet.save_revision()
+        get_response = self.get()
+        text_from_revision = get_response.context["form"].initial["text"]
+
+        post_response = self.post(
+            post_data={
+                "text": text_from_revision + " reverted",
+                "revision": self.initial_revision.pk,
+            }
+        )
+        self.assertRedirects(post_response, self.get_url("list", args=[]))
+
+        self.snippet.refresh_from_db()
+        latest_revision = self.snippet.get_latest_revision()
+        log_entry = ModelLogEntry.objects.filter(revision=latest_revision).first()
+        publish_log_entries = ModelLogEntry.objects.filter(
+            content_type=ContentType.objects.get_for_model(DraftStateModel),
+            action="wagtail.publish",
+            object_id=self.snippet.pk,
+        )
+
+        # The instance should not be updated
+        self.assertEqual(self.snippet.text, "Draft-enabled Foo")
+        # The initial revision, edited revision, and revert revision
+        self.assertEqual(self.snippet.revisions.count(), 3)
+        # The latest revision should be the revert revision
+        self.assertEqual(latest_revision.content["text"], "Draft-enabled Foo reverted")
+
+        # A new log entry with "wagtail.revert" action should be created
+        self.assertIsNotNone(log_entry)
+        self.assertEqual(log_entry.action, "wagtail.revert")
+
+        # There should be no log entries for the publish action
+        self.assertEqual(publish_log_entries.count(), 0)
+
+        # The instance should still be a draft
+        self.assertFalse(self.snippet.live)
+        self.assertTrue(self.snippet.has_unpublished_changes)
+        self.assertIsNone(self.snippet.first_published_at)
+        self.assertIsNone(self.snippet.last_published_at)
+        self.assertIsNone(self.snippet.live_revision)
+
+    def test_replace_publish(self):
+        self.snippet = DraftStateModel.objects.create(text="Draft-enabled Foo")
+        self.initial_revision = self.snippet.save_revision()
+        self.snippet.text = "Draft-enabled Foo edited"
+        self.edit_revision = self.snippet.save_revision()
+        get_response = self.get()
+        text_from_revision = get_response.context["form"].initial["text"]
+
+        timestamp = now()
+        with freeze_time(timestamp):
+            post_response = self.post(
+                post_data={
+                    "text": text_from_revision + " reverted",
+                    "revision": self.initial_revision.pk,
+                    "action-publish": "action-publish",
+                }
+            )
+
+        self.assertRedirects(post_response, self.get_url("list", args=[]))
+
+        self.snippet.refresh_from_db()
+        latest_revision = self.snippet.get_latest_revision()
+        log_entry = ModelLogEntry.objects.filter(revision=latest_revision).first()
+        revert_log_entries = ModelLogEntry.objects.filter(
+            content_type=ContentType.objects.get_for_model(DraftStateModel),
+            action="wagtail.revert",
+            object_id=self.snippet.pk,
+        )
+
+        # The instance should be updated
+        self.assertEqual(self.snippet.text, "Draft-enabled Foo reverted")
+        # The initial revision, edited revision, and revert revision
+        self.assertEqual(self.snippet.revisions.count(), 3)
+        # The latest revision should be the revert revision
+        self.assertEqual(latest_revision.content["text"], "Draft-enabled Foo reverted")
+
+        # The latest log entry should use the "wagtail.publish" action
+        self.assertIsNotNone(log_entry)
+        self.assertEqual(log_entry.action, "wagtail.publish")
+
+        # There should be a log entry for the revert action
+        self.assertEqual(revert_log_entries.count(), 1)
+
+        # The instance should be live
+        self.assertTrue(self.snippet.live)
+        self.assertFalse(self.snippet.has_unpublished_changes)
+        self.assertEqual(self.snippet.first_published_at, timestamp)
+        self.assertEqual(self.snippet.last_published_at, timestamp)
+        self.assertEqual(self.snippet.live_revision, self.snippet.latest_revision)
+
 
 class TestCompareRevisions(TestCase, WagtailTestUtils):
     # Actual tests for the comparison classes can be found in test_compare.py

wagtail/test/testapp/migrations/0003_draftstatemodel.py

@@ -0,0 +1,104 @@
+# Generated by Django 4.0.5 on 2022-07-07 04:13
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("wagtailcore", "0076_modellogentry_revision"),
+        ("tests", "0002_recreate_setting_models"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="DraftStateModel",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                (
+                    "live",
+                    models.BooleanField(
+                        default=True, editable=False, verbose_name="live"
+                    ),
+                ),
+                (
+                    "has_unpublished_changes",
+                    models.BooleanField(
+                        default=False,
+                        editable=False,
+                        verbose_name="has unpublished changes",
+                    ),
+                ),
+                (
+                    "first_published_at",
+                    models.DateTimeField(
+                        blank=True,
+                        db_index=True,
+                        null=True,
+                        verbose_name="first published at",
+                    ),
+                ),
+                (
+                    "last_published_at",
+                    models.DateTimeField(
+                        editable=False, null=True, verbose_name="last published at"
+                    ),
+                ),
+                (
+                    "go_live_at",
+                    models.DateTimeField(
+                        blank=True, null=True, verbose_name="go live date/time"
+                    ),
+                ),
+                (
+                    "expire_at",
+                    models.DateTimeField(
+                        blank=True, null=True, verbose_name="expiry date/time"
+                    ),
+                ),
+                (
+                    "expired",
+                    models.BooleanField(
+                        default=False, editable=False, verbose_name="expired"
+                    ),
+                ),
+                ("text", models.TextField()),
+                (
+                    "latest_revision",
+                    models.ForeignKey(
+                        blank=True,
+                        editable=False,
+                        null=True,
+                        on_delete=django.db.models.deletion.SET_NULL,
+                        related_name="+",
+                        to="wagtailcore.revision",
+                        verbose_name="latest revision",
+                    ),
+                ),
+                (
+                    "live_revision",
+                    models.ForeignKey(
+                        blank=True,
+                        editable=False,
+                        null=True,
+                        on_delete=django.db.models.deletion.SET_NULL,
+                        related_name="+",
+                        to="wagtailcore.revision",
+                        verbose_name="live revision",
+                    ),
+                ),
+            ],
+            options={
+                "abstract": False,
+            },
+        ),
+    ]

wagtail/test/testapp/models.py

@@ -58,6 +58,7 @@ from wagtail.images import get_image_model
 from wagtail.images.blocks import ImageChooserBlock
 from wagtail.images.models import AbstractImage, AbstractRendition, Image
 from wagtail.models import (
+    DraftStateMixin,
     Orderable,
     Page,
     PageManager,
@@ -967,6 +968,21 @@ class RevisableGrandChildModel(RevisableChildModel):
     pass
 
 
+# Models with DraftStateMixin
+class DraftStateModel(DraftStateMixin, RevisionMixin, models.Model):
+    text = models.TextField()
+
+    panels = [
+        FieldPanel("text"),
+    ]
+
+    def __str__(self):
+        return self.text
+
+
+register_snippet(DraftStateModel)
+
+
 class StandardIndex(Page):
     """Index for the site"""