fix(ui): issue with refresh for loggedInUser return 401 (#20990)

* fix(ui): issue with refresh for loggedInUser return 401

* fix playwright

openmetadata-ui/src/main/resources/ui/src/components/Auth/AuthProviders/AuthProvider.test.tsx

@@ -10,16 +10,14 @@
  *  See the License for the specific language governing permissions and
  *  limitations under the License.
  */
-import {
-  act,
-  render,
-  screen,
-  waitForElementToBeRemoved,
-} from '@testing-library/react';
+import { act, render, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
+import { AxiosResponse } from 'axios';
 import React from 'react';
 import { AuthProvider as AuthProviderProps } from '../../../generated/configuration/authenticationConfiguration';
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
+import axiosClient from '../../../rest';
+import TokenService from '../../../utils/Auth/TokenService/TokenServiceUtil';
 import AuthProvider from './AuthProvider';
 
 const localStorageMock = {
@@ -57,6 +55,66 @@ jest.mock('../../../rest/userAPI', () => ({
   updateUser: jest.fn().mockImplementation(() => Promise.resolve()),
 }));
 
+jest.mock('../../../utils/ToastUtils', () => ({
+  showErrorToast: jest.fn(),
+  showInfoToast: jest.fn(),
+}));
+
+const mockRefreshToken = jest
+  .fn()
+  .mockImplementation(() => Promise.resolve('newToken'));
+
+jest.mock('../../../utils/Auth/TokenService/TokenServiceUtil', () => {
+  return {
+    getInstance: jest.fn().mockImplementation(() => ({
+      refreshToken: mockRefreshToken,
+      isTokenUpdateInProgress: jest.fn().mockImplementation(() => false),
+      getToken: jest.fn().mockImplementation(() => Promise.resolve()),
+      clearRefreshInProgress: jest
+        .fn()
+        .mockImplementation(() => Promise.resolve()),
+      renewToken: jest.fn(),
+      refreshSuccessCallback: jest.fn(),
+      handleTokenUpdate: jest.fn(),
+      updateRenewToken: jest.fn(),
+      updateRefreshSuccessCallback: jest.fn(),
+      isTokenExpired: jest.fn(),
+      getTokenExpiry: jest.fn(),
+      fetchNewToken: jest.fn(),
+      setRefreshInProgress: jest.fn(),
+    })),
+  };
+});
+
+jest.mock('../../../hooks/useApplicationStore', () => ({
+  useApplicationStore: jest.fn().mockImplementation(() => ({
+    setHelperFunctionsRef: jest.fn(),
+    setCurrentUser: jest.fn(),
+    updateNewUser: jest.fn(),
+    setIsAuthenticated: jest.fn(),
+    setAuthConfig: jest.fn(),
+    setAuthorizerConfig: jest.fn(),
+    setIsSigningUp: jest.fn(),
+    authorizerConfig: {},
+    jwtPrincipalClaims: {},
+    jwtPrincipalClaimsMapping: {},
+    setJwtPrincipalClaims: jest.fn(),
+    setJwtPrincipalClaimsMapping: jest.fn(),
+    isApplicationLoading: false,
+    setApplicationLoading: jest.fn(),
+    authConfig: {
+      provider: AuthProviderProps.Basic,
+      providerName: 'Basic',
+      clientId: 'test',
+      authority: 'test',
+      callbackUrl: 'test',
+      jwtPrincipalClaims: [],
+      publicKeyUrls: [],
+      scope: 'openid',
+    },
+  })),
+}));
+
 describe('Test auth provider', () => {
   it('Logout handler should call the "updateUserDetails" method', async () => {
     const ConsumerComponent = () => {
@@ -75,8 +133,6 @@ describe('Test auth provider', () => {
       </AuthProvider>
     );
 
-    await waitForElementToBeRemoved(() => screen.getByTestId('loader'));
-
     const logoutButton = screen.getByTestId('logout-button');
 
     expect(logoutButton).toBeInTheDocument();
@@ -108,3 +164,253 @@ describe('Test auth provider', () => {
     expect(mockOnLogoutHandler).toHaveBeenCalled();
   });
 });
+
+describe('Test axios response interceptor', () => {
+  const ConsumerComponent = () => {
+    return <div>ConsumerComponent</div>;
+  };
+
+  const WrapperComponent = () => {
+    return (
+      <AuthProvider childComponentType={ConsumerComponent}>
+        <ConsumerComponent />
+      </AuthProvider>
+    );
+  };
+
+  beforeEach(() => {
+    jest.restoreAllMocks();
+  });
+
+  it('should set up response interceptor with correct signature', () => {
+    // Mock axios client
+    const mockUse = jest.spyOn(axiosClient.interceptors.response, 'use');
+    const mockAxios = jest.fn().mockResolvedValue({ data: 'success' });
+
+    jest.spyOn(axiosClient, 'request').mockImplementation(mockAxios);
+
+    render(<WrapperComponent />);
+
+    // Verify the interceptor was set up
+    expect(mockUse).toHaveBeenCalled();
+
+    // Get the arguments passed to use()
+    const [successHandler, errorHandler] = mockUse.mock.calls[0];
+
+    // Verify success handler signature
+    expect(typeof successHandler).toBe('function');
+    expect(successHandler).toHaveLength(1); // Takes one argument (response)
+
+    // Verify error handler signature
+    expect(typeof errorHandler).toBe('function');
+    expect(errorHandler).toHaveLength(1); // Takes one argument (error)
+
+    // Test success handler
+    const mockResponse = { data: 'test' } as AxiosResponse;
+
+    expect(successHandler?.(mockResponse)).toBe(mockResponse);
+
+    // Test error handler with 401 error
+    const mockError = {
+      response: {
+        status: 401,
+        data: { message: 'Token expired' },
+      },
+      config: { url: '/api/test' },
+    };
+
+    // The error handler should return a Promise
+    const result = errorHandler?.(mockError);
+
+    expect(result).toBeInstanceOf(Promise);
+    expect(mockRefreshToken).toHaveBeenCalled();
+  });
+
+  it('should handle 401 error when refresh is not in progress and refresh succeeds', async () => {
+    const mockUse = jest.spyOn(axiosClient.interceptors.response, 'use');
+    const mockAxios = jest.fn().mockResolvedValue({ data: 'success' });
+    jest.spyOn(axiosClient, 'request').mockImplementation(mockAxios);
+
+    await act(async () => {
+      render(<WrapperComponent />);
+    });
+
+    const [, errorHandler] = mockUse.mock.calls[0];
+    const mockError = {
+      response: {
+        status: 401,
+        data: { message: 'Token expired' },
+      },
+      config: { url: '/api/test' },
+    };
+
+    const result = await errorHandler?.(mockError);
+
+    expect(result).toEqual({ data: 'success' });
+    expect(mockRefreshToken).toHaveBeenCalled();
+    expect(mockAxios).toHaveBeenCalledWith(mockError.config);
+  });
+
+  it('should queue request when refresh is already in progress', async () => {
+    const mockUse = jest.spyOn(axiosClient.interceptors.response, 'use');
+    const mockAxios = jest.fn().mockResolvedValue({ data: 'success' });
+
+    jest.spyOn(axiosClient, 'request').mockImplementation(mockAxios);
+
+    // Mock isTokenUpdateInProgress to return true for this test
+    jest
+      .spyOn(TokenService.getInstance(), 'isTokenUpdateInProgress')
+      .mockReturnValue(true);
+
+    await act(async () => {
+      render(<WrapperComponent />);
+    });
+
+    const [, errorHandler] = mockUse.mock.calls[0];
+    const mockError = {
+      response: {
+        status: 401,
+        data: { message: 'Token expired' },
+      },
+      config: {
+        url: '/api/test',
+        headers: {},
+        baseURL: '',
+      },
+    };
+
+    const result = await errorHandler?.(mockError);
+
+    expect(mockRefreshToken).toHaveBeenCalled();
+    expect(mockAxios).toHaveBeenCalledWith(
+      expect.objectContaining(mockError.config)
+    );
+    expect(await result).toEqual({ data: 'success' });
+  });
+
+  it('should not call refresh for login api', async () => {
+    const mockUse = jest.spyOn(axiosClient.interceptors.response, 'use');
+    const mockAxios = jest.fn().mockResolvedValue({ data: 'success' });
+
+    jest.spyOn(axiosClient, 'request').mockImplementation(mockAxios);
+    await act(async () => {
+      render(<WrapperComponent />);
+    });
+
+    const [, errorHandler] = mockUse.mock.calls[0];
+    const mockError = {
+      response: {
+        status: 401,
+        data: { message: 'Token expired' },
+      },
+      config: {
+        url: '/users/login',
+        headers: {},
+        baseURL: '',
+      },
+    };
+
+    try {
+      await errorHandler?.(mockError);
+    } catch (error) {
+      // eslint-disable-next-line jest/no-conditional-expect, jest/no-try-expect
+      expect(error).toEqual(mockError);
+    }
+  });
+
+  it('should not call refresh for refresh api', async () => {
+    const mockUse = jest.spyOn(axiosClient.interceptors.response, 'use');
+    const mockAxios = jest.fn().mockResolvedValue({ data: 'success' });
+
+    jest.spyOn(axiosClient, 'request').mockImplementation(mockAxios);
+
+    await act(async () => {
+      render(<WrapperComponent />);
+    });
+
+    const [, errorHandler] = mockUse.mock.calls[0];
+    const mockError = {
+      response: {
+        status: 401,
+        data: { message: 'Token expired' },
+      },
+      config: {
+        url: '/users/refresh',
+        headers: {},
+        baseURL: '',
+      },
+    };
+
+    try {
+      await errorHandler?.(mockError);
+    } catch (error) {
+      // eslint-disable-next-line jest/no-conditional-expect, jest/no-try-expect
+      expect(error).toEqual(mockError);
+    }
+  });
+
+  it('should not call refresh for loggedInUser api if error is Token expired', async () => {
+    const mockUse = jest.spyOn(axiosClient.interceptors.response, 'use');
+    const mockAxios = jest.fn().mockResolvedValue({ data: 'success' });
+
+    jest.spyOn(axiosClient, 'request').mockImplementation(mockAxios);
+
+    await act(async () => {
+      render(<WrapperComponent />);
+    });
+
+    const [, errorHandler] = mockUse.mock.calls[0];
+    const mockError = {
+      response: {
+        status: 401,
+        data: { message: 'Token expired' },
+      },
+      config: {
+        url: '/users/loggedInUser',
+        headers: {},
+        baseURL: '',
+      },
+    };
+
+    try {
+      await errorHandler?.(mockError);
+    } catch (error) {
+      // eslint-disable-next-line jest/no-conditional-expect, jest/no-try-expect
+      expect(error).toEqual(mockError);
+    }
+  });
+
+  it('should call refresh for loggedInUser api if error other then Token expired', async () => {
+    const mockUse = jest.spyOn(axiosClient.interceptors.response, 'use');
+    const mockAxios = jest.fn().mockResolvedValue({ data: 'success' });
+    mockRefreshToken.mockImplementationOnce(() => Promise.resolve());
+
+    jest.spyOn(axiosClient, 'request').mockImplementation(mockAxios);
+
+    await act(async () => {
+      render(<WrapperComponent />);
+    });
+
+    const [, errorHandler] = mockUse.mock.calls[0];
+    const mockError = {
+      response: {
+        status: 401,
+        data: { message: 'token not valid' },
+      },
+      config: {
+        url: '/users/loggedInUser',
+        headers: {},
+        baseURL: '',
+      },
+    };
+
+    try {
+      await errorHandler?.(mockError);
+    } catch (error) {
+      // eslint-disable-next-line jest/no-conditional-expect, jest/no-try-expect
+      expect(error).toEqual(mockError);
+      // eslint-disable-next-line jest/no-conditional-expect, jest/no-try-expect
+      expect(mockRefreshToken).toHaveBeenCalledTimes(0);
+    }
+  });
+});

openmetadata-ui/src/main/resources/ui/src/components/Auth/AuthProviders/AuthProvider.tsx

@@ -526,8 +526,12 @@ export const AuthProvider = ({
           if (status === ClientErrors.UNAUTHORIZED) {
             // For login or refresh we don't want to fire another refresh req
             // Hence rejecting it
-            if (UN_AUTHORIZED_EXCLUDED_PATHS.includes(error.config.url)) {
-              return Promise.reject(error as Error);
+            if (
+              UN_AUTHORIZED_EXCLUDED_PATHS.includes(error.config.url) ||
+              (error.config.url === '/users/loggedInUser' &&
+                !error.response.data.message.includes('Expired token!'))
+            ) {
+              return Promise.reject(error);
             }
             handleStoreProtectedRedirectPath();
 
@@ -548,7 +552,7 @@ export const AuthProvider = ({
                     // Retry the pending requests
                     initializeAxiosInterceptors();
                     pendingRequests.forEach(({ resolve, reject, config }) => {
-                      axiosClient(config).then(resolve).catch(reject);
+                      axiosClient.request(config).then(resolve).catch(reject);
                     });
 
                     // Clear the queue after retrying

openmetadata-ui/src/main/resources/ui/src/components/GlobalSearchBar/GlobalSearchBar.tsx

@@ -13,7 +13,7 @@
 import Icon from '@ant-design/icons';
 import { Button, Divider, Input, Popover, Select, Tooltip } from 'antd';
 import classNames from 'classnames';
-import { debounce, isString } from 'lodash';
+import { debounce, isEmpty, isString } from 'lodash';
 import Qs from 'qs';
 import React, {
   useCallback,
@@ -61,7 +61,7 @@ export const GlobalSearchBar = () => {
   const [isSearchBoxOpen, setIsSearchBoxOpen] = useState<boolean>(false);
   const history = useHistory();
   const { isTourOpen, updateTourPage, updateTourSearch } = useTourProvider();
-
+  const { currentUser } = useApplicationStore();
   const parsedQueryString = Qs.parse(
     location.search.startsWith('?')
       ? location.search.substring(1)
@@ -166,14 +166,16 @@ export const GlobalSearchBar = () => {
   };
 
   const fetchNLPEnabledStatus = useCallback(() => {
-    getNLPEnabledStatus().then((enabled) => {
-      setNLPEnabled(enabled);
-    });
-  }, [setNLPEnabled]);
+    if (!isEmpty(currentUser)) {
+      getNLPEnabledStatus().then((enabled) => {
+        setNLPEnabled(enabled);
+      });
+    }
+  }, [setNLPEnabled, currentUser]);
 
   useEffect(() => {
     fetchNLPEnabledStatus();
-  }, []);
+  }, [fetchNLPEnabledStatus]);
 
   return (
     <div

openmetadata-ui/src/main/resources/ui/src/constants/Auth.constants.ts

@@ -33,8 +33,4 @@ export const HTTP_STATUS_CODE = {
   LIMIT_REACHED: 429, // Entity creation limit reached
 };
 
-export const UN_AUTHORIZED_EXCLUDED_PATHS = [
-  '/users/refresh',
-  '/users/login',
-  '/users/loggedInUser',
-];
+export const UN_AUTHORIZED_EXCLUDED_PATHS = ['/users/refresh', '/users/login'];