1 files changed, 5 insertions, 1 deletions
diff --git a/jstests/aggregation/sources/setWindowFields/memory_limit.js b/jstests/aggregation/sources/setWindowFields/memory_limit.js
index e64781d3eca..a88717d0331 100644
--- a/jstests/aggregation/sources/setWindowFields/memory_limit.js
+++ b/jstests/aggregation/sources/setWindowFields/memory_limit.js
@@ -61,6 +61,10 @@ assert.commandWorked(coll.runCommand({
     allowDiskUse: false
 }));
 
+setParameterOnAllHosts(nonConfigNodes,
+                       "internalDocumentSourceSetWindowFieldsMaxMemoryBytes",
+                       (perDocSize * docsPerPartition) + 1024);
+
 // Test that the query fails with a window function that stores documents.
 assert.commandFailedWithCode(coll.runCommand({
     aggregate: coll.getName(),
@@ -74,7 +78,7 @@ assert.commandFailedWithCode(coll.runCommand({
     cursor: {},
     allowDiskUse: false
 }),
-                             5414201);
+                             [5643011, 5414201]);
 // Reset limit for other tests.
 setParameterOnAllHosts(
     nonConfigNodes, "internalDocumentSourceSetWindowFieldsMaxMemoryBytes", 100 * 1024 * 1024);