Allow override of badges (#3191)

Follows up on #3189, fixes
[#1564](datagouv/data.gouv.fr#1564)

---------

Co-authored-by: maudetes <[email protected]>

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current (in progress)
 
 - Add more comments and types in the `api_field.py` "lib" [#3174](https://github.com/opendatateam/udata/pull/3174)
+- Allow overriding of badges (for example in plugins like udata-front) [#3191](https://github.com/opendatateam/udata/pull/3191)
 
 ## 10.0.0 (2024-11-07)
 

udata/api_fields.py

@@ -94,6 +94,8 @@ def convert_db_to_field(key, field, info) -> tuple[Callable | None, Callable | N
         params["min_length"] = field.min_length
         params["max_length"] = field.max_length
         params["enum"] = field.choices
+        if field.validation:
+            params["validation"] = validation_to_type(field.validation)
     elif isinstance(field, mongo_fields.ObjectIdField):
         constructor = restx_fields.String
     elif isinstance(field, mongo_fields.FloatField):
@@ -287,7 +289,7 @@ def wrapper(cls) -> Callable:
                 if not isinstance(
                     field, mongo_fields.ReferenceField | mongo_fields.LazyReferenceField
                 ):
-                    raise Exception("Cannot use additional_filters on not a ref.")
+                    raise Exception("Cannot use additional_filters on a field that is not a ref.")
 
                 ref_model: db.Document = field.document_type
 
@@ -678,9 +680,34 @@ def compute_filter(column: str, field, info, filterable) -> dict:
             filterable["type"] = boolean
         else:
             filterable["type"] = str
+    if field.validation:
+        filterable["type"] = validation_to_type(field.validation)
 
     filterable["choices"] = info.get("choices", None)
     if hasattr(field, "choices") and field.choices:
         filterable["choices"] = field.choices
 
     return filterable
+
+
+def validation_to_type(validation: Callable) -> Callable:
+    """Convert a mongo field's validation function to a ReqParser's type.
+
+    In flask_restx.ReqParser, validation is done by setting the param's type to
+    a callable that will either raise, or return the parsed value.
+
+    In mongo, a field's validation function cannot return anything, so this
+    helper wraps the mongo field's validation to return the value if it validated.
+    """
+    from udata.models import db
+
+    def wrapper(value: str) -> str:
+        try:
+            validation(value)
+        except db.ValidationError:
+            raise
+        return value
+
+    wrapper.__schema__ = {"type": "string", "format": "my-custom-format"}
+
+    return wrapper

udata/core/badges/tests/test_model.py

@@ -15,8 +15,13 @@
 }
 
 
+def validate_badge(value):
+    if value not in Fake.__badges__.keys():
+        raise db.ValidationError("Unknown badge type")
+
+
 class FakeBadge(Badge):
-    kind = db.StringField(required=True, choices=list(BADGES.keys()))
+    kind = db.StringField(required=True, validation=validate_badge)
 
 
 class FakeBadgeMixin(BadgeMixin):
@@ -34,12 +39,6 @@ def test_attributes(self):
         fake = Fake.objects.create()
         self.assertIsInstance(fake.badges, (list, tuple))
 
-    def test_choices(self):
-        """It should have a choice list on the badge field."""
-        self.assertEqual(
-            Fake._fields["badges"].field.document_type.kind.choices, list(Fake.__badges__.keys())
-        )
-
     def test_get_badge_found(self):
         """It allow to get a badge by kind if present"""
         fake = Fake.objects.create()
@@ -155,3 +154,19 @@ def test_create_disallow_unknown_badges(self):
         with self.assertRaises(db.ValidationError):
             fake = Fake.objects.create()
             fake.add_badge("unknown")
+
+    def test_validation(self):
+        """It should validate default badges as well as extended ones"""
+        # Model badges can be extended in plugins, for example in udata-front
+        # for french only badges.
+        Fake.__badges__["new"] = "new"
+
+        fake = FakeBadge(kind="test")
+        fake.validate()
+
+        fake = FakeBadge(kind="new")
+        fake.validate()
+
+        with self.assertRaises(db.ValidationError):
+            fake = FakeBadge(kind="doesnotexist")
+            fake.validate()

udata/core/dataset/models.py

@@ -521,8 +521,13 @@ def save(self, *args, **kwargs):
         self.dataset.save(*args, **kwargs)
 
 
+def validate_badge(value):
+    if value not in Dataset.__badges__.keys():
+        raise db.ValidationError("Unknown badge type")
+
+
 class DatasetBadge(Badge):
-    kind = db.StringField(required=True, choices=list(BADGES.keys()))
+    kind = db.StringField(required=True, validation=validate_badge)
 
 
 class DatasetBadgeMixin(BadgeMixin):

udata/core/organization/models.py

@@ -95,8 +95,13 @@ def with_badge(self, kind):
         return self(badges__kind=kind)
 
 
+def validate_badge(value):
+    if value not in Organization.__badges__.keys():
+        raise db.ValidationError("Unknown badge type")
+
+
 class OrganizationBadge(Badge):
-    kind = db.StringField(required=True, choices=list(BADGES.keys()))
+    kind = db.StringField(required=True, validation=validate_badge)
 
 
 class OrganizationBadgeMixin(BadgeMixin):

udata/core/reuse/models.py

@@ -35,8 +35,13 @@ def check_url_does_not_exists(url):
         raise FieldValidationError(_("This URL is already registered"), field="url")
 
 
+def validate_badge(value):
+    if value not in Reuse.__badges__.keys():
+        raise db.ValidationError("Unknown badge type")
+
+
 class ReuseBadge(Badge):
-    kind = db.StringField(required=True, choices=list(BADGES.keys()))
+    kind = db.StringField(required=True, validation=validate_badge)
 
 
 class ReuseBadgeMixin(BadgeMixin):

udata/tests/organization/test_organization_model.py

@@ -4,10 +4,10 @@
 from udata.core.dataset.factories import DatasetFactory, HiddenDatasetFactory
 from udata.core.followers.signals import on_follow, on_unfollow
 from udata.core.organization.factories import OrganizationFactory
-from udata.core.organization.models import Organization
+from udata.core.organization.models import Organization, OrganizationBadge
 from udata.core.reuse.factories import ReuseFactory, VisibleReuseFactory
 from udata.core.user.factories import UserFactory
-from udata.models import Dataset, Follow, Member, Reuse
+from udata.models import Dataset, Follow, Member, Reuse, db
 from udata.tests.helpers import assert_emit
 
 from .. import DBTestMixin, TestCase
@@ -71,3 +71,21 @@ def test_organization_queryset_with_badge(self):
         associations = list(Organization.objects.with_badge(org_constants.ASSOCIATION))
         assert len(associations) == 1
         assert org_certified_association in associations
+
+
+class OrganizationBadgeTest(DBTestMixin, TestCase):
+    # Model badges can be extended in plugins, for example in udata-front
+    # for french only badges.
+    Organization.__badges__["new"] = "new"
+
+    def test_validation(self):
+        """It should validate default badges as well as extended ones"""
+        badge = OrganizationBadge(kind="public-service")
+        badge.validate()
+
+        badge = OrganizationBadge(kind="new")
+        badge.validate()
+
+        with self.assertRaises(db.ValidationError):
+            badge = OrganizationBadge(kind="doesnotexist")
+            badge.validate()

udata/tests/reuse/test_reuse_model.py

@@ -5,9 +5,10 @@
 from udata.core.discussions.factories import DiscussionFactory
 from udata.core.organization.factories import OrganizationFactory
 from udata.core.reuse.factories import ReuseFactory, VisibleReuseFactory
+from udata.core.reuse.models import Reuse, ReuseBadge
 from udata.core.user.factories import UserFactory
 from udata.i18n import gettext as _
-from udata.models import Reuse
+from udata.models import db
 from udata.tests.helpers import assert_emit
 
 from .. import DBTestMixin, TestCase
@@ -124,3 +125,18 @@ def test_reuse_without_private(self):
         reuse.private = True
         reuse.save()
         self.assertEqual(reuse.private, True)
+
+
+class ReuseBadgeTest(DBTestMixin, TestCase):
+    # Model badges can be extended in plugins, for example in udata-front
+    # for french only badges.
+    Reuse.__badges__["new"] = "new"
+
+    def test_validation(self):
+        """It should validate default badges as well as extended ones"""
+        badge = ReuseBadge(kind="new")
+        badge.validate()
+
+        with self.assertRaises(db.ValidationError):
+            badge = ReuseBadge(kind="doesnotexist")
+            badge.validate()