test: add balances delay (more balance changes) flow

[arad-starkware]

This change is 

---

Note

Adds new balance-delay upgrade/attest test flows (incl. 2-pool trace entries), renames the existing flow, and wires up corresponding fork tests; minor ideas doc cleanup.

• Flows/Tests:
  • Add BalancesDelayUpgradeAttestFlow to verify staking power and rewards across epochs around an upgrade and attestations.
  • Add BalancesDelay2PoolTraceEntriesFlow to cover multi-epoch, multi-increment scenarios (STRK+BTC) pre/post-upgrade.
  • Rename and refactor BalancesDelayFlow -> BalancesDelay1PoolTraceEntryFlow (setup/test updated accordingly).
• Fork Tests:
  • Add balances_delay_2_pool_trace_entries_flow_test and balances_delay_upgrade_attest_flow_test.
  • Replace balances_delay_flow_test with balances_delay_1_pool_trace_entry_flow_test.
• Docs:
  • Trim obsolete migration notes from src/flow_test/flow_ideas.md.

^{Written by Cursor Bugbot for commit 91fd529. This will update automatically on new commits. Configure here.}

[arad-starkware]

• test: add balances delay (more balance changes) flow #84 👈 (View in Graphite)
• test: add accrue_rewards_across_versions flow #82 : 1 other dependent PR (#83 )
• test: add attest after upgrade test #71
• test: add token state after upgrade flow #59
• test: add update_rewards toggle tokens flow test #57
• test: add balances delay flow #65
• chore: update flow ideas #61
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[noa-starkware]

Reviewable status: 0 of 3 files reviewed, 3 unresolved discussions (waiting on @arad-starkware)

---

-- commits line 2 at r2:
Rename the PR

---

src/flow_test/flow_ideas.md line 38 at r2 (raw file):

- staker enter in V2, advance epoch, upgrade to V3, advance epoch, attest,
- staker in V2, update balance staker+update balance pool, upgrade, attest in current epoch, attest in next epoch, attest in next next epoch
- staker in V2, advance epoch, update balance staker+update balance pool, advance epoch, update balance staker+update balance pool, upgrade, update balance staker+update balance pool, attest in current epoch, attest in next epoch, attest in next next epoch

Do you implemented 2 tests or one for both lines?

[noa-starkware]

@noa-starkware reviewed 1 of 4 files at r6.
Reviewable status: 1 of 4 files reviewed, 2 unresolved discussions (waiting on @arad-starkware)

---

src/flow_test/flow_ideas.md line 19 at r6 (raw file):

## k=1 -> k=2 Migration Staker
- staker in V2, update balance staker+update balance pool, upgrade, attest in current epoch, attest in next epoch, attest in next next epoch

I think this part is missing

Code quote:

 attest in next next epoch

---

src/flow_test/flows.cairo line 10240 at r6 (raw file):

        let prev_strk_pool_rewards = expected_strk_pool_rewards;
        let prev_btc_commission_rewards = expected_btc_commission_rewards;
        let prev_btc_pool_rewards = expected_btc_pool_rewards;

Why here you check all vars and in the previous epoch you checked only for prev_staker_rewards? and why did you check is non zero in the prev epoch but not here?

[arad-starkware]

Reviewable status: 0 of 4 files reviewed, 2 unresolved discussions (waiting on @arad-starkware and @noa-starkware)

---

src/flow_test/flow_ideas.md line 19 at r6 (raw file):

Previously, noa-starkware wrote…

I think this part is missing

There are 3 attests after upgrading (current, next, next next)

---

src/flow_test/flows.cairo line 10240 at r6 (raw file):

Previously, noa-starkware wrote…

Why here you check all vars and in the previous epoch you checked only for prev_staker_rewards? and why did you check is non zero in the prev epoch but not here?

missed it with the iterations of the test...

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 95.81%. Comparing base (bd3cec8) to head (91fd529).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main      #84   +/-   ##
=======================================
  Coverage   95.81%   95.81%           
=======================================
  Files          42       42           
  Lines       10194    10194           
=======================================
  Hits         9767     9767           
  Misses        427      427           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[noa-starkware]

Reviewable status: 0 of 4 files reviewed, 1 unresolved discussion (waiting on @arad-starkware)

---

src/flow_test/flow_ideas.md line 19 at r6 (raw file):

Previously, arad-starkware wrote…

There are 3 attests after upgrading (current, next, next next)

So maybe its missing in the test doc?

[arad-starkware]

Reviewable status: 0 of 4 files reviewed, 1 unresolved discussion (waiting on @noa-starkware)

---

src/flow_test/flow_ideas.md line 19 at r6 (raw file):

Previously, noa-starkware wrote…

So maybe its missing in the test doc?

added

[noa-starkware]

Reviewable status: 0 of 4 files reviewed, 1 unresolved discussion (waiting on @arad-starkware)

---

src/flow_test/flows.cairo line 8789 at r8 (raw file):

    }

    #[feature("safe_dispatcher")]

Why?

Code quote:

#[feature("safe_dispatcher")]

[noa-starkware]

Reviewable status: 0 of 4 files reviewed, 1 unresolved discussion (waiting on @arad-starkware)

[arad-starkware]

Merge activity

• Dec 18, 8:35 AM UTC: @arad-starkware merged this pull request with Graphite.

[cursor]

Bug: Tests assume incorrect epoch delay for balance changes

The newly added test flows BalancesDelayUpgradeAttestFlow, BalancesDelay1PoolTraceEntryFlow, and BalancesDelay2PoolTraceEntriesFlow assume a one-epoch delay (K=1) for balance changes to become effective. However, the system's K constant is defined as 2 in src/constants.cairo, which implies a two-epoch delay. This inconsistency causes the tests to assert incorrect staking power and reward values for each epoch.

src/flow_test/flows.cairo#L8718-L9170

starknet-staking/src/flow_test/flows.cairo

Lines 8718 to 9170 in 91fd529

	/// Flow:
	/// Epoch 0:
	/// Staker stake
	/// Epoch 1:
	/// Staker increase stake
	/// Delegators delegate STRK and BTC
	/// Upgrade
	/// Test total staking power - according to Epoch 0
	/// Test staker balance with attestation rewards - according to Epoch 0
	/// Epoch 2:
	/// Test delegator balances with attestation rewards - according to Epoch 0
	/// Test total staking power - according to Epoch 1
	/// Test staker balance with attestation rewards - according to Epoch 1
	/// Epoch 3:
	/// Test delegator balances with attestation rewards - according to Epoch 1
	/// Test total staking power - according to Epoch 2
	/// Test staker balance with attestation rewards - according to Epoch 2
	/// Epoch 4:
	/// Test delegator balances with attestation rewards - according to Epoch 2
	#[derive(Drop, Copy)]
	pub(crate) struct BalancesDelayUpgradeAttestFlow {
	pub(crate) staker: Option<Staker>,
	pub(crate) stake_amount: Option<Amount>,
	pub(crate) commission: Option<Commission>,
	pub(crate) strk_delegated_amount: Option<Amount>,
	pub(crate) btc_delegated_amount: Option<Amount>,
	pub(crate) strk_delegator: Option<Delegator>,
	pub(crate) btc_delegator: Option<Delegator>,
	pub(crate) strk_pool: Option<ContractAddress>,
	pub(crate) btc_pool: Option<ContractAddress>,
	}
	pub(crate) impl BalancesDelayUpgradeAttestFlowImpl of FlowTrait<BalancesDelayUpgradeAttestFlow> {
	fn setup_v2(ref self: BalancesDelayUpgradeAttestFlow, ref system: SystemState) {
	let stake_amount = system.staking.get_min_stake();
	let strk_delegated_amount = STRK_CONFIG.min_for_rewards;
	let btc_delegated_amount = TEST_MIN_BTC_FOR_REWARDS;
	let staker = system.new_staker(amount: stake_amount * 3);
	let commission = 200;
	let strk_delegator = system.new_delegator(amount: strk_delegated_amount * 2);
	let btc_delegator = system
	.new_btc_delegator(amount: btc_delegated_amount * 2, token: system.btc_token);
	system.stake(:staker, amount: stake_amount, pool_enabled: true, :commission);
	let strk_pool = system.staking.get_pool(:staker);
	let btc_pool = system
	.set_open_for_delegation(:staker, token_address: system.btc_token.contract_address());
	system.advance_epoch(); // Epoch 0 - > 1
	system.increase_stake(:staker, amount: stake_amount);
	system.delegate(delegator: strk_delegator, pool: strk_pool, amount: strk_delegated_amount);
	system
	.delegate_btc(
	delegator: btc_delegator,
	pool: btc_pool,
	amount: btc_delegated_amount,
	token: system.btc_token,
	);
	system.set_staker_for_migration(staker_address: staker.staker.address);
	self.staker = Option::Some(staker);
	self.stake_amount = Option::Some(stake_amount);
	self.commission = Option::Some(commission);
	self.strk_delegated_amount = Option::Some(strk_delegated_amount);
	self.btc_delegated_amount = Option::Some(btc_delegated_amount);
	self.strk_delegator = Option::Some(strk_delegator);
	self.btc_delegator = Option::Some(btc_delegator);
	self.strk_pool = Option::Some(strk_pool);
	self.btc_pool = Option::Some(btc_pool);
	}
	fn test(self: BalancesDelayUpgradeAttestFlow, ref system: SystemState) {
	let staker = self.staker.unwrap();
	let stake_amount = self.stake_amount.unwrap();
	let commission = self.commission.unwrap();
	let strk_delegated_amount = self.strk_delegated_amount.unwrap();
	let btc_delegated_amount = self.btc_delegated_amount.unwrap();
	let strk_delegator = self.strk_delegator.unwrap();
	let btc_delegator = self.btc_delegator.unwrap();
	let strk_pool = self.strk_pool.unwrap();
	let btc_pool = self.btc_pool.unwrap();
	let staking_contract = system.staking.address;
	let minting_curve_contract = system.minting_curve.address;
	// Test total staking power - according to Epoch 0.
	let mut staker_stake = stake_amount;
	let total_staking_power = system.staking.get_current_total_staking_power_v2();
	let expected_total_staking_power = (
	NormalizedAmountTrait::from_strk_native_amount(staker_stake), Zero::zero(),
	);
	assert!(total_staking_power == expected_total_staking_power);
	// Calculate expected rewards - according to Epoch 0.
	let (expected_staker_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
	amount_own: staker_stake,
	pool_amount: Zero::zero(),
	:commission,
	:staking_contract,
	:minting_curve_contract,
	);
	assert!(expected_staker_rewards.is_non_zero());
	// Test staker balance with attestation rewards - according to Epoch 0.
	system
	.advance_block_into_attestation_window_custom_stake(
	staker_address: staker.staker.address, stake: staker_stake,
	);
	system.attest(:staker);
	let staker_rewards = system.staker_claim_rewards(:staker);
	assert!(staker_rewards == expected_staker_rewards);
	// Advance epoch - test delegator balances with attestation rewards - according to Epoch 0.
	system.advance_epoch(); // Epoch 1 - > 2
	let strk_delegator_rewards = system
	.delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
	let btc_delegator_rewards = system
	.delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
	assert!(strk_delegator_rewards == Zero::zero());
	assert!(btc_delegator_rewards == Zero::zero());
	// Update variables for Epoch 1.
	staker_stake += stake_amount;
	let mut strk_pool_balance = strk_delegated_amount;
	let mut btc_pool_balance = btc_delegated_amount;
	// Test total staking power - according to Epoch 1.
	let total_staking_power = system.staking.get_current_total_staking_power_v2();
	let expected_total_staking_power = (
	NormalizedAmountTrait::from_strk_native_amount(staker_stake + strk_pool_balance),
	NormalizedAmountTrait::from_native_amount(btc_pool_balance, TEST_BTC_DECIMALS),
	);
	assert!(total_staking_power == expected_total_staking_power);
	// Calculate expected rewards - according to Epoch 1.
	let prev_staker_rewards = expected_staker_rewards;
	let (expected_staker_rewards, expected_strk_pool_rewards) =
	calculate_staker_strk_rewards_with_balances_v2(
	amount_own: staker_stake,
	pool_amount: strk_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	);
	assert!(expected_staker_rewards != prev_staker_rewards);
	assert!(expected_strk_pool_rewards.is_non_zero());
	let (expected_btc_commission_rewards, expected_btc_pool_rewards) =
	calculate_staker_btc_pool_rewards_v2(
	pool_balance: btc_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	token_address: system.btc_token.contract_address(),
	);
	assert!(expected_btc_commission_rewards.is_non_zero());
	assert!(expected_btc_pool_rewards.is_non_zero());
	// Test staker balance with attestation rewards - according to Epoch 1.
	system
	.advance_block_into_attestation_window_custom_stake(
	staker_address: staker.staker.address, stake: staker_stake + strk_pool_balance,
	);
	system.attest(:staker);
	let staker_rewards = system.staker_claim_rewards(:staker);
	assert!(staker_rewards == expected_staker_rewards + expected_btc_commission_rewards);
	// Advance epoch - test delegator balances with attestation rewards - according to Epoch 1.
	system.advance_epoch(); // Epoch 2 - > 3
	let strk_delegator_rewards = system
	.delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
	let btc_delegator_rewards = system
	.delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
	assert!(strk_delegator_rewards == expected_strk_pool_rewards);
	assert!(btc_delegator_rewards == expected_btc_pool_rewards);
	// Test total staking power - according to Epoch 2 (same as Epoch 1).
	let total_staking_power = system.staking.get_current_total_staking_power_v2();
	assert!(total_staking_power == expected_total_staking_power);
	// Test staker balance with attestation rewards - according to Epoch 2 (same as Epoch 1).
	system
	.advance_block_into_attestation_window_custom_stake(
	staker_address: staker.staker.address, stake: staker_stake + strk_pool_balance,
	);
	system.attest(:staker);
	let staker_rewards = system.staker_claim_rewards(:staker);
	assert!(staker_rewards == expected_staker_rewards + expected_btc_commission_rewards);
	// Advance epoch - test delegator balances with attestation rewards - according to Epoch 2
	// (same as Epoch 1).
	system.advance_epoch(); // Epoch 3 - > 4
	let strk_delegator_rewards = system
	.delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
	let btc_delegator_rewards = system
	.delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
	assert!(strk_delegator_rewards == expected_strk_pool_rewards);
	assert!(btc_delegator_rewards == expected_btc_pool_rewards);
	}
	}
	/// Flow:
	/// Epoch 0:
	/// Staker stake
	/// Epoch 1:
	/// Staker increase stake
	/// Delegators delegate STRK and BTC
	/// Upgrade
	/// Staker increase stake
	/// Delegators increase delegation STRK and BTC
	/// Test total staking power - according to Epoch 0
	/// Test staker balance with attestation rewards - according to Epoch 0
	/// Epoch 2:
	/// Test delegator balances with attestation rewards - according to Epoch 0
	/// Test total staking power - according to Epoch 1 (pre-upgrade)
	/// Test staker balance with attestation rewards - according to Epoch 1 (pre-upgrade)
	/// Epoch 3:
	/// Test delegator balances with attestation rewards - according to Epoch 1 (pre-upgrade)
	/// Test total staking power - according to Epoch 1 (post-upgrade)
	/// Test staker balance with attestation rewards - according to Epoch 1 (post-upgrade)
	/// Epoch 4:
	/// Test delegator balances with attestation rewards - according to Epoch 1 (post-upgrade)
	#[derive(Drop, Copy)]
	pub(crate) struct BalancesDelay1PoolTraceEntryFlow {
	pub(crate) staker: Option<Staker>,
	pub(crate) stake_amount: Option<Amount>,
	pub(crate) commission: Option<Commission>,
	pub(crate) strk_delegated_amount: Option<Amount>,
	pub(crate) btc_delegated_amount: Option<Amount>,
	pub(crate) strk_delegator: Option<Delegator>,
	pub(crate) btc_delegator: Option<Delegator>,
	pub(crate) strk_pool: Option<ContractAddress>,
	pub(crate) btc_pool: Option<ContractAddress>,
	}
	pub(crate) impl BalancesDelay1PoolTraceEntryFlowImpl of FlowTrait<
	BalancesDelay1PoolTraceEntryFlow,
	> {
	fn setup_v2(ref self: BalancesDelay1PoolTraceEntryFlow, ref system: SystemState) {
	let stake_amount = system.staking.get_min_stake();
	let strk_delegated_amount = STRK_CONFIG.min_for_rewards;
	let btc_delegated_amount = TEST_MIN_BTC_FOR_REWARDS;
	let staker = system.new_staker(amount: stake_amount * 3);
	let commission = 200;
	let strk_delegator = system.new_delegator(amount: strk_delegated_amount * 2);
	let btc_delegator = system
	.new_btc_delegator(amount: btc_delegated_amount * 2, token: system.btc_token);
	system.stake(:staker, amount: stake_amount, pool_enabled: true, :commission);
	let strk_pool = system.staking.get_pool(:staker);
	let btc_pool = system
	.set_open_for_delegation(:staker, token_address: system.btc_token.contract_address());
	system.advance_epoch(); // Epoch 0 - > 1
	system.increase_stake(:staker, amount: stake_amount);
	system.delegate(delegator: strk_delegator, pool: strk_pool, amount: strk_delegated_amount);
	system
	.delegate_btc(
	delegator: btc_delegator,
	pool: btc_pool,
	amount: btc_delegated_amount,
	token: system.btc_token,
	);
	system.set_staker_for_migration(staker_address: staker.staker.address);
	self.staker = Option::Some(staker);
	self.stake_amount = Option::Some(stake_amount);
	self.commission = Option::Some(commission);
	self.strk_delegated_amount = Option::Some(strk_delegated_amount);
	self.btc_delegated_amount = Option::Some(btc_delegated_amount);
	self.strk_delegator = Option::Some(strk_delegator);
	self.btc_delegator = Option::Some(btc_delegator);
	self.strk_pool = Option::Some(strk_pool);
	self.btc_pool = Option::Some(btc_pool);
	}
	#[feature("safe_dispatcher")]
	fn test(self: BalancesDelay1PoolTraceEntryFlow, ref system: SystemState) {
	let staker = self.staker.unwrap();
	let stake_amount = self.stake_amount.unwrap();
	let commission = self.commission.unwrap();
	let strk_delegated_amount = self.strk_delegated_amount.unwrap();
	let btc_delegated_amount = self.btc_delegated_amount.unwrap();
	let strk_delegator = self.strk_delegator.unwrap();
	let btc_delegator = self.btc_delegator.unwrap();
	let strk_pool = self.strk_pool.unwrap();
	let btc_pool = self.btc_pool.unwrap();
	let staking_contract = system.staking.address;
	let minting_curve_contract = system.minting_curve.address;
	// Increase stake and delegations.
	system.increase_stake(:staker, amount: stake_amount);
	system
	.increase_delegate(
	delegator: strk_delegator, pool: strk_pool, amount: strk_delegated_amount,
	);
	system
	.increase_delegate_btc(
	delegator: btc_delegator,
	pool: btc_pool,
	amount: btc_delegated_amount,
	token: system.btc_token,
	);
	// Test total staking power - according to Epoch 0.
	let mut staker_stake = stake_amount;
	let total_staking_power = system.staking.get_current_total_staking_power_v2();
	let expected_total_staking_power = (
	NormalizedAmountTrait::from_strk_native_amount(staker_stake), Zero::zero(),
	);
	assert!(total_staking_power == expected_total_staking_power);
	// Calculate expected rewards - according to Epoch 0.
	let (expected_staker_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
	amount_own: staker_stake,
	pool_amount: Zero::zero(),
	:commission,
	:staking_contract,
	:minting_curve_contract,
	);
	assert!(expected_staker_rewards.is_non_zero());
	// Test staker balance with attestation rewards - according to Epoch 0.
	system
	.advance_block_into_attestation_window_custom_stake(
	staker_address: staker.staker.address, stake: staker_stake,
	);
	system.attest(:staker);
	let staker_rewards = system.staker_claim_rewards(:staker);
	assert!(staker_rewards == expected_staker_rewards);
	// Advance epoch - test delegator balances with attestation rewards - according to Epoch 0.
	system.advance_epoch(); // Epoch 1 - > 2
	let strk_delegator_rewards = system
	.delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
	let btc_delegator_rewards = system
	.delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
	assert!(strk_delegator_rewards == Zero::zero());
	assert!(btc_delegator_rewards == Zero::zero());
	// Update variables for Epoch 1 (pre-upgrade).
	staker_stake += stake_amount;
	let mut strk_pool_balance = strk_delegated_amount;
	let mut btc_pool_balance = btc_delegated_amount;
	// Test total staking power - according to Epoch 1 (pre-upgrade).
	let total_staking_power = system.staking.get_current_total_staking_power_v2();
	let expected_total_staking_power = (
	NormalizedAmountTrait::from_strk_native_amount(staker_stake + strk_pool_balance),
	NormalizedAmountTrait::from_native_amount(btc_pool_balance, TEST_BTC_DECIMALS),
	);
	assert!(total_staking_power == expected_total_staking_power);
	// Calculate expected rewards - according to Epoch 1 (pre-upgrade).
	let prev_staker_rewards = expected_staker_rewards;
	let (expected_staker_rewards, expected_strk_pool_rewards) =
	calculate_staker_strk_rewards_with_balances_v2(
	amount_own: staker_stake,
	pool_amount: strk_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	);
	assert!(expected_staker_rewards != prev_staker_rewards);
	assert!(expected_strk_pool_rewards.is_non_zero());
	let (expected_btc_commission_rewards, expected_btc_pool_rewards) =
	calculate_staker_btc_pool_rewards_v2(
	pool_balance: btc_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	token_address: system.btc_token.contract_address(),
	);
	assert!(expected_btc_commission_rewards.is_non_zero());
	assert!(expected_btc_pool_rewards.is_non_zero());
	// Test staker balance with attestation rewards - according to Epoch 1 (pre-upgrade).
	system
	.advance_block_into_attestation_window_custom_stake(
	staker_address: staker.staker.address, stake: staker_stake + strk_pool_balance,
	);
	system.attest(:staker);
	let staker_rewards = system.staker_claim_rewards(:staker);
	assert!(staker_rewards == expected_staker_rewards + expected_btc_commission_rewards);
	// Advance epoch - test delegator balances with attestation rewards - according to Epoch 1
	// (pre-upgrade).
	system.advance_epoch(); // Epoch 2 - > 3
	let strk_delegator_rewards = system
	.delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
	let btc_delegator_rewards = system
	.delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
	assert!(strk_delegator_rewards == expected_strk_pool_rewards);
	assert!(btc_delegator_rewards == expected_btc_pool_rewards);
	// Update variables for Epoch 1 (post-upgrade).
	staker_stake += stake_amount;
	strk_pool_balance += strk_delegated_amount;
	btc_pool_balance += btc_delegated_amount;
	// Test total staking power - according to Epoch 1 (post-upgrade).
	let total_staking_power = system.staking.get_current_total_staking_power_v2();
	let expected_total_staking_power = (
	NormalizedAmountTrait::from_strk_native_amount(staker_stake + strk_pool_balance),
	NormalizedAmountTrait::from_native_amount(btc_pool_balance, TEST_BTC_DECIMALS),
	);
	assert!(total_staking_power == expected_total_staking_power);
	// Calculate expected rewards - according to Epoch 1 (post-upgrade).
	let prev_staker_rewards = expected_staker_rewards;
	let prev_strk_pool_rewards = expected_strk_pool_rewards;
	let prev_btc_commission_rewards = expected_btc_commission_rewards;
	let prev_btc_pool_rewards = expected_btc_pool_rewards;
	let (expected_staker_rewards, expected_strk_pool_rewards) =
	calculate_staker_strk_rewards_with_balances_v2(
	amount_own: staker_stake,
	pool_amount: strk_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	);
	assert!(expected_staker_rewards != prev_staker_rewards);
	assert!(expected_strk_pool_rewards != prev_strk_pool_rewards);
	let (expected_btc_commission_rewards, expected_btc_pool_rewards) =
	calculate_staker_btc_pool_rewards_v2(
	pool_balance: btc_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	token_address: system.btc_token.contract_address(),
	);
	assert!(expected_btc_commission_rewards != prev_btc_commission_rewards);
	assert!(expected_btc_pool_rewards != prev_btc_pool_rewards);
	// Test staker balance with attestation rewards - according to Epoch 1 (post-upgrade).
	system
	.advance_block_into_attestation_window_custom_stake(
	staker_address: staker.staker.address, stake: staker_stake + strk_pool_balance,
	);
	system.attest(:staker);
	let staker_rewards = system.staker_claim_rewards(:staker);
	assert!(staker_rewards == expected_staker_rewards + expected_btc_commission_rewards);
	// Advance epoch - test delegator balances with attestation rewards - according to Epoch 1
	// (post-upgrade).
	system.advance_epoch(); // Epoch 3 - > 4
	let strk_delegator_rewards = system
	.delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
	let btc_delegator_rewards = system
	.delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
	assert!(strk_delegator_rewards == expected_strk_pool_rewards);
	assert!(btc_delegator_rewards == expected_btc_pool_rewards);
	}
	}

src/flow_test/flows.cairo#L10000-L10297

starknet-staking/src/flow_test/flows.cairo

Lines 10000 to 10297 in 91fd529

	/// Flow:
	/// Epoch 0:
	/// Staker stake
	/// Epoch 1:
	/// Staker increase stake
	/// Delegators delegate STRK and BTC
	/// Epoch 2:
	/// Staker increase stake
	/// Delegators increase delegation STRK and BTC
	/// Upgrade
	/// Staker increase stake
	/// Delegators increase delegation STRK and BTC
	/// Test total staking power - according to Epoch 1
	/// Test staker balance with attestation rewards - according to Epoch 1
	/// Epoch 3:
	/// Test delegator balances with attestation rewards - according to Epoch 1
	/// Test total staking power - according to Epoch 2 (pre-upgrade)
	/// Test staker balance with attestation rewards - according to Epoch 2 (pre-upgrade)
	/// Epoch 4:
	/// Test delegator balances with attestation rewards - according to Epoch 2 (pre-upgrade)
	/// Test total staking power - according to Epoch 2 (post-upgrade)
	/// Test staker balance with attestation rewards - according to Epoch 2 (post-upgrade)
	/// Epoch 5:
	/// Test delegator balances with attestation rewards - according to Epoch 2 (post-upgrade)
	#[derive(Drop, Copy)]
	pub(crate) struct BalancesDelay2PoolTraceEntriesFlow {
	pub(crate) staker: Option<Staker>,
	pub(crate) stake_amount: Option<Amount>,
	pub(crate) commission: Option<Commission>,
	pub(crate) strk_delegated_amount: Option<Amount>,
	pub(crate) btc_delegated_amount: Option<Amount>,
	pub(crate) strk_delegator: Option<Delegator>,
	pub(crate) btc_delegator: Option<Delegator>,
	pub(crate) strk_pool: Option<ContractAddress>,
	pub(crate) btc_pool: Option<ContractAddress>,
	}
	pub(crate) impl BalancesDelay2PoolTraceEntriesFlowImpl of FlowTrait<
	BalancesDelay2PoolTraceEntriesFlow,
	> {
	fn setup_v2(ref self: BalancesDelay2PoolTraceEntriesFlow, ref system: SystemState) {
	let stake_amount = system.staking.get_min_stake();
	let strk_delegated_amount = STRK_CONFIG.min_for_rewards;
	let btc_delegated_amount = TEST_MIN_BTC_FOR_REWARDS;
	let staker = system.new_staker(amount: stake_amount * 4);
	let commission = 200;
	let strk_delegator = system.new_delegator(amount: strk_delegated_amount * 3);
	let btc_delegator = system
	.new_btc_delegator(amount: btc_delegated_amount * 3, token: system.btc_token);
	system.stake(:staker, amount: stake_amount, pool_enabled: true, :commission);
	let strk_pool = system.staking.get_pool(:staker);
	let btc_pool = system
	.set_open_for_delegation(:staker, token_address: system.btc_token.contract_address());
	system.advance_epoch(); // Epoch 0 - > 1
	system.increase_stake(:staker, amount: stake_amount);
	system.delegate(delegator: strk_delegator, pool: strk_pool, amount: strk_delegated_amount);
	system
	.delegate_btc(
	delegator: btc_delegator,
	pool: btc_pool,
	amount: btc_delegated_amount,
	token: system.btc_token,
	);
	system.advance_epoch(); // Epoch 1 - > 2
	system.increase_stake(:staker, amount: stake_amount);
	system
	.increase_delegate(
	delegator: strk_delegator, pool: strk_pool, amount: strk_delegated_amount,
	);
	system
	.increase_delegate_btc(
	delegator: btc_delegator,
	pool: btc_pool,
	amount: btc_delegated_amount,
	token: system.btc_token,
	);
	system.set_staker_for_migration(staker_address: staker.staker.address);
	self.staker = Option::Some(staker);
	self.stake_amount = Option::Some(stake_amount);
	self.commission = Option::Some(commission);
	self.strk_delegated_amount = Option::Some(strk_delegated_amount);
	self.btc_delegated_amount = Option::Some(btc_delegated_amount);
	self.strk_delegator = Option::Some(strk_delegator);
	self.btc_delegator = Option::Some(btc_delegator);
	self.strk_pool = Option::Some(strk_pool);
	self.btc_pool = Option::Some(btc_pool);
	}
	fn test(self: BalancesDelay2PoolTraceEntriesFlow, ref system: SystemState) {
	let staker = self.staker.unwrap();
	let stake_amount = self.stake_amount.unwrap();
	let commission = self.commission.unwrap();
	let strk_delegated_amount = self.strk_delegated_amount.unwrap();
	let btc_delegated_amount = self.btc_delegated_amount.unwrap();
	let strk_delegator = self.strk_delegator.unwrap();
	let btc_delegator = self.btc_delegator.unwrap();
	let strk_pool = self.strk_pool.unwrap();
	let btc_pool = self.btc_pool.unwrap();
	let staking_contract = system.staking.address;
	let minting_curve_contract = system.minting_curve.address;
	// Increase stake and delegations.
	system.increase_stake(:staker, amount: stake_amount);
	system
	.increase_delegate(
	delegator: strk_delegator, pool: strk_pool, amount: strk_delegated_amount,
	);
	system
	.increase_delegate_btc(
	delegator: btc_delegator,
	pool: btc_pool,
	amount: btc_delegated_amount,
	token: system.btc_token,
	);
	// Initialize variables for Epoch 1.
	let mut staker_stake = stake_amount * 2;
	let mut strk_pool_balance = strk_delegated_amount;
	let mut btc_pool_balance = btc_delegated_amount;
	// Test total staking power - according to Epoch 1.
	let total_staking_power = system.staking.get_current_total_staking_power_v2();
	let expected_total_staking_power = (
	NormalizedAmountTrait::from_strk_native_amount(staker_stake + strk_pool_balance),
	NormalizedAmountTrait::from_native_amount(btc_pool_balance, TEST_BTC_DECIMALS),
	);
	assert!(total_staking_power == expected_total_staking_power);
	// Calculate expected rewards - according to Epoch 1.
	let (expected_staker_rewards, expected_strk_pool_rewards) =
	calculate_staker_strk_rewards_with_balances_v2(
	amount_own: staker_stake,
	pool_amount: strk_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	);
	assert!(expected_staker_rewards.is_non_zero());
	assert!(expected_strk_pool_rewards.is_non_zero());
	let (expected_btc_commission_rewards, expected_btc_pool_rewards) =
	calculate_staker_btc_pool_rewards_v2(
	pool_balance: btc_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	token_address: system.btc_token.contract_address(),
	);
	assert!(expected_btc_commission_rewards.is_non_zero());
	assert!(expected_btc_pool_rewards.is_non_zero());
	// Test staker balance with attestation rewards - according to Epoch 1.
	system
	.advance_block_into_attestation_window_custom_stake(
	staker_address: staker.staker.address, stake: staker_stake + strk_pool_balance,
	);
	system.attest(:staker);
	let staker_rewards = system.staker_claim_rewards(:staker);
	assert!(staker_rewards == expected_staker_rewards + expected_btc_commission_rewards);
	// Advance epoch - test delegator balances with attestation rewards - according to Epoch 1.
	system.advance_epoch(); // Epoch 2 - > 3
	let strk_delegator_rewards = system
	.delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
	let btc_delegator_rewards = system
	.delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
	assert!(strk_delegator_rewards == expected_strk_pool_rewards);
	assert!(btc_delegator_rewards == expected_btc_pool_rewards);
	// Update variables for Epoch 2 (pre-upgrade).
	staker_stake += stake_amount;
	strk_pool_balance += strk_delegated_amount;
	btc_pool_balance += btc_delegated_amount;
	// Test total staking power - according to Epoch 2 (pre-upgrade).
	let total_staking_power = system.staking.get_current_total_staking_power_v2();
	let expected_total_staking_power = (
	NormalizedAmountTrait::from_strk_native_amount(staker_stake + strk_pool_balance),
	NormalizedAmountTrait::from_native_amount(btc_pool_balance, TEST_BTC_DECIMALS),
	);
	assert!(total_staking_power == expected_total_staking_power);
	// Calculate expected rewards - according to Epoch 2 (pre-upgrade).
	let prev_staker_rewards = expected_staker_rewards;
	let prev_strk_pool_rewards = expected_strk_pool_rewards;
	let prev_btc_commission_rewards = expected_btc_commission_rewards;
	let prev_btc_pool_rewards = expected_btc_pool_rewards;
	let (expected_staker_rewards, expected_strk_pool_rewards) =
	calculate_staker_strk_rewards_with_balances_v2(
	amount_own: staker_stake,
	pool_amount: strk_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	);
	assert!(expected_staker_rewards != prev_staker_rewards);
	assert!(expected_staker_rewards.is_non_zero());
	assert!(expected_strk_pool_rewards != prev_strk_pool_rewards);
	assert!(expected_strk_pool_rewards.is_non_zero());
	let (expected_btc_commission_rewards, expected_btc_pool_rewards) =
	calculate_staker_btc_pool_rewards_v2(
	pool_balance: btc_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	token_address: system.btc_token.contract_address(),
	);
	assert!(expected_btc_commission_rewards != prev_btc_commission_rewards);
	assert!(expected_btc_commission_rewards.is_non_zero());
	assert!(expected_btc_pool_rewards != prev_btc_pool_rewards);
	assert!(expected_btc_pool_rewards.is_non_zero());
	// Test staker balance with attestation rewards - according to Epoch 2 (pre-upgrade).
	system
	.advance_block_into_attestation_window_custom_stake(
	staker_address: staker.staker.address, stake: staker_stake + strk_pool_balance,
	);
	system.attest(:staker);
	let staker_rewards = system.staker_claim_rewards(:staker);
	assert!(staker_rewards == expected_staker_rewards + expected_btc_commission_rewards);
	// Advance epoch - test delegator balances with attestation rewards - according to Epoch 2
	// (pre-upgrade).
	system.advance_epoch(); // Epoch 3 - > 4
	let strk_delegator_rewards = system
	.delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
	let btc_delegator_rewards = system
	.delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
	assert!(strk_delegator_rewards == expected_strk_pool_rewards);
	assert!(btc_delegator_rewards == expected_btc_pool_rewards);
	// Update variables for Epoch 2 (post-upgrade).
	staker_stake += stake_amount;
	strk_pool_balance += strk_delegated_amount;
	btc_pool_balance += btc_delegated_amount;
	// Test total staking power - according to Epoch 2 (post-upgrade).
	let total_staking_power = system.staking.get_current_total_staking_power_v2();
	let expected_total_staking_power = (
	NormalizedAmountTrait::from_strk_native_amount(staker_stake + strk_pool_balance),
	NormalizedAmountTrait::from_native_amount(btc_pool_balance, TEST_BTC_DECIMALS),
	);
	assert!(total_staking_power == expected_total_staking_power);
	// Calculate expected rewards - according to Epoch 2 (post-upgrade).
	let prev_staker_rewards = expected_staker_rewards;
	let prev_strk_pool_rewards = expected_strk_pool_rewards;
	let prev_btc_commission_rewards = expected_btc_commission_rewards;
	let prev_btc_pool_rewards = expected_btc_pool_rewards;
	let (expected_staker_rewards, expected_strk_pool_rewards) =
	calculate_staker_strk_rewards_with_balances_v2(
	amount_own: staker_stake,
	pool_amount: strk_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	);
	assert!(expected_staker_rewards != prev_staker_rewards);
	assert!(expected_staker_rewards.is_non_zero());
	assert!(expected_strk_pool_rewards != prev_strk_pool_rewards);
	assert!(expected_strk_pool_rewards.is_non_zero());
	let (expected_btc_commission_rewards, expected_btc_pool_rewards) =
	calculate_staker_btc_pool_rewards_v2(
	pool_balance: btc_pool_balance,
	:commission,
	:staking_contract,
	:minting_curve_contract,
	token_address: system.btc_token.contract_address(),
	);
	assert!(expected_btc_commission_rewards != prev_btc_commission_rewards);
	assert!(expected_btc_commission_rewards.is_non_zero());
	assert!(expected_btc_pool_rewards != prev_btc_pool_rewards);
	assert!(expected_btc_pool_rewards.is_non_zero());
	// Test staker balance with attestation rewards - according to Epoch 2 (post-upgrade).
	system
	.advance_block_into_attestation_window_custom_stake(
	staker_address: staker.staker.address, stake: staker_stake + strk_pool_balance,
	);
	system.attest(:staker);
	let staker_rewards = system.staker_claim_rewards(:staker);
	assert!(staker_rewards == expected_staker_rewards + expected_btc_commission_rewards);
	// Advance epoch - test delegator balances with attestation rewards - according to Epoch 2
	// (post-upgrade).
	system.advance_epoch(); // Epoch 4 - > 5
	let strk_delegator_rewards = system
	.delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
	let btc_delegator_rewards = system
	.delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
	assert!(wide_abs_diff(strk_delegator_rewards, expected_strk_pool_rewards) <= 1);
	assert!(btc_delegator_rewards == expected_btc_pool_rewards);
	}
	}

 

---