@@ -463,6 +463,14 @@ describe("completeAuth2Exchange()", () => {
463463 getOrCreateWalletWithTag : jest . fn ( ) . mockResolvedValue ( { walletId : "wallet-1" , tags : [ ] } ) ,
464464 } ;
465465
466+ /** Access token whose JWT `aud` includes a wallet URN so `Auth2Token.wallet` is set (migration runs only in that case). */
467+ const accessTokenWithWalletClaim = makeJwt ( {
468+ sub : "user-1" ,
469+ client_id : "test-client" ,
470+ ext : { a2t : DEFAULT_A2T } ,
471+ aud : [ "urn:phantom:wallet:wallet-1:0" ] ,
472+ } ) ;
473+
466474 beforeEach ( ( ) => {
467475 mockKms . getOrCreatePhantomOrganization . mockResolvedValue ( { organizationId : "org-1" } ) ;
468476 mockKms . listPendingMigrations . mockResolvedValue ( { pendingMigrations : [ ] } ) ;
@@ -542,7 +550,14 @@ describe("completeAuth2Exchange()", () => {
542550 expect ( mockKms . getOrCreatePhantomOrganization ) . toHaveBeenCalledWith ( "encoded-pub-key" ) ;
543551 } ) ;
544552
545- it ( "calls kms.listPendingMigrations with organizationId" , async ( ) => {
553+ it ( "calls kms.listPendingMigrations with organizationId when the token has a wallet claim" , async ( ) => {
554+ mockExchangeAuthCode . mockResolvedValueOnce ( {
555+ accessToken : accessTokenWithWalletClaim ,
556+ idType : "Bearer" ,
557+ expiresInMs : 3_600_000 ,
558+ refreshToken : "refresh-tok" ,
559+ } ) ;
560+
546561 await completeAuth2Exchange ( {
547562 stamper : makeStamper ( ) ,
548563 kms : mockKms as any ,
@@ -638,6 +653,12 @@ describe("completeAuth2Exchange()", () => {
638653 } ) ;
639654
640655 it ( "calls completeWalletTransfer for each pending migration" , async ( ) => {
656+ mockExchangeAuthCode . mockResolvedValueOnce ( {
657+ accessToken : accessTokenWithWalletClaim ,
658+ idType : "Bearer" ,
659+ expiresInMs : 3_600_000 ,
660+ refreshToken : "refresh-tok" ,
661+ } ) ;
641662 mockKms . listPendingMigrations . mockResolvedValueOnce ( {
642663 pendingMigrations : [ { migrationId : "migration-1" } , { migrationId : "migration-2" } ] ,
643664 } ) ;
@@ -663,6 +684,12 @@ describe("completeAuth2Exchange()", () => {
663684 } ) ;
664685
665686 it ( "skips migration entries that have no migrationId" , async ( ) => {
687+ mockExchangeAuthCode . mockResolvedValueOnce ( {
688+ accessToken : accessTokenWithWalletClaim ,
689+ idType : "Bearer" ,
690+ expiresInMs : 3_600_000 ,
691+ refreshToken : "refresh-tok" ,
692+ } ) ;
666693 mockKms . listPendingMigrations . mockResolvedValueOnce ( {
667694 pendingMigrations : [ { migrationId : "migration-1" } , { } , { migrationId : "migration-3" } ] ,
668695 } ) ;
@@ -688,6 +715,12 @@ describe("completeAuth2Exchange()", () => {
688715 } ) ;
689716
690717 it ( "does not call completeWalletTransfer when pendingMigrations key is absent" , async ( ) => {
718+ mockExchangeAuthCode . mockResolvedValueOnce ( {
719+ accessToken : accessTokenWithWalletClaim ,
720+ idType : "Bearer" ,
721+ expiresInMs : 3_600_000 ,
722+ refreshToken : "refresh-tok" ,
723+ } ) ;
691724 mockKms . listPendingMigrations . mockResolvedValueOnce ( { } ) ;
692725
693726 await completeAuth2Exchange ( {
@@ -703,6 +736,12 @@ describe("completeAuth2Exchange()", () => {
703736 } ) ;
704737
705738 it ( "does not call completeWalletTransfer when there are no pending migrations" , async ( ) => {
739+ mockExchangeAuthCode . mockResolvedValueOnce ( {
740+ accessToken : accessTokenWithWalletClaim ,
741+ idType : "Bearer" ,
742+ expiresInMs : 3_600_000 ,
743+ refreshToken : "refresh-tok" ,
744+ } ) ;
706745 mockKms . listPendingMigrations . mockResolvedValueOnce ( { pendingMigrations : [ ] } ) ;
707746
708747 await completeAuth2Exchange ( {
@@ -718,6 +757,12 @@ describe("completeAuth2Exchange()", () => {
718757 } ) ;
719758
720759 it ( "continues and resolves even when a completeWalletTransfer call throws" , async ( ) => {
760+ mockExchangeAuthCode . mockResolvedValueOnce ( {
761+ accessToken : accessTokenWithWalletClaim ,
762+ idType : "Bearer" ,
763+ expiresInMs : 3_600_000 ,
764+ refreshToken : "refresh-tok" ,
765+ } ) ;
721766 mockKms . listPendingMigrations . mockResolvedValueOnce ( {
722767 pendingMigrations : [ { migrationId : "migration-fail" } , { migrationId : "migration-ok" } ] ,
723768 } ) ;
@@ -875,8 +920,12 @@ describe("_getOrMigrateWallet()", () => {
875920 expect ( mockKms . getOrCreatePhantomOrganization ) . toHaveBeenCalledWith ( "encoded-pub-key" ) ;
876921 } ) ;
877922
878- it ( "calls listPendingMigrations with organizationId" , async ( ) => {
879- await _getOrMigrateWallet ( { stamper : makeStamper ( ) , kms : mockKms as any , auth2Token : makeToken ( ) } ) ;
923+ it ( "calls listPendingMigrations with organizationId when the token has a wallet claim" , async ( ) => {
924+ await _getOrMigrateWallet ( {
925+ stamper : makeStamper ( ) ,
926+ kms : mockKms as any ,
927+ auth2Token : makeToken ( { walletUrn : "urn:phantom:wallet:w:0" } ) ,
928+ } ) ;
880929
881930 expect ( mockKms . listPendingMigrations ) . toHaveBeenCalledWith ( "org-1" ) ;
882931 } ) ;
@@ -886,7 +935,11 @@ describe("_getOrMigrateWallet()", () => {
886935 pendingMigrations : [ { migrationId : "m-1" } , { migrationId : "m-2" } ] ,
887936 } ) ;
888937
889- await _getOrMigrateWallet ( { stamper : makeStamper ( ) , kms : mockKms as any , auth2Token : makeToken ( ) } ) ;
938+ await _getOrMigrateWallet ( {
939+ stamper : makeStamper ( ) ,
940+ kms : mockKms as any ,
941+ auth2Token : makeToken ( { walletUrn : "urn:phantom:wallet:wallet-1:0" } ) ,
942+ } ) ;
890943
891944 expect ( mockKms . completeWalletTransfer ) . toHaveBeenCalledTimes ( 2 ) ;
892945 expect ( mockKms . completeWalletTransfer ) . toHaveBeenCalledWith ( { organizationId : "org-1" , migrationId : "m-1" } ) ;
@@ -898,7 +951,11 @@ describe("_getOrMigrateWallet()", () => {
898951 pendingMigrations : [ { migrationId : "m-1" } , { } , { migrationId : "m-3" } ] ,
899952 } ) ;
900953
901- await _getOrMigrateWallet ( { stamper : makeStamper ( ) , kms : mockKms as any , auth2Token : makeToken ( ) } ) ;
954+ await _getOrMigrateWallet ( {
955+ stamper : makeStamper ( ) ,
956+ kms : mockKms as any ,
957+ auth2Token : makeToken ( { walletUrn : "urn:phantom:wallet:wallet-1:0" } ) ,
958+ } ) ;
902959
903960 expect ( mockKms . completeWalletTransfer ) . toHaveBeenCalledTimes ( 2 ) ;
904961 expect ( mockKms . completeWalletTransfer ) . not . toHaveBeenCalledWith (
@@ -909,15 +966,23 @@ describe("_getOrMigrateWallet()", () => {
909966 it ( "does not call completeWalletTransfer when pendingMigrations key is absent" , async ( ) => {
910967 mockKms . listPendingMigrations . mockResolvedValueOnce ( { } ) ;
911968
912- await _getOrMigrateWallet ( { stamper : makeStamper ( ) , kms : mockKms as any , auth2Token : makeToken ( ) } ) ;
969+ await _getOrMigrateWallet ( {
970+ stamper : makeStamper ( ) ,
971+ kms : mockKms as any ,
972+ auth2Token : makeToken ( { walletUrn : "urn:phantom:wallet:wallet-1:0" } ) ,
973+ } ) ;
913974
914975 expect ( mockKms . completeWalletTransfer ) . not . toHaveBeenCalled ( ) ;
915976 } ) ;
916977
917978 it ( "does not call completeWalletTransfer when pendingMigrations is empty" , async ( ) => {
918979 mockKms . listPendingMigrations . mockResolvedValueOnce ( { pendingMigrations : [ ] } ) ;
919980
920- await _getOrMigrateWallet ( { stamper : makeStamper ( ) , kms : mockKms as any , auth2Token : makeToken ( ) } ) ;
981+ await _getOrMigrateWallet ( {
982+ stamper : makeStamper ( ) ,
983+ kms : mockKms as any ,
984+ auth2Token : makeToken ( { walletUrn : "urn:phantom:wallet:wallet-1:0" } ) ,
985+ } ) ;
921986
922987 expect ( mockKms . completeWalletTransfer ) . not . toHaveBeenCalled ( ) ;
923988 } ) ;
@@ -931,7 +996,7 @@ describe("_getOrMigrateWallet()", () => {
931996 const result = await _getOrMigrateWallet ( {
932997 stamper : makeStamper ( ) ,
933998 kms : mockKms as any ,
934- auth2Token : makeToken ( ) ,
999+ auth2Token : makeToken ( { walletUrn : "urn:phantom:wallet:wallet-1:0" } ) ,
9351000 } ) ;
9361001
9371002 expect ( mockKms . completeWalletTransfer ) . toHaveBeenCalledTimes ( 2 ) ;
0 commit comments