Skip to content

Commit 6a65594

Browse files
vasyafromrussiavasyafromrussia
committed
restore restake_all tests
1 parent 2d5146e commit 6a65594

File tree

4 files changed

+125
-69
lines changed

4 files changed

+125
-69
lines changed

‎contract/src/claim/api.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ impl Contract {
7777
account_rollback.score = Some(account.score);
7878
account_rollback.jars = Some(rollback_jars);
7979

80-
account.score.claim_score();
80+
account.score.try_claim_score();
8181

8282
if accumulator.get_total().0 > 0 {
8383
self.claim_interest(

‎contract/src/jar/tests/restake_all.rs

+112-68
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
1-
use near_sdk::test_utils::test_env::alice;
1+
use near_sdk::{
2+
json_types::{U128, U64},
3+
test_utils::test_env::alice,
4+
};
25
use sweat_jar_model::{
3-
api::{JarApi, ProductApi},
4-
MS_IN_DAY, MS_IN_YEAR,
6+
api::{ClaimApi, JarApi, ProductApi},
7+
jar::JarView,
8+
MS_IN_DAY, MS_IN_MINUTE, MS_IN_YEAR,
59
};
610

711
use crate::{
@@ -94,68 +98,108 @@ fn restake_all() {
9498
assert!(all_jar_ids.contains(&create_synthetic_jar_id(long_term_product_to_disable.id, 0)));
9599
}
96100

97-
// #[test]
98-
// fn restake_all_after_maturity_for_restakable_product_one_jar() {
99-
// let alice = alice();
100-
// let admin = admin();
101-
//
102-
// let product = Product::new().with_allows_restaking(true);
103-
// let jar = Jar::new(0).principal(PRINCIPAL);
104-
// let mut context = Context::new(admin).with_products(&[product]).with_jars(&[jar.clone()]);
105-
//
106-
// context.set_block_timestamp_in_days(366);
107-
//
108-
// context.switch_account(&alice);
109-
// let restaked = context.contract().restake_all(None).into_iter().next().unwrap();
110-
//
111-
// assert_eq!(restaked.account_id, alice);
112-
// assert_eq!(restaked.principal.0, PRINCIPAL);
113-
// assert_eq!(restaked.claimed_balance.0, 0);
114-
//
115-
// let alice_jars = context.contract().get_jars_for_account(alice);
116-
//
117-
// assert_eq!(2, alice_jars.len());
118-
// assert_eq!(0, alice_jars.iter().find(|item| item.id.0 == 0).unwrap().principal.0);
119-
// assert_eq!(
120-
// PRINCIPAL,
121-
// alice_jars.iter().find(|item| item.id.0 == 1).unwrap().principal.0
122-
// );
123-
// }
124-
//
125-
// #[test]
126-
// fn batch_restake_all() {
127-
// let alice = alice();
128-
// let admin = admin();
129-
//
130-
// let product = Product::new().with_allows_restaking(true).lockup_term(MS_IN_YEAR);
131-
//
132-
// let jars: Vec<_> = (0..8)
133-
// .map(|id| Jar::new(id).principal(PRINCIPAL + id as u128))
134-
// .collect();
135-
//
136-
// let mut context = Context::new(admin).with_products(&[product]).with_jars(&jars);
137-
//
138-
// context.set_block_timestamp_in_days(366);
139-
//
140-
// context.switch_account(&alice);
141-
//
142-
// context.contract().claim_total(None);
143-
//
144-
// let restaked: Vec<_> = context
145-
// .contract()
146-
// .restake_all(Some(vec![1.into(), 2.into(), 5.into()]))
147-
// .into_iter()
148-
// .map(|j| j.id.0)
149-
// .collect();
150-
//
151-
// assert_eq!(restaked, [8, 9, 10]);
152-
//
153-
// let jars: Vec<_> = context
154-
// .contract()
155-
// .get_jars_for_account(alice)
156-
// .into_iter()
157-
// .map(|j| j.id.0)
158-
// .collect();
159-
//
160-
// assert_eq!(jars, [0, 7, 8, 3, 4, 9, 6, 10,]);
161-
// }
101+
#[test]
102+
fn restake_all_after_maturity_for_restakable_product_one_jar() {
103+
let alice = alice();
104+
let admin = admin();
105+
106+
let product = ProductV2::new();
107+
let jar = JarV2::new().with_deposit(0, PRINCIPAL);
108+
let mut context = Context::new(admin)
109+
.with_products(&[product.clone()])
110+
.with_jars(&alice, &[(product.id.clone(), jar.clone())]);
111+
112+
let restake_time = MS_IN_YEAR + MS_IN_MINUTE;
113+
context.set_block_timestamp_in_ms(restake_time);
114+
115+
context.switch_account(&alice);
116+
let result = context.contract().restake_all(None);
117+
118+
assert_eq!(1, result.len());
119+
assert_eq!((product.id.clone(), PRINCIPAL), result.first().unwrap().clone());
120+
121+
let alice_jars = context.contract().get_jars_for_account(alice);
122+
assert_eq!(1, alice_jars.len());
123+
assert_eq!(PRINCIPAL, alice_jars.first().unwrap().principal.0);
124+
assert_eq!(restake_time, alice_jars.first().unwrap().created_at.0);
125+
}
126+
127+
#[test]
128+
fn batch_restake_all() {
129+
let alice = alice();
130+
let admin = admin();
131+
132+
let product_one = ProductV2::new().id("product_one");
133+
let product_two = ProductV2::new().id("product_two");
134+
let product_three = ProductV2::new().id("product_three");
135+
136+
let product_one_jar_deposit_first = (0, PRINCIPAL);
137+
let product_one_jar_deposit_second = (MS_IN_DAY, 5 * PRINCIPAL);
138+
let product_two_jar_deposit = (0, 2 * PRINCIPAL);
139+
let product_three_jar_deposit = (0, 3 * PRINCIPAL);
140+
141+
let product_one_jar = JarV2::new()
142+
.with_deposit(product_one_jar_deposit_first.0, product_one_jar_deposit_first.1)
143+
.with_deposit(product_one_jar_deposit_second.0, product_one_jar_deposit_second.1);
144+
let product_two_jar = JarV2::new().with_deposit(product_two_jar_deposit.0, product_two_jar_deposit.1);
145+
let product_three_jar = JarV2::new().with_deposit(product_three_jar_deposit.0, product_three_jar_deposit.1);
146+
147+
let mut context = Context::new(admin)
148+
.with_products(&[product_one.clone(), product_two.clone(), product_three.clone()])
149+
.with_jars(
150+
&alice,
151+
&[
152+
(product_one.id.clone(), product_one_jar),
153+
(product_two.id.clone(), product_two_jar),
154+
(product_three.id.clone(), product_three_jar),
155+
],
156+
);
157+
158+
let restake_time = MS_IN_YEAR + MS_IN_DAY;
159+
context.set_block_timestamp_in_ms(restake_time);
160+
161+
context.switch_account(&alice);
162+
163+
context.contract().claim_total(None);
164+
165+
let result: Vec<_> = context
166+
.contract()
167+
.restake_all(Some(vec![product_one.id.clone(), product_two.id.clone()]));
168+
169+
assert_eq!(2, result.len());
170+
assert!(result.contains(&(product_one.id.clone(), product_one_jar_deposit_first.1)));
171+
assert!(result.contains(&(product_two.id.clone(), product_two_jar_deposit.1)));
172+
173+
let mut jars: Vec<_> = context.contract().get_jars_for_account(alice);
174+
jars.sort_by(|a, b| b.id.cmp(&a.id));
175+
assert_eq!(4, jars.len());
176+
177+
let mut expected_jars = vec![
178+
JarView {
179+
id: create_synthetic_jar_id(product_one.id.clone(), restake_time),
180+
product_id: product_one.id.clone(),
181+
created_at: U64(restake_time),
182+
principal: U128(product_one_jar_deposit_first.1),
183+
},
184+
JarView {
185+
id: create_synthetic_jar_id(product_one.id.clone(), product_one_jar_deposit_second.0),
186+
product_id: product_one.id.clone(),
187+
created_at: U64(product_one_jar_deposit_second.0),
188+
principal: U128(product_one_jar_deposit_second.1),
189+
},
190+
JarView {
191+
id: create_synthetic_jar_id(product_two.id.clone(), restake_time),
192+
product_id: product_two.id.clone(),
193+
created_at: U64(restake_time),
194+
principal: U128(product_two_jar_deposit.1),
195+
},
196+
JarView {
197+
id: create_synthetic_jar_id(product_three.id.clone(), 0),
198+
product_id: product_three.id.clone(),
199+
created_at: U64(0),
200+
principal: U128(product_three_jar_deposit.1),
201+
},
202+
];
203+
expected_jars.sort_by(|a, b| b.id.cmp(&a.id));
204+
assert_eq!(jars, expected_jars);
205+
}

‎contract/src/score/account_score.rs

+8
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,14 @@ impl AccountScore {
4545
}
4646
}
4747

48+
pub fn try_claim_score(&mut self) -> Option<Vec<Score>> {
49+
if self.is_valid() {
50+
Some(self.claim_score())
51+
} else {
52+
None
53+
}
54+
}
55+
4856
/// On claim we need to clear active scores so they aren't claimed twice or more.
4957
// TODO: at least rename
5058
pub fn claim_score(&mut self) -> Vec<Score> {

‎contract/src/test_utils.rs

+4
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,10 @@ impl JarV2 {
3838
}
3939
}
4040

41+
pub(crate) fn total_principal(&self) -> TokenAmount {
42+
self.deposits.iter().map(|deposit| deposit.principal).sum()
43+
}
44+
4145
pub(crate) fn with_deposit(mut self, created_at: Timestamp, principal: TokenAmount) -> Self {
4246
self.deposits.push(Deposit::new(created_at, principal));
4347
self

0 commit comments

Comments
 (0)