2 files changed, 26 insertions, 4 deletions
diff --git a/mediagoblin/messages.py b/mediagoblin/messages.py
index 80d8ece7..d58f13d4 100644
--- a/mediagoblin/messages.py
+++ b/mediagoblin/messages.py
@@ -14,16 +14,24 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+from mediagoblin.tools import common
+
 DEBUG = 'debug'
 INFO = 'info'
 SUCCESS = 'success'
 WARNING = 'warning'
 ERROR = 'error'
 
+ADD_MESSAGE_TEST = []
+
 
 def add_message(request, level, text):
     messages = request.session.setdefault('messages', [])
     messages.append({'level': level, 'text': text})
+
+    if common.TESTS_ENABLED:
+        ADD_MESSAGE_TEST.append(messages)
+
     request.session.save()
 
 
@@ -33,4 +41,10 @@ def fetch_messages(request, clear_from_session=True):
         # Save that we removed the messages from the session
         request.session['messages'] = []
         request.session.save()
+
     return messages
+
+
+def clear_add_message():
+    global ADD_MESSAGE_TEST
+    ADD_MESSAGE_TEST = []
diff --git a/mediagoblin/tests/test_messages.py b/mediagoblin/tests/test_messages.py
index 3ac917b0..22f9e800 100644
--- a/mediagoblin/tests/test_messages.py
+++ b/mediagoblin/tests/test_messages.py
@@ -14,7 +14,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-from mediagoblin.messages import fetch_messages, add_message
+from mediagoblin import messages
 from mediagoblin.tools import template
 
 
@@ -32,11 +32,19 @@ def test_messages(test_app):
     # The message queue should be empty
     assert request.session.get('messages', []) == []
 
+    # First of all, we should clear the messages queue
+    messages.clear_add_message()
     # Adding a message should modify the session accordingly
-    add_message(request, 'herp_derp', 'First!')
+    messages.add_message(request, 'herp_derp', 'First!')
     test_msg_queue = [{'text': 'First!', 'level': 'herp_derp'}]
-    assert request.session['messages'] == test_msg_queue
+
+    # Alternative tests to the following, test divided in two steps:
+    # assert request.session['messages'] == test_msg_queue
+    # 1. Tests if add_message worked
+    assert messages.ADD_MESSAGE_TEST[-1] == test_msg_queue
+    # 2. Tests if add_message updated session information
+    assert messages.ADD_MESSAGE_TEST[-1] == request.session['messages']
 
     # fetch_messages should return and empty the queue
-    assert fetch_messages(request) == test_msg_queue
+    assert messages.fetch_messages(request) == test_msg_queue
     assert request.session.get('messages') == []