From 398576b3509bbaed80d5cd2a078594385d17069e Mon Sep 17 00:00:00 2001
From: brentstone <brent@heliax.dev>
Date: Thu, 27 Jul 2023 15:16:51 -0400
Subject: [PATCH] WIP fix more tests

---
 proof_of_stake/src/lib.rs   | 35 ++++++++++++++++++++++++++++++-----
 proof_of_stake/src/tests.rs |  6 +++---
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/proof_of_stake/src/lib.rs b/proof_of_stake/src/lib.rs
index c7ef4053076..1deeb8564c6 100644
--- a/proof_of_stake/src/lib.rs
+++ b/proof_of_stake/src/lib.rs
@@ -4800,17 +4800,23 @@ where
     }
 
     // Update the validator stakes
+    // TODO: need to generalize for arbitrary pipeline_len!!
     for (validator, slash_amounts) in map_validator_slash {
         // Use the `slash_amounts` to deduct from the deltas for the current and
         // the next epochs (no adjusting at pipeline)
+
+        let next_stake_pre = read_validator_stake(
+            storage,
+            &params,
+            &validator,
+            current_epoch.next(),
+        )?
+        .unwrap_or_default()
+        .change();
         let delta_cur = slash_amounts
             .get(&current_epoch)
             .cloned()
-            .expect("Expected a map element for the current epoch");
-        let delta_next = slash_amounts
-            .get(&current_epoch.next())
-            .cloned()
-            .expect("Expected a map element for the next epoch");
+            .unwrap_or_default();
         update_validator_deltas(
             storage,
             &validator,
@@ -4818,6 +4824,24 @@ where
             current_epoch,
             0,
         )?;
+        update_total_deltas(storage, -delta_cur, current_epoch, 0)?;
+
+        let next_stake_post = next_stake_pre
+            - slash_amounts
+                .get(&current_epoch.next())
+                .cloned()
+                .unwrap_or_default();
+
+        let next_stake_read = read_validator_stake(
+            storage,
+            &params,
+            &validator,
+            current_epoch.next(),
+        )?
+        .unwrap_or_default()
+        .change();
+        let delta_next = next_stake_post - next_stake_read;
+
         update_validator_deltas(
             storage,
             &validator,
@@ -4825,6 +4849,7 @@ where
             current_epoch.next(),
             0,
         )?;
+        update_total_deltas(storage, -delta_next, current_epoch.next(), 0)?;
     }
 
     Ok(())
diff --git a/proof_of_stake/src/tests.rs b/proof_of_stake/src/tests.rs
index 61a982dadc4..0a402e4daa7 100644
--- a/proof_of_stake/src/tests.rs
+++ b/proof_of_stake/src/tests.rs
@@ -3811,8 +3811,8 @@ fn test_slash_validator() {
     let infraction_stake = token::Change::from(23);
 
     let initial_stakes = BTreeMap::from_iter([
+        (Epoch(10), infraction_stake),
         (Epoch(11), infraction_stake),
-        (Epoch(12), infraction_stake),
     ]);
 
     let current_epoch = Epoch(10);
@@ -4134,8 +4134,8 @@ fn test_slash_validator() {
     assert_eq!(
         res,
         BTreeMap::from_iter([
-            (current_epoch.next(), token::Change::zero()),
-            (current_epoch.next().next(), token::Change::zero())
+            (current_epoch, token::Change::zero()),
+            (current_epoch.next(), token::Change::zero())
         ])
     );
 }