Refactor mkPState' and add tests

Author: aniketd

libs/cardano-ledger-api/src/Cardano/Ledger/Api/State/Query.hs

@@ -70,6 +70,7 @@ module Cardano.Ledger.Api.State.Query (
   -- * @GetPoolState@
   queryPoolState,
   PState' (..),
+  mkPState',
 
   -- * For testing
   getNextEpochCommitteeMembers,
@@ -463,21 +464,25 @@ instance DecShareCBOR PState' where
 instance DecShareCBOR PState' => DecCBOR PState' where
   decCBOR = decNoShareCBOR
 
+mkPState' ::
+  (forall b. Map.Map (KeyHash 'StakePool) b -> Map.Map (KeyHash 'StakePool) b) ->
+  PState era ->
+  PState'
+mkPState' f ps =
+  PState'
+    { psStakePoolParams = Map.mapWithKey stakePoolStateToPoolParams $ f $ psStakePools ps
+    , psFutureStakePoolParams = Map.mapWithKey stakePoolStateToPoolParams $ f $ psFutureStakePools ps
+    , psRetiring' = f $ psRetiring ps
+    , psDeposits = Map.map (fromCompact . spsDeposit) $ f $ psStakePools ps
+    }
+
 -- | Query the PState'. This is slightly different from the internal
 -- representation used by Ledger and is intended to resemble how the internal
 -- representation used to be.
 queryPoolState ::
   EraCertState era => NewEpochState era -> Maybe (Set (KeyHash 'StakePool)) -> PState'
 queryPoolState nes mPoolKeys =
   let pstate = nes ^. nesEsL . esLStateL . lsCertStateL . certPStateL
-      mkPState' :: (forall b. Map.Map (KeyHash 'StakePool) b -> Map.Map (KeyHash 'StakePool) b) -> PState'
-      mkPState' f =
-        PState'
-          { psStakePoolParams = Map.mapWithKey stakePoolStateToPoolParams $ f $ psStakePools pstate
-          , psFutureStakePoolParams = Map.mapWithKey stakePoolStateToPoolParams $ f $ psFutureStakePools pstate
-          , psRetiring' = f $ psRetiring pstate
-          , psDeposits = Map.map (fromCompact . spsDeposit) $ f $ psStakePools pstate
-          }
    in case mPoolKeys of
-        Nothing -> mkPState' id
-        Just keys -> mkPState' (`Map.restrictKeys` keys)
+        Nothing -> mkPState' id pstate
+        Just keys -> mkPState' (`Map.restrictKeys` keys) pstate

libs/cardano-ledger-api/test/Test/Cardano/Ledger/Api/State/QuerySpec.hs

@@ -16,8 +16,11 @@ import Cardano.Ledger.Api.State.Query (
   HotCredAuthStatus (..),
   MemberStatus (..),
   NextEpochChange (..),
+  PState',
   getNextEpochCommitteeMembers,
+  mkPState',
   queryCommitteeMembersState,
+  queryPoolState,
   queryStakePoolDelegsAndRewards,
  )
 import Cardano.Ledger.BaseTypes
@@ -54,12 +57,30 @@ import Test.Cardano.Ledger.Common
 import Test.Cardano.Ledger.Conway.Arbitrary ()
 import Test.Cardano.Ledger.Conway.Era (ShelleyEraTest)
 import Test.Cardano.Ledger.Core.Arbitrary (genValidUMapWithCreds)
+import Test.Cardano.Ledger.Core.Binary.RoundTrip (roundTripShareEraExpectation)
 import Test.Cardano.Ledger.Era (accountsFromUMap)
 import Test.Cardano.Ledger.Shelley.Arbitrary ()
 import Test.Cardano.Slotting.Numeric ()
 
 spec :: Spec
 spec = do
+  describe "API Types" $ do
+    describe "Roundtrip" $ do
+      prop "Shelley" $ roundTripShareEraExpectation @ShelleyEra @PState'
+      prop "Allegra" $ roundTripShareEraExpectation @AllegraEra @PState'
+      prop "Mary" $ roundTripShareEraExpectation @MaryEra @PState'
+      prop "Alonzo" $ roundTripShareEraExpectation @AlonzoEra @PState'
+      prop "Babbage" $ roundTripShareEraExpectation @BabbageEra @PState'
+      prop "Conway" $ roundTripShareEraExpectation @ConwayEra @PState'
+      prop "Dijkstra" $ roundTripShareEraExpectation @DijkstraEra @PState'
+    describe "QuerySpec" $ do
+      describe "queryPoolState" $ do
+        queryPoolStateSpec @ShelleyEra
+        queryPoolStateSpec @AllegraEra
+        queryPoolStateSpec @MaryEra
+        queryPoolStateSpec @AlonzoEra
+        queryPoolStateSpec @BabbageEra
+        queryPoolStateSpec @ConwayEra
   queryStakePoolDelegsAndRewardsSpec @ShelleyEra
   queryStakePoolDelegsAndRewardsSpec @AllegraEra
   queryStakePoolDelegsAndRewardsSpec @MaryEra
@@ -70,6 +91,12 @@ spec = do
     committeeMembersStateSpec @ConwayEra
     committeeMembersStateSpec @DijkstraEra
 
+queryPoolStateSpec :: forall era. ShelleyEraTest era => Spec
+queryPoolStateSpec =
+  prop ("mkPState' works as expected - " <> eraName @era) $ \(ps :: PState era) ->
+    let nes = def & nesEsL . esLStateL . lsCertStateL . certPStateL .~ ps
+     in queryPoolState nes Nothing `shouldBe` mkPState' id ps
+
 queryStakePoolDelegsAndRewardsSpec :: forall era. ShelleyEraTest era => Spec
 queryStakePoolDelegsAndRewardsSpec =
   describe (eraName @era) $ do

libs/cardano-ledger-api/testlib/Test/Cardano/Ledger/Api/Arbitrary.hs

@@ -2,9 +2,12 @@
 
 module Test.Cardano.Ledger.Api.Arbitrary () where
 
-import Cardano.Ledger.Api.State.Query (MemberStatus)
+import Cardano.Ledger.Api.State.Query (MemberStatus, PState', mkPState')
 import Test.Cardano.Ledger.Common
 import Test.Cardano.Ledger.Dijkstra.Arbitrary ()
 
 instance Arbitrary MemberStatus where
   arbitrary = arbitraryBoundedEnum
+
+instance Arbitrary PState' where
+  arbitrary = mkPState' id <$> arbitrary

libs/cardano-ledger-api/testlib/Test/Cardano/Ledger/Api/Upgrade.hs

@@ -17,6 +17,7 @@ import Cardano.Ledger.Core
 import Cardano.Ledger.MemoBytes (EqRaw (eqRaw))
 import Data.Default (Default (def))
 import qualified Prettyprinter as Pretty
+import Test.Cardano.Ledger.Api.Arbitrary ()
 import Test.Cardano.Ledger.Binary.RoundTrip
 import Test.Cardano.Ledger.Common
 import Test.Cardano.Ledger.Core.Arbitrary ()