From 5037f3ef0a457aa7b7f15f9048a1d5087466e488 Mon Sep 17 00:00:00 2001
From: "Nicholas H.Tollervey" <ntoll@ntoll.org>
Date: Thu, 23 Jan 2020 13:32:43 +0000
Subject: [PATCH] Fixes and test for @creviera's feedback.

---
 securedrop_client/gui/widgets.py |  6 +++--
 tests/factory.py                 | 14 ++++++++++++
 tests/gui/test_widgets.py        | 39 +++++++++++++++++++++++++++++++-
 3 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/securedrop_client/gui/widgets.py b/securedrop_client/gui/widgets.py
index eec2575034..d4fe5bfa90 100644
--- a/securedrop_client/gui/widgets.py
+++ b/securedrop_client/gui/widgets.py
@@ -31,7 +31,8 @@
     QPushButton, QVBoxLayout, QLineEdit, QScrollArea, QDialog, QAction, QMenu, QMessageBox, \
     QToolButton, QSizePolicy, QPlainTextEdit, QStatusBar, QGraphicsDropShadowEffect
 
-from securedrop_client.db import DraftReply, Source, Message, File, Reply, User
+from securedrop_client.db import (DraftReply, Source, Message, File, Reply, User,
+                                  ReplySendStatusCodes)
 from securedrop_client.storage import source_exists
 from securedrop_client.export import ExportStatus, ExportError
 from securedrop_client.gui import SecureQLabel, SvgLabel, SvgPushButton, SvgToggleButton
@@ -2469,7 +2470,8 @@ def update_conversation(self, collection: list) -> None:
                         item_widget.message.setText(conversation_item.content)
                 # Check if this is a draft reply then ensure it's removed.
                 if isinstance(conversation_item, DraftReply):
-                    self.conversation_layout.removeWidget(item_widget)
+                    if conversation_item.send_status.name == ReplySendStatusCodes.PENDING.value:
+                        self.conversation_layout.removeWidget(item_widget)
             else:
                 # add a new item to be displayed.
                 if isinstance(conversation_item, Message):
diff --git a/tests/factory.py b/tests/factory.py
index 308e9f6ec3..2ef97f5b5b 100644
--- a/tests/factory.py
+++ b/tests/factory.py
@@ -13,6 +13,7 @@
 FILE_COUNT = 0
 REPLY_COUNT = 0
 DRAFT_REPLY_COUNT = 0
+REPLY_SEND_STATUS_COUNT = 0
 USER_COUNT = 0
 
 
@@ -95,6 +96,7 @@ def DraftReply(**attrs):
         file_counter=1,
         uuid='draft-reply-uuid-{}'.format(REPLY_COUNT),
         content='content',
+        send_status_id=1,
     )
 
     defaults.update(attrs)
@@ -102,6 +104,18 @@ def DraftReply(**attrs):
     return db.DraftReply(**defaults)
 
 
+def ReplySendStatus(**attrs):
+    global REPLY_SEND_STATUS_COUNT
+    REPLY_SEND_STATUS_COUNT += 1
+    defaults = dict(
+        name=db.ReplySendStatusCodes.PENDING.value,
+    )
+
+    defaults.update(attrs)
+
+    return db.ReplySendStatus(**defaults)
+
+
 def File(**attrs):
     global FILE_COUNT
     FILE_COUNT += 1
diff --git a/tests/gui/test_widgets.py b/tests/gui/test_widgets.py
index 3923a538d5..8db1dc3338 100644
--- a/tests/gui/test_widgets.py
+++ b/tests/gui/test_widgets.py
@@ -2769,13 +2769,15 @@ def test_update_conversation_removes_draft_items(mocker, session):
     """
     source = factory.Source()
     session.add(source)
+    send_status = factory.ReplySendStatus()
+    session.add(send_status)
     session.commit()
 
     file_ = factory.File(filename='1-source-doc.gpg', source=source)
     session.add(file_)
     message = factory.Message(filename='2-source-msg.gpg', source=source)
     session.add(message)
-    draft_reply = factory.DraftReply(source=source)
+    draft_reply = factory.DraftReply(source=source, send_status=send_status)
     session.add(draft_reply)
     session.commit()
 
@@ -2795,6 +2797,41 @@ def test_update_conversation_removes_draft_items(mocker, session):
     assert cv.conversation_layout.count() == 3
 
 
+def test_update_conversation_keeps_failed_draft_items(mocker, session):
+    """
+    Calling update_conversation keeps items that were added as drafts but which
+    have failed.
+    """
+    source = factory.Source()
+    session.add(source)
+    send_status = factory.ReplySendStatus(name="FAILED")
+    session.add(send_status)
+    session.commit()
+
+    file_ = factory.File(filename='1-source-doc.gpg', source=source)
+    session.add(file_)
+    message = factory.Message(filename='2-source-msg.gpg', source=source)
+    session.add(message)
+    draft_reply = factory.DraftReply(source=source, send_status=send_status)
+    session.add(draft_reply)
+    session.commit()
+
+    mock_get_file = mocker.MagicMock(return_value=file_)
+    mock_controller = mocker.MagicMock(get_file=mock_get_file)
+
+    cv = ConversationView(source, mock_controller)
+    assert cv.conversation_layout.count() == 3  # precondition with draft
+
+    # add the new message and persist
+    new_message = factory.Message(filename='4-source-msg.gpg', source=source)
+    session.add(new_message)
+    session.commit()
+
+    # New message added, draft message retained.
+    cv.update_conversation(cv.source.collection)
+    assert cv.conversation_layout.count() == 4
+
+
 def test_update_conversation_adds_new_items(mocker, session):
     """
     Calling update_conversation adds new items to layout