SERVER-33255 clean up TransportLayer API

1 files changed, 47 insertions, 6 deletions

diff --git a/src/mongo/transport/mock_session.h b/src/mongo/transport/mock_session.h
index 147e13e4f7e..2f091854193 100644
--- a/src/mongo/transport/mock_session.h
+++ b/src/mongo/transport/mock_session.h
@@ -28,14 +28,14 @@
 
 #pragma once
 
+#include "mongo/base/checked_cast.h"
 #include "mongo/transport/session.h"
+#include "mongo/transport/transport_layer_mock.h"
 #include "mongo/util/net/hostandport.h"
 
 namespace mongo {
 namespace transport {
 
-class TransportLayer;
-
 class MockSession : public Session {
     MONGO_DISALLOW_COPYING(MockSession);
 
@@ -65,12 +65,53 @@ public:
         return _local;
     }
 
-protected:
-    explicit MockSession(TransportLayer* tl) : _tl(tl), _remote(), _local() {}
+    void end() override {
+        if (!_tl->owns(id()))
+            return;
+        _tl->_sessions[id()].ended = true;
+    }
+
+    StatusWith<Message> sourceMessage() override {
+        if (_tl->inShutdown()) {
+            return TransportLayer::ShutdownStatus;
+        } else if (!_tl->owns(id())) {
+            return TransportLayer::SessionUnknownStatus;
+        } else if (_tl->_sessions[id()].ended) {
+            return TransportLayer::TicketSessionClosedStatus;
+        }
+
+        return Message();  // Subclasses can do something different.
+    }
+
+    void asyncSourceMessage(std::function<void(StatusWith<Message>)> cb) override {
+        cb(sourceMessage());
+    }
+
+    Status sinkMessage(Message message) override {
+        if (_tl->inShutdown()) {
+            return TransportLayer::ShutdownStatus;
+        } else if (!_tl->owns(id())) {
+            return TransportLayer::SessionUnknownStatus;
+        } else if (_tl->_sessions[id()].ended) {
+            return TransportLayer::TicketSessionClosedStatus;
+        }
+
+        return Status::OK();
+    }
+
+    void asyncSinkMessage(Message message, std::function<void(Status)> cb) override {
+        cb(sinkMessage(message));
+    }
+
+    explicit MockSession(TransportLayer* tl)
+        : _tl(checked_cast<TransportLayerMock*>(tl)), _remote(), _local() {}
     explicit MockSession(HostAndPort remote, HostAndPort local, TransportLayer* tl)
-        : _tl(tl), _remote(std::move(remote)), _local(std::move(local)) {}
+        : _tl(checked_cast<TransportLayerMock*>(tl)),
+          _remote(std::move(remote)),
+          _local(std::move(local)) {}
 
-    TransportLayer* _tl;
+protected:
+    TransportLayerMock* _tl;
 
     HostAndPort _remote;
     HostAndPort _local;