SERVER-23913 Implement TransportLayer test suite and mock

19 files changed, 593 insertions, 58 deletions

diff --git a/src/mongo/base/error_codes.err b/src/mongo/base/error_codes.err
index 7c8f58de4fd..cee5d694d4d 100644
--- a/src/mongo/base/error_codes.err
+++ b/src/mongo/base/error_codes.err
@@ -170,15 +170,17 @@ error_code("InvalidPipelineOperator", 168)
 error_code("CommandOnShardedViewNotSupportedOnMongod", 169)
 error_code("TooManyMatchingDocuments", 170)
 error_code("CannotIndexParallelArrays", 171)
-error_code("TransportSessionNotFound", 172)
-error_code("QueryPlanKilled", 173)
-error_code("FileOpenFailed", 174)
-error_code("ZoneNotFound", 175)
-error_code("RangeOverlapConflict", 176)
-error_code("WindowsPdhError", 177)
-error_code("BadPerfCounterPath", 178)
-error_code("AmbiguousIndexKeyPattern", 179)
-error_code("InvalidViewDefinition", 180);
+error_code("TransportSessionClosed", 172)
+error_code("TransportSessionNotFound", 173)
+error_code("TransportSessionUnknown", 174)
+error_code("QueryPlanKilled", 175)
+error_code("FileOpenFailed", 176)
+error_code("ZoneNotFound", 177)
+error_code("RangeOverlapConflict", 178)
+error_code("WindowsPdhError", 179)
+error_code("BadPerfCounterPath", 180)
+error_code("AmbiguousIndexKeyPattern", 181)
+error_code("InvalidViewDefinition", 182);
 
 # Non-sequential error codes (for compatibility only)
 error_code("SocketException", 9001)
diff --git a/src/mongo/transport/SConscript b/src/mongo/transport/SConscript
index cbc2d3b9432..38ef3410619 100644
--- a/src/mongo/transport/SConscript
+++ b/src/mongo/transport/SConscript
@@ -17,6 +17,7 @@ env.Library(
     source=[
         'session.cpp',
         'ticket.cpp',
+        'transport_layer.cpp',
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/unittest/unittest',
@@ -89,3 +90,13 @@ env.CppUnitTest(
         'service_entry_point_test_suite',
     ],
 )
+
+env.CppUnitTest(
+    target='transport_layer_mock_test',
+    source=[
+        'transport_layer_mock_test.cpp',
+    ],
+    LIBDEPS=[
+        'transport_layer_mock',
+    ],
+)
diff --git a/src/mongo/transport/service_entry_point_mock_test.cpp b/src/mongo/transport/service_entry_point_mock_test.cpp
index c996c446b37..66806111980 100644
--- a/src/mongo/transport/service_entry_point_mock_test.cpp
+++ b/src/mongo/transport/service_entry_point_mock_test.cpp
@@ -46,19 +46,33 @@ std::unique_ptr<ServiceEntryPoint> mockSEPFactory(transport::TransportLayer* tl)
 
 }  // namespace
 
-TEST_F(ServiceEntryPointTestSuite, ServiceEntryPointMockTest) {
+TEST_F(ServiceEntryPointTestSuite, NoLifeCycleTest) {
     setServiceEntryPoint(&mockSEPFactory);
-
-    // Lifecycle tests
     noLifeCycleTest();
+}
+
+TEST_F(ServiceEntryPointTestSuite, HalfLifeCycleTest) {
+    setServiceEntryPoint(&mockSEPFactory);
     halfLifeCycleTest();
+}
+
+TEST_F(ServiceEntryPointTestSuite, FullLifeCycleTest) {
+    setServiceEntryPoint(&mockSEPFactory);
     fullLifeCycleTest();
+}
 
-    // Concurrent session tests
+TEST_F(ServiceEntryPointTestSuite, InterruptingSessionTest) {
+    setServiceEntryPoint(&mockSEPFactory);
     interruptingSessionTest();
+}
 
-    // Stress tests
+TEST_F(ServiceEntryPointTestSuite, BurstStressTest) {
+    setServiceEntryPoint(&mockSEPFactory);
     burstStressTest();
+}
+
+TEST_F(ServiceEntryPointTestSuite, LongSessionStressTest) {
+    setServiceEntryPoint(&mockSEPFactory);
     longSessionStressTest();
 }
 
diff --git a/src/mongo/transport/service_entry_point_test_suite.cpp b/src/mongo/transport/service_entry_point_test_suite.cpp
index e2a9c7f336c..efd43fd9a9b 100644
--- a/src/mongo/transport/service_entry_point_test_suite.cpp
+++ b/src/mongo/transport/service_entry_point_test_suite.cpp
@@ -158,7 +158,7 @@ TransportLayer::Stats ServiceEntryPointTestSuite::MockTLHarness::sessionStats()
     return Stats();
 }
 
-void ServiceEntryPointTestSuite::MockTLHarness::end(const Session& session) {
+void ServiceEntryPointTestSuite::MockTLHarness::end(Session& session) {
     return _end(session);
 }
 
diff --git a/src/mongo/transport/service_entry_point_test_suite.h b/src/mongo/transport/service_entry_point_test_suite.h
index 828fee5fe3c..623f0533267 100644
--- a/src/mongo/transport/service_entry_point_test_suite.h
+++ b/src/mongo/transport/service_entry_point_test_suite.h
@@ -132,7 +132,7 @@ private:
         std::string getX509SubjectName(const transport::Session& session) override;
         void registerTags(const transport::Session& session) override;
         Stats sessionStats() override;
-        void end(const transport::Session& session) override;
+        void end(transport::Session& session) override;
         void endAllSessions(
             transport::Session::TagMask tags = transport::Session::kEmptyTagMask) override;
         Status start() override;
diff --git a/src/mongo/transport/session.cpp b/src/mongo/transport/session.cpp
index ebf237f8bd4..b341cd97630 100644
--- a/src/mongo/transport/session.cpp
+++ b/src/mongo/transport/session.cpp
@@ -42,15 +42,19 @@ AtomicUInt64 sessionIdCounter(0);
 
 }  // namespace
 
+const Status Session::ClosedStatus =
+    Status(ErrorCodes::TransportSessionClosed, "Session is closed.");
+
 Session::Session(HostAndPort remote, HostAndPort local, TransportLayer* tl)
-    : _id(sessionIdCounter.addAndFetch(1)),
+    : _ended(false),
+      _id(sessionIdCounter.addAndFetch(1)),
       _remote(std::move(remote)),
       _local(std::move(local)),
       _tags(kEmptyTagMask),
       _tl(tl) {}
 
 Session::~Session() {
-    if (_tl != nullptr) {
+    if (_tl != nullptr && !_ended) {
         _tl->end(*this);
     }
 }
@@ -95,5 +99,12 @@ std::string Session::getX509SubjectName() const {
     return _tl->getX509SubjectName(*this);
 }
 
+void Session::end() {
+    if (!_ended) {
+        _ended = true;
+        _tl->end(*this);
+    }
+}
+
 }  // namespace transport
 }  // namespace mongo
diff --git a/src/mongo/transport/session.h b/src/mongo/transport/session.h
index b340b0fecf8..8551f607037 100644
--- a/src/mongo/transport/session.h
+++ b/src/mongo/transport/session.h
@@ -33,6 +33,7 @@
 #include "mongo/transport/ticket.h"
 #include "mongo/util/net/hostandport.h"
 #include "mongo/util/net/message.h"
+#include "mongo/util/time_support.h"
 
 namespace mongo {
 namespace transport {
@@ -56,6 +57,9 @@ public:
      * Tags for groups of connections.
      */
     using TagMask = uint32_t;
+
+    static const Status ClosedStatus;
+
     static constexpr TagMask kEmptyTagMask = 0;
     static constexpr TagMask kKeepOpen = 1;
 
@@ -136,7 +140,21 @@ public:
         return _tl;
     }
 
+    /*
+     * End the session.
+     */
+    void end();
+
+    /*
+     * Return true if the session ended, false otherwise.
+     */
+    bool ended() const {
+        return _ended;
+    }
+
 private:
+    bool _ended = false;
+
     Id _id;
 
     HostAndPort _remote;
diff --git a/src/mongo/transport/ticket.cpp b/src/mongo/transport/ticket.cpp
index 77a290087dd..f7ee14778c4 100644
--- a/src/mongo/transport/ticket.cpp
+++ b/src/mongo/transport/ticket.cpp
@@ -28,6 +28,7 @@
 
 #include "mongo/platform/basic.h"
 
+#include "mongo/base/status.h"
 #include "mongo/transport/ticket.h"
 #include "mongo/transport/ticket_impl.h"
 #include "mongo/transport/transport_layer.h"
@@ -37,9 +38,16 @@ namespace transport {
 
 const Date_t Ticket::kNoExpirationDate{Date_t::max()};
 
+Status Ticket::ExpiredStatus = Status(ErrorCodes::ExceededTimeLimit, "Ticket has expired.");
+
+Status Ticket::SessionClosedStatus =
+    Status(ErrorCodes::TransportSessionClosed, "Ticket's Session is closed.");
+
 Ticket::Ticket(TransportLayer* tl, std::unique_ptr<TicketImpl> ticket)
     : _tl(tl), _ticket(std::move(ticket)) {}
 
+Ticket::Ticket(Status status) : _status(status) {}
+
 Ticket::~Ticket() = default;
 
 Ticket::Ticket(Ticket&&) = default;
@@ -53,5 +61,21 @@ void Ticket::asyncWait(TicketCallback cb)&& {
     return _tl->asyncWait(std::move(*this), std::move(cb));
 }
 
+bool Ticket::valid() {
+    return _status == Status::OK() && !expired();
+}
+
+Status Ticket::status() const {
+    return _status;
+}
+
+bool Ticket::expired() {
+    bool expired = expiration() <= Date_t::now();
+    if (_status == Status::OK() && expired) {
+        _status = Status(ErrorCodes::ExceededTimeLimit, "Ticket has expired.");
+    }
+    return expired;
+}
+
 }  // namespace transport
 }  // namespace mongo
diff --git a/src/mongo/transport/ticket.h b/src/mongo/transport/ticket.h
index dd8927085aa..63b064d104b 100644
--- a/src/mongo/transport/ticket.h
+++ b/src/mongo/transport/ticket.h
@@ -31,14 +31,15 @@
 #include <memory>
 
 #include "mongo/base/disallow_copying.h"
+#include "mongo/base/status.h"
 #include "mongo/stdx/functional.h"
-#include "mongo/transport/session_id.h"
 #include "mongo/transport/ticket_impl.h"
 #include "mongo/util/time_support.h"
 
 namespace mongo {
 namespace transport {
 
+class Session;
 class TransportLayer;
 
 /**
@@ -59,7 +60,16 @@ public:
      */
     static const Date_t kNoExpirationDate;
 
+    static Status ExpiredStatus;
+    static Status SessionClosedStatus;
+
     Ticket(TransportLayer* tl, std::unique_ptr<TicketImpl> ticket);
+
+    /**
+     * An invalid ticket and a Status for why it's invalid.
+     */
+    Ticket(Status status);
+
     ~Ticket();
 
     /**
@@ -96,7 +106,27 @@ public:
      */
     void asyncWait(TicketCallback cb) &&;
 
-protected:
+    /*
+     * Return's the Status of the ticket.
+     */
+    Status status() const;
+
+    /*
+     * Returns true if the ticket has expired, false otherwise.
+     *
+     * If the ticket has expired, changes the Status of the ticket to TicketExpired but
+     * only if the current Status is Status::OK().
+     */
+    bool expired();
+
+    /*
+     * Return true if the ticket is usable, false otherwise.
+     *
+     * If the ticket has expired, changes the Status of the ticket to TicketExpired but
+     * only if the current Status is Status::OK().
+     */
+    bool valid();
+
     /**
      * Return a non-owning pointer to the underlying TicketImpl type
      */
@@ -106,6 +136,7 @@ protected:
 
 private:
     TransportLayer* _tl;
+    Status _status = Status::OK();
     std::unique_ptr<TicketImpl> _ticket;
 };
 
diff --git a/src/mongo/transport/ticket_impl.h b/src/mongo/transport/ticket_impl.h
index 4cbbc0a22d0..22adf9414e7 100644
--- a/src/mongo/transport/ticket_impl.h
+++ b/src/mongo/transport/ticket_impl.h
@@ -30,6 +30,7 @@
 
 #include "mongo/base/disallow_copying.h"
 #include "mongo/transport/session_id.h"
+#include "mongo/util/net/message.h"
 #include "mongo/util/time_support.h"
 
 namespace mongo {
diff --git a/src/mongo/transport/transport_layer.cpp b/src/mongo/transport/transport_layer.cpp
new file mode 100644
index 00000000000..877070191b4
--- /dev/null
+++ b/src/mongo/transport/transport_layer.cpp
@@ -0,0 +1,47 @@
+/**
+ *    Copyright (C) 2016 MongoDB Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    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/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/base/status.h"
+#include "mongo/transport/transport_layer.h"
+
+namespace mongo {
+namespace transport {
+
+const Status TransportLayer::SessionUnknownStatus =
+    Status(ErrorCodes::TransportSessionUnknown, "TransportLayer does not own the Session.");
+
+const Status TransportLayer::ShutdownStatus =
+    Status(ErrorCodes::ShutdownInProgress, "TransportLayer is in shutdown.");
+
+const Status TransportLayer::TicketSessionUnknownStatus = Status(
+    ErrorCodes::TransportSessionUnknown, "TransportLayer does not own the Ticket's Session.");
+
+}  // namespace transport
+}  // namespace mongo
diff --git a/src/mongo/transport/transport_layer.h b/src/mongo/transport/transport_layer.h
index dde6248eaf5..44995fc375b 100644
--- a/src/mongo/transport/transport_layer.h
+++ b/src/mongo/transport/transport_layer.h
@@ -57,6 +57,10 @@ class TransportLayer {
     MONGO_DISALLOW_COPYING(TransportLayer);
 
 public:
+    static const Status SessionUnknownStatus;
+    static const Status ShutdownStatus;
+    static const Status TicketSessionUnknownStatus;
+
     /**
      * Stats for sessions open in the Transport Layer.
      */
@@ -171,7 +175,7 @@ public:
      *
      * This method is idempotent and synchronous.
      */
-    virtual void end(const Session& session) = 0;
+    virtual void end(Session& session) = 0;
 
     /**
      * End all active sessions in the TransportLayer. Tickets that have already been started via
diff --git a/src/mongo/transport/transport_layer_legacy.cpp b/src/mongo/transport/transport_layer_legacy.cpp
index b45d6a6ad07..8d1f95b6a23 100644
--- a/src/mongo/transport/transport_layer_legacy.cpp
+++ b/src/mongo/transport/transport_layer_legacy.cpp
@@ -163,7 +163,7 @@ void TransportLayerLegacy::asyncWait(Ticket&& ticket, TicketCallback callback) {
     MONGO_UNREACHABLE;
 }
 
-void TransportLayerLegacy::end(const Session& session) {
+void TransportLayerLegacy::end(Session& session) {
     stdx::lock_guard<stdx::mutex> lk(_connectionsMutex);
     auto conn = _connections.find(session.id());
     if (conn != _connections.end()) {
@@ -225,11 +225,11 @@ void TransportLayerLegacy::shutdown() {
 
 Status TransportLayerLegacy::_runTicket(Ticket ticket) {
     if (!_running.load()) {
-        return {ErrorCodes::ShutdownInProgress, "TransportLayer in shutdown"};
+        return TransportLayer::ShutdownStatus;
     }
 
     if (ticket.expiration() < Date_t::now()) {
-        return {ErrorCodes::ExceededTimeLimit, "Ticket has expired"};
+        return Ticket::ExpiredStatus;
     }
 
     AbstractMessagingPort* amp;
@@ -239,7 +239,7 @@ Status TransportLayerLegacy::_runTicket(Ticket ticket) {
 
         auto conn = _connections.find(ticket.sessionId());
         if (conn == _connections.end()) {
-            return {ErrorCodes::TransportSessionNotFound, "No such session in TransportLayer"};
+            return TransportLayer::TicketSessionUnknownStatus;
         }
 
         // "check out" the port
diff --git a/src/mongo/transport/transport_layer_legacy.h b/src/mongo/transport/transport_layer_legacy.h
index ec17c76ca1e..1ef193b4754 100644
--- a/src/mongo/transport/transport_layer_legacy.h
+++ b/src/mongo/transport/transport_layer_legacy.h
@@ -83,8 +83,9 @@ public:
 
     Stats sessionStats() override;
 
-    void end(const Session& session) override;
+    void end(Session& session) override;
     void endAllSessions(transport::Session::TagMask tags = Session::kKeepOpen) override;
+
     void shutdown() override;
 
 private:
diff --git a/src/mongo/transport/transport_layer_manager.cpp b/src/mongo/transport/transport_layer_manager.cpp
index c64023b2067..7ba1797a21f 100644
--- a/src/mongo/transport/transport_layer_manager.cpp
+++ b/src/mongo/transport/transport_layer_manager.cpp
@@ -100,7 +100,7 @@ void TransportLayerManager::registerTags(const Session& session) {
     session.getTransportLayer()->registerTags(session);
 }
 
-void TransportLayerManager::end(const Session& session) {
+void TransportLayerManager::end(Session& session) {
     session.getTransportLayer()->end(session);
 }
 
diff --git a/src/mongo/transport/transport_layer_manager.h b/src/mongo/transport/transport_layer_manager.h
index 391bdf4ebea..aeed86edcfe 100644
--- a/src/mongo/transport/transport_layer_manager.h
+++ b/src/mongo/transport/transport_layer_manager.h
@@ -69,11 +69,10 @@ public:
 
     Stats sessionStats() override;
 
-    void end(const Session& session) override;
+    void end(Session& session) override;
     void endAllSessions(Session::TagMask tags = Session::kEmptyTagMask) override;
 
     Status start() override;
-
     void shutdown() override;
 
     Status addAndStartTransportLayer(std::unique_ptr<TransportLayer> tl);
diff --git a/src/mongo/transport/transport_layer_mock.cpp b/src/mongo/transport/transport_layer_mock.cpp
index c550570b401..3f71b5d16e0 100644
--- a/src/mongo/transport/transport_layer_mock.cpp
+++ b/src/mongo/transport/transport_layer_mock.cpp
@@ -28,42 +28,82 @@
 
 #include "mongo/platform/basic.h"
 
-#include <memory>
+#include "mongo/transport/transport_layer_mock.h"
 
 #include "mongo/base/status.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/transport/session.h"
 #include "mongo/transport/ticket.h"
 #include "mongo/transport/ticket_impl.h"
-#include "mongo/transport/transport_layer_mock.h"
+#include "mongo/transport/transport_layer.h"
+#include "mongo/util/net/message.h"
 #include "mongo/util/time_support.h"
 
-#include "mongo/stdx/memory.h"
-
 namespace mongo {
 namespace transport {
 
-Session::Id TransportLayerMock::MockTicket::sessionId() const {
-    return Session::Id{};
+TransportLayerMock::TicketMock::TicketMock(const Session* session,
+                                           Message* message,
+                                           Date_t expiration)
+    : _session(session), _message(message), _expiration(expiration) {}
+
+TransportLayerMock::TicketMock::TicketMock(const Session* session, Date_t expiration)
+    : _session(session), _expiration(expiration) {}
+
+Session::Id TransportLayerMock::TicketMock::sessionId() const {
+    return _session->id();
 }
 
-Date_t TransportLayerMock::MockTicket::expiration() const {
-    return Date_t::now();
+Date_t TransportLayerMock::TicketMock::expiration() const {
+    return _expiration;
 }
 
+boost::optional<Message*> TransportLayerMock::TicketMock::msg() const {
+    return _message;
+}
+
+TransportLayerMock::TransportLayerMock() : _shutdown(false) {}
+
 Ticket TransportLayerMock::sourceMessage(const Session& session,
                                          Message* message,
                                          Date_t expiration) {
-    return Ticket(this, stdx::make_unique<MockTicket>());
+    if (inShutdown()) {
+        return Ticket(TransportLayer::ShutdownStatus);
+    } else if (!owns(session.id())) {
+        return Ticket(TransportLayer::SessionUnknownStatus);
+    } else if (session.ended()) {
+        return Ticket(Session::ClosedStatus);
+    }
+
+    return Ticket(this,
+                  stdx::make_unique<TransportLayerMock::TicketMock>(&session, message, expiration));
 }
 
 Ticket TransportLayerMock::sinkMessage(const Session& session,
                                        const Message& message,
                                        Date_t expiration) {
-    return Ticket(this, stdx::make_unique<MockTicket>());
+    if (inShutdown()) {
+        return Ticket(TransportLayer::ShutdownStatus);
+    } else if (!owns(session.id())) {
+        return Ticket(TransportLayer::SessionUnknownStatus);
+    } else if (session.ended()) {
+        return Ticket(Session::ClosedStatus);
+    }
+
+    return Ticket(this, stdx::make_unique<TransportLayerMock::TicketMock>(&session, expiration));
 }
 
 Status TransportLayerMock::wait(Ticket&& ticket) {
+    if (inShutdown()) {
+        return ShutdownStatus;
+    } else if (!ticket.valid()) {
+        return ticket.status();
+    } else if (!owns(ticket.sessionId())) {
+        return TicketSessionUnknownStatus;
+    } else if (get(ticket.sessionId())->ended()) {
+        return Ticket::SessionClosedStatus;
+    }
+
     return Status::OK();
 }
 
@@ -81,15 +121,57 @@ TransportLayer::Stats TransportLayerMock::sessionStats() {
 
 void TransportLayerMock::registerTags(const Session& session) {}
 
-void TransportLayerMock::end(const Session& session) {}
+Session* TransportLayerMock::createSession() {
+    std::unique_ptr<Session> session =
+        stdx::make_unique<Session>(HostAndPort(), HostAndPort(), this);
+    Session::Id sessionId = session->id();
+
+    _sessions[sessionId] = std::move(session);
+
+    return _sessions[sessionId].get();
+}
+
+Session* TransportLayerMock::get(Session::Id id) {
+    if (!owns(id))
+        return nullptr;
+
+    return _sessions[id].get();
+}
 
-void TransportLayerMock::endAllSessions(Session::TagMask tags) {}
+bool TransportLayerMock::owns(Session::Id id) {
+    return _sessions.count(id) > 0;
+}
+
+void TransportLayerMock::end(Session& session) {
+    session.end();
+}
+
+void TransportLayerMock::endAllSessions(Session::TagMask tags) {
+    auto it = _sessions.begin();
+    while (it != _sessions.end()) {
+        end(*it->second.get());
+        it++;
+    }
+}
 
 Status TransportLayerMock::start() {
     return Status::OK();
 }
 
-void TransportLayerMock::shutdown() {}
+void TransportLayerMock::shutdown() {
+    if (!inShutdown()) {
+        _shutdown = true;
+        endAllSessions();
+    }
+}
+
+bool TransportLayerMock::inShutdown() const {
+    return _shutdown;
+}
+
+TransportLayerMock::~TransportLayerMock() {
+    shutdown();
+}
 
 }  // namespace transport
 }  // namespace mongo
diff --git a/src/mongo/transport/transport_layer_mock.h b/src/mongo/transport/transport_layer_mock.h
index 400d6b02258..f519713e9bc 100644
--- a/src/mongo/transport/transport_layer_mock.h
+++ b/src/mongo/transport/transport_layer_mock.h
@@ -28,6 +28,8 @@
 
 #pragma once
 
+#include <unordered_map>
+
 #include "mongo/base/status.h"
 #include "mongo/transport/session.h"
 #include "mongo/transport/ticket.h"
@@ -46,7 +48,33 @@ class TransportLayerMock : public TransportLayer {
     MONGO_DISALLOW_COPYING(TransportLayerMock);
 
 public:
-    TransportLayerMock() = default;
+    class TicketMock : public TicketImpl {
+    public:
+        // Source constructor
+        TicketMock(const Session* session,
+                   Message* message,
+                   Date_t expiration = Ticket::kNoExpirationDate);
+
+        // Sink constructor
+        TicketMock(const Session* session, Date_t expiration = Ticket::kNoExpirationDate);
+
+        TicketMock(TicketMock&&) = default;
+        TicketMock& operator=(TicketMock&&) = default;
+
+        SessionId sessionId() const override;
+
+        Date_t expiration() const override;
+
+        boost::optional<Message*> msg() const;
+
+    private:
+        const Session* _session;
+        boost::optional<Message*> _message;
+        Date_t _expiration;
+    };
+
+    TransportLayerMock();
+    ~TransportLayerMock();
 
     Ticket sourceMessage(const Session& session,
                          Message* message,
@@ -63,27 +91,19 @@ public:
 
     Stats sessionStats() override;
 
-    void end(const Session& session) override;
+    Session* createSession();
+    Session* get(Session::Id id);
+    bool owns(Session::Id id);
+    void end(Session& session) override;
     void endAllSessions(Session::TagMask tags = Session::kEmptyTagMask) override;
 
     Status start() override;
     void shutdown() override;
+    bool inShutdown() const;
 
 private:
-    /**
-     * A class for Tickets issued from the TransportLayerMock. These will
-     * route to the appropriate TransportLayer when run with wait() or
-     * asyncWait().
-     */
-    class MockTicket : public TicketImpl {
-        MONGO_DISALLOW_COPYING(MockTicket);
-
-    public:
-        MockTicket() = default;
-
-        Session::Id sessionId() const override;
-        Date_t expiration() const override;
-    };
+    std::unordered_map<Session::Id, std::unique_ptr<Session>> _sessions;
+    bool _shutdown;
 };
 
 }  // namespace transport
diff --git a/src/mongo/transport/transport_layer_mock_test.cpp b/src/mongo/transport/transport_layer_mock_test.cpp
new file mode 100644
index 00000000000..22a1da12d88
--- /dev/null
+++ b/src/mongo/transport/transport_layer_mock_test.cpp
@@ -0,0 +1,270 @@
+/**
+ *    Copyright (C) 2016 MongoDB Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    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/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/stdx/memory.h"
+#include "mongo/transport/session.h"
+#include "mongo/transport/transport_layer.h"
+#include "mongo/transport/transport_layer_mock.h"
+#include "mongo/unittest/unittest.h"
+
+namespace mongo {
+namespace transport {
+
+class TransportLayerMockTest : public mongo::unittest::Test {
+public:
+    void setUp() {
+        _transportLayer = stdx::make_unique<TransportLayerMock>();
+    }
+
+    TransportLayerMock* tl() {
+        return _transportLayer.get();
+    }
+
+private:
+    std::unique_ptr<TransportLayerMock> _transportLayer;
+};
+
+// sinkMessage() generates a valid Ticket
+TEST_F(TransportLayerMockTest, SinkMessageGeneratesTicket) {
+    Message* msg = nullptr;
+    Session* session = tl()->createSession();
+
+    // call sinkMessage() with no expiration
+    Ticket ticket = tl()->sinkMessage(*session, *msg);
+    ASSERT(ticket.valid());
+    ASSERT_OK(ticket.status());
+    ASSERT_EQUALS(ticket.sessionId(), session->id());
+    ASSERT_EQUALS(ticket.expiration(), Ticket::kNoExpirationDate);
+
+    // call sinkMessage() with an expiration
+    Date_t expiration = Date_t::now() + Hours(1);
+    ticket = tl()->sinkMessage(*session, *msg, expiration);
+    ASSERT(ticket.valid());
+    ASSERT_OK(ticket.status());
+    ASSERT_EQUALS(ticket.sessionId(), session->id());
+    ASSERT_EQUALS(ticket.expiration(), expiration);
+}
+
+// sinkMessage() generates an invalid Ticket if the Session is closed
+TEST_F(TransportLayerMockTest, SinkMessageSessionClosed) {
+    Message* msg = nullptr;
+    Session* session = tl()->createSession();
+
+    tl()->end(*session);
+
+    Ticket ticket = tl()->sinkMessage(*session, *msg);
+    ASSERT_FALSE(ticket.valid());
+    ASSERT_EQUALS(ticket.status().code(), ErrorCodes::TransportSessionClosed);
+}
+
+// sinkMessage() generates an invalid Ticket if the TransportLayer does not own the Session
+TEST_F(TransportLayerMockTest, SinkMessageSessionUnknown) {
+    Message* msg = nullptr;
+
+    std::unique_ptr<TransportLayerMock> anotherTL = stdx::make_unique<TransportLayerMock>();
+    Session* session = anotherTL->createSession();
+
+    Ticket ticket = tl()->sinkMessage(*session, *msg);
+    ASSERT_FALSE(ticket.valid());
+    ASSERT_EQUALS(ticket.status().code(), ErrorCodes::TransportSessionUnknown);
+}
+
+// sinkMessage() generates an invalid Ticket if the TransportLayer is in shutdown
+TEST_F(TransportLayerMockTest, SinkMessageTLShutdown) {
+    Message* msg = nullptr;
+    Session* session = tl()->createSession();
+
+    tl()->shutdown();
+
+    Ticket ticket = tl()->sinkMessage(*session, *msg);
+    ASSERT_FALSE(ticket.valid());
+    ASSERT_EQUALS(ticket.status().code(), ErrorCodes::ShutdownInProgress);
+}
+
+// sourceMessage() generates a valid ticket
+TEST_F(TransportLayerMockTest, SourceMessageGeneratesTicket) {
+    Message* msg = nullptr;
+    Session* session = tl()->createSession();
+
+    // call sourceMessage() with no expiration
+    Ticket ticket = tl()->sourceMessage(*session, msg);
+    ASSERT(ticket.valid());
+    ASSERT_OK(ticket.status());
+    ASSERT_EQUALS(ticket.sessionId(), session->id());
+    ASSERT_EQUALS(*static_cast<TransportLayerMock::TicketMock*>(ticket.impl())->msg(), msg);
+    ASSERT_EQUALS(ticket.expiration(), Ticket::kNoExpirationDate);
+
+    // call sourceMessage() with an expiration
+    Date_t expiration = Date_t::now() + Hours(1);
+    ticket = tl()->sourceMessage(*session, msg, expiration);
+    ASSERT(ticket.valid());
+    ASSERT_OK(ticket.status());
+    ASSERT_EQUALS(ticket.sessionId(), session->id());
+    ASSERT_EQUALS(*static_cast<TransportLayerMock::TicketMock*>(ticket.impl())->msg(), msg);
+    ASSERT_EQUALS(ticket.expiration(), expiration);
+}
+
+// sourceMessage() generates an invalid ticket if the Session is closed
+TEST_F(TransportLayerMockTest, SourceMessageSessionClosed) {
+    Message* msg = nullptr;
+    Session* session = tl()->createSession();
+
+    tl()->end(*session);
+
+    Ticket ticket = tl()->sourceMessage(*session, msg);
+    ASSERT_FALSE(ticket.valid());
+    ASSERT_EQUALS(ticket.status().code(), ErrorCodes::TransportSessionClosed);
+}
+
+// sourceMessage() generates an invalid ticket if the TransportLayer does not own the Session
+TEST_F(TransportLayerMockTest, SourceMessageSessionUnknown) {
+    Message* msg = nullptr;
+
+    std::unique_ptr<TransportLayerMock> anotherTL = stdx::make_unique<TransportLayerMock>();
+    Session* session = anotherTL->createSession();
+
+    Ticket ticket = tl()->sourceMessage(*session, msg);
+    ASSERT_FALSE(ticket.valid());
+    ASSERT_EQUALS(ticket.status().code(), ErrorCodes::TransportSessionUnknown);
+}
+
+// sourceMessage() generates an invalid ticket if the TransportLayer is in shutdown
+TEST_F(TransportLayerMockTest, SourceMessageTLShutdown) {
+    Message* msg = nullptr;
+    Session* session = tl()->createSession();
+
+    tl()->shutdown();
+
+    Ticket ticket = tl()->sourceMessage(*session, msg);
+    ASSERT_FALSE(ticket.valid());
+    ASSERT_EQUALS(ticket.status().code(), ErrorCodes::ShutdownInProgress);
+}
+
+// wait() returns an OK status
+TEST_F(TransportLayerMockTest, Wait) {
+    Session* session = tl()->createSession();
+    Ticket ticket = Ticket(tl(), stdx::make_unique<TransportLayerMock::TicketMock>(session));
+
+    Status status = tl()->wait(std::move(ticket));
+    ASSERT_OK(status);
+}
+
+// wait() returns an TicketExpired error status if the Ticket expired
+TEST_F(TransportLayerMockTest, WaitExpiredTicket) {
+    Session* session = tl()->createSession();
+    Ticket expiredTicket =
+        Ticket(tl(), stdx::make_unique<TransportLayerMock::TicketMock>(session, Date_t::now()));
+
+    Status status = tl()->wait(std::move(expiredTicket));
+    ASSERT_EQUALS(status.code(), ErrorCodes::ExceededTimeLimit);
+}
+
+// wait() returns the invalid Ticket's Status
+TEST_F(TransportLayerMockTest, WaitInvalidTicket) {
+    Ticket invalidTicket = Ticket(Status(ErrorCodes::UnknownError, ""));
+    ASSERT_FALSE(invalidTicket.valid());
+
+    Status status = tl()->wait(std::move(invalidTicket));
+    ASSERT_EQUALS(status.code(), ErrorCodes::UnknownError);
+}
+
+// wait() returns a SessionClosed error status if the Ticket's Session is closed
+TEST_F(TransportLayerMockTest, WaitSessionClosed) {
+    Session* session = tl()->createSession();
+    Ticket ticket = Ticket(tl(), stdx::make_unique<TransportLayerMock::TicketMock>(session));
+
+    tl()->end(*session);
+
+    Status status = tl()->wait(std::move(ticket));
+    ASSERT_EQUALS(status.code(), ErrorCodes::TransportSessionClosed);
+}
+
+// wait() returns a SessionUnknown error status if the TransportLayer does not own the Ticket's
+// Session
+TEST_F(TransportLayerMockTest, WaitSessionUnknown) {
+    std::unique_ptr<TransportLayerMock> anotherTL = stdx::make_unique<TransportLayerMock>();
+    Session* session = anotherTL->createSession();
+    Ticket ticket = Ticket(tl(), stdx::make_unique<TransportLayerMock::TicketMock>(session));
+
+    Status status = tl()->wait(std::move(ticket));
+    ASSERT_EQUALS(status.code(), ErrorCodes::TransportSessionUnknown);
+}
+
+// wait() returns a ShutdownInProgress status if the TransportLayer is in shutdown
+TEST_F(TransportLayerMockTest, WaitTLShutdown) {
+    Session* session = tl()->createSession();
+    Ticket ticket = Ticket(tl(), stdx::make_unique<TransportLayerMock::TicketMock>(session));
+
+    tl()->shutdown();
+
+    Status status = tl()->wait(std::move(ticket));
+    ASSERT_EQUALS(status.code(), ErrorCodes::ShutdownInProgress);
+}
+
+// end() closes the session
+TEST_F(TransportLayerMockTest, EndSession) {
+    Session* session = tl()->createSession();
+    tl()->end(*session);
+    ASSERT(session->ended());
+}
+
+std::vector<Session*> createSessions(TransportLayerMock* tl) {
+    int numSessions = 10;
+    std::vector<Session*> sessions;
+    for (int i = 0; i < numSessions; i++) {
+        Session* session = tl->createSession();
+        sessions.push_back(session);
+    }
+    return sessions;
+}
+
+void assertEnded(std::vector<Session*> sessions) {
+    for (auto session : sessions) {
+        ASSERT(session->ended());
+    }
+}
+
+// endAllSessions() ends all sessions
+TEST_F(TransportLayerMockTest, EndAllSessions) {
+    std::vector<Session*> sessions = createSessions(tl());
+    tl()->endAllSessions();
+    assertEnded(sessions);
+}
+
+// shutdown() ends all sessions and shuts down
+TEST_F(TransportLayerMockTest, Shutdown) {
+    std::vector<Session*> sessions = createSessions(tl());
+    tl()->shutdown();
+    assertEnded(sessions);
+    ASSERT(tl()->inShutdown());
+}
+
+}  // namespace transport
+}  // namespace mongo