1 files changed, 50 insertions, 33 deletions
diff --git a/Swiften/TLS/UnitTest/ClientServerTest.cpp b/Swiften/TLS/UnitTest/ClientServerTest.cpp
index 24bd7c5..a356dcf 100644
--- a/Swiften/TLS/UnitTest/ClientServerTest.cpp
+++ b/Swiften/TLS/UnitTest/ClientServerTest.cpp
@@ -1,3 +1,3 @@
 /*
- * Copyright (c) 2010-2018 Isode Limited.
+ * Copyright (c) 2010-2019 Isode Limited.
  * All rights reserved.
@@ -410,3 +410,2 @@ class TLSEventToSafeByteArrayVisitor : public boost::static_visitor<SafeByteArra
         }
-
 };
@@ -461,2 +460,19 @@ class TLSClientServerEventHistory {
 
+        template<class TLSEventType>
+        boost::optional<TLSEventType> getEvent(const std::string& peer, size_t number = 0) {
+            for (const auto& pair : events) {
+                if (pair.first == peer) {
+                    if (pair.second.type() == typeid(TLSEventType)) {
+                        if (number == 0) {
+                            return boost::optional<TLSEventType>(boost::get<TLSEventType>(pair.second));
+                        }
+                        else {
+                            number--;
+                        }
+                    }
+                }
+            }
+            return {};
+        }
+
     private:
@@ -604,8 +620,8 @@ TEST(ClientServerTest, testClientServerBasicCommunication) {
 
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
 }
@@ -634,8 +650,8 @@ TEST(ClientServerTest, testClientServerBasicCommunicationEncryptedPrivateKeyRigh
 
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
 }
@@ -741,10 +757,11 @@ TEST(ClientServerTest, testClientServerSNIRequestedHostAvailable) {
     serverContext->handleDataFromApplication(createSafeByteArray("This is a test message from the server."));
-    ASSERT_EQ("This is a test message from the client.", safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ("This is a test message from the server.", safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
 
-    ASSERT_EQ("/CN=montague.example", boost::get<TLSConnected>(events.events[5].second).chain[0]->getSubjectName());
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
+
+    ASSERT_EQ("/CN=montague.example", events.getEvent<TLSConnected>("client")->chain[0]->getSubjectName());
 }
@@ -827,8 +844,8 @@ TEST(ClientServerTest, testClientServerBasicCommunicationWith2048BitDHParams) {
 
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
 }
@@ -859,8 +876,8 @@ TEST(ClientServerTest, testClientServerBasicCommunicationWith1024BitDHParams) {
 
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
 }