Skip to content
Advertisement

Django model saves when it shouldn’t

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")
Advertisement