I’m trying to set custom validation rules on my Django models.
I have the following code:
Entity.py
from django.core.exceptions import ObjectDoesNotExist, PermissionDenied
from django.db import models
class Entity(models.Model):
TYPE_CHOICES = [
('asset', 'asset'),
('company', 'company')
]
entity_type = models.CharField(max_length=25, choices=TYPE_CHOICES)
class Asset(Entity):
class Company(Entity):
Security.py
from django.core.exceptions import ObjectDoesNotExist, PermissionDenied, ValidationError
from django.db import models
from .Entity import Entity
class Security(models.Model):
entities = models.ManyToManyField(Entity)
def clean(self, *args, **kwargs) -> None:
try:
# Check that object has already been saved
Security.objects.get(name=self.name)
print("All entities:", self.entities.all())
if len([e for e in self.entities.all() if e.entity_type == 'company']) > 1:
raise ValidationError("A security can only be related to one company")
except ObjectDoesNotExist:
pass
return super(Security, self).clean(*args, **kwargs)
def save(self, *args, **kwargs):
try:
self.full_clean()
return super(Security, self).save(*args, **kwargs)
except ValidationError:
raise ValidationError("A security can only be related to one company")
class Bond(Security):
class Equity(Security):
What I’m trying to do is to prevent a Security
from being saved if its field entities
holds more than 1 company. The entities
field is a ManyToMany relationship with Entity
, as it can hold multiple Asset
objects but only one Company
.
This is my test file
class EntityTestCase(TestCase):
def test_multiple_companies_fail(self):
b = Bond.objects.get(name='Bond 1')
c1 = Company.objects.get(name='Company 1')
c2 = Company.objects.get(name='Company 2')
a1 = Asset.objects.get(name='Asset 1')
a2 = Asset.objects.get(name='Asset 2')
b.entities.add(c1)
b.save()
b.entities.add(a1)
b.entities.add(a2)
b.save()
try:
b.entities.add(c2)
b.save()
self.fail("Should not be able to link more than one company to a security")
except ValidationError:
b = Bond.objects.get(name='Bond 1')
# Exception gets caught, but following assert fails with `AssertionError: 2 != 1`.
self.assertEquals(len([e for e in b.entities.all() if e.entity_type == 'company']), 1)
I tried with the Django shell and I have the same behaviour. The record gets saved even though I raise an exception in the save
method, which doesn’t seem logical.
Would anyone know why ? What am I doing wrong, and how I could I avoid saving based on conditions?
Advertisement
Answer
I actually found the solution, thanks to Williem’s comment.
The m2m relation isn’t managed by save()
, add(...)
does it. So what did the trick is this function:
@receiver(m2m_changed, sender=Security.entities.through)
def check_unique_company(sender, instance, **kwargs):
if len([e for e in instance.entities.all() if e.entity_type == 'company']) > 1:
raise ValidationError("A security can only be related to one company")