1 files changed, 3 insertions, 1 deletions
diff --git a/Swift/Controllers/UnitTest/MockChatWindow.h b/Swift/Controllers/UnitTest/MockChatWindow.h
index 56f118d..389d787 100644
--- a/Swift/Controllers/UnitTest/MockChatWindow.h
+++ b/Swift/Controllers/UnitTest/MockChatWindow.h
@@ -1,53 +1,54 @@
 /*
  * Copyright (c) 2010-2017 Isode Limited.
  * All rights reserved.
  * See the COPYING file for more information.
  */
 
 #pragma once
 
 #include <memory>
 #include <string>
 
 #include <Swift/Controllers/UIInterfaces/ChatWindow.h>
 
 namespace Swift {
     class MockChatWindow : public ChatWindow {
         public:
             MockChatWindow() {}
             virtual ~MockChatWindow();
 
-            virtual std::string addMessage(const ChatMessage& message, const std::string& senderName, bool senderIsSelf, std::shared_ptr<SecurityLabel> /*label*/, const std::string& /*avatarPath*/, const boost::posix_time::ptime& /*time*/) {
+            virtual std::string addMessage(const ChatMessage& message, const std::string& senderName, bool senderIsSelf, std::shared_ptr<SecurityLabel> label, const std::string& /*avatarPath*/, const boost::posix_time::ptime& /*time*/) {
                 lastAddedMessage_ = message;
                 lastAddedMessageSenderName_ = senderName;
                 lastAddedMessageSenderIsSelf_ = senderIsSelf;
+                lastAddedMessageSecurityLabel_ = label;
                 return "id";
             }
 
             virtual std::string addAction(const ChatMessage& message, const std::string& senderName, bool senderIsSelf, std::shared_ptr<SecurityLabel> /*label*/, const std::string& /*avatarPath*/, const boost::posix_time::ptime& /*time*/) {
                 lastAddedAction_ = message;
                 lastAddedActionSenderName_ = senderName;
                 lastAddedActionSenderIsSelf_ = senderIsSelf;
                 return "id";
             }
 
             virtual std::string addSystemMessage(const ChatMessage& message, Direction /*direction*/) {
                 lastAddedSystemMessage_ = message;
                 return "id";
             }
 
             virtual void addPresenceMessage(const ChatMessage& message, Direction /*direction*/) {
                 lastAddedPresence_ = message;
             }
 
             virtual void addErrorMessage(const ChatMessage& message) {
                 lastAddedErrorMessage_ = message;
             }
 
             virtual void replaceMessage(const ChatMessage& message, const std::string& /*id*/, const boost::posix_time::ptime& /*time*/) {
                 lastReplacedMessage_ = message;
             }
 
             virtual void replaceWithAction(const ChatMessage& /*message*/, const std::string& /*id*/, const boost::posix_time::ptime& /*time*/) {}
             virtual void replaceLastMessage(const ChatMessage& message, const TimestampBehaviour /*timestampBehaviour*/) {
                 lastReplacedLastMessage_ = message;
@@ -118,53 +119,54 @@ namespace Swift {
                 std::shared_ptr<ChatHighlightingMessagePart> highlight;
                 for (auto &&part : message.getParts()) {
                     if ((text = std::dynamic_pointer_cast<ChatTextMessagePart>(part))) {
                         body += text->text;
                     }
                     else if ((highlight = std::dynamic_pointer_cast<ChatHighlightingMessagePart>(part))) {
                         body += highlight->text;
                     }
                 }
                 return body;
             }
 
             void resetLastMessages() {
                 lastAddedMessage_ = lastAddedAction_ = lastAddedPresence_ = lastReplacedLastMessage_ = lastAddedSystemMessage_ = lastReplacedSystemMessage_ = ChatMessage();
                 lastAddedMessageSenderName_ = lastAddedActionSenderName_ = "";
                 lastAddedMessageSenderIsSelf_ = lastAddedActionSenderIsSelf_ = false;
             }
 
             void setChatSecurityMarking(const std::string& markingValue, const std::string& markingForegroundColorValue, const std::string& markingBackgroundColorValue) {
                 markingValue_ = markingValue;
                 markingForegroundColorValue_ = markingForegroundColorValue;
                 markingBackgroundColorValue_ = markingBackgroundColorValue;
             }
 
             void removeChatSecurityMarking() {}
 
             std::string name_;
             ChatMessage lastAddedMessage_;
             std::string lastAddedMessageSenderName_;
             bool lastAddedMessageSenderIsSelf_ = false;
+            std::shared_ptr<SecurityLabel> lastAddedMessageSecurityLabel_ = nullptr;
             ChatMessage lastAddedAction_;
             std::string lastAddedActionSenderName_;
             bool lastAddedActionSenderIsSelf_ = false;
             ChatMessage lastAddedPresence_;
             ChatMessage lastReplacedMessage_;
             ChatMessage lastReplacedLastMessage_;
             ChatMessage lastAddedSystemMessage_;
             ChatMessage lastReplacedSystemMessage_;
             ChatMessage lastAddedErrorMessage_;
             JID lastMUCInvitationJID_;
             std::vector<SecurityLabelsCatalog::Item> labels_;
             bool labelsEnabled_ = false;
             bool impromptuMUCSupported_ = false;
             SecurityLabelsCatalog::Item label_;
             Roster* roster_ = nullptr;
             std::vector<std::pair<std::string, ReceiptState>> receiptChanges_;
             boost::optional<MUCType> mucType_;
             std::string markingValue_;
             std::string markingForegroundColorValue_;
             std::string markingBackgroundColorValue_;
     };
 }