From bf6843f803a5516f99eb3cb89ea548ad8a0d2a3c Mon Sep 17 00:00:00 2001
From: Kyrylo Stepanov <k.stepanov@tutanota.com>
Date: Thu, 27 Jun 2024 17:32:23 +0300
Subject: [PATCH] add lockup_period_ends_at for Flex LockupPeriod

---
 programs/rewards/src/utils.rs           |  3 +-
 programs/rewards/tests/rewards/claim.rs | 77 +++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/programs/rewards/src/utils.rs b/programs/rewards/src/utils.rs
index bd612bcb..a7d0dabb 100644
--- a/programs/rewards/src/utils.rs
+++ b/programs/rewards/src/utils.rs
@@ -271,10 +271,11 @@ impl LockupPeriod {
         let beginning_of_the_day = start_ts - (start_ts % SECONDS_PER_DAY);
 
         match self {
-            LockupPeriod::None | LockupPeriod::Flex => Err(MplxRewardsError::InvalidLockupPeriod),
+            LockupPeriod::None => Err(MplxRewardsError::InvalidLockupPeriod),
             LockupPeriod::ThreeMonths => Ok(beginning_of_the_day + SECONDS_PER_DAY * 90),
             LockupPeriod::SixMonths => Ok(beginning_of_the_day + SECONDS_PER_DAY * 180),
             LockupPeriod::OneYear => Ok(beginning_of_the_day + SECONDS_PER_DAY * 365),
+            LockupPeriod::Flex => Ok(0)
         }
     }
 
diff --git a/programs/rewards/tests/rewards/claim.rs b/programs/rewards/tests/rewards/claim.rs
index d84974b5..ac0efa4f 100644
--- a/programs/rewards/tests/rewards/claim.rs
+++ b/programs/rewards/tests/rewards/claim.rs
@@ -778,3 +778,80 @@ async fn claim_after_withdraw_is_correct() {
     )
     .await;
 }
+
+#[tokio::test]
+async fn with_two_users_with_flex() {
+    let (mut context, test_rewards, rewarder) = setup().await;
+
+    let (user_a, user_rewards_a, user_mining_a) =
+        create_end_user(&mut context, &test_rewards).await;
+    test_rewards
+        .deposit_mining(
+            &mut context,
+            &user_mining_a,
+            100,
+            LockupPeriod::Flex,
+            &user_a.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    let (user_b, user_rewards_b, user_mining_b) =
+        create_end_user(&mut context, &test_rewards).await;
+    test_rewards
+        .deposit_mining(
+            &mut context,
+            &user_mining_b,
+            100,
+            LockupPeriod::Flex,
+            &user_b.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;
+
+    test_rewards
+        .fill_vault(&mut context, &rewarder, 100, distribution_ends_at)
+        .await
+        .unwrap();
+
+    test_rewards.distribute_rewards(&mut context).await.unwrap();
+
+    test_rewards
+        .claim(
+            &mut context,
+            &user_a,
+            &user_mining_a,
+            &user_rewards_a.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    test_rewards
+        .claim(
+            &mut context,
+            &user_b,
+            &user_mining_b,
+            &user_rewards_b.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    let user_reward_account_a = get_account(&mut context, &user_rewards_a.pubkey()).await;
+    let user_rewards_a = Account::unpack(user_reward_account_a.data.borrow()).unwrap();
+
+    assert_eq!(user_rewards_a.amount, 50);
+
+    let user_reward_account_b = get_account(&mut context, &user_rewards_b.pubkey()).await;
+    let user_rewards_b = Account::unpack(user_reward_account_b.data.borrow()).unwrap();
+
+    assert_eq!(user_rewards_b.amount, 50);
+}
\ No newline at end of file