Add integration tests for migration of reward balances

lib/core-integration/src/Test/Integration/Scenario/API/Shelley/Migrations.hs

@@ -57,14 +57,10 @@ import Data.Functor
     ( (<&>) )
 import Data.Generics.Internal.VL.Lens
     ( view, (^.) )
-import Data.Maybe
-    ( mapMaybe )
 import Data.Proxy
     ( Proxy )
 import Data.Quantity
     ( Quantity (..) )
-import Data.Word
-    ( Word64 )
 import Numeric.Natural
     ( Natural )
 import Test.Hspec
@@ -96,6 +92,7 @@ import Test.Integration.Framework.DSL
     , postWallet
     , randomAddresses
     , request
+    , rewardWallet
     , unsafeRequest
     , unsafeResponse
     , verify
@@ -114,6 +111,7 @@ import qualified Cardano.Wallet.Api.Link as Link
 import qualified Cardano.Wallet.Api.Types as ApiTypes
 import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
 import qualified Data.Foldable as F
+import qualified Data.List.NonEmpty as NE
 import qualified Data.Map.Strict as Map
 import qualified Data.Set as Set
 import qualified Network.HTTP.Types.Status as HTTP
@@ -196,7 +194,6 @@ spec = describe "SHELLEY_MIGRATIONS" $ do
                         (errMsg404NoWallet $ sourceWallet ^. walletId)
                     ]
 
-
     it "SHELLEY_CREATE_MIGRATION_PLAN_04 - \
         \Cannot create a plan for a wallet that only contains freeriders."
         $ \ctx -> runResourceT $ do
@@ -262,6 +259,49 @@ spec = describe "SHELLEY_MIGRATIONS" $ do
                 _ ->
                     error "Unable to compare plans."
 
+    it "SHELLEY_CREATE_MIGRATION_PLAN_06 - \
+        \Can create a migration plan for a wallet that has rewards."
+        $ \ctx -> runResourceT $ do
+            (sourceWallet, _sourceWalletMnemonic) <- rewardWallet ctx
+            -- Check that the source wallet has the expected balance.
+            request @ApiWallet ctx
+                (Link.getWallet @'Shelley sourceWallet) Default Empty
+                >>= flip verify
+                [ expectField (#balance . #reward . #getQuantity)
+                    (`shouldBe` 1_000_000_000_000)
+                , expectField (#balance . #available . #getQuantity)
+                    (`shouldBe`   100_000_000_000)
+                , expectField (#balance . #total . #getQuantity)
+                    (`shouldBe` 1_100_000_000_000)
+                ]
+            targetWallet <- emptyWallet ctx
+            targetAddresses <- listAddresses @n ctx targetWallet
+            let targetAddressIds = targetAddresses <&>
+                    (\(ApiTypes.ApiAddress addrId _ _) -> addrId)
+            let ep = Link.createMigrationPlan @'Shelley sourceWallet
+            response <- request @(ApiWalletMigrationPlan n) ctx ep Default
+                (Json [json|{addresses: #{targetAddressIds}}|])
+            verify response
+                [ expectResponseCode HTTP.status202
+                , expectField (#totalFee . #getQuantity)
+                    (`shouldBe` 139_300)
+                , expectField (#selections)
+                    ((`shouldBe` 1) . length)
+                , expectField (#selections)
+                    ((`shouldBe` 1) . length . view #withdrawals . NE.head)
+                , expectField (#selections)
+                    ((`shouldBe` 1_000_000_000_000)
+                        . view #getQuantity
+                        . view #amount
+                        . head
+                        . view #withdrawals
+                        . NE.head)
+                , expectField (#balanceSelected . #ada . #getQuantity)
+                    (`shouldBe` 1_100_000_000_000)
+                , expectField (#balanceLeftover . #ada . #getQuantity)
+                    (`shouldBe` 0)
+                ]
+
     describe "SHELLEY_MIGRATE_01 - \
         \After a migration operation successfully completes, the correct \
         \amounts eventually become available in the target wallet for an \
@@ -614,6 +654,81 @@ spec = describe "SHELLEY_MIGRATIONS" $ do
                     )
                 ]
 
+    Hspec.it "SHELLEY_MIGRATE_09 - \
+        \Can migrate a wallet that has rewards."
+        $ \ctx -> runResourceT @IO $ do
+
+            -- Create a source wallet with rewards:
+            (sourceWallet, _sourceWalletMnemonic) <- rewardWallet ctx
+
+            -- Check that the source wallet has the expected balance:
+            let expectedAdaBalanceAvailable =   100_000_000_000
+            let expectedAdaBalanceReward    = 1_000_000_000_000
+            let expectedAdaBalanceTotal     = 1_100_000_000_000
+            request @ApiWallet ctx
+                (Link.getWallet @'Shelley sourceWallet) Default Empty
+                >>= flip verify
+                [ expectField (#balance . #available . #getQuantity)
+                    (`shouldBe` expectedAdaBalanceAvailable)
+                , expectField (#balance . #reward . #getQuantity)
+                    (`shouldBe` expectedAdaBalanceReward)
+                , expectField (#balance . #total . #getQuantity)
+                    (`shouldBe` expectedAdaBalanceTotal)
+                ]
+
+            -- Create an empty target wallet:
+            targetWallet <- emptyWallet ctx
+            targetAddresses <- listAddresses @n ctx targetWallet
+            let targetAddressIds = targetAddresses <&>
+                    (\(ApiTypes.ApiAddress addrId _ _) -> addrId)
+
+            -- Perform a migration:
+            let ep = Link.migrateWallet @'Shelley sourceWallet
+            responseMigrate <- request @[ApiTransaction n] ctx ep Default $
+                Json [json|
+                    { passphrase: #{fixturePassphrase}
+                    , addresses: #{targetAddressIds}
+                    }|]
+
+            -- Verify the fee is as expected:
+            let expectedFee = 139_300
+            verify responseMigrate
+                [ expectResponseCode HTTP.status202
+                , expectField id ((`shouldBe` 1) . length)
+                , expectField id
+                    $ (`shouldBe` expectedFee)
+                    . sum
+                    . fmap apiTransactionFee
+                ]
+
+            -- Check that funds become available in the target wallet:
+            let expectedTargetBalance = expectedAdaBalanceTotal - expectedFee
+            eventually "Target wallet balance reaches the expected amount." $ do
+                request @ApiWallet ctx
+                    (Link.getWallet @'Shelley targetWallet) Default Empty
+                    >>= flip verify
+                    [ expectField
+                        (#balance . #available)
+                        (`shouldBe` Quantity expectedTargetBalance)
+                    , expectField
+                        (#balance . #total)
+                        (`shouldBe` Quantity expectedTargetBalance)
+                    ]
+
+            -- Check that the source wallet has been depleted:
+            eventually "Source wallet balance is depleted." $ do
+                request @ApiWallet ctx
+                    (Link.getWallet @'Shelley sourceWallet) Default Empty
+                    >>= flip verify
+                    [ expectResponseCode HTTP.status200
+                    , expectField
+                        (#balance . #available)
+                        (`shouldBe` Quantity 0)
+                    , expectField
+                        (#balance . #total)
+                        (`shouldBe` Quantity 0)
+                    ]
+
     Hspec.it "SHELLEY_MIGRATE_MULTI_ASSET_01 - \
         \Can migrate a multi-asset wallet."
         $ \ctx -> runResourceT @IO $ do
@@ -738,19 +853,8 @@ spec = describe "SHELLEY_MIGRATIONS" $ do
                 ]
   where
     -- Compute the fee associated with an API transaction.
-    apiTransactionFee :: ApiTransaction n -> Word64
-    apiTransactionFee t =
-        inputBalance t - outputBalance t
-      where
-        inputBalance = fromIntegral
-            . sum
-            . fmap (view (#amount . #getQuantity))
-            . mapMaybe ApiTypes.source
-            . view #inputs
-        outputBalance = fromIntegral
-            . sum
-            . fmap (view (#amount . #getQuantity))
-            . view #outputs
+    apiTransactionFee :: ApiTransaction n -> Natural
+    apiTransactionFee = view (#fee . #getQuantity)
 
     migrateWallet
         :: Context

lib/core/src/Cardano/Wallet.hs

@@ -1873,10 +1873,30 @@ migrationPlanToSelectionWithdrawals plan rewardWithdrawal outputAddressesToCycle
             { inputsSelected = view #inputIds migrationSelection
             , outputsCovered
             , utxoRemaining = UTxOIndex.empty
-            , extraCoinSource = Nothing
+            , extraCoinSource
             , changeGenerated = []
             }
 
+        -- NOTE:
+        --
+        -- Due to a quirk of history, we need to populate the 'extraCoinSource'
+        -- field with the reward withdrawal amount, since the transaction layer
+        -- uses the 'selectionDelta' function to calculate the final fee, and
+        -- that particular function doesn't know about reward withdrawals.
+        --
+        -- This is non-ideal, because we're returning the reward withdrawal
+        -- amount in two places in the output of this function.
+        --
+        -- In future, it would be better to return a single record whose fields
+        -- more closely resemble exactly what is needed to build a transaction,
+        -- and have the transaction layer calculate the actual fee based only
+        -- on the contents of that record.
+        --
+        extraCoinSource =
+            if (view #rewardWithdrawal migrationSelection) > Coin 0
+            then Just (view #rewardWithdrawal migrationSelection)
+            else Nothing
+
         withdrawal =
             if (view #rewardWithdrawal migrationSelection) > Coin 0
             then rewardWithdrawal

lib/core/src/Cardano/Wallet/Api/Server.hs

@@ -472,8 +472,6 @@ import Data.Time
     ( UTCTime )
 import Data.Type.Equality
     ( (:~:) (..), type (==), testEquality )
-import Data.Void
-    ( Void )
 import Data.Word
     ( Word32 )
 import Fmt
@@ -2040,13 +2038,16 @@ createMigrationPlan
         , WalletKey k
         )
     => ctx
+    -> Maybe ApiWithdrawalPostData
+        -- ^ What type of reward withdrawal to attempt
     -> ApiT WalletId
         -- ^ Source wallet
     -> ApiWalletMigrationPlanPostData n
         -- ^ Target addresses
     -> Handler (ApiWalletMigrationPlan n)
-createMigrationPlan ctx (ApiT wid) postData = do
-    (rewardWithdrawal, _) <- mkRewardAccountBuilder @_ @s @_ @n ctx wid Nothing
+createMigrationPlan ctx withdrawalType (ApiT wid) postData = do
+    (rewardWithdrawal, _) <-
+        mkRewardAccountBuilder @_ @s @_ @n ctx wid withdrawalType
     withWorkerCtx ctx wid liftE liftE $ \wrk -> liftHandler $ do
         (wallet, _, _) <- withExceptT ErrCreateMigrationPlanNoSuchWallet $
             W.readWallet wrk wid
@@ -2079,7 +2080,7 @@ mkApiWalletMigrationPlan s addresses rewardWithdrawal plan =
     maybeSelections = fmap mkApiCoinSelectionForMigration <$> maybeUnsignedTxs
 
     maybeSelectionWithdrawals
-        :: Maybe (NonEmpty (SelectionResult Void, Withdrawal))
+        :: Maybe (NonEmpty (W.SelectionResultWithoutChange, Withdrawal))
     maybeSelectionWithdrawals
         = W.migrationPlanToSelectionWithdrawals plan rewardWithdrawal
         $ getApiT . fst <$> addresses
@@ -2099,10 +2100,15 @@ mkApiWalletMigrationPlan s addresses rewardWithdrawal plan =
         & mkApiWalletMigrationBalance
 
     balanceSelected :: ApiWalletMigrationBalance
-    balanceSelected = plan
-        & view #selections
-        & F.foldMap (view #inputBalance)
-        & mkApiWalletMigrationBalance
+    balanceSelected = mkApiWalletMigrationBalance $
+        TokenBundle.fromCoin balanceRewardWithdrawal <> balanceUTxO
+      where
+        balanceUTxO = plan
+            & view #selections
+            & F.foldMap (view #inputBalance)
+        balanceRewardWithdrawal = plan
+            & view #selections
+            & F.foldMap (view #rewardWithdrawal)
 
     mkApiCoinSelectionForMigration unsignedTx =
         mkApiCoinSelection [] Nothing Nothing unsignedTx
@@ -2125,12 +2131,14 @@ migrateWallet
         , WalletKey k
         )
     => ctx
+    -> Maybe ApiWithdrawalPostData
+        -- ^ What type of reward withdrawal to attempt
     -> ApiT WalletId
     -> ApiWalletMigrationPostData n p
     -> Handler (NonEmpty (ApiTransaction n))
-migrateWallet ctx (ApiT wid) postData = do
+migrateWallet ctx withdrawalType (ApiT wid) postData = do
     (rewardWithdrawal, mkRewardAccount) <-
-        mkRewardAccountBuilder @_ @s @_ @n ctx wid Nothing
+        mkRewardAccountBuilder @_ @s @_ @n ctx wid withdrawalType
     withWorkerCtx ctx wid liftE liftE $ \wrk -> do
         plan <- liftHandler $ W.createMigrationPlan wrk wid rewardWithdrawal
         txTimeToLive <- liftIO $ W.getTxExpiry ti Nothing

lib/shelley/src/Cardano/Wallet/Shelley/Api/Server.hs

@@ -132,6 +132,7 @@ import Cardano.Wallet.Api.Types
     , ApiSelectCoinsData (..)
     , ApiStakePool
     , ApiT (..)
+    , ApiWithdrawalPostData (..)
     , HealthCheckSMASH (..)
     , MaintenanceAction (..)
     , SettingsPutData (..)
@@ -295,8 +296,8 @@ server byron icarus shelley multisig spl ntp =
 
     shelleyMigrations :: Server (ShelleyMigrations n)
     shelleyMigrations =
-             createMigrationPlan @_ @_ shelley
-        :<|> migrateWallet shelley
+             createMigrationPlan @_ @_ shelley (Just SelfWithdrawal)
+        :<|> migrateWallet shelley (Just SelfWithdrawal)
 
     stakePools :: Server (StakePools n ApiStakePool)
     stakePools =
@@ -449,12 +450,12 @@ server byron icarus shelley multisig spl ntp =
     byronMigrations :: Server (ByronMigrations n)
     byronMigrations =
              (\wid postData -> withLegacyLayer wid
-                (byron , createMigrationPlan @_ @_ byron wid postData)
-                (icarus, createMigrationPlan @_ @_ icarus wid postData)
+                (byron , createMigrationPlan @_ @_ byron Nothing wid postData)
+                (icarus, createMigrationPlan @_ @_ icarus Nothing wid postData)
              )
         :<|> (\wid m -> withLegacyLayer wid
-                (byron , migrateWallet byron wid m)
-                (icarus, migrateWallet icarus wid m)
+                (byron , migrateWallet byron Nothing wid m)
+                (icarus, migrateWallet icarus Nothing wid m)
              )
 
     network' :: Server Network