diff --git a/__tests__/useAccessToken.spec.tsx b/__tests__/useAccessToken.spec.tsx index 64ba744..0bdc7c8 100644 --- a/__tests__/useAccessToken.spec.tsx +++ b/__tests__/useAccessToken.spec.tsx @@ -57,15 +57,15 @@ describe('useAccessToken', () => { ); }; - it('should fetch an access token on mount without showing loading state', async () => { + it('should fetch an access token on mount and show loading state initially', async () => { const mockToken = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwic2lkIjoic2Vzc2lvbl8xMjMiLCJleHAiOjk5OTk5OTk5OTl9.mock-signature'; (getAccessTokenAction as jest.Mock).mockResolvedValueOnce(mockToken); const { getByTestId } = render(); - // Loading should remain false for background fetches - expect(getByTestId('loading')).toHaveTextContent('false'); + // Loading should be true during initial fetch + expect(getByTestId('loading')).toHaveTextContent('true'); await waitFor(() => { expect(getAccessTokenAction).toHaveBeenCalledTimes(1); @@ -77,7 +77,7 @@ describe('useAccessToken', () => { }); }); - it('should handle token refresh when an expiring token is received without showing loading', async () => { + it('should handle token refresh when an expiring token is received', async () => { // Create a token that's about to expire (exp is very close to current time) const currentTimeInSeconds = Math.floor(Date.now() / 1000); // Use 25 seconds to ensure it's within the 30-second buffer for short-lived tokens @@ -97,8 +97,8 @@ describe('useAccessToken', () => { const { getByTestId } = render(); - // Loading should remain false throughout - expect(getByTestId('loading')).toHaveTextContent('false'); + // Loading should be true initially during token fetch + expect(getByTestId('loading')).toHaveTextContent('true'); await waitFor(() => { expect(getByTestId('loading')).toHaveTextContent('false'); @@ -155,14 +155,14 @@ describe('useAccessToken', () => { }); }); - it('should handle errors during token fetch without showing loading', async () => { + it('should handle errors during token fetch', async () => { const error = new Error('Failed to fetch token'); (getAccessTokenAction as jest.Mock).mockRejectedValueOnce(error); const { getByTestId } = render(); - // Loading should remain false even when there's an error - expect(getByTestId('loading')).toHaveTextContent('false'); + // Loading should be true initially + expect(getByTestId('loading')).toHaveTextContent('true'); await waitFor(() => { expect(getByTestId('loading')).toHaveTextContent('false'); @@ -381,8 +381,8 @@ describe('useAccessToken', () => { const { getByTestId } = render(); - // Loading should remain false for background operations - expect(getByTestId('loading')).toHaveTextContent('false'); + // Loading should be true initially during fetch + expect(getByTestId('loading')).toHaveTextContent('true'); await waitFor(() => { expect(fetchCalls).toBe(1); @@ -459,6 +459,44 @@ describe('useAccessToken', () => { }); }); + it('should show loading state immediately on first render when user exists but no token', () => { + // Mock user with no token initially + (useAuth as jest.Mock).mockImplementation(() => ({ + user: { id: 'user_123' }, + sessionId: 'session_123', + refreshAuth: jest.fn().mockResolvedValue({}), + })); + + (getAccessTokenAction as jest.Mock).mockImplementation( + () => new Promise((resolve) => setTimeout(() => resolve('token'), 100)), + ); + + const { getByTestId } = render(); + + expect(getByTestId('loading')).toHaveTextContent('true'); + expect(getByTestId('token')).toHaveTextContent('no-token'); + }); + + it('should not show loading when a valid token already exists', async () => { + const existingToken = + 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJleGlzdGluZyIsInNpZCI6InNlc3Npb24xMjMiLCJleHAiOjk5OTk5OTk5OTl9.existing'; + + await act(async () => { + (getAccessTokenAction as jest.Mock).mockResolvedValueOnce(existingToken); + await tokenStore.getAccessTokenSilently(); + }); + + // Reset the mock to track new calls + (getAccessTokenAction as jest.Mock).mockClear(); + + const { getByTestId } = render(); + + expect(getByTestId('loading')).toHaveTextContent('false'); + expect(getByTestId('token')).toHaveTextContent(existingToken); + + expect(getAccessTokenAction).not.toHaveBeenCalled(); + }); + // Additional test cases to increase coverage it('should handle concurrent manual refresh attempts', async () => { const initialToken = diff --git a/src/components/useAccessToken.ts b/src/components/useAccessToken.ts index dfd3f05..e3b622d 100644 --- a/src/components/useAccessToken.ts +++ b/src/components/useAccessToken.ts @@ -1,4 +1,4 @@ -import { useCallback, useEffect, useRef, useSyncExternalStore } from 'react'; +import { useCallback, useEffect, useRef, useState, useSyncExternalStore } from 'react'; import { useAuth } from './authkit-provider.js'; import { tokenStore } from './tokenStore.js'; @@ -42,9 +42,17 @@ export function useAccessToken(): UseAccessTokenReturn { const tokenState = useSyncExternalStore(tokenStore.subscribe, tokenStore.getSnapshot, tokenStore.getServerSnapshot); + // Track if we're waiting for the initial token fetch for the current user + // Initialize synchronously to prevent first-paint flash + const [isInitialTokenLoading, setIsInitialTokenLoading] = useState(() => { + // Only show loading if we have a user but no token yet + return Boolean(user && !tokenState.token && !tokenState.error); + }); + useEffect(() => { if (!user) { tokenStore.clearToken(); + setIsInitialTokenLoading(false); return; } @@ -59,10 +67,28 @@ export function useAccessToken(): UseAccessTokenReturn { prevSessionRef.current = sessionId; prevUserIdRef.current = userId; + // Check if getAccessTokenSilently will actually fetch (not just return cached) + const currentToken = tokenStore.getSnapshot().token; + const tokenData = currentToken ? tokenStore.parseToken(currentToken) : null; + const willActuallyFetch = !currentToken || (tokenData && tokenData.isExpiring); + + // Only show loading if we're actually going to fetch + if (willActuallyFetch) { + setIsInitialTokenLoading(true); + } + /* istanbul ignore next */ - tokenStore.getAccessTokenSilently().catch(() => { - // Error is handled in the store - }); + tokenStore + .getAccessTokenSilently() + .catch(() => { + // Error is handled in the store + }) + .finally(() => { + // Only clear loading if we were actually loading + if (willActuallyFetch) { + setIsInitialTokenLoading(false); + } + }); }, [userId, sessionId]); useEffect(() => { @@ -112,9 +138,12 @@ export function useAccessToken(): UseAccessTokenReturn { return tokenStore.refreshToken(); }, []); + // Combine loading states: initial token fetch OR token store is loading + const isLoading = isInitialTokenLoading || tokenState.loading; + return { accessToken: tokenState.token, - loading: tokenState.loading, + loading: isLoading, error: tokenState.error, refresh, getAccessToken,