From 72c6c9573ddf4e8950ee05d5650682473ff7c051 Mon Sep 17 00:00:00 2001
From: Aniket Deshpande <git@aniketd.mailer.me>
Date: Fri, 10 Jan 2025 14:30:17 +0530
Subject: [PATCH] Add randomised reg-unreg certs for Imptests.

* `genRegTxCert` to choose between
  `RegTxCert` and `RegDepositTxCert`
* `genUnRegTxCert` to choose between
  `UnRegTxCert` and `UnRegDepositTxCert`
---
 eras/conway/impl/CHANGELOG.md                 |  1 +
 .../Cardano/Ledger/Conway/Imp/DelegSpec.hs    | 12 +++----
 .../Test/Cardano/Ledger/Conway/ImpTest.hs     | 35 +++++++++++++++++++
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/eras/conway/impl/CHANGELOG.md b/eras/conway/impl/CHANGELOG.md
index 937acea2750..c486e983407 100644
--- a/eras/conway/impl/CHANGELOG.md
+++ b/eras/conway/impl/CHANGELOG.md
@@ -9,6 +9,7 @@
 
 ### `testlib`
 
+* Add `genRegTxCert` and `genUnRegTxCert`. #4830
 * Add `Arbitrary` instance for `ConwayBbodyPredFailure` and `ConwayMempoolPredFailure`
 
 ## 1.18.0.0
diff --git a/eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/Imp/DelegSpec.hs b/eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/Imp/DelegSpec.hs
index 1571fb2eb4a..4bcc0712cac 100644
--- a/eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/Imp/DelegSpec.hs
+++ b/eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/Imp/DelegSpec.hs
@@ -49,10 +49,12 @@ spec = do
       expectedDeposit <- getsNES $ nesEsL . curPParamsEpochStateL . ppKeyDepositL
 
       freshKeyHash >>= \kh -> do
+        let cred = KeyHashObj kh
+        regTxCert <- genRegTxCert cred
         submitTx_ $
           mkBasicTx mkBasicTxBody
-            & bodyTxL . certsTxBodyL .~ [RegTxCert (KeyHashObj kh)]
-        expectRegistered (KeyHashObj kh)
+            & bodyTxL . certsTxBodyL .~ [regTxCert]
+        expectRegistered cred
 
       freshKeyHash >>= \kh -> do
         submitTx_ $
@@ -207,13 +209,11 @@ spec = do
       expectRegisteredRewardAddress otherRewardAccount
       submitAndExpireProposalToMakeReward otherStakeCred
       lookupReward otherStakeCred `shouldReturn` govActionDeposit
+      unRegTxCert <- genUnRegTxCert stakeCred
       submitTx_ . mkBasicTx $
         mkBasicTxBody
           & certsTxBodyL
-            .~ SSeq.fromList
-              -- https://github.com/IntersectMBO/formal-ledger-specifications/issues/636
-              -- we use this inplace of UnRegTxCert to make conformance-spec happy
-              [UnRegDepositTxCert stakeCred keyDeposit]
+            .~ SSeq.fromList [unRegTxCert]
           & withdrawalsTxBodyL
             .~ Withdrawals
               ( Map.fromList
diff --git a/eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/ImpTest.hs b/eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/ImpTest.hs
index 9c449bd19a9..0910850c9aa 100644
--- a/eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/ImpTest.hs
+++ b/eras/conway/impl/testlib/Test/Cardano/Ledger/Conway/ImpTest.hs
@@ -43,6 +43,8 @@ module Test.Cardano.Ledger.Conway.ImpTest (
   submitYesVote_,
   submitFailingVote,
   trySubmitVote,
+  genRegTxCert,
+  genUnRegTxCert,
   registerDRep,
   unRegisterDRep,
   updateDRep,
@@ -215,6 +217,7 @@ import Cardano.Ledger.Shelley.LedgerState (
  )
 import Cardano.Ledger.TxIn (TxId (..))
 import Cardano.Ledger.UMap (dRepMap)
+import qualified Cardano.Ledger.UMap as UMap
 import Cardano.Ledger.UTxO (EraUTxO, UTxO, balance, sumAllValue, txInsFilter)
 import Cardano.Ledger.Val (Val (..), (<->))
 import Control.Monad (forM)
@@ -390,6 +393,38 @@ unRegisterDRep drep = do
       & bodyTxL . certsTxBodyL
         .~ SSeq.singleton (UnRegDRepTxCert drep refund)
 
+genUnRegTxCert ::
+  forall era.
+  ( ShelleyEraImp era
+  , ConwayEraTxCert era
+  ) =>
+  Credential 'Staking ->
+  ImpTestM era (TxCert era)
+genUnRegTxCert stakingCredential = do
+  umap <- getsNES unifiedL
+  let mumapDeposit = UMap.rdDepositCoin <$> UMap.lookup stakingCredential (UMap.RewDepUView umap)
+  case mumapDeposit of
+    Nothing -> pure $ UnRegTxCert stakingCredential
+    Just umapDeposit ->
+      elements
+        [ UnRegTxCert stakingCredential
+        , UnRegDepositTxCert stakingCredential umapDeposit
+        ]
+
+genRegTxCert ::
+  forall era.
+  ( ShelleyEraImp era
+  , ConwayEraTxCert era
+  ) =>
+  Credential 'Staking ->
+  ImpTestM era (TxCert era)
+genRegTxCert stakingCredential =
+  oneof
+    [ pure $ RegTxCert stakingCredential
+    , RegDepositTxCert stakingCredential
+        <$> getsNES (nesEsL . curPParamsEpochStateL . ppKeyDepositL)
+    ]
+
 -- | Submit a transaction that updates a given DRep
 updateDRep ::
   forall era.