From 3636188e260e06f1a126403cd5d6a9b4d3b27b73 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 20 Apr 2021 05:35:22 +0000
Subject: [PATCH] Add property test for `categorizeUTxOEntries`.

We demonstrate that `uncategorize . categorize == id`.

In response to review feedback:

https://github.com/input-output-hk/cardano-wallet/pull/2618#issuecomment-822413404
---
 .../Primitive/Migration/PlanningSpec.hs       | 61 ++++++++++++++++++-
 1 file changed, 60 insertions(+), 1 deletion(-)

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 0ce53848517..d9de5e24236 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Migration/PlanningSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Migration/PlanningSpec.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DuplicateRecordFields #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE NamedFieldPuns #-}
@@ -26,6 +27,7 @@ import Cardano.Wallet.Primitive.Migration.Selection
 import Cardano.Wallet.Primitive.Migration.SelectionSpec
     ( MockInputId
     , MockTxConstraints (..)
+    , Pretty (..)
     , conjoinMap
     , counterexampleMap
     , genCoinRange
@@ -51,6 +53,8 @@ import Data.Set
     ( Set )
 import Fmt
     ( padLeftF, pretty )
+import GHC.Generics
+    ( Generic )
 import Test.Hspec
     ( Spec, describe, it )
 import Test.Hspec.Core.QuickCheck
@@ -69,7 +73,9 @@ import Test.QuickCheck
     , label
     , oneof
     , property
+    , shrinkList
     , withMaxSuccess
+    , (===)
     )
 
 import qualified Cardano.Wallet.Primitive.Migration.Selection as Selection
@@ -99,6 +105,8 @@ spec = describe "Cardano.Wallet.Primitive.Migration.PlanningSpec" $
 
     parallel $ describe "Categorizing UTxO entries" $ do
 
+        it "prop_categorizeUTxOEntries" $
+            property prop_categorizeUTxOEntries
         it "prop_categorizeUTxOEntry" $
             property prop_categorizeUTxOEntry
 
@@ -319,7 +327,58 @@ prop_createPlan mockArgs =
         ]
 
 --------------------------------------------------------------------------------
--- Categorizing UTxO entries
+-- Categorizing multiple UTxO entries
+--------------------------------------------------------------------------------
+
+data ArgsForCategorizeUTxOEntries = ArgsForCategorizeUTxOEntries
+    { mockConstraints :: MockTxConstraints
+    , mockEntries :: [(MockInputId, TokenBundle)]
+    }
+    deriving (Eq, Generic, Show)
+
+instance Arbitrary ArgsForCategorizeUTxOEntries where
+    arbitrary = genArgsForCategorizeUTxOEntries
+    shrink = shrinkArgsForCategorizeUTxOEntries
+
+genArgsForCategorizeUTxOEntries :: Gen ArgsForCategorizeUTxOEntries
+genArgsForCategorizeUTxOEntries = do
+    mockConstraints <- genMockTxConstraints
+    mockEntryCount <- choose (0, 100)
+    mockEntries <- replicateM mockEntryCount (genMockInput mockConstraints)
+    pure ArgsForCategorizeUTxOEntries {mockConstraints, mockEntries}
+
+shrinkArgsForCategorizeUTxOEntries
+    :: ArgsForCategorizeUTxOEntries -> [ArgsForCategorizeUTxOEntries]
+shrinkArgsForCategorizeUTxOEntries args = do
+    mockEntriesShrunk <- shrinkList shrinkInput (mockEntries args)
+    pure ArgsForCategorizeUTxOEntries
+        { mockConstraints = view #mockConstraints args
+        , mockEntries = mockEntriesShrunk
+        }
+  where
+    shrinkInput :: (MockInputId, TokenBundle) -> [(MockInputId, TokenBundle)]
+    shrinkInput (inputId, TokenBundle c m)
+        | c /= Coin 0, m /= mempty =
+            [(inputId, TokenBundle c mempty)]
+        | c /= Coin 0, m == mempty =
+            [(inputId, TokenBundle (Coin 0) mempty)]
+        | otherwise =
+            []
+
+prop_categorizeUTxOEntries :: Pretty ArgsForCategorizeUTxOEntries -> Property
+prop_categorizeUTxOEntries args =
+    Pretty (L.sortOn fst (uncategorizeUTxOEntries categorizedEntries))
+        === Pretty (L.sortOn fst mockEntries)
+  where
+    categorizedEntries = categorizeUTxOEntries constraints mockEntries
+    Pretty ArgsForCategorizeUTxOEntries
+        { mockConstraints
+        , mockEntries
+        } = args
+    constraints = unMockTxConstraints mockConstraints
+
+--------------------------------------------------------------------------------
+-- Categorizing individual UTxO entries
 --------------------------------------------------------------------------------
 
 data ArgsForCategorizeUTxOEntry = ArgsForCategorizeUTxOEntry