From 22b3312324dad0ff8ad7904d0fef00c0df70bc91 Mon Sep 17 00:00:00 2001
From: brentstone <brentstone01@gmail.com>
Date: Thu, 8 Feb 2024 19:51:18 -0500
Subject: [PATCH] fix testing

---
 .../lib/node/ledger/shell/finalize_block.rs   | 12 +++++++++---
 .../proof_of_stake/src/tests/state_machine.rs | 19 +++++++++++++------
 .../src/tests/state_machine_v2.rs             | 19 +++++++++++++------
 3 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/crates/apps/src/lib/node/ledger/shell/finalize_block.rs b/crates/apps/src/lib/node/ledger/shell/finalize_block.rs
index 5c92d3f8db..2fb4e7869b 100644
--- a/crates/apps/src/lib/node/ledger/shell/finalize_block.rs
+++ b/crates/apps/src/lib/node/ledger/shell/finalize_block.rs
@@ -3953,6 +3953,7 @@ mod test_finalize_block {
         let enqueued_slash = enqueued_slashes_handle()
             .at(&processing_epoch)
             .at(&val1.address)
+            .at(&height.0)
             .front(&shell.state)
             .unwrap()
             .unwrap();
@@ -3990,7 +3991,7 @@ mod test_finalize_block {
                     address: pkh1,
                     power: Default::default(),
                 },
-                height: height.try_into().unwrap(),
+                height: height.next_height().try_into().unwrap(),
                 time: tendermint::Time::unix_epoch(),
                 total_voting_power: Default::default(),
             },
@@ -4023,8 +4024,13 @@ mod test_finalize_block {
             .at(&processing_epoch.next())
             .at(&val1.address);
 
-        assert_eq!(enqueued_slashes_8.len(&shell.state).unwrap(), 2_u64);
-        assert_eq!(enqueued_slashes_9.len(&shell.state).unwrap(), 1_u64);
+        let num_enqueued_8 =
+            enqueued_slashes_8.iter(&shell.state).unwrap().count();
+        let num_enqueued_9 =
+            enqueued_slashes_9.iter(&shell.state).unwrap().count();
+
+        assert_eq!(num_enqueued_8, 2);
+        assert_eq!(num_enqueued_9, 1);
         let last_slash =
             namada_proof_of_stake::storage::read_validator_last_slash_epoch(
                 &shell.state,
diff --git a/crates/proof_of_stake/src/tests/state_machine.rs b/crates/proof_of_stake/src/tests/state_machine.rs
index 7bf14ce6d5..ed5c47c78f 100644
--- a/crates/proof_of_stake/src/tests/state_machine.rs
+++ b/crates/proof_of_stake/src/tests/state_machine.rs
@@ -737,6 +737,7 @@ impl StateMachineTest for ConcretePosState {
                     &params,
                     current_epoch,
                     infraction_epoch,
+                    height,
                     slash_type,
                     &address,
                 );
@@ -1355,6 +1356,7 @@ impl ConcretePosState {
         params: &PosParams,
         current_epoch: Epoch,
         infraction_epoch: Epoch,
+        infraction_height: u64,
         slash_type: SlashType,
         validator: &Address,
     ) {
@@ -1391,6 +1393,7 @@ impl ConcretePosState {
         let slash = enqueued_slashes_handle()
             .at(&processing_epoch)
             .at(validator)
+            .at(&infraction_height)
             .back(&self.s)
             .unwrap();
         if let Some(slash) = slash {
@@ -5211,12 +5214,16 @@ fn arb_slash(state: &AbstractPosState) -> impl Strategy<Value = Transition> {
         .checked_sub(state.params.unbonding_len)
         .unwrap_or_default()..=current_epoch)
         .prop_map(Epoch::from);
-    (arb_validator, arb_type, arb_epoch).prop_map(
-        |(validator, slash_type, infraction_epoch)| Transition::Misbehavior {
-            address: validator,
-            slash_type,
-            infraction_epoch,
-            height: 0,
+    let arb_height = 0_u64..10_000_u64;
+
+    (arb_validator, arb_type, arb_epoch, arb_height).prop_map(
+        |(validator, slash_type, infraction_epoch, height)| {
+            Transition::Misbehavior {
+                address: validator,
+                slash_type,
+                infraction_epoch,
+                height,
+            }
         },
     )
 }
diff --git a/crates/proof_of_stake/src/tests/state_machine_v2.rs b/crates/proof_of_stake/src/tests/state_machine_v2.rs
index 9a2d1234d7..45a51389c8 100644
--- a/crates/proof_of_stake/src/tests/state_machine_v2.rs
+++ b/crates/proof_of_stake/src/tests/state_machine_v2.rs
@@ -2690,6 +2690,7 @@ impl StateMachineTest for ConcretePosState {
                     &params,
                     current_epoch,
                     infraction_epoch,
+                    height,
                     slash_type,
                     &address,
                 );
@@ -3057,6 +3058,7 @@ impl ConcretePosState {
         params: &PosParams,
         current_epoch: Epoch,
         infraction_epoch: Epoch,
+        infraction_height: u64,
         slash_type: SlashType,
         validator: &Address,
     ) {
@@ -3101,6 +3103,7 @@ impl ConcretePosState {
         let slash = enqueued_slashes_handle()
             .at(&processing_epoch)
             .at(validator)
+            .at(&infraction_height)
             .back(&self.s)
             .unwrap();
         if let Some(slash) = slash {
@@ -4588,12 +4591,16 @@ fn arb_slash(state: &AbstractPosState) -> impl Strategy<Value = Transition> {
         .checked_sub(state.params.unbonding_len)
         .unwrap_or_default()..=current_epoch)
         .prop_map(Epoch::from);
-    (arb_validator, arb_type, arb_epoch).prop_map(
-        |(validator, slash_type, infraction_epoch)| Transition::Misbehavior {
-            address: validator,
-            slash_type,
-            infraction_epoch,
-            height: 0,
+    let arb_height = 0_u64..10_000_u64;
+
+    (arb_validator, arb_type, arb_epoch, arb_height).prop_map(
+        |(validator, slash_type, infraction_epoch, height)| {
+            Transition::Misbehavior {
+                address: validator,
+                slash_type,
+                infraction_epoch,
+                height,
+            }
         },
     )
 }