pyra-margin 0.4.2

Margin weight, balance, and price calculations for Drift spot positions
Documentation
use pyra_types::Vault;

/// Returns the effective remaining spend limit for the current timeframe in base units.
///
/// - If `timeframe_in_seconds == 0`, spending is disabled — returns 0.
/// - If the timeframe has expired (`next_timeframe_reset_timestamp < now_unix`),
///   returns the full `spend_limit_per_timeframe`.
/// - Otherwise, returns `remaining_spend_limit_per_timeframe`.
///
/// `now_unix` is the current Unix timestamp in seconds (e.g. `Utc::now().timestamp() as u64`).
/// It is passed as a parameter instead of calling `Utc::now()` directly to keep this
/// function pure and testable without time mocking.
pub fn get_remaining_timeframe_limit(vault: &Vault, now_unix: u64) -> u64 {
    if vault.timeframe_in_seconds == 0 {
        return 0;
    }

    if vault.next_timeframe_reset_timestamp < now_unix {
        vault.spend_limit_per_timeframe
    } else {
        vault.remaining_spend_limit_per_timeframe
    }
}

#[cfg(test)]
#[allow(
    clippy::allow_attributes,
    clippy::allow_attributes_without_reason,
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::arithmetic_side_effects,
    reason = "test code"
)]
mod tests {
    use super::*;

    fn make_vault(
        timeframe_in_seconds: u64,
        next_timeframe_reset_timestamp: u64,
        spend_limit_per_timeframe: u64,
        remaining_spend_limit_per_timeframe: u64,
    ) -> Vault {
        Vault {
            owner: vec![0; 32],
            bump: 0,
            spend_limit_per_transaction: 0,
            spend_limit_per_timeframe,
            remaining_spend_limit_per_timeframe,
            next_timeframe_reset_timestamp,
            timeframe_in_seconds,
        }
    }

    const NOW: u64 = 1_700_000_000;

    #[test]
    fn zero_timeframe_blocks_spending() {
        let vault = make_vault(0, 0, 1_000_000, 500_000);
        assert_eq!(get_remaining_timeframe_limit(&vault, NOW), 0);
    }

    #[test]
    fn zero_timeframe_blocks_even_with_future_reset() {
        let vault = make_vault(0, NOW + 9999, 1_000_000, 500_000);
        assert_eq!(get_remaining_timeframe_limit(&vault, NOW), 0);
    }

    #[test]
    fn expired_timeframe_returns_full_limit() {
        let vault = make_vault(86_400, NOW - 100, 1_000_000, 200_000);
        assert_eq!(get_remaining_timeframe_limit(&vault, NOW), 1_000_000);
    }

    #[test]
    fn active_timeframe_returns_remaining() {
        let vault = make_vault(86_400, NOW + 3600, 1_000_000, 300_000);
        assert_eq!(get_remaining_timeframe_limit(&vault, NOW), 300_000);
    }

    #[test]
    fn reset_at_exact_now_returns_remaining() {
        // next_reset == now means the reset hasn't happened yet (< vs <=)
        let vault = make_vault(86_400, NOW, 1_000_000, 400_000);
        assert_eq!(get_remaining_timeframe_limit(&vault, NOW), 400_000);
    }

    #[test]
    fn reset_one_second_before_now_returns_full() {
        let vault = make_vault(86_400, NOW - 1, 1_000_000, 100_000);
        assert_eq!(get_remaining_timeframe_limit(&vault, NOW), 1_000_000);
    }
}