From c30bc3174b7476f1e1832ba57a59b5cb506402d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Lebleu?=
 <jeromelebleu@users.noreply.github.com>
Date: Tue, 7 Jan 2020 18:19:15 +0100
Subject: [PATCH] core: catch errors in load_initial_data command

---
 .../management/commands/load_initial_data.py  | 15 ++++++++++---
 .../stupid_extension_1/modo_extension.py      |  4 +++-
 .../stupid_extension_2/modo_extension.py      |  3 +++
 modoboa/core/tests/test_extensions.py         | 21 ++++++++++++++++---
 4 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/modoboa/core/management/commands/load_initial_data.py b/modoboa/core/management/commands/load_initial_data.py
index d43f27bfc..7811c13e8 100644
--- a/modoboa/core/management/commands/load_initial_data.py
+++ b/modoboa/core/management/commands/load_initial_data.py
@@ -61,9 +61,18 @@ class Command(BaseCommand):
 
         for extname in list(extensions.exts_pool.extensions.keys()):
             extension = extensions.exts_pool.get_extension(extname)
-            extension.load_initial_data()
-            signals.initial_data_loaded.send(
-                sender=self.__class__, extname=extname)
+            try:
+                extension.load_initial_data()
+            except Exception as e:
+                self.stderr.write(
+                    "Unable to load initial data for '{}' ({}).".format(
+                        extname, str(e)
+                    )
+                )
+            else:
+                signals.initial_data_loaded.send(
+                    sender=self.__class__, extname=extname
+                )
 
         if options["extra_fixtures"]:
             from modoboa.admin import factories
diff --git a/modoboa/core/tests/stupid_extension_1/modo_extension.py b/modoboa/core/tests/stupid_extension_1/modo_extension.py
index 5fdf23318..a645634d5 100644
--- a/modoboa/core/tests/stupid_extension_1/modo_extension.py
+++ b/modoboa/core/tests/stupid_extension_1/modo_extension.py
@@ -19,7 +19,9 @@ class StupidExtension1(extensions.ModoExtension):
         pass
 
     def load_initial_data(self):
-        raise RuntimeError
+        from modoboa.admin.factories import DomainFactory
+
+        DomainFactory(name="stupid_1.com")
 
 
 extensions.exts_pool.register_extension(StupidExtension1)
diff --git a/modoboa/core/tests/stupid_extension_2/modo_extension.py b/modoboa/core/tests/stupid_extension_2/modo_extension.py
index c7a58d1cc..fd42cb5cb 100644
--- a/modoboa/core/tests/stupid_extension_2/modo_extension.py
+++ b/modoboa/core/tests/stupid_extension_2/modo_extension.py
@@ -15,5 +15,8 @@ class StupidExtension2(extensions.ModoExtension):
     version = "1.0.0"
     description = "A stupid extension"
 
+    def load_initial_data(self):
+        raise RuntimeError
+
 
 extensions.exts_pool.register_extension(StupidExtension2, show=False)
diff --git a/modoboa/core/tests/test_extensions.py b/modoboa/core/tests/test_extensions.py
index af420f526..e9d4e8416 100644
--- a/modoboa/core/tests/test_extensions.py
+++ b/modoboa/core/tests/test_extensions.py
@@ -6,13 +6,15 @@ from __future__ import unicode_literals
 
 import os
 import sys
+from io import StringIO
 
 from testfixtures import compare
 
 from django.core import management
 from django.test import TestCase, override_settings
 
-from .. import extensions
+from modoboa.admin.models import Domain
+from .. import extensions, signals
 
 sys.path.append(os.path.dirname(__file__))
 
@@ -74,8 +76,21 @@ class ExtensionTestCase(TestCase):
 
     def test_load_initial_data(self):
         """Check if method is called."""
-        with self.assertRaises(RuntimeError):
-            management.call_command("load_initial_data")
+        self.signal_called = 0
+
+        def handler(sender, extname, **kwargs):
+            self.assertEqual(extname, "stupid_extension_1")
+            self.signal_called += 1
+
+        signals.initial_data_loaded.connect(handler)
+
+        stderr_out = StringIO()
+        management.call_command("load_initial_data", stderr=stderr_out)
+        self.assertTrue(Domain.objects.filter(name="stupid_1.com").exists())
+        self.assertIn("stupid_extension_2", stderr_out.getvalue())
+        self.assertEqual(self.signal_called, 1)
+
+        signals.initial_data_loaded.disconnect(handler)
 
     def test_get_urls(self):
         """Load extensions urls."""