From 2445551d17a2ea7d448ca1cdb41a5e76f35fe22b Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Thu, 22 Apr 2021 11:52:19 +0000
Subject: [PATCH] Check that `createPlan` preserves the total non-ada token
 balance.

---
 .../Wallet/Primitive/Migration/PlanningSpec.hs  | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/Migration/PlanningSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/Migration/PlanningSpec.hs
index 9464d6c0c34..ccc0107387b 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Migration/PlanningSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Migration/PlanningSpec.hs
@@ -38,6 +38,8 @@ import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..) )
 import Cardano.Wallet.Primitive.Types.TokenBundle
     ( TokenBundle (..) )
+import Cardano.Wallet.Primitive.Types.TokenMap
+    ( TokenMap )
 import Control.Monad
     ( replicateM )
 import Data.Either
@@ -162,6 +164,9 @@ prop_createPlan_inner mockConstraints inputs reward =
         . verify
             (totalInputAda >= totalOutputAda)
             "ada is consumed and not created"
+        . verify
+            (totalInputTokenBalance == totalOutputTokenBalance)
+            "balance of non-ada tokens is preserved"
         . verify
             (totalFee result == totalFeeExpected)
             "total fee is correct"
@@ -200,6 +205,10 @@ prop_createPlan_inner mockConstraints inputs reward =
             "total fee expected"
         . report (totalFee result)
             "total fee actual"
+        . report totalInputTokenBalance
+            "total input token balance"
+        . report totalOutputTokenBalance
+            "total output token balance"
 
     makeStatistics
         = tabulate "Number of transactions required"
@@ -322,6 +331,14 @@ prop_createPlan_inner mockConstraints inputs reward =
     totalOutputAda =
         F.foldMap (view #coin . F.fold . view #outputs) (selections result)
 
+    totalInputTokenBalance :: TokenMap
+    totalInputTokenBalance =
+        F.foldMap (view #tokens . view #inputBalance) (selections result)
+
+    totalOutputTokenBalance :: TokenMap
+    totalOutputTokenBalance =
+        F.foldMap (view #tokens . F.fold . view #outputs) (selections result)
+
 --------------------------------------------------------------------------------
 -- Categorizing multiple UTxO entries
 --------------------------------------------------------------------------------