exted deposit test with LockupPeriod::Flex

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);
+}

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;
+}