From 118fd90c74e3f7cdb71babd0b2760441d25e7839 Mon Sep 17 00:00:00 2001
From: Jesse Abramowitz <jesse@entropy.xyz>
Date: Fri, 18 Oct 2024 13:32:07 -0400
Subject: [PATCH] correct truncation

---
 pallets/staking/src/lib.rs   |  9 +++------
 pallets/staking/src/tests.rs | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/pallets/staking/src/lib.rs b/pallets/staking/src/lib.rs
index 64e0a7e83..2500ec229 100644
--- a/pallets/staking/src/lib.rs
+++ b/pallets/staking/src/lib.rs
@@ -735,12 +735,9 @@ pub mod pallet {
                 } else {
                     remove_index_len = remove_indexs.len();
                     let remove_indexs_reversed: Vec<_> = remove_indexs.iter().rev().collect();
-                    // TODO: Only remove up to threhsold https://github.com/entropyxyz/entropy-core/issues/1114
-                    let truncated = if remove_index_len > signers_info.threshold as usize {
-                        remove_indexs_reversed[..signers_info.threshold as usize].to_vec()
-                    } else {
-                        remove_indexs_reversed
-                    };
+                    let truncated = remove_indexs_reversed
+                        [..(signers_info.total_signers as usize - signers_info.threshold as usize)]
+                        .to_vec();
                     for remove_index in truncated {
                         current_signers.remove(*remove_index);
                     }
diff --git a/pallets/staking/src/tests.rs b/pallets/staking/src/tests.rs
index 5493e758d..4b64807c8 100644
--- a/pallets/staking/src/tests.rs
+++ b/pallets/staking/src/tests.rs
@@ -476,10 +476,9 @@ fn it_tests_new_session_handler() {
             last_session_change: 0,
         });
 
-        assert_ok!(Staking::new_session_handler(&[1, 2, 3]));
-        // takes signers original (5,6) not in validators pops off both, adds (fake randomness in mock so adds 1 and 3)
-        assert_eq!(Staking::next_signers().unwrap().next_signers, vec![1, 3]);
-
+        assert_ok!(Staking::new_session_handler(&[1, 5, 6]));
+        // takes signers original (5,6) pops off one and adds in new validator
+        assert_eq!(Staking::next_signers().unwrap().next_signers, vec![6, 1]);
         assert_eq!(
             Staking::reshare_data().block_number,
             101,
@@ -487,7 +486,7 @@ fn it_tests_new_session_handler() {
         );
         assert_eq!(
             Staking::reshare_data().new_signers,
-            vec![1u64.encode(), 3u64.encode()],
+            vec![1u64.encode()],
             "Check reshare next signer up is 3"
         );
         assert_eq!(
@@ -522,6 +521,32 @@ fn it_tests_new_session_handler() {
     });
 }
 
+#[test]
+fn it_tests_new_session_handler_truncating() {
+    new_test_ext().execute_with(|| {
+        // Start with current validators as 5 and 6 based off the Mock `GenesisConfig`.
+        Signers::<Test>::put(vec![7, 8]);
+        System::set_block_number(100);
+        pallet_parameters::SignersInfo::<Test>::put(SignersSize {
+            total_signers: 2,
+            threshold: 2,
+            last_session_change: 0,
+        });
+        // test truncates none if t and n = 0
+        assert_ok!(Staking::new_session_handler(&[1, 2, 3]));
+        assert_eq!(Staking::next_signers().unwrap().next_signers, vec![7, 8]);
+
+        pallet_parameters::SignersInfo::<Test>::put(SignersSize {
+            total_signers: 2,
+            threshold: 1,
+            last_session_change: 0,
+        });
+        // test truncates 1 if n - t = 1
+        assert_ok!(Staking::new_session_handler(&[1, 2, 3]));
+        assert_eq!(Staking::next_signers().unwrap().next_signers, vec![7, 1]);
+    });
+}
+
 #[test]
 fn it_tests_new_session_handler_signer_size_changes() {
     new_test_ext().execute_with(|| {