From 7eca5caf85a527abee2e4bbc9fe565d45cb8a3f7 Mon Sep 17 00:00:00 2001
From: Jules Robichaud-Gagnon <j.robichaudg@gmail.com>
Date: Fri, 15 Sep 2023 12:52:35 -0400
Subject: [PATCH] Call post_command when the command raises an unhandled
 exception

---
 django_extensions/management/utils.py |  9 +++++--
 tests/test_management_command.py      | 35 ++++++++++++++++++++++++++-
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/django_extensions/management/utils.py b/django_extensions/management/utils.py
index 3a5ed51d9..9bdeb327c 100644
--- a/django_extensions/management/utils.py
+++ b/django_extensions/management/utils.py
@@ -59,8 +59,13 @@ def signalcommand(func):
 
     def inner(self, *args, **kwargs):
         pre_command.send(self.__class__, args=args, kwargs=kwargs)
-        ret = func(self, *args, **kwargs)
-        post_command.send(self.__class__, args=args, kwargs=kwargs, outcome=ret)
+        try:
+            ret = func(self, *args, **kwargs)
+        except Exception as e:
+            post_command.send(self.__class__, args=args, kwargs=kwargs, outcome=e)
+            raise
+        else:
+            post_command.send(self.__class__, args=args, kwargs=kwargs, outcome=ret)
         return ret
     return inner
 
diff --git a/tests/test_management_command.py b/tests/test_management_command.py
index b0b4a7719..28cbf31ef 100644
--- a/tests/test_management_command.py
+++ b/tests/test_management_command.py
@@ -5,12 +5,13 @@
 import importlib
 
 from django.conf import settings
-from django.core.management import call_command, find_commands, load_command_class
+from django.core.management import call_command, find_commands, load_command_class, BaseCommand
 from django.test import TestCase
 from io import StringIO
 
 from django_extensions.management.modelviz import use_model, generate_graph_data
 from django_extensions.management.commands.merge_model_instances import get_model_to_deduplicate, get_field_names, keep_first_or_last_instance
+from django_extensions.management.utils import signalcommand
 from . import force_color_support
 from .testapp.models import Person, Name, Note, Personality, Club, Membership, Permission
 from .testapp.jobs.hourly.test_hourly_job import HOURLY_JOB_MOCK
@@ -105,6 +106,38 @@ def post(sender, **kwargs):
         self.assertIn('kwargs', CommandSignalTests.post)
         self.assertIn('outcome', CommandSignalTests.post)
 
+    def test_exception(self):
+        from django_extensions.management.signals import post_command, pre_command
+
+        exception = Exception()
+
+        class TestCommand(BaseCommand):
+            @signalcommand
+            def handle(self, *args, **options):
+                raise exception
+
+        def pre(sender, **kwargs):
+            CommandSignalTests.pre = dict(**kwargs)
+
+        def post(sender, **kwargs):
+            CommandSignalTests.post = dict(**kwargs)
+
+        pre_command.connect(pre, TestCommand)
+        post_command.connect(post, TestCommand)
+
+        out = StringIO()
+
+        with self.assertRaises(Exception) as e_info:
+            call_command(TestCommand(), stdout=out)
+        self.assertEqual(exception, e_info.exception)
+        self.assertIn('args', CommandSignalTests.pre)
+        self.assertIn('kwargs', CommandSignalTests.pre)
+
+        self.assertIn('args', CommandSignalTests.post)
+        self.assertIn('kwargs', CommandSignalTests.post)
+        self.assertIn('outcome', CommandSignalTests.post)
+        self.assertEqual(e_info.exception, CommandSignalTests.post["outcome"])
+
 
 class CommandClassTests(TestCase):
     def setUp(self):