SERVER-59853 Simplify $lookup and $unionWith tests to avoid needing manual disposal

3 files changed, 94 insertions, 90 deletions

diff --git a/src/mongo/db/pipeline/aggregation_context_fixture.h b/src/mongo/db/pipeline/aggregation_context_fixture.h
index 76cc01a40c4..e7595382094 100644
--- a/src/mongo/db/pipeline/aggregation_context_fixture.h
+++ b/src/mongo/db/pipeline/aggregation_context_fixture.h
@@ -33,6 +33,7 @@
 #include <memory>
 
 #include "mongo/db/concurrency/locker_noop_client_observer.h"
+#include "mongo/db/pipeline/document_source.h"
 #include "mongo/db/pipeline/expression_context_for_test.h"
 #include "mongo/db/service_context_test_fixture.h"
 #include "mongo/unittest/temp_dir.h"
@@ -76,6 +77,14 @@ private:
     boost::intrusive_ptr<ExpressionContextForTest> _expCtx;
 };
 
+// A custom-deleter which disposes a DocumentSource when it goes out of scope.
+struct DocumentSourceDeleter {
+    void operator()(DocumentSource* docSource) {
+        docSource->dispose();
+        delete docSource;
+    }
+};
+
 class ServerlessAggregationContextFixture : public AggregationContextFixture {
 public:
     ServerlessAggregationContextFixture()
diff --git a/src/mongo/db/pipeline/document_source_lookup_test.cpp b/src/mongo/db/pipeline/document_source_lookup_test.cpp
index 82a0b6bbd61..aae4d7beef5 100644
--- a/src/mongo/db/pipeline/document_source_lookup_test.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup_test.cpp
@@ -82,6 +82,13 @@ public:
     }
 };
 
+auto makeLookUpFromBson(BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& expCtx) {
+    auto docSource = DocumentSourceLookUp::createFromBson(elem, expCtx);
+    auto lookup = static_cast<DocumentSourceLookUp*>(docSource.detach());
+    return std::unique_ptr<DocumentSourceLookUp, DocumentSourceDeleter>(lookup,
+                                                                        DocumentSourceDeleter());
+}
+
 // A 'let' variable defined in a $lookup stage is expected to be available to all sub-pipelines. For
 // sub-pipelines below the immediate one, they are passed to via ExpressionContext. This test
 // confirms that variables defined in the ExpressionContext are captured by the $lookup stage.
@@ -869,9 +876,7 @@ TEST_F(DocumentSourceLookUpTest, ShouldPropagatePauses) {
                                          {"foreignField", "_id"_sd},
                                          {"as", "foreignDocs"_sd}}}}
                           .toBson();
-    auto parsed = DocumentSourceLookUp::createFromBson(lookupSpec.firstElement(), expCtx);
-    auto lookup = static_cast<DocumentSourceLookUp*>(parsed.get());
-
+    auto lookup = makeLookUpFromBson(lookupSpec.firstElement(), expCtx);
     lookup->setSource(mockLocalSource.get());
 
     auto next = lookup->getNext();
@@ -890,7 +895,6 @@ TEST_F(DocumentSourceLookUpTest, ShouldPropagatePauses) {
 
     ASSERT_TRUE(lookup->getNext().isEOF());
     ASSERT_TRUE(lookup->getNext().isEOF());
-    lookup->dispose();
 }
 
 TEST_F(DocumentSourceLookUpTest, ShouldPropagatePausesWhileUnwinding) {
@@ -905,6 +909,14 @@ TEST_F(DocumentSourceLookUpTest, ShouldPropagatePausesWhileUnwinding) {
     expCtx->mongoProcessInterface =
         std::make_shared<MockMongoInterface>(std::move(mockForeignContents));
 
+    // Mock its input, pausing every other result.
+    auto mockLocalSource =
+        DocumentSourceMock::createForTest({Document{{"foreignId", 0}},
+                                           DocumentSource::GetNextResult::makePauseExecution(),
+                                           Document{{"foreignId", 1}},
+                                           DocumentSource::GetNextResult::makePauseExecution()},
+                                          expCtx);
+
     // Set up the $lookup stage.
     auto lookupSpec = Document{{"$lookup",
                                 Document{{"from", fromNs.coll()},
@@ -912,21 +924,13 @@ TEST_F(DocumentSourceLookUpTest, ShouldPropagatePausesWhileUnwinding) {
                                          {"foreignField", "_id"_sd},
                                          {"as", "foreignDoc"_sd}}}}
                           .toBson();
-    auto parsed = DocumentSourceLookUp::createFromBson(lookupSpec.firstElement(), expCtx);
-    auto lookup = static_cast<DocumentSourceLookUp*>(parsed.get());
+    auto lookup = makeLookUpFromBson(lookupSpec.firstElement(), expCtx);
 
     const bool preserveNullAndEmptyArrays = false;
     const boost::optional<std::string> includeArrayIndex = boost::none;
     lookup->setUnwindStage(DocumentSourceUnwind::create(
         expCtx, "foreignDoc", preserveNullAndEmptyArrays, includeArrayIndex));
 
-    // Mock its input, pausing every other result.
-    auto mockLocalSource =
-        DocumentSourceMock::createForTest({Document{{"foreignId", 0}},
-                                           DocumentSource::GetNextResult::makePauseExecution(),
-                                           Document{{"foreignId", 1}},
-                                           DocumentSource::GetNextResult::makePauseExecution()},
-                                          expCtx);
     lookup->setSource(mockLocalSource.get());
 
     auto next = lookup->getNext();
@@ -945,7 +949,6 @@ TEST_F(DocumentSourceLookUpTest, ShouldPropagatePausesWhileUnwinding) {
 
     ASSERT_TRUE(lookup->getNext().isEOF());
     ASSERT_TRUE(lookup->getNext().isEOF());
-    lookup->dispose();
 }
 
 TEST_F(DocumentSourceLookUpTest, LookupReportsAsFieldIsModified) {
@@ -961,14 +964,12 @@ TEST_F(DocumentSourceLookUpTest, LookupReportsAsFieldIsModified) {
                                          {"foreignField", "_id"_sd},
                                          {"as", "foreignDocs"_sd}}}}
                           .toBson();
-    auto parsed = DocumentSourceLookUp::createFromBson(lookupSpec.firstElement(), expCtx);
-    auto lookup = static_cast<DocumentSourceLookUp*>(parsed.get());
+    auto lookup = makeLookUpFromBson(lookupSpec.firstElement(), expCtx);
 
     auto modifiedPaths = lookup->getModifiedPaths();
     ASSERT(modifiedPaths.type == DocumentSource::GetModPathsReturn::Type::kFiniteSet);
     ASSERT_EQ(1U, modifiedPaths.paths.size());
     ASSERT_EQ(1U, modifiedPaths.paths.count("foreignDocs"));
-    lookup->dispose();
 }
 
 TEST_F(DocumentSourceLookUpTest, LookupReportsFieldsModifiedByAbsorbedUnwind) {
@@ -984,8 +985,7 @@ TEST_F(DocumentSourceLookUpTest, LookupReportsFieldsModifiedByAbsorbedUnwind) {
                                          {"foreignField", "_id"_sd},
                                          {"as", "foreignDoc"_sd}}}}
                           .toBson();
-    auto parsed = DocumentSourceLookUp::createFromBson(lookupSpec.firstElement(), expCtx);
-    auto lookup = static_cast<DocumentSourceLookUp*>(parsed.get());
+    auto lookup = makeLookUpFromBson(lookupSpec.firstElement(), expCtx);
 
     const bool preserveNullAndEmptyArrays = false;
     const boost::optional<std::string> includeArrayIndex = std::string("arrIndex");
@@ -997,7 +997,6 @@ TEST_F(DocumentSourceLookUpTest, LookupReportsFieldsModifiedByAbsorbedUnwind) {
     ASSERT_EQ(2U, modifiedPaths.paths.size());
     ASSERT_EQ(1U, modifiedPaths.paths.count("foreignDoc"));
     ASSERT_EQ(1U, modifiedPaths.paths.count("arrIndex"));
-    lookup->dispose();
 }
 
 BSONObj sequentialCacheStageObj(const StringData status = "kBuilding"_sd,
diff --git a/src/mongo/db/pipeline/document_source_union_with_test.cpp b/src/mongo/db/pipeline/document_source_union_with_test.cpp
index 04f440fa91a..05e0feb7baa 100644
--- a/src/mongo/db/pipeline/document_source_union_with_test.cpp
+++ b/src/mongo/db/pipeline/document_source_union_with_test.cpp
@@ -60,6 +60,19 @@ using MockMongoInterface = StubLookupSingleDocumentProcessInterface;
 // This provides access to getExpCtx(), but we'll use a different name for this test suite.
 using DocumentSourceUnionWithTest = AggregationContextFixture;
 
+auto makeUnion(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+               std::unique_ptr<Pipeline, PipelineDeleter> pipeline) {
+    return std::unique_ptr<DocumentSourceUnionWith, DocumentSourceDeleter>(
+        new DocumentSourceUnionWith(expCtx, std::move(pipeline)), DocumentSourceDeleter());
+}
+
+auto makeUnionFromBson(BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& expCtx) {
+    auto docSource = DocumentSourceUnionWith::createFromBson(elem, expCtx);
+    auto unionWith = static_cast<DocumentSourceUnionWith*>(docSource.detach());
+    return std::unique_ptr<DocumentSourceUnionWith, DocumentSourceDeleter>(unionWith,
+                                                                           DocumentSourceDeleter());
+}
+
 TEST_F(DocumentSourceUnionWithTest, BasicSerialUnions) {
     const auto docs = std::array{Document{{"a", 1}}, Document{{"b", 1}}, Document{{"c", 1}}};
     const auto mock = DocumentSourceMock::createForTest(docs[0], getExpCtx());
@@ -69,19 +82,19 @@ TEST_F(DocumentSourceUnionWithTest, BasicSerialUnions) {
     mockCtxOne->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeOne);
     const auto mockCtxTwo = getExpCtx()->copyWith({});
     mockCtxTwo->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeTwo);
-    auto unionWithOne = DocumentSourceUnionWith(
-        mockCtxOne,
-        Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
-    auto unionWithTwo = DocumentSourceUnionWith(
-        mockCtxTwo,
-        Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
-    unionWithOne.setSource(mock.get());
-    unionWithTwo.setSource(&unionWithOne);
+    auto unionWithOne =
+        makeUnion(mockCtxOne,
+                  Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
+    auto unionWithTwo =
+        makeUnion(mockCtxTwo,
+                  Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
+    unionWithOne->setSource(mock.get());
+    unionWithTwo->setSource(unionWithOne.get());
 
     auto comparator = DocumentComparator();
     auto results = comparator.makeUnorderedDocumentSet();
     for (auto& doc [[maybe_unused]] : docs) {
-        auto next = unionWithTwo.getNext();
+        auto next = unionWithTwo->getNext();
         ASSERT_TRUE(next.isAdvanced());
         const auto [ignored, inserted] = results.insert(next.releaseDocument());
         ASSERT_TRUE(inserted);
@@ -89,12 +102,9 @@ TEST_F(DocumentSourceUnionWithTest, BasicSerialUnions) {
     for (const auto& doc : docs)
         ASSERT_TRUE(results.find(doc) != results.end());
 
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-
-    unionWithOne.dispose();
-    unionWithTwo.dispose();
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
 }
 
 TEST_F(DocumentSourceUnionWithTest, BasicNestedUnions) {
@@ -109,16 +119,16 @@ TEST_F(DocumentSourceUnionWithTest, BasicNestedUnions) {
     auto unionWithOne = make_intrusive<DocumentSourceUnionWith>(
         mockCtxOne,
         Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
-    auto unionWithTwo = DocumentSourceUnionWith(
-        mockCtxTwo,
-        Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{unionWithOne},
-                         getExpCtx()));
-    unionWithTwo.setSource(mock.get());
+    auto unionWithTwo =
+        makeUnion(mockCtxTwo,
+                  Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{unionWithOne},
+                                   getExpCtx()));
+    unionWithTwo->setSource(mock.get());
 
     auto comparator = DocumentComparator();
     auto results = comparator.makeUnorderedDocumentSet();
     for (auto& doc [[maybe_unused]] : docs) {
-        auto next = unionWithTwo.getNext();
+        auto next = unionWithTwo->getNext();
         ASSERT_TRUE(next.isAdvanced());
         const auto [ignored, inserted] = results.insert(next.releaseDocument());
         ASSERT_TRUE(inserted);
@@ -126,11 +136,9 @@ TEST_F(DocumentSourceUnionWithTest, BasicNestedUnions) {
     for (const auto& doc : docs)
         ASSERT_TRUE(results.find(doc) != results.end());
 
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-
-    unionWithTwo.dispose();
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
 }
 
 TEST_F(DocumentSourceUnionWithTest, UnionsWithNonEmptySubPipelines) {
@@ -145,19 +153,19 @@ TEST_F(DocumentSourceUnionWithTest, UnionsWithNonEmptySubPipelines) {
     mockCtxTwo->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeTwo);
     const auto filter = DocumentSourceMatch::create(BSON("d" << 1), mockCtxOne);
     const auto proj = DocumentSourceAddFields::create(BSON("d" << 1), mockCtxTwo);
-    auto unionWithOne = DocumentSourceUnionWith(
+    auto unionWithOne = makeUnion(
         mockCtxOne,
         Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{filter}, getExpCtx()));
-    auto unionWithTwo = DocumentSourceUnionWith(
+    auto unionWithTwo = makeUnion(
         mockCtxTwo,
         Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{proj}, getExpCtx()));
-    unionWithOne.setSource(mock.get());
-    unionWithTwo.setSource(&unionWithOne);
+    unionWithOne->setSource(mock.get());
+    unionWithTwo->setSource(unionWithOne.get());
 
     auto comparator = DocumentComparator();
     auto results = comparator.makeUnorderedDocumentSet();
     for (auto& doc [[maybe_unused]] : outputDocs) {
-        auto next = unionWithTwo.getNext();
+        auto next = unionWithTwo->getNext();
         ASSERT_TRUE(next.isAdvanced());
         const auto [ignored, inserted] = results.insert(next.releaseDocument());
         ASSERT_TRUE(inserted);
@@ -165,12 +173,9 @@ TEST_F(DocumentSourceUnionWithTest, UnionsWithNonEmptySubPipelines) {
     for (const auto& doc : outputDocs)
         ASSERT_TRUE(results.find(doc) != results.end());
 
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-
-    unionWithOne.dispose();
-    unionWithTwo.dispose();
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
 }
 
 TEST_F(DocumentSourceUnionWithTest, SerializeAndParseWithPipeline) {
@@ -315,26 +320,23 @@ TEST_F(DocumentSourceUnionWithTest, PropagatePauses) {
     mockCtxOne->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeOne);
     const auto mockCtxTwo = getExpCtx()->copyWith({});
     mockCtxTwo->mongoProcessInterface = std::make_unique<MockMongoInterface>(mockDequeTwo);
-    auto unionWithOne = DocumentSourceUnionWith(
-        mockCtxOne,
-        Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
-    auto unionWithTwo = DocumentSourceUnionWith(
-        mockCtxTwo,
-        Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
-    unionWithOne.setSource(mock.get());
-    unionWithTwo.setSource(&unionWithOne);
-
-    ASSERT_TRUE(unionWithTwo.getNext().isAdvanced());
-    ASSERT_TRUE(unionWithTwo.getNext().isPaused());
-    ASSERT_TRUE(unionWithTwo.getNext().isAdvanced());
-    ASSERT_TRUE(unionWithTwo.getNext().isPaused());
-
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-    ASSERT_TRUE(unionWithTwo.getNext().isEOF());
-
-    unionWithOne.dispose();
-    unionWithTwo.dispose();
+    auto unionWithOne =
+        makeUnion(mockCtxOne,
+                  Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
+    auto unionWithTwo =
+        makeUnion(mockCtxTwo,
+                  Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
+    unionWithOne->setSource(mock.get());
+    unionWithTwo->setSource(unionWithOne.get());
+
+    ASSERT_TRUE(unionWithTwo->getNext().isAdvanced());
+    ASSERT_TRUE(unionWithTwo->getNext().isPaused());
+    ASSERT_TRUE(unionWithTwo->getNext().isAdvanced());
+    ASSERT_TRUE(unionWithTwo->getNext().isPaused());
+
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
+    ASSERT_TRUE(unionWithTwo->getNext().isEOF());
 }
 
 TEST_F(DocumentSourceUnionWithTest, ReturnEOFAfterBeingDisposed) {
@@ -406,10 +408,10 @@ TEST_F(DocumentSourceUnionWithTest, RespectsViewDefinition) {
     expCtx->mongoProcessInterface =
         std::make_shared<MockMongoInterface>(std::move(mockForeignContents));
 
-    auto bson = BSON("$unionWith" << nsToUnionWith.coll());
-    auto unionWith = DocumentSourceUnionWith::createFromBson(bson.firstElement(), expCtx);
     const auto localMock =
         DocumentSourceMock::createForTest({Document{{"_id"_sd, "local"_sd}}}, getExpCtx());
+    auto bson = BSON("$unionWith" << nsToUnionWith.coll());
+    auto unionWith = makeUnionFromBson(bson.firstElement(), expCtx);
     unionWith->setSource(localMock.get());
 
     auto result = unionWith->getNext();
@@ -421,8 +423,6 @@ TEST_F(DocumentSourceUnionWithTest, RespectsViewDefinition) {
     ASSERT_DOCUMENT_EQ(result.getDocument(), (Document{{"_id"_sd, 2}}));
 
     ASSERT_TRUE(unionWith->getNext().isEOF());
-
-    unionWith->dispose();
 }
 
 TEST_F(DocumentSourceUnionWithTest, ConcatenatesViewDefinitionToPipeline) {
@@ -445,7 +445,7 @@ TEST_F(DocumentSourceUnionWithTest, ConcatenatesViewDefinitionToPipeline) {
                          "coll" << viewNsToUnionWith.coll() << "pipeline"
                                 << BSON_ARRAY(fromjson(
                                        "{$set: {originalId: '$_id', _id: {$add: [1, '$_id']}}}"))));
-    auto unionWith = DocumentSourceUnionWith::createFromBson(bson.firstElement(), expCtx);
+    auto unionWith = makeUnionFromBson(bson.firstElement(), expCtx);
     unionWith->setSource(localMock.get());
 
     auto result = unionWith->getNext();
@@ -459,8 +459,6 @@ TEST_F(DocumentSourceUnionWithTest, ConcatenatesViewDefinitionToPipeline) {
     ASSERT_DOCUMENT_EQ(result.getDocument(), (Document{{"_id"_sd, 3}, {"originalId"_sd, 2}}));
 
     ASSERT_TRUE(unionWith->getNext().isEOF());
-
-    unionWith->dispose();
 }
 
 TEST_F(DocumentSourceUnionWithTest, RejectUnionWhenDepthLimitIsExceeded) {
@@ -482,9 +480,9 @@ TEST_F(DocumentSourceUnionWithTest, RejectUnionWhenDepthLimitIsExceeded) {
 }
 
 TEST_F(DocumentSourceUnionWithTest, ConstraintsWithoutPipelineAreCorrect) {
-    auto emptyUnion = DocumentSourceUnionWith(
-        getExpCtx(),
-        Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
+    auto emptyUnion =
+        makeUnion(getExpCtx(),
+                  Pipeline::create(std::list<boost::intrusive_ptr<DocumentSource>>{}, getExpCtx()));
     StageConstraints defaultConstraints(StageConstraints::StreamType::kStreaming,
                                         StageConstraints::PositionRequirement::kNone,
                                         StageConstraints::HostTypeRequirement::kAnyShard,
@@ -493,9 +491,7 @@ TEST_F(DocumentSourceUnionWithTest, ConstraintsWithoutPipelineAreCorrect) {
                                         StageConstraints::TransactionRequirement::kNotAllowed,
                                         StageConstraints::LookupRequirement::kAllowed,
                                         StageConstraints::UnionRequirement::kAllowed);
-    ASSERT_TRUE(emptyUnion.constraints(Pipeline::SplitState::kUnsplit) == defaultConstraints);
-
-    emptyUnion.dispose();
+    ASSERT_TRUE(emptyUnion->constraints(Pipeline::SplitState::kUnsplit) == defaultConstraints);
 }
 
 TEST_F(DocumentSourceUnionWithTest, ConstraintsWithMixedSubPipelineAreCorrect) {