Skip to content

Commit 08c02fa

Browse files
nicknisinicknisi
authored
fix: improve useAccessToken timer management and prevent background flashing (#280)
1 parent dfe4759 commit 08c02fa

File tree

2 files changed

+107
-25
lines changed

2 files changed

+107
-25
lines changed

‎__tests__/useAccessToken.spec.tsx‎

Lines changed: 71 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -49,24 +49,27 @@ describe('useAccessToken', () => {
4949
);
5050
};
5151

52-
it('should fetch an access token on mount', async () => {
52+
it('should fetch an access token on mount without showing loading state', async () => {
5353
const mockToken =
5454
'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwic2lkIjoic2Vzc2lvbl8xMjMiLCJleHAiOjk5OTk5OTk5OTl9.mock-signature';
5555
(getAccessTokenAction as jest.Mock).mockResolvedValueOnce(mockToken);
5656

5757
const { getByTestId } = render(<TestComponent />);
5858

59-
expect(getByTestId('loading')).toHaveTextContent('true');
59+
// Loading should remain false for background fetches
60+
expect(getByTestId('loading')).toHaveTextContent('false');
61+
62+
await waitFor(() => {
63+
expect(getAccessTokenAction).toHaveBeenCalledTimes(1);
64+
});
6065

6166
await waitFor(() => {
6267
expect(getByTestId('loading')).toHaveTextContent('false');
6368
expect(getByTestId('token')).toHaveTextContent(mockToken);
6469
});
65-
66-
expect(getAccessTokenAction).toHaveBeenCalledTimes(1);
6770
});
6871

69-
it('should handle token refresh when an expiring token is received', async () => {
72+
it('should handle token refresh when an expiring token is received without showing loading', async () => {
7073
// Create a token that's about to expire (exp is very close to current time)
7174
const currentTimeInSeconds = Math.floor(Date.now() / 1000);
7275
const payload = { sub: '1234567890', sid: 'session_123', exp: currentTimeInSeconds + 30 };
@@ -80,6 +83,9 @@ describe('useAccessToken', () => {
8083

8184
const { getByTestId } = render(<TestComponent />);
8285

86+
// Loading should remain false throughout
87+
expect(getByTestId('loading')).toHaveTextContent('false');
88+
8389
await waitFor(() => {
8490
expect(getByTestId('loading')).toHaveTextContent('false');
8591
expect(getByTestId('token')).toHaveTextContent(refreshedToken);
@@ -89,7 +95,7 @@ describe('useAccessToken', () => {
8995
expect(refreshAccessTokenAction).toHaveBeenCalledTimes(1);
9096
});
9197

92-
it('should handle token refresh on manual refresh', async () => {
98+
it('should handle token refresh on manual refresh and show loading state', async () => {
9399
const initialToken =
94100
'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwic2lkIjoic2Vzc2lvbl8xMjMiLCJleHAiOjk5OTk5OTk5OTl9.mock-signature';
95101
const refreshedToken =
@@ -108,16 +114,21 @@ describe('useAccessToken', () => {
108114

109115
await waitFor(() => {
110116
expect(getByTestId('token')).toHaveTextContent(initialToken);
117+
expect(getByTestId('loading')).toHaveTextContent('false');
111118
});
112119

113120
act(() => {
114121
getByTestId('refresh').click();
115122
});
116123

124+
// Should show loading for user-initiated refresh
125+
expect(getByTestId('loading')).toHaveTextContent('true');
126+
117127
await waitFor(() => {
118128
expect(mockRefreshAuth).toHaveBeenCalledTimes(1);
119129
expect(getAccessTokenAction).toHaveBeenCalledTimes(2);
120130
expect(getByTestId('token')).toHaveTextContent(refreshedToken);
131+
expect(getByTestId('loading')).toHaveTextContent('false');
121132
});
122133
});
123134

@@ -164,12 +175,15 @@ describe('useAccessToken', () => {
164175
});
165176
});
166177

167-
it('should handle errors during token fetch', async () => {
178+
it('should handle errors during token fetch without showing loading', async () => {
168179
const error = new Error('Failed to fetch token');
169180
(getAccessTokenAction as jest.Mock).mockRejectedValueOnce(error);
170181

171182
const { getByTestId } = render(<TestComponent />);
172183

184+
// Loading should remain false even when there's an error
185+
expect(getByTestId('loading')).toHaveTextContent('false');
186+
173187
await waitFor(() => {
174188
expect(getByTestId('loading')).toHaveTextContent('false');
175189
expect(getByTestId('error')).toHaveTextContent('Failed to fetch token');
@@ -250,7 +264,7 @@ describe('useAccessToken', () => {
250264
});
251265
});
252266

253-
it('should retry fetching when an error occurs', async () => {
267+
it('should retry fetching when an error occurs without showing loading', async () => {
254268
const error = new Error('Failed to fetch token');
255269
const mockToken =
256270
'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwic2lkIjoic2Vzc2lvbl8xMjMiLCJleHAiOjk5OTk5OTk5OTl9.mock-signature';
@@ -261,16 +275,21 @@ describe('useAccessToken', () => {
261275

262276
await waitFor(() => {
263277
expect(getByTestId('error')).toHaveTextContent('Failed to fetch token');
278+
expect(getByTestId('loading')).toHaveTextContent('false');
279+
expect(getByTestId('token')).toHaveTextContent('no-token');
264280
});
265281

266282
act(() => {
267283
jest.advanceTimersByTime(5 * 60 * 1000); // RETRY_DELAY
268284
});
269285

286+
// Loading should remain false during retry
287+
expect(getByTestId('loading')).toHaveTextContent('false');
288+
270289
await waitFor(() => {
271290
expect(getAccessTokenAction).toHaveBeenCalledTimes(2);
272291
expect(getByTestId('token')).toHaveTextContent(mockToken);
273-
expect(getByTestId('error')).toHaveTextContent('no-error');
292+
expect(getByTestId('loading')).toHaveTextContent('false');
274293
});
275294
});
276295

@@ -376,7 +395,8 @@ describe('useAccessToken', () => {
376395

377396
const { getByTestId } = render(<TestComponent />);
378397

379-
expect(getByTestId('loading')).toHaveTextContent('true');
398+
// Loading should remain false for background operations
399+
expect(getByTestId('loading')).toHaveTextContent('false');
380400

381401
await waitFor(() => {
382402
expect(fetchCalls).toBe(1);
@@ -441,6 +461,42 @@ describe('useAccessToken', () => {
441461
expect(refreshAuthCalls).toBe(1);
442462
});
443463

464+
it('should not show loading state for background refresh when token is unchanged', async () => {
465+
// Create a token that expires in 2 minutes
466+
const currentTimeInSeconds = Math.floor(Date.now() / 1000);
467+
const expTimeInSeconds = currentTimeInSeconds + 120;
468+
const payload = { sub: '1234567890', sid: 'session_123', exp: expTimeInSeconds };
469+
const sameToken = `eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.${btoa(JSON.stringify(payload))}.mock-signature`;
470+
471+
// Mock to return the same token twice
472+
(getAccessTokenAction as jest.Mock).mockResolvedValue(sameToken);
473+
474+
const { getByTestId } = render(<TestComponent />);
475+
476+
// Wait for initial token to be set
477+
await waitFor(() => {
478+
expect(getByTestId('token')).toHaveTextContent(sameToken);
479+
expect(getByTestId('loading')).toHaveTextContent('false');
480+
});
481+
482+
// Clear previous mock calls
483+
(getAccessTokenAction as jest.Mock).mockClear();
484+
485+
// Trigger scheduled refresh by advancing time to just before expiry
486+
act(() => {
487+
jest.advanceTimersByTime((120 - 60) * 1000); // Advance to trigger refresh (2 min - buffer)
488+
});
489+
490+
await waitFor(() => {
491+
expect(getAccessTokenAction).toHaveBeenCalledTimes(1); // Called once for the refresh
492+
});
493+
494+
// Should not show loading during background refresh
495+
expect(getByTestId('loading')).toHaveTextContent('false');
496+
// Token should remain the same
497+
expect(getByTestId('token')).toHaveTextContent(sameToken);
498+
});
499+
444500
it('should handle non-Error objects thrown during token fetch', async () => {
445501
// Simulate a string error being thrown
446502
(getAccessTokenAction as jest.Mock).mockImplementation(() => {
@@ -452,6 +508,7 @@ describe('useAccessToken', () => {
452508
await waitFor(() => {
453509
expect(getByTestId('loading')).toHaveTextContent('false');
454510
expect(getByTestId('error')).toHaveTextContent('String error message');
511+
expect(getByTestId('token')).toHaveTextContent('no-token');
455512
});
456513
});
457514

@@ -476,14 +533,17 @@ describe('useAccessToken', () => {
476533
refreshAuth: mockRefreshAuth,
477534
}));
478535

479-
(getAccessTokenAction as jest.Mock).mockResolvedValueOnce(initialToken).mockResolvedValueOnce(refreshedToken);
536+
(getAccessTokenAction as jest.Mock).mockResolvedValue(initialToken);
480537

481538
const { getByTestId } = render(<TestComponent />);
482539

483540
await waitFor(() => {
484541
expect(getByTestId('token')).toHaveTextContent(initialToken);
485542
});
486543

544+
// Now update the mock for the refresh
545+
(getAccessTokenAction as jest.Mock).mockResolvedValue(refreshedToken);
546+
487547
act(() => {
488548
getByTestId('refresh').click();
489549
});

‎src/components/useAccessToken.ts‎

Lines changed: 36 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ function tokenReducer(state: TokenState, action: TokenAction): TokenState {
2525
case 'FETCH_START':
2626
return { ...state, loading: true, error: null };
2727
case 'FETCH_SUCCESS':
28-
return { ...state, loading: false, token: action.token };
28+
return { ...state, loading: false, token: action.token, error: null };
2929
case 'FETCH_ERROR':
3030
return { ...state, loading: false, error: action.error };
3131
case 'RESET':
@@ -90,14 +90,34 @@ export function useAccessToken() {
9090
}
9191
}, []);
9292

93+
// Store the current token in a ref to avoid stale closures
94+
const currentTokenRef = useRef<string | undefined>(state.token);
95+
currentTokenRef.current = state.token;
96+
97+
// Store updateToken in a ref to break circular dependency
98+
const updateTokenRef = useRef<() => Promise<string | undefined>>();
99+
100+
// Centralized timer scheduling function
101+
const scheduleNextRefresh = useCallback(
102+
(delay: number) => {
103+
clearRefreshTimeout();
104+
refreshTimeoutRef.current = setTimeout(() => {
105+
if (updateTokenRef.current) {
106+
updateTokenRef.current();
107+
}
108+
}, delay);
109+
},
110+
[clearRefreshTimeout],
111+
);
112+
93113
const updateToken = useCallback(async () => {
94114
// istanbul ignore next - safety guard against concurrent fetches
95115
if (fetchingRef.current) {
96116
return;
97117
}
98118

99119
fetchingRef.current = true;
100-
dispatch({ type: 'FETCH_START' });
120+
101121
try {
102122
let token = await getAccessTokenAction();
103123
if (token) {
@@ -107,26 +127,30 @@ export function useAccessToken() {
107127
}
108128
}
109129

110-
dispatch({ type: 'FETCH_SUCCESS', token });
130+
// Only update state if token has changed
131+
if (token !== currentTokenRef.current) {
132+
dispatch({ type: 'FETCH_SUCCESS', token });
133+
}
111134

112135
if (token) {
113136
const tokenData = parseTokenPayload(token);
114137
if (tokenData) {
115138
const delay = getRefreshDelay(tokenData.timeUntilExpiry);
116-
clearRefreshTimeout();
117-
refreshTimeoutRef.current = setTimeout(updateToken, delay);
139+
scheduleNextRefresh(delay);
118140
}
119141
}
120142

121143
return token;
122144
} catch (error) {
123145
dispatch({ type: 'FETCH_ERROR', error: error instanceof Error ? error : new Error(String(error)) });
124-
clearRefreshTimeout();
125-
refreshTimeoutRef.current = setTimeout(updateToken, RETRY_DELAY_SECONDS * 1000);
146+
scheduleNextRefresh(RETRY_DELAY_SECONDS * 1000);
126147
} finally {
127148
fetchingRef.current = false;
128149
}
129-
}, [clearRefreshTimeout]);
150+
}, [scheduleNextRefresh]);
151+
152+
// Assign updateToken to ref for use in scheduleNextRefresh
153+
updateTokenRef.current = updateToken;
130154

131155
const refresh = useCallback(async () => {
132156
if (fetchingRef.current) {
@@ -146,21 +170,19 @@ export function useAccessToken() {
146170
const tokenData = parseTokenPayload(token);
147171
if (tokenData) {
148172
const delay = getRefreshDelay(tokenData.timeUntilExpiry);
149-
clearRefreshTimeout();
150-
refreshTimeoutRef.current = setTimeout(updateToken, delay);
173+
scheduleNextRefresh(delay);
151174
}
152175
}
153176

154177
return token;
155178
} catch (error) {
156179
const typedError = error instanceof Error ? error : new Error(String(error));
157180
dispatch({ type: 'FETCH_ERROR', error: typedError });
158-
clearRefreshTimeout();
159-
refreshTimeoutRef.current = setTimeout(updateToken, RETRY_DELAY_SECONDS * 1000);
181+
scheduleNextRefresh(RETRY_DELAY_SECONDS * 1000);
160182
} finally {
161183
fetchingRef.current = false;
162184
}
163-
}, [refreshAuth, clearRefreshTimeout, updateToken]);
185+
}, [refreshAuth, scheduleNextRefresh, updateToken]);
164186

165187
useEffect(() => {
166188
if (!user) {
@@ -171,7 +193,7 @@ export function useAccessToken() {
171193
updateToken();
172194

173195
return clearRefreshTimeout;
174-
}, [userId, sessionId, updateToken, clearRefreshTimeout]);
196+
}, [userId, sessionId, clearRefreshTimeout]);
175197

176198
return {
177199
accessToken: state.token,

0 commit comments

Comments
 (0)