1 files changed, 28 insertions, 33 deletions
diff --git a/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp b/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp
index e261c5ac537..3286896f3f6 100644
--- a/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp
+++ b/src/mongo/db/query/optimizer/utils/unit_test_utils.cpp
@@ -254,66 +254,61 @@ std::unique_ptr<CardinalityEstimator> makeHintedCE(ce::PartialSchemaSelHints hin
     return std::make_unique<ce::HintedEstimator>(std::move(hints));
 }
 
+cost_model::CostModelCoefficients getTestCostModel() {
+    return cost_model::CostModelManager::getDefaultCoefficients();
+}
+
 std::unique_ptr<CostEstimator> makeCostEstimator() {
-    return std::make_unique<cost_model::CostEstimatorImpl>(
-        cost_model::CostModelManager::getDefaultCoefficients());
+    return makeCostEstimator(getTestCostModel());
 }
 
-OptPhaseManager makePhaseManager(OptPhaseManager::PhaseSet phaseSet,
-                                 PrefixId& prefixId,
-                                 Metadata metadata,
-                                 DebugInfo debugInfo,
-                                 QueryHints queryHints) {
+std::unique_ptr<CostEstimator> makeCostEstimator(
+    const cost_model::CostModelCoefficients& costModel) {
+    return std::make_unique<cost_model::CostEstimatorImpl>(costModel);
+}
+
+
+OptPhaseManager makePhaseManager(
+    OptPhaseManager::PhaseSet phaseSet,
+    PrefixId& prefixId,
+    Metadata metadata,
+    const boost::optional<cost_model::CostModelCoefficients>& costModel,
+    DebugInfo debugInfo,
+    QueryHints queryHints) {
     return OptPhaseManager{std::move(phaseSet),
                            prefixId,
                            false /*requireRID*/,
                            std::move(metadata),
                            makeHeuristicCE(),  // primary CE
                            makeHeuristicCE(),  // substitution phase CE, same as primary
-                           makeCostEstimator(),
+                           makeCostEstimator(costModel ? *costModel : getTestCostModel()),
                            defaultConvertPathToInterval,
                            ConstEval::constFold,
                            std::move(debugInfo),
                            std::move(queryHints)};
 }
 
-OptPhaseManager makePhaseManager(OptPhaseManager::PhaseSet phaseSet,
-                                 PrefixId& prefixId,
-                                 Metadata metadata,
-                                 std::unique_ptr<CardinalityEstimator> ce,
-                                 DebugInfo debugInfo,
-                                 QueryHints queryHints) {
+OptPhaseManager makePhaseManager(
+    OptPhaseManager::PhaseSet phaseSet,
+    PrefixId& prefixId,
+    Metadata metadata,
+    std::unique_ptr<CardinalityEstimator> ce,
+    const boost::optional<cost_model::CostModelCoefficients>& costModel,
+    DebugInfo debugInfo,
+    QueryHints queryHints) {
     return OptPhaseManager{std::move(phaseSet),
                            prefixId,
                            false /*requireRID*/,
                            std::move(metadata),
                            std::move(ce),      // primary CE
                            makeHeuristicCE(),  // substitution phase CE
-                           makeCostEstimator(),
+                           makeCostEstimator(costModel ? *costModel : getTestCostModel()),
                            defaultConvertPathToInterval,
                            ConstEval::constFold,
                            std::move(debugInfo),
                            std::move(queryHints)};
 }
 
-OptPhaseManager makePhaseManager(OptPhaseManager::PhaseSet phaseSet,
-                                 PrefixId& prefixId,
-                                 Metadata metadata,
-                                 DebugInfo debugInfo,
-                                 mongo::cost_model::CostModelCoefficients coefs,
-                                 QueryHints queryHints) {
-    return OptPhaseManager{std::move(phaseSet),
-                           prefixId,
-                           false /*requireRID*/,
-                           std::move(metadata),
-                           makeHeuristicCE(),  // primary CE
-                           makeHeuristicCE(),  // substitution phase CE, same as primary
-                           std::make_unique<cost_model::CostEstimatorImpl>(coefs),
-                           defaultConvertPathToInterval,
-                           ConstEval::constFold,
-                           std::move(debugInfo),
-                           std::move(queryHints)};
-}
 
 OptPhaseManager makePhaseManagerRequireRID(OptPhaseManager::PhaseSet phaseSet,
                                            PrefixId& prefixId,