From c4c34f6e17a948b88f83c2487362a9c9dd84909e Mon Sep 17 00:00:00 2001
From: Ethan Frey <ethanfrey@noreply.github.com>
Date: Fri, 10 May 2024 11:24:11 +0200
Subject: [PATCH] Native-staking* and vault tests pass

---
 .../native-staking-proxy/src/multitest.rs     | 37 +++++++++++--------
 .../provider/native-staking/src/multitest.rs  | 10 +++++
 contracts/provider/vault/src/multitest.rs     | 13 +++++--
 3 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/contracts/provider/native-staking-proxy/src/multitest.rs b/contracts/provider/native-staking-proxy/src/multitest.rs
index a7243a9d..410f067c 100644
--- a/contracts/provider/native-staking-proxy/src/multitest.rs
+++ b/contracts/provider/native-staking-proxy/src/multitest.rs
@@ -281,10 +281,7 @@ fn unstaking() {
     );
 
     // Advance time until the unbonding period is over
-    app.update_block(|block| {
-        block.height += 1234;
-        block.time = block.time.plus_seconds(UNBONDING_PERIOD);
-    });
+    process_staking_unbondings(&app);
 
     // Check that the contract now has the funds
     assert_eq!(
@@ -296,6 +293,7 @@ fn unstaking() {
     );
 }
 
+
 #[test]
 fn burning() {
     let owner = "vault_admin";
@@ -338,11 +336,9 @@ fn burning() {
         coin(0, OSMO)
     );
 
+
     // Advance time until the unbonding period is over
-    app.update_block(|block| {
-        block.height += 1234;
-        block.time = block.time.plus_seconds(UNBONDING_PERIOD);
-    });
+    process_staking_unbondings(&app);
 
     // Check that the contract now has the funds
     assert_eq!(
@@ -424,10 +420,7 @@ fn burning_multiple_delegations() {
     );
 
     // Advance time until the unbonding period is over
-    app.update_block(|block| {
-        block.height += 1234;
-        block.time = block.time.plus_seconds(UNBONDING_PERIOD);
-    });
+    process_staking_unbondings(&app);
 
     // Check that the contract now has the funds
     assert_eq!(
@@ -482,10 +475,7 @@ fn releasing_unbonded() {
     assert!(delegation.is_none());
 
     // Advance time until the unbonding period is over
-    app.update_block(|block| {
-        block.height += 12345;
-        block.time = block.time.plus_seconds(UNBONDING_PERIOD + 1);
-    });
+    process_staking_unbondings(&app);
 
     // Release the unbonded funds
     staking_proxy.release_unbonded().call(user).unwrap();
@@ -551,3 +541,18 @@ fn withdrawing_rewards() {
         .unwrap();
     assert_eq!(original_vault_funds, vault_funds);
 }
+
+fn process_staking_unbondings(app: &App<MtApp>) {
+    // Advance unbonding period
+    app.app_mut().update_block(|block| {
+        block.time = block.time.plus_seconds(UNBONDING_PERIOD);
+        block.height += UNBONDING_PERIOD / 5;
+    });
+    // This is deprecated as unneeded, but tests fail if it isn't here. What's up???
+    app.app_mut()
+        .sudo(cw_multi_test::SudoMsg::Staking(
+            cw_multi_test::StakingSudo::ProcessQueue {},
+        ))
+        .unwrap();
+}
+
diff --git a/contracts/provider/native-staking/src/multitest.rs b/contracts/provider/native-staking/src/multitest.rs
index 140098cb..0bd18b4f 100644
--- a/contracts/provider/native-staking/src/multitest.rs
+++ b/contracts/provider/native-staking/src/multitest.rs
@@ -340,6 +340,8 @@ fn releasing_proxy_stake() {
         .unstake(validator.to_string(), coin(100, OSMO))
         .call(user)
         .unwrap();
+    // Important: we need to wait the unbonding period until this is released
+    app.update_block(advance_unbonding_period);
     staking_proxy.release_unbonded().call(user).unwrap();
 
     // Check that the vault has the funds again
@@ -351,3 +353,11 @@ fn releasing_proxy_stake() {
     let claims = vault.account_claims(user.to_owned(), None, None).unwrap();
     assert_eq!(claims.claims, []);
 }
+
+pub fn advance_unbonding_period(block: &mut cosmwasm_std::BlockInfo) {
+    // Default unbonding time in cw_multi_test is 60, from looking at the code...
+    // Wish I could find this somewhere in this setup somewhere.
+    const UNBONDING_TIME: u64 = 60;
+    block.time = block.time.plus_seconds(5 * UNBONDING_TIME);
+    block.height += UNBONDING_TIME;
+}
diff --git a/contracts/provider/vault/src/multitest.rs b/contracts/provider/vault/src/multitest.rs
index 61657925..7b448b89 100644
--- a/contracts/provider/vault/src/multitest.rs
+++ b/contracts/provider/vault/src/multitest.rs
@@ -273,9 +273,16 @@ fn proxy_for_user<'a>(
 }
 
 fn process_staking_unbondings(app: &App<MtApp>) {
-    let mut block_info = app.block_info();
-    block_info.time = block_info.time.plus_seconds(61);
-    app.set_block(block_info);
+    app.app_mut().update_block(|block| {
+        block.time = block.time.plus_seconds(61);
+        block.height += 13;
+    });
+    // This is deprecated as unneeded, but tests fail if it isn't here. What's up???
+    app.app_mut()
+        .sudo(cw_multi_test::SudoMsg::Staking(
+            cw_multi_test::StakingSudo::ProcessQueue {},
+        ))
+        .unwrap();
 }
 
 #[track_caller]