SERVER-24153 Add unit tests for DocumentSource pausing behavior.

1 files changed, 6 insertions, 7 deletions

diff --git a/src/mongo/db/pipeline/tee_buffer.cpp b/src/mongo/db/pipeline/tee_buffer.cpp
index 51ed90241e9..f16a6ffd289 100644
--- a/src/mongo/db/pipeline/tee_buffer.cpp
+++ b/src/mongo/db/pipeline/tee_buffer.cpp
@@ -81,13 +81,6 @@ void TeeBuffer::loadNextBatch() {
 
     auto input = _source->getNext();
     for (; input.isAdvanced(); input = _source->getNext()) {
-
-        // For the following reasons, we invariant that we never get a paused input:
-        //   - TeeBuffer is the only place where a paused GetNextReturn will be returned.
-        //   - The $facet stage is the only stage that uses TeeBuffer.
-        //   - We currently disallow nested $facet stages.
-        invariant(!input.isPaused());
-
         bytesInBuffer += input.getDocument().getApproximateSize();
         _buffer.push_back(std::move(input));
 
@@ -96,6 +89,12 @@ void TeeBuffer::loadNextBatch() {
         }
     }
 
+    // For the following reasons, we invariant that we never get a paused input:
+    //   - TeeBuffer is the only place where a paused GetNextReturn will be returned.
+    //   - The $facet stage is the only stage that uses TeeBuffer.
+    //   - We currently disallow nested $facet stages.
+    invariant(!input.isPaused());
+
     // Populate the pending returns.
     for (size_t consumerId = 0; consumerId < _consumers.size(); ++consumerId) {
         if (_consumers[consumerId].stillInUse) {