From dc51a6c4b198e924b38328971056a2d342cafbca Mon Sep 17 00:00:00 2001
From: PelleK <elfjes@users.noreply.github.com>
Date: Wed, 11 Dec 2024 15:01:39 +0100
Subject: [PATCH] Keep django messages in queue on htmx redirects or refreshes

---
 changelog.d/1071.added.md     |  1 +
 src/argus/htmx/middleware.py  |  4 +++
 tests/htmx/test_middleware.py | 51 +++++++++++++++++++++++++++++++++--
 3 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 changelog.d/1071.added.md

diff --git a/changelog.d/1071.added.md b/changelog.d/1071.added.md
new file mode 100644
index 000000000..b2151faa7
--- /dev/null
+++ b/changelog.d/1071.added.md
@@ -0,0 +1 @@
+argus.htmx: Keep django messages in queue on htmx redirects or refreshes
diff --git a/src/argus/htmx/middleware.py b/src/argus/htmx/middleware.py
index e6df5dfb7..cec287c91 100644
--- a/src/argus/htmx/middleware.py
+++ b/src/argus/htmx/middleware.py
@@ -69,6 +69,10 @@ def process_response(self, request: HtmxHttpRequest, response: HttpResponse) ->
         if 300 <= response.status_code < 400:
             return response
 
+        # do not add messages to hx redirects/refreshes
+        if response.headers.get("HX-Refresh") == "true" or {"HX-Redirect", "HX-Location"} & response.headers.keys():
+            return response
+
         if not response.writable():
             return response
 
diff --git a/tests/htmx/test_middleware.py b/tests/htmx/test_middleware.py
index 64783a4a7..4655921cc 100644
--- a/tests/htmx/test_middleware.py
+++ b/tests/htmx/test_middleware.py
@@ -1,10 +1,18 @@
 from unittest.mock import Mock
 
 from django import test
-from django.http import HttpResponseRedirect
+from django.contrib import messages
+from django.contrib.messages.middleware import MessageMiddleware
+from django.contrib.sessions.middleware import SessionMiddleware
+from django.http import HttpResponse, HttpResponseRedirect
 from django.test.client import RequestFactory
+from django_htmx.http import (
+    HttpResponseClientRedirect,
+    HttpResponseClientRefresh,
+    HttpResponseLocation,
+)
 
-from argus.htmx.middleware import LoginRequiredMiddleware
+from argus.htmx.middleware import HtmxMessageMiddleware, LoginRequiredMiddleware
 
 
 class TestLoginRequiredMiddleware(test.TestCase):
@@ -48,3 +56,42 @@ def view_func():
         result = LoginRequiredMiddleware(lambda x: x).process_view(self.request, view_func, None, {})
         self.assertIsNotNone(result)
         self.assertIsInstance(result, HttpResponseRedirect)
+
+
+class TestHtmxMessageMiddleware(test.TestCase):
+    def setUp(self):
+        request = RequestFactory().get("/foo")
+        request.htmx = True
+        request.user = Mock()
+        self.request = request
+
+        SessionMiddleware(lambda x: x).process_request(self.request)
+        MessageMiddleware(lambda x: x).process_request(self.request)
+        messages.info(self.request, "a message")
+
+        self.middleware = HtmxMessageMiddleware(lambda x: x)
+
+    def process_response(self, response: HttpResponse):
+        return self.middleware.process_response(self.request, response).content.decode()
+
+    def tearDown(self):
+        # expire current messages
+        messages.get_messages(self.request)
+
+    def test_adds_message_to_response(self):
+        self.assertIn("a message", self.process_response(HttpResponse()))
+
+    def test_doesnt_add_message_to_response_if_not_htmx(self):
+        self.request.htmx = False
+        self.assertNotIn("a message", self.process_response(HttpResponse()))
+
+    def test_doesnt_add_message_on_redirect_response(self):
+        responses = [
+            ("redirect", HttpResponseRedirect("/")),
+            ("hx-redirect", HttpResponseClientRedirect("/")),
+            ("hx-location", HttpResponseLocation("/")),
+            ("hx-refresh", HttpResponseClientRefresh()),
+        ]
+        for name, response in responses:
+            with self.subTest(name):
+                self.assertNotIn("a message", self.process_response(response))