SERVER-10026 update dbtests for keeping invalidated docs in query

2 files changed, 12 insertions, 15 deletions

diff --git a/src/mongo/dbtests/query_stage_fetch.cpp b/src/mongo/dbtests/query_stage_fetch.cpp
index 2c187542351..8892ddce892 100644
--- a/src/mongo/dbtests/query_stage_fetch.cpp
+++ b/src/mongo/dbtests/query_stage_fetch.cpp
@@ -273,18 +273,15 @@ namespace QueryStageFetch {
             // Invalidate the DL.
             fetchStage->invalidate(member->loc, INVALIDATION_DELETION);
 
-            bool fetchReturnsInvalidated = false;
-            if (fetchReturnsInvalidated) {
-                // Next call to work() should give us the OWNED obj as it was invalidated mid-page-in.
-                state = fetchStage->work(&id);
-                ASSERT_EQUALS(PlanStage::ADVANCED, state);
-                ASSERT_EQUALS(WorkingSetMember::OWNED_OBJ, member->state);
-
-                // We should be able to get data from the obj now.
-                BSONElement elt;
-                ASSERT_TRUE(member->getFieldDotted("foo", &elt));
-                ASSERT_EQUALS(elt.numberInt(), 5);
-            }
+            // Next call to work() should give us the OWNED obj as it was invalidated mid-page-in.
+            state = fetchStage->work(&id);
+            ASSERT_EQUALS(PlanStage::ADVANCED, state);
+            ASSERT_EQUALS(WorkingSetMember::OWNED_OBJ, member->state);
+
+            // We should be able to get data from the obj now.
+            BSONElement elt;
+            ASSERT_TRUE(member->getFieldDotted("foo", &elt));
+            ASSERT_EQUALS(elt.numberInt(), 5);
 
             // Mock stage is EOF so fetch should be too.
             ASSERT_TRUE(fetchStage->isEOF());
diff --git a/src/mongo/dbtests/query_stage_sort.cpp b/src/mongo/dbtests/query_stage_sort.cpp
index 23842445255..2d432215612 100644
--- a/src/mongo/dbtests/query_stage_sort.cpp
+++ b/src/mongo/dbtests/query_stage_sort.cpp
@@ -294,7 +294,7 @@ namespace QueryStageSortTests {
             }
             ss->recoverFromYield();
 
-            // After invalidating all our data, we have nothing left to sort.
+            // Invalidation of data in the sort stage fetches it but passes it through.
             int count = 0;
             while (!ss->isEOF()) {
                 WorkingSetID id;
@@ -306,8 +306,8 @@ namespace QueryStageSortTests {
                 ++count;
             }
 
-            // Therefore, we expect an empty result set from running the sort stage to completion.
-            ASSERT_EQUALS(0, count);
+            // Returns all docs.
+            ASSERT_EQUALS(limit() ? limit() : numObj(), count);
         }
     };