From ed67be2711a7b4f5fa92e853fc6a726ccb3157ec Mon Sep 17 00:00:00 2001
From: Kyrylo Stepanov <k.stepanov@tutanota.com>
Date: Thu, 27 Jun 2024 17:38:53 +0300
Subject: [PATCH] exted deposit test with LockupPeriod::Flex

---
 .../rewards/tests/rewards/deposit_mining.rs   | 19 ++++++++
 .../tests/rewards/distribute_rewards.rs       | 46 +++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/programs/rewards/tests/rewards/deposit_mining.rs b/programs/rewards/tests/rewards/deposit_mining.rs
index 311fd58e..2bcb96ed 100644
--- a/programs/rewards/tests/rewards/deposit_mining.rs
+++ b/programs/rewards/tests/rewards/deposit_mining.rs
@@ -55,3 +55,22 @@ async fn success() {
     let mining = Mining::unpack(mining_account.data.borrow()).unwrap();
     assert_eq!(mining.share, 200);
 }
+
+#[tokio::test]
+async fn success_with_flex() {
+    let (mut context, test_rewards, user, mining) = setup().await;
+
+    test_rewards
+        .deposit_mining(&mut context, &mining, 100, LockupPeriod::Flex, &user)
+        .await
+        .unwrap();
+
+    let reward_pool_account = get_account(&mut context, &test_rewards.mining_reward_pool).await;
+    let reward_pool = RewardPool::unpack(reward_pool_account.data.borrow()).unwrap();
+
+    assert_eq!(reward_pool.total_share, 100);
+
+    let mining_account = get_account(&mut context, &mining).await;
+    let mining = Mining::unpack(mining_account.data.borrow()).unwrap();
+    assert_eq!(mining.share, 100);
+}
\ No newline at end of file
diff --git a/programs/rewards/tests/rewards/distribute_rewards.rs b/programs/rewards/tests/rewards/distribute_rewards.rs
index 2eabbcd9..23a16606 100644
--- a/programs/rewards/tests/rewards/distribute_rewards.rs
+++ b/programs/rewards/tests/rewards/distribute_rewards.rs
@@ -90,3 +90,49 @@ async fn happy_path() {
 
     assert_tokens(&mut context, &user_rewards.pubkey(), 1).await;
 }
+
+#[tokio::test]
+async fn happy_path_with_flex() {
+    let (mut context, test_rewards, rewarder) = setup().await;
+
+    let (user, user_rewards, user_mining_addr) = create_end_user(&mut context, &test_rewards).await;
+    test_rewards
+        .deposit_mining(
+            &mut context,
+            &user_mining_addr,
+            100,
+            LockupPeriod::Flex,
+            &user.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    // fill vault with tokens
+    let distribution_ends_at = context
+        .banks_client
+        .get_sysvar::<solana_program::clock::Clock>()
+        .await
+        .unwrap()
+        .unix_timestamp as u64
+        + SECONDS_PER_DAY * 100;
+
+    test_rewards
+        .fill_vault(&mut context, &rewarder, 100, distribution_ends_at)
+        .await
+        .unwrap();
+    // distribute rewards to users
+    test_rewards.distribute_rewards(&mut context).await.unwrap();
+
+    // // user claims their rewards
+    test_rewards
+        .claim(
+            &mut context,
+            &user,
+            &user_mining_addr,
+            &user_rewards.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    assert_tokens(&mut context, &user_rewards.pubkey(), 1).await;
+}
\ No newline at end of file