Repository: wagtail/wagtail
File: wagtail/admin/tests/test_workflows.py (Python)

Code before:
        notifications = ["approved", "rejected", "submitted"]
        for notification in notifications:
            with self.subTest(f"Testing with {notification}_notifications"):
                notifier.notification = notification
                self.assertSetEqual(
                    notifier.get_valid_recipients(self.object), {self.submitter}
                )

        # remove notifications and re-test
        userprofile = UserProfile.get_for_user(self.submitter)
        updated_fields = []
        for notification in notifications:
            attribute = f"{notification}_notifications"
            setattr(userprofile, attribute, False)
            updated_fields.append(attribute)
        userprofile.save(update_fields=updated_fields)

        for notification in notifications:
            with self.subTest(f"Testing with {notification}_notifications"):
                notifier.notification = notification
                self.assertSetEqual(notifier.get_valid_recipients(self.object), set())


class TestCustomWorkflowLockOnTask(BasePageWorkflowTests):
    model_name = "page"

    def setup_workflow_and_tasks(self):
        self.workflow = Workflow.objects.create(name="test_workflow")
        self.task_1 = CustomLockTask.objects.create(name="test_task_1")
        WorkflowTask.objects.create(
            workflow=self.workflow, task=self.task_1, sort_order=1
        )

    def test_custom_lock_class(self):
        self.post("submit")
        response = self.client.get(self.get_url("edit"))
        self.assertContains(response, "If there is a door, there must be a key")
        self.assertIsInstance(self.object.get_lock(), CustomWorkflowLock)

    @mock.patch.object(CustomLockTask, "lock_class", new_callable=mock.PropertyMock)
    def test_typeerror_if_custom_lock_class_inherits_basic_locks(self, mock_property):
        mock_property.return_value = BasicLock

        self.post("submit")

        with self.assertRaises(TypeError), self.assertLogs(level="WARNING"):
            self.client.get(self.get_url("edit"))


Review comment (suggestion):
```suggestion
```

and we should make a subclass of this test with `BaseSnippetWorkflowTests`

Code after:
        notifications = ["approved", "rejected", "submitted"]
        for notification in notifications:
            with self.subTest(f"Testing with {notification}_notifications"):
                notifier.notification = notification
                self.assertSetEqual(
                    notifier.get_valid_recipients(self.object), {self.submitter}
                )

        # remove notifications and re-test
        userprofile = UserProfile.get_for_user(self.submitter)
        updated_fields = []
        for notification in notifications:
            attribute = f"{notification}_notifications"
            setattr(userprofile, attribute, False)
            updated_fields.append(attribute)
        userprofile.save(update_fields=updated_fields)

        for notification in notifications:
            with self.subTest(f"Testing with {notification}_notifications"):
                notifier.notification = notification
                self.assertSetEqual(notifier.get_valid_recipients(self.object), set())


class TestCustomWorkflowLockOnTask(BasePageWorkflowTests):
    def setup_workflow_and_tasks(self):
        self.workflow = Workflow.objects.create(name="test_workflow")
        self.task_1 = CustomLockTask.objects.create(name="test_task_1")
        WorkflowTask.objects.create(
            workflow=self.workflow, task=self.task_1, sort_order=1
        )

    def test_custom_lock_class(self):
        self.post("submit")
        response = self.client.get(self.get_url("edit"))
        self.assertContains(response, "If there is a door, there must be a key")
        self.assertIsInstance(self.object.get_lock(), CustomWorkflowLock)

    @mock.patch.object(CustomLockTask, "lock_class", new_callable=mock.PropertyMock)
    def test_typeerror_if_custom_lock_class_inherits_basic_locks(self, mock_property):
        mock_property.return_value = BasicLock

        self.post("submit")

        with self.assertRaises(TypeError):
            self.client.get(self.get_url("edit"))


class TestCustomWorkflowLockOnTaskWithSnippets(
    TestCustomWorkflowLockOnTask,
    BaseSnippetWorkflowTests,
):