From 8c70de3b58e9f9d0605e93255ce962d1221671df Mon Sep 17 00:00:00 2001
From: Rhys Howell <rhys.howell@mongodb.com>
Date: Tue, 25 Jun 2024 13:05:50 -0400
Subject: [PATCH] fix(aggregations): use maxTimeMS default on preview docs
 COMPASS-7798

---
 .../src/modules/input-documents.spec.ts       | 59 ++++++++++++++++++-
 .../src/modules/input-documents.ts            |  8 ++-
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/packages/compass-aggregations/src/modules/input-documents.spec.ts b/packages/compass-aggregations/src/modules/input-documents.spec.ts
index faf305c8842..d59862c0fbf 100644
--- a/packages/compass-aggregations/src/modules/input-documents.spec.ts
+++ b/packages/compass-aggregations/src/modules/input-documents.spec.ts
@@ -1,12 +1,22 @@
+import { expect } from 'chai';
+import { defaultPreferencesInstance } from 'compass-preferences-model';
+import sinon from 'sinon';
+
 import reducer, {
   toggleInputDocumentsCollapsed,
   updateInputDocuments,
   loadingInputDocuments,
+  refreshInputDocuments,
   ActionTypes,
 } from './input-documents';
-import { expect } from 'chai';
+import type { RootState } from '.';
+import type { DataService } from './data-service';
 
 describe('input documents module', function () {
+  afterEach(function () {
+    sinon.restore();
+  });
+
   describe('#toggleInputDocumentsCollapsed', function () {
     it('returns the ActionTypes.CollapseToggled action', function () {
       expect(toggleInputDocumentsCollapsed()).to.deep.equal({
@@ -33,6 +43,53 @@ describe('input documents module', function () {
     });
   });
 
+  describe('#refreshInputDocuments', function () {
+    it('should apply maxTimeMS to the aggregation when it is set', async function () {
+      const refreshInputDocumentsThunk = refreshInputDocuments();
+
+      const mockAggregate = sinon.stub().resolves([]);
+      const mockState: Partial<RootState> = {
+        dataService: {
+          dataService: {
+            aggregate: mockAggregate,
+          } as unknown as DataService,
+        },
+        namespace: 'test.namespace',
+        maxTimeMS: undefined,
+        settings: {
+          isExpanded: false,
+          isCommentMode: false,
+          isDirty: false,
+          limit: 10,
+          sampleSize: 10,
+        },
+      };
+
+      await refreshInputDocumentsThunk(
+        sinon.stub(),
+        () => mockState as RootState,
+        { preferences: defaultPreferencesInstance } as any
+      );
+
+      expect(mockAggregate.calledOnce).to.be.true;
+      expect(mockAggregate.firstCall.args[2]).to.deep.equal({
+        maxTimeMS: 60_000,
+      });
+
+      mockState.maxTimeMS = 1000;
+      await refreshInputDocumentsThunk(
+        sinon.stub(),
+        () => mockState as RootState,
+        { preferences: defaultPreferencesInstance } as any
+      );
+
+      expect(mockAggregate.calledTwice).to.be.true;
+      expect(mockAggregate.secondCall.args[2]).to.deep.equal({
+        maxTimeMS: 1000,
+      });
+    });
+  });
+
   describe('#reducer', function () {
     context(
       'when the action is not toggle input documents collapsed',
diff --git a/packages/compass-aggregations/src/modules/input-documents.ts b/packages/compass-aggregations/src/modules/input-documents.ts
index 592300c9085..7876e985702 100644
--- a/packages/compass-aggregations/src/modules/input-documents.ts
+++ b/packages/compass-aggregations/src/modules/input-documents.ts
@@ -3,6 +3,7 @@ import { capMaxTimeMSAtPreferenceLimit } from 'compass-preferences-model/provide
 import type { PipelineBuilderThunkAction } from '.';
 import type { AnyAction } from 'redux';
 import { isAction } from '../utils/is-action';
+import { DEFAULT_MAX_TIME_MS } from '../constants';
 
 export enum ActionTypes {
   CollapseToggled = 'aggregations/input-documents/CollapseToggled',
@@ -111,9 +112,10 @@ export const refreshInputDocuments = (): PipelineBuilderThunkAction<
     }
 
     const options = {
-      maxTimeMS: capMaxTimeMSAtPreferenceLimit(preferences, maxTimeMS) as
-        | number
-        | undefined,
+      maxTimeMS: capMaxTimeMSAtPreferenceLimit(
+        preferences,
+        maxTimeMS ?? DEFAULT_MAX_TIME_MS
+      ),
     };
 
     const aggregateOptions = { ...options };