improve coverage

modoboa/admin/management/commands/subcommands/_mx.py

@@ -64,19 +64,23 @@ class CheckMXRecords(BaseCommand):
     def get_mx_records_for_domain(self, domain, ttl=7200):
         """Return one or more `models.MXRecord` for `domain`. DNS queries are
         not performed while `ttl` (in seconds) is still valid"""
-        try:
-            answers = dns.resolver.query(domain.name, "MX")
-        except dns.resolver.NoAnswer:
-            return
         now = datetime.now()
-        delta = timedelta(seconds=ttl)
         records = models.MXRecord.objects.filter(domain=domain,
-                                                 updated__gte=now - delta)
+                                                 updated__gt=now)
         if records.count():
             for record in records:
                 yield record
             raise StopIteration()
-        models.MXRecord.objects.filter(domain=domain).delete()
+        else:
+            models.MXRecord.objects.filter(domain=domain).delete()
+
+        try:
+            answers = dns.resolver.query(domain.name, "MX")
+        except dns.resolver.NoAnswer:
+            raise StopIteration()
+
+        delta = timedelta(seconds=ttl)
+
         for answer in answers:
             address = None
             try:
@@ -94,7 +98,7 @@ class CheckMXRecords(BaseCommand):
                         domain=domain,
                         name=str(answer.exchange).strip('.'),
                         address=address,
-                        updated=now)
+                        updated=now + delta)
                     yield record
 
     def query_dnsbl(self, mx_list, provider):

modoboa/admin/tests/test_mx.py

@@ -10,6 +10,41 @@ from .. import factories
 from .. import models
 
 
+class MXTestCase(ModoTestCase):
+    """TestCase for DNSBL related features."""
+
+    @classmethod
+    def setUpTestData(cls):
+        """Create some data."""
+        super(MXTestCase, cls).setUpTestData()
+        cls.domain = factories.DomainFactory(name="modoboa.org")
+        factories.DomainFactory(name="pouet.com")  # should not exist
+
+    def test_management_command(self):
+        """Check that command works fine."""
+        self.assertEqual(models.MXRecord.objects.count(), 0)
+        management.call_command("modo", "check_mx", "--no-dnsbl",
+                                "--valid-mx=127.0.0.1", "--ttl=0")
+        self.assertTrue(
+            models.MXRecord.objects.filter(domain=self.domain).exists())
+
+        # we passed a ttl to 0. this will reset the cache this time
+        qs = models.MXRecord.objects.filter(domain=self.domain)
+        id = qs[0].id
+        management.call_command("modo", "check_mx", "--no-dnsbl",
+                                "--valid-mx=127.0.0.1", "--ttl=7200")
+        qs = models.MXRecord.objects.filter(domain=self.domain)
+        self.assertNotEqual(id, qs[0].id)
+        id = qs[0].id
+
+        # assume that mxrecords ids are the same. means that we taking care of
+        # ttl
+        management.call_command("modo", "check_mx", "--no-dnsbl",
+                                "--valid-mx=127.0.0.1")
+        qs = models.MXRecord.objects.filter(domain=self.domain)
+        self.assertEqual(id, qs[0].id)
+
+
 class DNSBLTestCase(ModoTestCase):
     """TestCase for DNSBL related features."""