From 303305e05fa48eab9d02496aff91d52badde0fdc Mon Sep 17 00:00:00 2001
From: Oliver <oliver.henry.walters@gmail.com>
Date: Tue, 20 Jun 2023 16:57:20 +1000
Subject: [PATCH] Fix constraint for address model (#5076)

* Fix constraint for address model

- Do not handle at database level
- Add a "validate_unique" method to the address model
- Fixes https://github.com/inventree/InvenTree/issues/5070

* Remove unique constraint rule

* Unit test fix

- Adjust unit test around new code

* Further unit test updates
---
 .../migrations/0063_auto_20230502_1956.py        |  4 ----
 InvenTree/company/models.py                      | 13 ++++++++++---
 InvenTree/company/tests.py                       | 16 ++++++++++++----
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/InvenTree/company/migrations/0063_auto_20230502_1956.py b/InvenTree/company/migrations/0063_auto_20230502_1956.py
index 5daf2418e7..c0cce4e188 100644
--- a/InvenTree/company/migrations/0063_auto_20230502_1956.py
+++ b/InvenTree/company/migrations/0063_auto_20230502_1956.py
@@ -30,8 +30,4 @@ class Migration(migrations.Migration):
                 ('company', models.ForeignKey(help_text='Select company', on_delete=django.db.models.deletion.CASCADE, related_name='addresses', to='company.company', verbose_name='Company')),
             ],
         ),
-        migrations.AddConstraint(
-            model_name='address',
-            constraint=models.UniqueConstraint(condition=models.Q(('primary', True)), fields=('company',), name='one_primary_per_company'),
-        ),
     ]
diff --git a/InvenTree/company/models.py b/InvenTree/company/models.py
index 9a44ecf3dd..6334ce6dd9 100644
--- a/InvenTree/company/models.py
+++ b/InvenTree/company/models.py
@@ -306,9 +306,6 @@ class Address(models.Model):
 
     class Meta:
         """Metaclass defines extra model options"""
-        constraints = [
-            UniqueConstraint(fields=['company'], condition=Q(primary=True), name='one_primary_per_company')
-        ]
         verbose_name_plural = "Addresses"
 
     @staticmethod
@@ -316,6 +313,16 @@ class Address(models.Model):
         """Return the API URL associated with the Contcat model"""
         return reverse('api-address-list')
 
+    def validate_unique(self, exclude=None):
+        """Ensure that only one primary address exists per company"""
+
+        super().validate_unique(exclude=exclude)
+
+        if self.primary:
+            # Check that no other primary address exists for this company
+            if Address.objects.filter(company=self.company, primary=True).exclude(pk=self.pk).exists():
+                raise ValidationError({'primary': _('Company already has a primary address')})
+
     company = models.ForeignKey(Company, related_name='addresses',
                                 on_delete=models.CASCADE,
                                 verbose_name=_('Company'),
diff --git a/InvenTree/company/tests.py b/InvenTree/company/tests.py
index 807b78e866..b6c49a5eb6 100644
--- a/InvenTree/company/tests.py
+++ b/InvenTree/company/tests.py
@@ -5,7 +5,6 @@ from decimal import Decimal
 
 from django.core.exceptions import ValidationError
 from django.db import transaction
-from django.db.utils import IntegrityError
 from django.test import TestCase
 
 from part.models import Part
@@ -199,17 +198,23 @@ class AddressTest(TestCase):
         c2 = Company.objects.create(name='Test Corp2.', description='We make stuff good')
         Address.objects.create(company=self.c, primary=True)
         Address.objects.create(company=self.c, primary=False)
+
         self.assertEqual(Address.objects.count(), 2)
 
         # Testing the constraint itself
         # Intentionally throwing exceptions breaks unit tests unless performed in an atomic block
         with transaction.atomic():
-            self.assertRaises(IntegrityError, Address.objects.create, company=self.c, primary=True, confirm_primary=False)
+            with self.assertRaises(ValidationError):
+                addr = Address(company=self.c, primary=True, confirm_primary=False)
+                addr.validate_unique()
 
         Address.objects.create(company=c2, primary=True, line1="Hellothere", line2="generalkenobi")
 
         with transaction.atomic():
-            self.assertRaises(IntegrityError, Address.objects.create, company=c2, primary=True)
+            with self.assertRaises(ValidationError):
+                addr = Address(company=c2, primary=True, confirm_primary=False)
+                addr.validate_unique()
+
         self.assertEqual(Address.objects.count(), 3)
 
     def test_first_address_is_primary(self):
@@ -219,7 +224,10 @@ class AddressTest(TestCase):
 
         self.assertTrue(addr.primary)
 
-        self.assertRaises(IntegrityError, Address.objects.create, company=self.c, primary=True)
+        # Create another address, which should error out if primary is not set to False
+        with self.assertRaises(ValidationError):
+            addr = Address(company=self.c, primary=True)
+            addr.validate_unique()
 
     def test_model_str(self):
         """Test value of __str__"""