From 108f9500300007fdcc6ba767cbe4a271213daf37 Mon Sep 17 00:00:00 2001
From: Antoine Nguyen <tonio@ngyn.org>
Date: Fri, 27 Nov 2015 18:59:59 +0100
Subject: [PATCH] Database cleanup.

fix #792
---
 modoboa/relaydomains/handlers.py | 16 ++++++++++++++++
 modoboa/relaydomains/tests.py    | 30 ++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/modoboa/relaydomains/handlers.py b/modoboa/relaydomains/handlers.py
index 16129dfdc..3c8d4d857 100644
--- a/modoboa/relaydomains/handlers.py
+++ b/modoboa/relaydomains/handlers.py
@@ -1,5 +1,6 @@
 """Django signal handlers for relaydomains."""
 
+from django.db.models import signals
 from django.dispatch import receiver
 
 from modoboa.admin import models as admin_models
@@ -23,3 +24,18 @@ def check_relaydomain_alias(sender, **kwargs):
             domain__name=domain, address=localpart).exists()):
         return False
     return True
+
+
+@receiver(signals.post_save, sender=admin_models.Domain)
+def clean_domain(sender, instance, **kwargs):
+    """Remove or create RelayDomain record if needed."""
+    if kwargs.get("created"):
+        return
+    has_relaydom = hasattr(instance, "relaydomain")
+    if instance.type == "domain" and has_relaydom:
+        models.RelayDomain.objects.filter(domain=instance).delete()
+    elif instance.type == "relaydomain" and not has_relaydom:
+        # Make sure to create a RelayDomain instance since we can't do it
+        # at form level...
+        models.RelayDomain.objects.create(
+            domain=instance, service=models.Service.objects.first())
diff --git a/modoboa/relaydomains/tests.py b/modoboa/relaydomains/tests.py
index 288284ae9..8d9af675c 100644
--- a/modoboa/relaydomains/tests.py
+++ b/modoboa/relaydomains/tests.py
@@ -138,6 +138,36 @@ class RelayDomainsTestCase(ModoTestCase, Operations):
         )
         RelayDomain.objects.get(domain__name='relaydomain.org')
 
+    def test_relaydomain_domain_switch(self):
+        """Check domain <-> relaydomain transitions."""
+        domain = self.rdom.domain
+        values = {
+            "name": "relaydomain.tld",
+            "type": "domain",
+            "enabled": True,
+        }
+        self.ajax_post(
+            reverse('admin:domain_change', args=[domain.pk]), values)
+        self.assertFalse(
+            RelayDomain.objects.filter(
+                domain__name="relaydomain.tld").exists())
+        self.assertEqual(
+            admin_models.Domain.objects.get(name="relaydomain.tld").type,
+            "domain")
+        values = {
+            "name": "relaydomain.tld",
+            "type": "relaydomain",
+            "enabled": True,
+        }
+        self.ajax_post(
+            reverse("admin:domain_change", args=[domain.pk]), values)
+        self.assertEqual(
+            admin_models.Domain.objects.get(name="relaydomain.tld").type,
+            "relaydomain")
+        self.assertTrue(
+            RelayDomain.objects.filter(
+                domain__name="relaydomain.tld").exists())
+
     def test_edit_relaydomainalias(self):
         """Test the modification of a relay domain alias.