diff options
Diffstat (limited to 'src/mongo/executor/network_interface_mock.h')
| -rw-r--r-- | src/mongo/executor/network_interface_mock.h | 21 |
1 files changed, 19 insertions, 2 deletions
diff --git a/src/mongo/executor/network_interface_mock.h b/src/mongo/executor/network_interface_mock.h index e096589d640..b8610d58d9d 100644 --- a/src/mongo/executor/network_interface_mock.h +++ b/src/mongo/executor/network_interface_mock.h @@ -41,6 +41,7 @@ #include "mongo/stdx/condition_variable.h" #include "mongo/stdx/unordered_map.h" #include "mongo/stdx/unordered_set.h" +#include "mongo/transport/mock_session.h" #include "mongo/util/clock_source.h" #include "mongo/util/clock_source_mock.h" #include "mongo/util/time_support.h" @@ -148,6 +149,20 @@ public: void testEgress(const HostAndPort&, transport::ConnectSSLMode, Milliseconds, Status) override {} + using LeasedStreamMaker = + std::function<std::unique_ptr<NetworkInterface::LeasedStream>(HostAndPort)>; + + void setLeasedStreamMaker(LeasedStreamMaker lsm) { + _leasedStreamMaker = std::move(lsm); + } + + SemiFuture<std::unique_ptr<NetworkInterface::LeasedStream>> leaseStream( + const HostAndPort& hp, transport::ConnectSSLMode, Milliseconds) override { + invariant(_leasedStreamMaker, + "Tried to lease a stream from NetworkInterfaceMock without providing one"); + return (*_leasedStreamMaker)(hp); + } + //////////////////////////////////////////////////////////////////////////////// // // Methods for simulating network operations and the passage of time. @@ -273,9 +288,9 @@ public: void runReadyNetworkOperations(); /** - * Sets the reply of the 'isMaster' handshake for a specific host. This reply will only + * Sets the reply of the 'hello' handshake for a specific host. This reply will only * be given to the 'validateHost' method of the ConnectionHook set on this object - NOT - * to the completion handlers of any 'isMaster' commands scheduled with 'startCommand'. + * to the completion handlers of any 'hello' commands scheduled with 'startCommand'. * * This reply will persist until it is changed again using this method. * @@ -446,6 +461,8 @@ private: // The handshake replies set for each host. stdx::unordered_map<HostAndPort, RemoteCommandResponse> _handshakeReplies; // (M) + + boost::optional<LeasedStreamMaker> _leasedStreamMaker; }; /** |