fix: review changes

src/core_modules/capture-core/components/MetadataAutoSelectInitializer/hooks/useMetadataAutoSelect.js

@@ -1,9 +1,10 @@
 // @flow
 import { useCallback, useEffect, useState } from 'react';
 import { useHistory } from 'react-router-dom';
-import { useApiMetadataQuery, useIndexedDBQuery } from '../../../utils/reactQueryHelpers';
+import { useIndexedDBQuery } from '../../../utils/reactQueryHelpers';
 import { getUserStorageController, userStores } from '../../../storageControllers';
 import { buildUrlQueryString, useLocationQuery } from '../../../utils/routing';
+import { useOrgUnitAutoSelect } from '../../../dataQueries';
 
 const getAllPrograms = () => {
     const userStorageController = getUserStorageController();
@@ -28,21 +29,8 @@ export const useMetadataAutoSelect = () => {
         },
     );
 
-    const { data: searchOrgUnits, isLoading: loadingOrgUnits } = useApiMetadataQuery(
-        ['searchOrgUnitsForAutoSelect'],
-        {
-            resource: 'organisationUnits',
-            params: {
-                fields: 'id',
-                withinUserSearchHierarchy: true,
-                pageSize: 2,
-            },
-        },
-        {
-            enabled: Object.keys(urlParams).length === 0 && !mounted,
-            select: ({ organisationUnits }) => organisationUnits,
-        },
-    );
+    const queryOptions = { enabled: Object.keys(urlParams).length === 0 && !mounted };
+    const { isLoading: loadingOrgUnits, data: searchOrgUnits } = useOrgUnitAutoSelect(queryOptions);
 
     const updateUrlIfApplicable = useCallback(() => {
         const paramsToAdd = {

src/core_modules/capture-core/components/WidgetRelatedStages/WidgetRelatedStages.component.js

@@ -1,7 +1,7 @@
 // @flow
 import React, { forwardRef, useCallback, useEffect, useImperativeHandle, useState } from 'react';
 import { useRelatedStages } from './useRelatedStages';
-import { useOrgUnitAutoSelect } from './hooks/useAutoSelctOrgUnitRelatedStage';
+import { useOrgUnitAutoSelect } from '../../dataQueries';
 import type { Props, RelatedStageDataValueStates } from './WidgetRelatedStages.types';
 import { RelatedStagesActions } from './RelatedStagesActions';
 import { relatedStageStatus } from './constants';
@@ -37,7 +37,15 @@ const WidgetRelatedStagesPlain = ({
         orgUnit: undefined,
         linkedEventId: undefined,
     });
-    const { isLoading: orgUnitLoading } = useOrgUnitAutoSelect(setRelatedStageDataValues);
+    const { isLoading: orgUnitLoading, data } = useOrgUnitAutoSelect();
+    useEffect(() => {
+        if (!orgUnitLoading && data?.length === 1) {
+            setRelatedStageDataValues(prev => ({
+                ...prev,
+                orgUnit: data[0],
+            }));
+        }
+    }, [data, orgUnitLoading, setRelatedStageDataValues]);
 
     const addErrorMessage = (message: ErrorMessagesForRelatedStages) => {
         setErrorMessages((prevMessages: Object) => ({

src/core_modules/capture-core/dataQueries/index.js

@@ -1 +1,2 @@
 export { useOrganisationUnit } from './useOrganisationUnit';
+export { useOrgUnitAutoSelect } from './useOrgUnitsForAutoSelect';

src/core_modules/capture-core/dataQueries/useOrgUnitsForAutoSelect.js

@@ -0,0 +1,26 @@
+// @flow
+import { useApiMetadataQuery } from '../utils/reactQueryHelpers';
+
+export const useOrgUnitAutoSelect = (customQueryOptions: Object) => {
+    const queryKey = ['organisationUnits'];
+    const queryFn = {
+        resource: 'organisationUnits',
+        params: {
+            fields: ['id, displayName~rename(name), path'],
+            withinUserHierarchy: true,
+            pageSize: 2,
+        },
+    };
+    const defaultQueryOptions = {
+        select: ({ organisationUnits }) => organisationUnits,
+    };
+
+    const queryOptions = { ...defaultQueryOptions, ...customQueryOptions };
+
+    const { data, isLoading } = useApiMetadataQuery(queryKey, queryFn, queryOptions);
+
+    return {
+        isLoading,
+        data,
+    };
+};