WIP fix more tests

anoma · Jul 27, 2023 · 398576b · 398576b
1 parent e9b3aa6
commit 398576b
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 8 deletions.
diff --git a/proof_of_stake/src/lib.rs b/proof_of_stake/src/lib.rs
@@ -4800,31 +4800,56 @@ 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,
             -delta_cur,
             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,
             -delta_next,
             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
@@ -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())
         ])
     );
 }