11import React from 'react' ;
22import { render , screen , fireEvent , waitFor , act } from '@testing-library/react' ;
33import { WalletProvider , useWallet } from '@/context/WalletContext' ;
4+ import { getSessionItem } from '@/auth/session' ;
45
56jest . mock ( '@stellar/freighter-api' , ( ) => ( {
67 isConnected : jest . fn ( ) ,
@@ -202,7 +203,7 @@ describe('WalletContext', () => {
202203 await waitFor ( ( ) => expect ( screen . getByTestId ( 'is-loading' ) ) . toHaveTextContent ( 'Ready' ) ) ;
203204 fireEvent . click ( screen . getByTestId ( 'connect-btn' ) ) ;
204205
205- await waitFor ( ( ) => expect ( localStorage . getItem ( STORAGE_KEY ) ) . toBe ( PUBLIC_KEY ) ) ;
206+ await waitFor ( ( ) => expect ( getSessionItem ( STORAGE_KEY ) ) . resolves . toBe ( PUBLIC_KEY ) ) ;
206207 } ) ;
207208
208209 it ( 'should restore wallet from localStorage when it matches the current Freighter address' , async ( ) => {
@@ -218,7 +219,7 @@ describe('WalletContext', () => {
218219
219220 await waitFor ( ( ) => expect ( screen . getByTestId ( 'public-key' ) ) . toHaveTextContent ( PUBLIC_KEY ) ) ;
220221 expect ( screen . getByTestId ( 'is-connected' ) ) . toHaveTextContent ( 'Connected' ) ;
221- expect ( localStorage . getItem ( STORAGE_KEY ) ) . toBe ( PUBLIC_KEY ) ;
222+ await expect ( getSessionItem ( STORAGE_KEY ) ) . resolves . toBe ( PUBLIC_KEY ) ;
222223 } ) ;
223224
224225 it ( 'should clear stored wallet when Freighter is disconnected on restore' , async ( ) => {
@@ -340,7 +341,7 @@ describe('WalletContext', () => {
340341 </ WalletProvider >
341342 ) ;
342343
343- await waitFor ( ( ) => expect ( localStorage . getItem ( STORAGE_KEY ) ) . toBe ( PUBLIC_KEY ) ) ;
344+ await waitFor ( ( ) => expect ( getSessionItem ( STORAGE_KEY ) ) . resolves . toBe ( PUBLIC_KEY ) ) ;
344345
345346 fireEvent . click ( screen . getByTestId ( 'disconnect-btn' ) ) ;
346347
@@ -382,7 +383,7 @@ describe('WalletContext', () => {
382383 await waitFor ( ( ) => expect ( screen . getByTestId ( 'is-loading' ) ) . toHaveTextContent ( 'Ready' ) ) ;
383384 fireEvent . click ( screen . getByTestId ( 'connect-btn' ) ) ;
384385
385- await waitFor ( ( ) => expect ( localStorage . getItem ( STORAGE_KEY ) ) . toBe ( PUBLIC_KEY ) ) ;
386+ await waitFor ( ( ) => expect ( getSessionItem ( STORAGE_KEY ) ) . resolves . toBe ( PUBLIC_KEY ) ) ;
386387
387388 unmount ( ) ;
388389
0 commit comments