fix scenarios

Author: paula-stacho

packages/compass-global-writes/src/components/states/incomplete-sharding-setup.spec.tsx

@@ -10,7 +10,7 @@ import { renderWithStore } from '../../../tests/create-store';
 import { type ConnectionInfo } from '@mongodb-js/compass-connections/provider';
 import { type ShardZoneData } from '../../store/reducer';
 
-describe.only('IncompleteShardingSetup', function () {
+describe('IncompleteShardingSetup', function () {
   const shardZones: ShardZoneData[] = [
     {
       zoneId: '45893084',

packages/compass-global-writes/src/components/states/incomplete-sharding-setup.tsx

@@ -27,7 +27,7 @@ const containerStyles = css({
   marginBottom: spacing[400],
 });
 
-const unmanageBtnStyles = css({
+const manageBtnStyles = css({
   marginTop: spacing[100],
 });
 
@@ -68,7 +68,7 @@ export function IncompleteShardingSetup({
             onClick={onResume}
             variant={ButtonVariant.Default}
             isLoading={isSubmittingForSharding}
-            className={unmanageBtnStyles}
+            className={manageBtnStyles}
           >
             Enable Global Writes
           </Button>

packages/compass-global-writes/src/components/states/shard-key-correct.tsx

@@ -18,11 +18,8 @@ import {
 } from '../../store/reducer';
 import { ShardZonesTable } from '../shard-zones-table';
 import ShardKeyMarkup from '../shard-key-markup';
-import {
-  containerStyles,
-  bannerStyles,
-} from '../common-styles';
-import ExampleCommandsMarkup from '../example-commands-markup'
+import { containerStyles, bannerStyles } from '../common-styles';
+import ExampleCommandsMarkup from '../example-commands-markup';
 
 const nbsp = '\u00a0';
 

packages/compass-global-writes/src/store/index.spec.ts

@@ -6,6 +6,7 @@ import {
   type CreateShardKeyData,
   unmanageNamespace,
   cancelSharding,
+  resumeManagedNamespace,
   POLLING_INTERVAL,
   type ShardKey,
 } from './reducer';
@@ -179,9 +180,11 @@ describe('GlobalWritesStore Store', function () {
   context('scenarios', function () {
     it('not managed -> sharding -> valid shard key', async function () {
       let mockShardKey = false;
+      let mockManagedNamespace = false;
       // initial state === unsharded
       const store = createStore({
         hasShardKey: Sinon.fake(() => mockShardKey),
+        isNamespaceManaged: Sinon.fake(() => mockManagedNamespace),
       });
       await waitFor(() => {
         expect(store.getState().status).to.equal('UNSHARDED');
@@ -194,6 +197,7 @@ describe('GlobalWritesStore Store', function () {
       });
       const promise = store.dispatch(createShardKey(shardKeyData));
       expect(store.getState().status).to.equal('SUBMITTING_FOR_SHARDING');
+      mockManagedNamespace = true;
       await promise;
       expect(store.getState().status).to.equal('SHARDING');
 
@@ -309,6 +313,34 @@ describe('GlobalWritesStore Store', function () {
       });
     });
 
+    it('incomplete setup -> shard key correct', async function () {
+      // initial state -> incomplete shardingSetup
+      clock = sinon.useFakeTimers({
+        shouldAdvanceTime: true,
+      });
+      let mockManagedNamespace = false;
+      const store = createStore({
+        isNamespaceManaged: Sinon.fake(() => mockManagedNamespace),
+        hasShardKey: () => true,
+      });
+      await waitFor(() => {
+        expect(store.getState().status).to.equal('INCOMPLETE_SHARDING_SETUP');
+        expect(store.getState().managedNamespace).to.be.undefined;
+      });
+
+      // user asks to resume geosharding
+      const promise = store.dispatch(resumeManagedNamespace());
+      mockManagedNamespace = true;
+      expect(store.getState().status).to.equal(
+        'SUBMITTING_FOR_SHARDING_INCOMPLETE'
+      );
+      await promise;
+      clock.tick(POLLING_INTERVAL);
+      await waitFor(() => {
+        expect(store.getState().status).to.equal('SHARD_KEY_CORRECT');
+      });
+    });
+
     it('valid shard key -> not managed', async function () {
       // initial state === shard key correct
       const store = createStore({

packages/compass-global-writes/src/store/reducer.ts

@@ -321,6 +321,25 @@ const reducer: Reducer<RootState, Action> = (state = initialState, action) => {
     };
   }
 
+  if (
+    isAction<NamespaceShardKeyFetchedAction>(
+      action,
+      GlobalWritesActionTypes.NamespaceShardKeyFetched
+    ) &&
+    state.status === ShardingStatuses.NOT_READY &&
+    !action.shardKey &&
+    !state.managedNamespace
+  ) {
+    if (state.pollingTimeout) {
+      throw new Error('Polling was not stopped');
+    }
+    return {
+      ...state,
+      status: ShardingStatuses.UNSHARDED,
+      pollingTimeout: state.pollingTimeout,
+    };
+  }
+
   if (
     isAction<ShardZonesFetchedAction>(
       action,
@@ -759,15 +778,23 @@ export const fetchNamespaceShardKey = (): GlobalWritesThunkAction<
     getState,
     { atlasGlobalWritesService, logger, connectionInfoRef }
   ) => {
-    const { namespace, status } = getState();
+    const { namespace, status, managedNamespace } = getState();
 
     try {
       const [shardingError, shardKey] = await Promise.all([
         atlasGlobalWritesService.getShardingError(namespace),
         atlasGlobalWritesService.getShardingKeys(namespace),
       ]);
 
-      if (shardingError && !shardKey) {
+      if (status === ShardingStatuses.SHARDING && (shardKey || shardingError)) {
+        dispatch(stopPollingForShardKey());
+      }
+
+      if (managedNamespace && !shardKey) {
+        if (!shardingError) {
+          dispatch(setNamespaceBeingSharded());
+          return;
+        }
         // if there is an existing shard key and an error both,
         // means we have a key mismatch
         // this will be handled in NamespaceShardKeyFetched
@@ -781,14 +808,10 @@ export const fetchNamespaceShardKey = (): GlobalWritesThunkAction<
         return;
       }
 
-      if (status === ShardingStatuses.SHARDING) {
-        dispatch(stopPollingForShardKey());
-      }
       dispatch({
         type: GlobalWritesActionTypes.NamespaceShardKeyFetched,
         shardKey,
       });
-
       // if there is a key, we fetch sharding zones
       if (!shardKey) return;
       void dispatch(fetchShardingZones());
@@ -874,36 +897,33 @@ export function getStatusFromShardKeyAndManaged(
 ) {
   const [firstShardKey, secondShardKey] = shardKey.fields;
 
-  if (!shardKey && !managedNamespace) {
-    return ShardingStatuses.UNSHARDED;
-  }
-
   // For a shard key to be valid:
   // 1. The first key must be location and of type RANGE.
   // 2. The second key name must match managedNamespace.customShardKey and
   //   the type must match the managedNamespace.isCustomShardKeyHashed.
 
   const isLocatonKeyValid =
     firstShardKey.name === 'location' && firstShardKey.type === 'RANGE';
+
+  if (!isLocatonKeyValid || !secondShardKey) {
+    return ShardingStatuses.SHARD_KEY_INVALID;
+  }
+
+  if (!managedNamespace) {
+    return ShardingStatuses.INCOMPLETE_SHARDING_SETUP;
+  }
+
   const isCustomKeyValid =
     managedNamespace &&
     managedNamespace.isShardKeyUnique === shardKey.isUnique &&
     secondShardKey.name === managedNamespace.customShardKey &&
     secondShardKey.type ===
       (managedNamespace.isCustomShardKeyHashed ? 'HASHED' : 'RANGE');
 
-  if (!isLocatonKeyValid || !secondShardKey) {
-    return ShardingStatuses.SHARD_KEY_INVALID;
-  }
-
   if (!isCustomKeyValid) {
     return ShardingStatuses.SHARD_KEY_MISMATCH;
   }
 
-  if (!managedNamespace) {
-    return ShardingStatuses.INCOMPLETE_SHARDING_SETUP;
-  }
-
   return ShardingStatuses.SHARD_KEY_CORRECT;
 }