Fix testnet queries

Author: carbolymer

cardano-testnet/src/Testnet/Components/Query.hs

@@ -39,7 +39,7 @@ module Testnet.Components.Query
   , getProtocolParams
   , getGovActionLifetime
   , getKeyDeposit
-  , getDelegationState
+  , getAccountsStates
   , getTxIx
   ) where
 
@@ -50,12 +50,11 @@ import qualified Cardano.Api.UTxO as Utxo
 
 import           Cardano.Ledger.Api (ConwayGovState)
 import qualified Cardano.Ledger.Api as L
+import qualified Cardano.Ledger.Api.State.Query as SQ
 import qualified Cardano.Ledger.Conway.Governance as L
 import qualified Cardano.Ledger.Conway.PParams as L
 import qualified Cardano.Ledger.Shelley.LedgerState as L
-import qualified Cardano.Ledger.UMap as L
-import qualified Cardano.Ledger.Api.State.Query as SQ
-import qualified Data.Set as Set
+import qualified Cardano.Ledger.State as L
 
 import           Prelude
 
@@ -69,6 +68,7 @@ import qualified Data.Map as Map
 import           Data.Map.Strict (Map)
 import           Data.Maybe
 import           Data.Ord (Down (..))
+import qualified Data.Set as Set
 import           Data.Text (Text)
 import qualified Data.Text as T
 import           Data.Type.Equality
@@ -412,7 +412,7 @@ checkDRepState epochStateView@EpochStateView{nodeConfigPath, socketPath} sbe f =
       $ \(AnyNewEpochState actualEra newEpochState _) _slotNumber _blockNumber -> do
         Refl <- either error pure $ assertErasEqual sbe actualEra
         let dreps =
-              shelleyBasedEraConstraints sbe 
+              shelleyBasedEraConstraints sbe
                 $ SQ.queryDRepState newEpochState Set.empty
         case f dreps of
           Nothing -> pure ConditionNotMet
@@ -467,7 +467,7 @@ getTreasuryValue
   -> m L.Coin -- ^ The current value of the treasury
 getTreasuryValue epochStateView = withFrozenCallStack $ do
   AnyNewEpochState _ newEpochState _ <- getEpochState epochStateView
-  pure $ newEpochState ^. L.nesEpochStateL . L.epochStateTreasuryL
+  pure $ newEpochState ^. L.nesEpochStateL . L.treasuryL
 
 -- | Obtain minimum deposit amount for governance action from node
 getMinGovActionDeposit
@@ -590,19 +590,21 @@ getKeyDeposit epochStateView ceo = conwayEraOnwardsConstraints ceo $ do
                       . L.ppKeyDepositL
 
 
--- | Returns delegation state from the epoch state.
-getDelegationState :: (H.MonadAssertion m, MonadTest m, MonadIO m)
+-- | Returns staking accounts state
+getAccountsStates :: (H.MonadAssertion m, MonadTest m, MonadIO m)
   => EpochStateView
-  -> m L.StakeCredentials
-getDelegationState epochStateView = do
+  -> ShelleyBasedEra era
+  -> m (Map (L.Credential L.Staking) (L.AccountState (ShelleyLedgerEra era)))
+getAccountsStates epochStateView sbe' = shelleyBasedEraConstraints sbe' $ do
   AnyNewEpochState sbe newEpochState _ <- getEpochState epochStateView
-  let pools = shelleyBasedEraConstraints sbe $ newEpochState
-                ^. L.nesEsL
-                 . L.esLStateL
-                 . L.lsCertStateL
-                 . L.certDStateL
-
-  pure $ L.toStakeCredentials pools
+  Refl <- H.nothingFail $ testEquality sbe sbe'
+  pure $ newEpochState
+          ^. L.nesEsL
+           . L.esLStateL
+           . L.lsCertStateL
+           . L.certDStateL
+           . L.accountsL
+           . L.accountsMapL
 
 -- | Returns the transaction index of a transaction with a given amount and ID.
 getTxIx :: forall m era. HasCallStack

cardano-testnet/src/Testnet/Process/Cli/SPO.hs

@@ -18,8 +18,8 @@ module Testnet.Process.Cli.SPO
 import           Cardano.Api hiding (cardanoEra)
 import qualified Cardano.Api.Ledger as L
 
-import qualified Cardano.Ledger.Api.State.Query as L
 import qualified Cardano.Ledger.Shelley.LedgerState as L
+import qualified Cardano.Ledger.State as L
 import qualified Cardano.Ledger.UMap as L
 
 import           Control.Monad
@@ -130,30 +130,59 @@ checkStakeKeyRegistered tempAbsP nodeConfigFile sPath terminationEpoch execConfi
                     ]
  where
   handler :: StakeAddress -> AnyNewEpochState -> SlotNo -> BlockNo -> StateT DelegationsAndRewards IO ConditionResult
-  handler (StakeAddress network sCred) (AnyNewEpochState sbe newEpochState _) _ _ =
-    let umap = shelleyBasedEraConstraints sbe $ newEpochState ^. L.nesEsL . L.epochStateUMapL
-        dag = L.filterStakePoolDelegsAndRewards umap $ Set.singleton sCred
-        allStakeCredentials = umap ^. L.umElemsL -- This does not include pointer addresses
-        delegsAndRewards = shelleyBasedEraConstraints sbe $ toDelegationsAndRewards network sbe dag
-    in case Map.lookup sCred allStakeCredentials of
-         Nothing -> return ConditionNotMet
-         Just _ -> StateT.put delegsAndRewards >> return ConditionMet
+  handler (StakeAddress network sCred) (AnyNewEpochState sbe newEpochState _) _ _ = shelleyBasedEraConstraints sbe $ do
+    let accountsMap = newEpochState
+                        ^. L.nesEsL
+                         . L.esLStateL
+                         . L.lsCertStateL
+                         . L.certDStateL
+                         . L.accountsL
+                         . L.accountsMapL
+
+
+    -- let umap = shelleyBasedEraConstraints sbe $ newEpochState ^. L.nesEsL . L.epochStateUMapL
+    --     dag = L.filterStakePoolDelegsAndRewards umap $ Set.singleton sCred
+    --     allStakeCredentials = umap ^. L.umElemsL -- This does not include pointer addresses
+    --     delegsAndRewards = shelleyBasedEraConstraints sbe $ toDelegationsAndRewards network sbe dag
+
+    case Map.lookup sCred accountsMap of
+      Nothing -> pure ConditionNotMet
+      Just _ -> do
+        StateT.put $ toDelegationsAndRewards sbe network accountsMap
+        pure ConditionMet
 
   toDelegationsAndRewards
-    :: L.Network
-    -> ShelleyBasedEra era
-    -> (Map (L.Credential L.Staking) (L.KeyHash L.StakePool), Map (L.Credential 'L.Staking) L.Coin)
+    :: ShelleyBasedEra era
+    -> L.Network
+    -> Map (L.Credential L.Staking) (L.AccountState (ShelleyLedgerEra era))
     -> DelegationsAndRewards
-  toDelegationsAndRewards n _ (delegationMap, rewardsMap) =
-    let apiDelegationMap = Map.map toApiPoolId $ Map.mapKeys (toApiStakeAddress n) delegationMap
-        apiRewardsMap = Map.mapKeys (toApiStakeAddress n) rewardsMap
-    in DelegationsAndRewards (apiRewardsMap, apiDelegationMap)
+  toDelegationsAndRewards sbe n accountsMap = do
+    let accountsMap' = Map.mapKeys (toApiStakeAddress n) accountsMap
+    let apiDelegationMap = Map.mapMaybe (toApiPoolId sbe) accountsMap'
+        apiRewardsMap = Map.map (toBalance sbe) accountsMap'
+    DelegationsAndRewards (apiRewardsMap, apiDelegationMap)
+
+  -- toApiPoolId ::  L.KeyHash L.StakePool -> PoolId
+  toApiPoolId :: ShelleyBasedEra era
+              -> L.AccountState (ShelleyLedgerEra era)
+              -> Maybe PoolId
+  toApiPoolId sbe accountState =
+    fmap StakePoolKeyHash $
+      shelleyBasedEraConstraints sbe $
+        accountState ^. L.stakePoolDelegationAccountStateL
+
+
+  toBalance :: ShelleyBasedEra era
+            -> L.AccountState (ShelleyLedgerEra era)
+            -> L.Coin
+  toBalance sbe accountState =
+    shelleyBasedEraConstraints sbe $
+      accountState ^. L.balanceAccountStateL . to L.fromCompact
+
 
 toApiStakeAddress :: L.Network -> L.Credential 'L.Staking -> StakeAddress
 toApiStakeAddress = StakeAddress
 
-toApiPoolId ::  L.KeyHash L.StakePool -> PoolId
-toApiPoolId = StakePoolKeyHash
 
 createStakeDelegationCertificate
   :: (MonadTest m, MonadCatch m, MonadIO m, HasCallStack)