1 files changed, 8 insertions, 0 deletions
diff --git a/src/mongo/s/query/router_stage_mock.h b/src/mongo/s/query/router_stage_mock.h
index 7cb31ce9745..b83e2879096 100644
--- a/src/mongo/s/query/router_stage_mock.h
+++ b/src/mongo/s/query/router_stage_mock.h
@@ -49,6 +49,8 @@ public:
 
     bool remotesExhausted() final;
 
+    Status setAwaitDataTimeout(Milliseconds awaitDataTimeout) final;
+
     /**
      * Queues a BSONObj to be returned.
      */
@@ -70,9 +72,15 @@ public:
      */
     void markRemotesExhausted();
 
+    /**
+     * Gets the timeout for awaitData, or an error if none was set.
+     */
+    StatusWith<Milliseconds> getAwaitDataTimeout();
+
 private:
     std::queue<StatusWith<boost::optional<BSONObj>>> _resultsQueue;
     bool _remotesExhausted = false;
+    boost::optional<Milliseconds> _awaitDataTimeout;
 };
 
 }  // namespace mongo