test.rs

#![cfg(test)]
extern crate std;

use crate::{StakingContract, StakingContractClient};
use soroban_sdk::token::{Client as TokenClient, StellarAssetClient as TokenAdminClient};
use soroban_sdk::{
    testutils::{Address as _, Ledger},
    Address, Env,
};

#[test]
fn test_staking_success() {
    let env = Env::default();
    env.mock_all_auths();

    let admin = Address::generate(&env);
    let staker = Address::generate(&env);

    // Deploy token
    let token_admin = Address::generate(&env);
    let token_id = env.register_stellar_asset_contract_v2(token_admin.clone());
    let token_address = token_id.address();
    let token_client = TokenClient::new(&env, &token_address);
    let token_admin_client = TokenAdminClient::new(&env, &token_address);

    // Deploy Staking Contract
    let staking_id = env.register_contract(None, StakingContract);
    let staking_client = StakingContractClient::new(&env, &staking_id);
    staking_client.initialize(&admin, &token_address);

    // Mint tokens
    token_admin_client.mint(&staker, &1000);

    // Stake
    staking_client.stake(&staker, &500);
    assert_eq!(token_client.balance(&staker), 500);
    assert_eq!(token_client.balance(&staking_id), 500);
    assert_eq!(staking_client.get_voting_power(&staker), 500);

    // Unstake
    staking_client.unstake(&staker, &200);
    assert_eq!(staking_client.get_voting_power(&staker), 300);

    // Try withdraw (should fail - cooldown)
    let result = staking_client.try_withdraw(&staker);
    assert!(result.is_err());

    // Jump 7 days
    env.ledger().set_timestamp(604801);
    staking_client.withdraw(&staker);
    assert_eq!(token_client.balance(&staker), 700);
}

#[test]
fn test_staking_rewards() {
    let env = Env::default();
    env.mock_all_auths();

    let admin = Address::generate(&env);
    let staker1 = Address::generate(&env);
    let staker2 = Address::generate(&env);

    // Deploy token
    let token_admin = Address::generate(&env);
    let token_id = env.register_stellar_asset_contract_v2(token_admin.clone());
    let token_address = token_id.address();
    let token_admin_client = TokenAdminClient::new(&env, &token_address);

    // Deploy Staking Contract
    let staking_id = env.register_contract(None, StakingContract);
    let staking_client = StakingContractClient::new(&env, &staking_id);
    staking_client.initialize(&admin, &token_address);

    // Mint tokens & Stake
    token_admin_client.mint(&staker1, &1000);
    token_admin_client.mint(&staker2, &1000);
    staking_client.stake(&staker1, &600);
    staking_client.stake(&staker2, &400);

    // Admin deposits rewards
    token_admin_client.mint(&admin, &1000);
    staking_client.deposit_rewards(&admin, &100); // 60 for staker1, 40 for staker2

    // Claim rewards
    let claimed1 = staking_client.claim_staking_rewards(&staker1);
    let claimed2 = staking_client.claim_staking_rewards(&staker2);

    assert_eq!(claimed1, 60);
    assert_eq!(claimed2, 40);
}

#[test]
fn test_delegation() {
    let env = Env::default();
    env.mock_all_auths();

    let admin = Address::generate(&env);
    let staker1 = Address::generate(&env);
    let staker2 = Address::generate(&env);

    // Deploy token
    let token_admin = Address::generate(&env);
    let token_id = env.register_stellar_asset_contract_v2(token_admin.clone());
    let token_address = token_id.address();
    let token_admin_client = TokenAdminClient::new(&env, &token_address);

    // Deploy Staking Contract
    let staking_id = env.register_contract(None, StakingContract);
    let staking_client = StakingContractClient::new(&env, &staking_id);
    staking_client.initialize(&admin, &token_address);

    // Stake
    token_admin_client.mint(&staker1, &1000);
    token_admin_client.mint(&staker2, &1000);
    staking_client.stake(&staker1, &600);
    staking_client.stake(&staker2, &400);

    assert_eq!(staking_client.get_voting_power(&staker1), 600);
    assert_eq!(staking_client.get_voting_power(&staker2), 400);

    // Delegate staker1 -> staker2
    staking_client.delegate_voting_power(&staker1, &Some(staker2.clone()));

    assert_eq!(staking_client.get_voting_power(&staker1), 600); // Still has own voting power?
                                                                // Wait, usually delegation means giving power to someone else.
                                                                // The requirement says "Delegate voting power".
                                                                // My implementation: get_voting_power = staked + delegated_to_me.
                                                                // So staker2 should have 400 + 600 = 1000.
    assert_eq!(staking_client.get_voting_power(&staker2), 1000);

    // Unstake staker1 partial
    staking_client.unstake(&staker1, &100);
    assert_eq!(staking_client.get_voting_power(&staker2), 900);

    // Remove delegation
    staking_client.delegate_voting_power(&staker1, &None);
    assert_eq!(staking_client.get_voting_power(&staker2), 400);
}

// ============================================================
// Property / invariant tests (proptest-style state transitions)
// ============================================================

mod proptests {
    use super::*;
    use proptest::prelude::*;

    const SCALING_FACTOR: i128 = 1_000_000_000_000i128;

    #[derive(Clone)]
    struct StakerModel {
        staked: i128,
        accumulated_rewards: i128,
        last_reward_index: i128,
        delegated_to: Option<usize>,
    }

    fn update_staker_rewards(model: &mut [StakerModel], idx: usize, reward_index: i128) {
        let staker = &mut model[idx];
        if staker.last_reward_index < reward_index {
            let delta = reward_index - staker.last_reward_index;
            let earned = (staker.staked * delta) / SCALING_FACTOR;
            staker.accumulated_rewards += earned;
            staker.last_reward_index = reward_index;
        }
    }

    fn voting_power(model: &[StakerModel], delegated_amounts: &[i128], idx: usize) -> i128 {
        model[idx].staked + delegated_amounts[idx]
    }

    proptest! {
        #![proptest_config(ProptestConfig { cases: 16, .. ProptestConfig::default() })]
        #[test]
        fn prop_staking_conserves_tokens_and_rewards_match_claim(
            steps in prop::collection::vec(0u8..=3u8, 1usize..=12),
            stake_idx in 0usize..3usize,
            amount_seed in 1u32..=500u32,
            deposit_seed in 1u32..=500u32
        ) {
            let env = Env::default();
            env.mock_all_auths();

            let admin = Address::generate(&env);

            // Deploy token
            let token_admin = Address::generate(&env);
            let token_id = env.register_stellar_asset_contract_v2(token_admin.clone());
            let token_address = token_id.address();
            let token_client = TokenClient::new(&env, &token_address);
            let token_admin_client = TokenAdminClient::new(&env, &token_address);

            // Deploy staking
            let staking_id = env.register_contract(None, StakingContract);
            let staking_client = StakingContractClient::new(&env, &staking_id);
            staking_client.initialize(&admin, &token_address);

            // Fixed staker set (3)
            let mut stakers: std::vec::Vec<Address> = std::vec::Vec::new();
            stakers.push(Address::generate(&env));
            stakers.push(Address::generate(&env));
            stakers.push(Address::generate(&env));
            for s in &stakers {
                token_admin_client.mint(s, &1_000_000);
            }
            token_admin_client.mint(&admin, &1_000_000);

            let mut model: std::vec::Vec<StakerModel> = std::vec::Vec::new();
            model.push(StakerModel { staked: 0, accumulated_rewards: 0, last_reward_index: 0, delegated_to: None });
            model.push(StakerModel { staked: 0, accumulated_rewards: 0, last_reward_index: 0, delegated_to: None });
            model.push(StakerModel { staked: 0, accumulated_rewards: 0, last_reward_index: 0, delegated_to: None });

            let mut delegated_amounts: std::vec::Vec<i128> = std::vec::Vec::new();
            delegated_amounts.push(0i128);
            delegated_amounts.push(0i128);
            delegated_amounts.push(0i128);
            let mut reward_index: i128 = 0;
            let mut reward_pool: i128 = 0; // tokens in staking contract beyond principal

            // Use seeds to get deterministic amounts per case.
            let mut stake_amount = amount_seed as i128;
            let mut deposit_amount = deposit_seed as i128;

            for (step_no, op) in steps.into_iter().enumerate() {
                // Capture balances for invariants.
                let before_contract_balance = token_client.balance(&staking_id);
                let before_total_tokens = {
                    let mut t = token_client.balance(&admin);
                    for s in &stakers {
                        t += token_client.balance(s);
                    }
                    t + token_client.balance(&staking_id)
                };

                // Deterministic pseudo-random indexing from step_no.
                let idx = (step_no + stake_idx) % 3;
                let delegatee_idx = (step_no + stake_idx + 1) % 3;

                match op {
                    // 0 => stake
                    0 => {
                        let amount = ((stake_amount + step_no as i128) % 500i128) + 1i128;
                        let staker = stakers[idx].clone();
                        let res = staking_client.try_stake(&staker, &amount);
                        if res.is_ok() {
                            update_staker_rewards(&mut model, idx, reward_index);
                            model[idx].staked += amount;
                            // If delegated, add to delegatee's voting power backing.
                            if let Some(d) = model[idx].delegated_to {
                                delegated_amounts[d] += amount;
                            }
                        }
                    }
                    // 1 => deposit rewards
                    1 => {
                        let amount = ((deposit_amount + step_no as i128 * 7) % 500i128) + 1i128;
                        let res = staking_client.try_deposit_rewards(&admin, &amount);
                        if res.is_ok() {
                            // deposit_rewards only works if total_staked > 0 (contract validates)
                            let total_staked: i128 = model.iter().map(|m| m.staked).sum();
                            if total_staked > 0 {
                                reward_pool += amount;
                                reward_index += (amount * SCALING_FACTOR) / total_staked;
                            }
                        }
                    }
                    // 2 => delegate / clear delegation
                    2 => {
                        // Alternate between Some(delegatee) and None.
                        let delegator = stakers[idx].clone();
                        let delegatee_addr = stakers[delegatee_idx].clone();
                        let delegate_opt = if step_no % 2 == 0 {
                            Some(delegatee_addr.clone())
                        } else {
                            None
                        };

                        let res = staking_client.try_delegate_voting_power(&delegator, &delegate_opt);
                        if res.is_ok() {
                            let old = model[idx].delegated_to;
                            if let Some(old_idx) = old {
                                delegated_amounts[old_idx] -= model[idx].staked;
                            }
                            model[idx].delegated_to = if step_no % 2 == 0 {
                                Some(delegatee_idx)
                            } else {
                                None
                            };
                            if let Some(new_idx) = model[idx].delegated_to {
                                delegated_amounts[new_idx] += model[idx].staked;
                            }
                        }
                    }
                    // 3 => claim rewards
                    _ => {
                        let staker = stakers[idx].clone();
                        let res = staking_client.try_claim_staking_rewards(&staker);
                        if let Ok(Ok(claimed)) = res {
                            update_staker_rewards(&mut model, idx, reward_index);
                            prop_assert_eq!(claimed, model[idx].accumulated_rewards);
                            model[idx].accumulated_rewards = 0;
                            reward_pool -= claimed;
                        }
                    }
                }

                // Token conservation must always hold.
                let after_total_tokens = {
                    let mut t = token_client.balance(&admin);
                    for s in &stakers {
                        t += token_client.balance(s);
                    }
                    t + token_client.balance(&staking_id)
                };
                prop_assert_eq!(after_total_tokens, before_total_tokens);

                // Contract balance must equal principal + unclaimed reward pool.
                let model_total_staked: i128 = model.iter().map(|m| m.staked).sum();
                prop_assert_eq!(
                    token_client.balance(&staking_id),
                    model_total_staked + reward_pool
                );

                // Voting power matches staked + delegated backing.
                for i in 0..3 {
                    let expected = voting_power(&model, &delegated_amounts, i);
                    let actual = staking_client.get_voting_power(&stakers[i]);
                    prop_assert_eq!(actual, expected);
                }

                // Ensure we didn't underflow the reward pool.
                prop_assert!(reward_pool >= 0);

                // Keep seeds bounded.
                stake_amount = (stake_amount + step_no as i128) % 1000i128;
                deposit_amount = (deposit_amount + step_no as i128 * 3) % 1000i128;
                let _ = before_contract_balance;
            }
        }
    }
}