test and reducer bug fixes

Author: ruchitharajaghatta

packages/compass-generative-ai/src/store/atlas-optin-reducer.spec.ts

@@ -19,20 +19,23 @@ describe('atlasOptInReducer', function () {
 
   beforeEach(async function () {
     mockPreferences = await createSandboxFromDefaultPreferences();
+    await mockPreferences.savePreferences({
+      optInDataExplorerGenAIFeatures: false,
+    });
   });
 
   afterEach(function () {
     sandbox.reset();
   });
 
   describe('optIn', function () {
-    it('should check authenticated state and set state to success if already authenticated', async function () {
-      const mockAtlasService = {
-        optIn: sandbox.stub().resolves({ sub: '1234' }),
+    it('should check state and set state to success if already opted in', async function () {
+      const mockAtlasAiService = {
+        optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }),
       };
       const store = configureStore({
-        atlasAuthService: mockAtlasService as any,
-        atlasAiService: mockAtlasService as any,
+        atlasAuthService: {} as any,
+        atlasAiService: mockAtlasAiService as any,
         preferences: mockPreferences,
       });
 
@@ -41,22 +44,22 @@ describe('atlasOptInReducer', function () {
         'initial'
       );
       void store.dispatch(atlasAiServiceOptedIn());
-      console.log(store.getState());
       await store.dispatch(optIn());
-      expect(mockAtlasService.optIn).not.to.have.been.called;
+      expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).not.to.have.been
+        .called;
       expect(store.getState().optIn).to.have.nested.property(
         'state',
         'optin-success'
       );
     });
 
     it('should start opt in, and set state to success', async function () {
-      const mockAtlasService = {
-        optIn: sandbox.stub().resolves({ sub: '1234' }),
+      const mockAtlasAiService = {
+        optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }),
       };
       const store = configureStore({
-        atlasAuthService: mockAtlasService as any,
-        atlasAiService: mockAtlasService as any,
+        atlasAuthService: {} as any,
+        atlasAiService: mockAtlasAiService as any,
         preferences: mockPreferences,
       });
 
@@ -66,20 +69,23 @@ describe('atlasOptInReducer', function () {
       );
       void store.dispatch(optIntoGenAIWithModalPrompt()).catch(() => {});
       await store.dispatch(optIn());
-      expect(mockAtlasService.optIn).to.have.been.calledOnce;
+      expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).to.have.been
+        .calledOnce;
       expect(store.getState().optIn).to.have.nested.property(
         'state',
         'optin-success'
       );
     });
 
     it('should fail opt in if opt in failed', async function () {
-      const mockAtlasService = {
-        optIn: sandbox.stub().rejects(new Error('Pineapples!')),
+      const mockAtlasAiService = {
+        optIntoGenAIFeaturesAtlas: sandbox
+          .stub()
+          .rejects(new Error('Whooops!')),
       };
       const store = configureStore({
-        atlasAuthService: mockAtlasService as any,
-        atlasAiService: mockAtlasService as any,
+        atlasAuthService: {} as any,
+        atlasAiService: mockAtlasAiService as any,
         preferences: mockPreferences,
       });
 
@@ -88,7 +94,8 @@ describe('atlasOptInReducer', function () {
       // Avoid unhandled rejections.
       AttemptStateMap.get(attemptId)?.promise.catch(() => {});
       await optInPromise;
-      expect(mockAtlasService.optIn).to.have.been.calledOnce;
+      expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).to.have.been
+        .calledOnce;
       expect(store.getState().optIn).to.have.nested.property('state', 'error');
     });
   });
@@ -112,8 +119,8 @@ describe('atlasOptInReducer', function () {
     });
 
     it('should cancel opt in if opt in is in progress', async function () {
-      const mockAtlasService = {
-        optIn: sandbox
+      const mockAtlasAiService = {
+        optIntoGenAIFeaturesAtlas: sandbox
           .stub()
           .callsFake(({ signal }: { signal: AbortSignal }) => {
             return new Promise((resolve, reject) => {
@@ -123,9 +130,10 @@ describe('atlasOptInReducer', function () {
             });
           }),
       };
+
       const store = configureStore({
-        atlasAuthService: mockAtlasService as any,
-        atlasAiService: mockAtlasService as any,
+        atlasAuthService: {} as any,
+        atlasAiService: mockAtlasAiService as any,
         preferences: mockPreferences,
       });
 
@@ -144,12 +152,12 @@ describe('atlasOptInReducer', function () {
 
   describe('optIntoAtlasWithModalPrompt', function () {
     it('should resolve when user finishes opt in with prompt flow', async function () {
-      const mockAtlasService = {
-        optIn: sandbox.stub().resolves({ sub: '1234' }),
+      const mockAtlasAiService = {
+        optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }),
       };
       const store = configureStore({
-        atlasAuthService: mockAtlasService as any,
-        atlasAiService: mockAtlasService as any,
+        atlasAuthService: {} as any,
+        atlasAiService: mockAtlasAiService as any,
         preferences: mockPreferences,
       });
 
@@ -161,12 +169,12 @@ describe('atlasOptInReducer', function () {
     });
 
     it('should reject if opt in flow fails', async function () {
-      const mockAtlasService = {
-        optIn: sandbox.stub().rejects(new Error('Whoops!')),
+      const mockAtlasAiService = {
+        optIntoGenAIFeaturesAtlas: sandbox.stub().rejects(new Error('Whoops!')),
       };
       const store = configureStore({
-        atlasAuthService: mockAtlasService as any,
-        atlasAiService: mockAtlasService as any,
+        atlasAuthService: {} as any,
+        atlasAiService: mockAtlasAiService as any,
         preferences: mockPreferences,
       });
 
@@ -184,12 +192,12 @@ describe('atlasOptInReducer', function () {
     });
 
     it('should reject if user dismissed the modal', async function () {
-      const mockAtlasService = {
-        optIn: sandbox.stub().resolves({ sub: '1234' }),
+      const mockAtlasAiService = {
+        optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }),
       };
       const store = configureStore({
-        atlasAuthService: mockAtlasService as any,
-        atlasAiService: mockAtlasService as any,
+        atlasAuthService: {} as any,
+        atlasAiService: mockAtlasAiService as any,
         preferences: mockPreferences,
       });
 
@@ -207,12 +215,12 @@ describe('atlasOptInReducer', function () {
     });
 
     it('should reject if provided signal was aborted', async function () {
-      const mockAtlasService = {
-        optIn: sandbox.stub().resolves({ sub: '1234' }),
+      const mockAtlasAiService = {
+        optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }),
       };
       const store = configureStore({
-        atlasAuthService: mockAtlasService as any,
-        atlasAiService: mockAtlasService as any,
+        atlasAuthService: {} as any,
+        atlasAiService: mockAtlasAiService as any,
         preferences: mockPreferences,
       });
 

packages/compass-generative-ai/src/store/atlas-optin-reducer.ts

@@ -199,7 +199,7 @@ const startAttempt = (
   fn: () => void
 ): GenAIAtlasOptInThunkAction<AttemptState> => {
   return (dispatch, getState) => {
-    if (getState().attemptId) {
+    if (getState().optIn.attemptId) {
       throw new Error(
         "Can't start opt in with prompt while another opt in attempt is in progress"
       );
@@ -230,7 +230,7 @@ export const optIntoGenAIWithModalPrompt = ({
 > => {
   return (dispatch, getState, { preferences }) => {
     // Nothing to do if we already opted in.
-    const { state } = getState();
+    const { state } = getState().optIn;
     if (
       state === 'optin-success' ||
       preferences.getPreferences().optInDataExplorerGenAIFeatures
@@ -251,23 +251,25 @@ export const optIntoGenAIWithModalPrompt = ({
 
 export const optIn = (): GenAIAtlasOptInThunkAction<Promise<void>> => {
   return async (dispatch, getState, { atlasAiService }) => {
-    const { attemptId } = getState();
+    if (['in-progress', 'optin-success'].includes(getState().optIn.state)) {
+      return;
+    }
+    const { attemptId } = getState().optIn;
     if (attemptId === null) {
       return;
     }
     const {
       controller: { signal },
       resolve,
       reject,
-    } = getAttempt(getState().attemptId);
+    } = getAttempt(getState().optIn.attemptId);
     dispatch({
       type: AtlasOptInActions.Start,
     });
 
     try {
       throwIfAborted(signal);
-      console.log(atlasAiService);
-      await atlasAiService.optIn.optIntoGenAIFeaturesAtlas();
+      await atlasAiService.optIntoGenAIFeaturesAtlas();
       dispatch(atlasAiServiceOptedIn());
       resolve();
     } catch (err) {
@@ -300,10 +302,10 @@ export const cancelOptIn = (reason?: any): GenAIAtlasOptInThunkAction<void> => {
   return (dispatch, getState) => {
     // Can't cancel opt in after the flow was finished indicated by current
     // attempt id being set to null.
-    if (getState().attemptId === null) {
+    if (getState().optIn.attemptId === null) {
       return;
     }
-    const attempt = getAttempt(getState().attemptId);
+    const attempt = getAttempt(getState().optIn.attemptId);
     attempt.controller.abort();
     attempt.reject(reason ?? attempt.controller.signal.reason);
     dispatch({ type: AtlasOptInActions.Cancel });

packages/compass-generative-ai/src/store/atlas-signin-reducer.spec.ts

@@ -227,7 +227,6 @@ describe('atlasSignInReducer', function () {
       } catch (err) {
         expect(err).to.have.property('message', 'Aborted from outside');
       }
-      expect(store.getState()).to.include('signIn');
       expect(store.getState().signIn).to.have.property('state', 'canceled');
     });
   });

packages/compass-generative-ai/src/store/atlas-signin-reducer.ts

@@ -227,7 +227,7 @@ const startAttempt = (
   fn: () => void
 ): GenAIAtlasSignInThunkAction<AttemptState> => {
   return (dispatch, getState) => {
-    if (getState().attemptId) {
+    if (getState().signIn.attemptId) {
       throw new Error(
         "Can't start sign in with prompt while another sign in attempt is in progress"
       );
@@ -257,7 +257,7 @@ export const signIntoAtlasWithModalPrompt = ({
 > => {
   return (dispatch, getState) => {
     // Nothing to do if we already signed in.
-    const { state } = getState();
+    const { state } = getState().signIn;
     if (state === 'success') {
       return Promise.resolve();
     }
@@ -275,18 +275,18 @@ export const signIntoAtlasWithModalPrompt = ({
 
 export const signIn = (): GenAIAtlasSignInThunkAction<Promise<void>> => {
   return async (dispatch, getState, { atlasAuthService }) => {
-    if (['in-progress', 'authenticated'].includes(getState().state)) {
+    if (['in-progress', 'authenticated'].includes(getState().signIn.state)) {
       return;
     }
-    const { attemptId } = getState();
+    const { attemptId } = getState().signIn;
     if (attemptId === null) {
       return;
     }
     const {
       controller: { signal },
       resolve,
       reject,
-    } = getAttempt(getState().attemptId);
+    } = getAttempt(getState().signIn.attemptId);
     dispatch({
       type: AtlasSignInActions.Start,
     });
@@ -332,10 +332,10 @@ export const cancelSignIn = (
   return (dispatch, getState) => {
     // Can't cancel sign in after the flow was finished indicated by current
     // attempt id being set to null.
-    if (getState().attemptId === null) {
+    if (getState().signIn.attemptId === null) {
       return;
     }
-    const attempt = getAttempt(getState().attemptId);
+    const attempt = getAttempt(getState().signIn.attemptId);
     attempt.controller.abort();
     attempt.reject(reason ?? attempt.controller.signal.reason);
     dispatch({ type: AtlasSignInActions.Cancel });