1 files changed, 32 insertions, 3 deletions
diff --git a/src/mongo/db/pipeline/window_function/window_function_exec_non_removable_test.cpp b/src/mongo/db/pipeline/window_function/window_function_exec_non_removable_test.cpp
index d8042421f95..2263d96ff1c 100644
--- a/src/mongo/db/pipeline/window_function/window_function_exec_non_removable_test.cpp
+++ b/src/mongo/db/pipeline/window_function/window_function_exec_non_removable_test.cpp
@@ -31,6 +31,7 @@
 
 #include "mongo/db/exec/document_value/document_value_test_util.h"
 #include "mongo/db/pipeline/accumulator.h"
+#include "mongo/db/pipeline/accumulator_for_window_functions.h"
 #include "mongo/db/pipeline/aggregation_context_fixture.h"
 #include "mongo/db/pipeline/document_source.h"
 #include "mongo/db/pipeline/document_source_mock.h"
@@ -53,14 +54,27 @@ public:
     WindowFunctionExecNonRemovable<AccumulatorState> createForFieldPath(
         std::deque<DocumentSource::GetNextResult> docs,
         const std::string& inputPath,
-        WindowBounds::Bound<int> upper) {
+        WindowBounds::Bound<int> upper,
+        boost::optional<std::string> sortByPath = boost::none) {
         _docSource = DocumentSourceMock::createForTest(std::move(docs), getExpCtx());
         _iter = std::make_unique<PartitionIterator>(
             getExpCtx().get(), _docSource.get(), boost::none, boost::none);
         auto input = ExpressionFieldPath::parse(
             getExpCtx().get(), inputPath, getExpCtx()->variablesParseState);
-        return WindowFunctionExecNonRemovable<AccumulatorState>(
-            _iter.get(), std::move(input), AccumulatorType::create(getExpCtx().get()), upper);
+        if (sortByPath) {
+            auto sortBy = ExpressionFieldPath::parse(
+                getExpCtx().get(), *sortByPath, getExpCtx()->variablesParseState);
+            return WindowFunctionExecNonRemovable<AccumulatorState>(
+                _iter.get(),
+                ExpressionArray::create(
+                    getExpCtx().get(),
+                    std::vector<boost::intrusive_ptr<Expression>>{sortBy, input}),
+                AccumulatorType::create(getExpCtx().get()),
+                upper);
+        } else {
+            return WindowFunctionExecNonRemovable<AccumulatorState>(
+                _iter.get(), std::move(input), AccumulatorType::create(getExpCtx().get()), upper);
+        }
     }
 
     auto advanceIterator() {
@@ -173,5 +187,20 @@ TEST_F(WindowFunctionExecNonRemovableTest, InputExpressionAllowedToCreateVariabl
     ASSERT_VALUE_EQ(Value(std::vector<Value>{Value(2), Value(3)}), exec.getNext());
 }
 
+TEST_F(WindowFunctionExecNonRemovableTest, CanReceiveSortByExpression) {
+    const auto docs = std::deque<DocumentSource::GetNextResult>{
+        Document{{"x", 1}, {"y", 0}}, Document{{"x", 3}, {"y", 2}}, Document{{"x", 5}, {"y", 4}}};
+    auto mgr = createForFieldPath<AccumulatorIntegral>(
+        docs, "$y" /* input */, 0, std::string("$x") /* sortBy */);
+    double expectedIntegral = 0;
+    ASSERT_VALUE_EQ(Value(expectedIntegral), mgr.getNext());
+    advanceIterator();
+    expectedIntegral += 2.0;  // (2 + 0) * (3 - 1) / 2.0 = 2.0
+    ASSERT_VALUE_EQ(Value(expectedIntegral), mgr.getNext());
+    advanceIterator();
+    expectedIntegral += 6.0;  // (4 + 2) * (5 - 3) / 2.0 = 6.0
+    ASSERT_VALUE_EQ(Value(expectedIntegral), mgr.getNext());
+}
+
 }  // namespace
 }  // namespace mongo