Use a form to clean the instance

src/wagtail_personalisation/forms.py

@@ -0,0 +1,18 @@
+from django.core.exceptions import ValidationError
+from wagtail.wagtailadmin.forms import WagtailAdminModelForm
+
+
+class SegmentAdminForm(WagtailAdminModelForm):
+    def clean(self):
+        cleaned_data = super(SegmentAdminForm, self).clean()
+
+        rules = [form for formset in self.formsets.values() for form in formset if form not in formset.deleted_forms]
+        consistent = rules and all(rule.instance.static for rule in rules)
+
+        from .models import Segment
+        if cleaned_data.get('type') == Segment.TYPE_STATIC and not cleaned_data.get('count') and not consistent:
+            raise ValidationError({
+                'count': ('Static segments with non-static compatible rules must include a count.'),
+            })
+
+        return cleaned_data

src/wagtail_personalisation/models.py

@@ -26,6 +26,8 @@ from wagtail.wagtailcore.models import Page
 from wagtail_personalisation.rules import AbstractBaseRule
 from wagtail_personalisation.utils import count_active_days
 
+from .forms import SegmentAdminForm
+
 
 SessionStore = import_module(settings.SESSION_ENGINE).SessionStore
 
@@ -103,6 +105,8 @@ class Segment(ClusterableModel):
 
     objects = SegmentQuerySet.as_manager()
 
+    base_form_class = SegmentAdminForm
+
     def __init__(self, *args, **kwargs):
         Segment.panels = [
             MultiFieldPanel([
@@ -131,12 +135,6 @@ class Segment(ClusterableModel):
     def __str__(self):
         return self.name
 
-    def clean(self):
-        if self.is_static and not self.is_consistent and not self.count:
-            raise ValidationError({
-                'count': _('Static segments with non-static compatible rules must include a count.'),
-            })
-
     @property
     def is_static(self):
         return self.type == self.TYPE_STATIC

tests/unit/test_static_dynamic_segments.py

@@ -5,7 +5,9 @@ import datetime
 import pytest
 
 from django.core.exceptions import ValidationError
+from django.forms.models import model_to_dict
 from tests.factories.segment import SegmentFactory
+from wagtail_personalisation.forms import SegmentAdminForm
 from wagtail_personalisation.models import Segment
 from wagtail_personalisation.rules import TimeRule, VisitCountRule
 
@@ -125,37 +127,54 @@ def test_does_not_calculate_the_segment_again(rf, site, client, mocker):
     assert mock_test_rule.call_count == 0
 
 
+def form_with_data(segment, rule):
+    model_fields = ['type', 'status', 'count', 'name']
+
+    class TestSegmentAdminForm(SegmentAdminForm):
+        class Meta:
+            model = Segment
+            fields = model_fields
+
+    data = model_to_dict(segment, model_fields)
+    for formset in TestSegmentAdminForm().formsets.values():
+        rule_data = {}
+        if isinstance(rule, formset.model):
+            rule_data = model_to_dict(rule)
+            for key, value in rule_data.items():
+                data['{}-0-{}'.format(formset.prefix, key)] = value
+        data['{}-INITIAL_FORMS'.format(formset.prefix)] = 0
+        data['{}-TOTAL_FORMS'.format(formset.prefix)] = 1 if rule_data else 0
+    return TestSegmentAdminForm(data)
+
+
+
 @pytest.mark.django_db
 def test_non_static_rules_have_a_count():
     segment = SegmentFactory(type=Segment.TYPE_STATIC, count=0)
-    TimeRule.objects.create(
+    rule = TimeRule.objects.create(
         start_time=datetime.time(0, 0, 0),
         end_time=datetime.time(23, 59, 59),
         segment=segment,
     )
-    with pytest.raises(ValidationError):
-        segment.clean()
+    form = form_with_data(segment, rule)
+    assert not form.is_valid()
 
 
 @pytest.mark.django_db
 def test_static_segment_with_static_rules_needs_no_count(site):
     segment = SegmentFactory(type=Segment.TYPE_STATIC, count=0)
-    VisitCountRule.objects.create(counted_page=site.root_page, segment=segment)
-    try:
-        segment.clean()
-    except ValidationError:
-        pytest.fail('Should not raise ValidationError.')
+    rule = VisitCountRule.objects.create(counted_page=site.root_page, segment=segment)
+    form = form_with_data(segment, rule)
+    assert form.is_valid()
 
 
 @pytest.mark.django_db
 def test_dynamic_segment_with_non_static_rules_have_a_count():
     segment = SegmentFactory(type=Segment.TYPE_DYNAMIC, count=0)
-    TimeRule.objects.create(
+    rule = TimeRule.objects.create(
         start_time=datetime.time(0, 0, 0),
         end_time=datetime.time(23, 59, 59),
         segment=segment,
     )
-    try:
-        segment.clean()
-    except ValidationError:
-        pytest.fail('Should not raise ValidationError.')
+    form = form_with_data(segment, rule)
+    assert form.is_valid(), form.errors