templar_common/vault/
lock.rs

1use super::*;
2
3/// Tracks which markets are currently locked for exclusive operation (e.g. during rebalance withdrawals).
4#[derive(Default)]
5#[near(serializers = [borsh, serde])]
6pub struct Locker {
7    to_lock: Vec<MarketId>,
8}
9
10impl Locker {
11    pub fn lock(&mut self, market: MarketId) {
12        if self.is_locked(market) {
13            env::panic_str("Market is locked");
14        }
15        Event::LockChange {
16            is_locked: true,
17            market,
18        }
19        .emit();
20        self.to_lock.push(market);
21    }
22
23    pub fn unlock(&mut self, market: MarketId) {
24        if !self.is_locked(market) {
25            return;
26        }
27        Event::LockChange {
28            is_locked: false,
29            market,
30        }
31        .emit();
32        self.to_lock.retain(|&x| x != market);
33    }
34
35    /// Clears the lock status for all markets.
36    /// This method should be used with caution as it will unlock all markets
37    pub fn clear(&mut self) {
38        for market in self.to_lock.iter().copied() {
39            Event::LockChange {
40                is_locked: false,
41                market,
42            }
43            .emit();
44        }
45        self.to_lock.clear();
46    }
47
48    pub fn is_locked(&self, market: MarketId) -> bool {
49        self.to_lock.contains(&market)
50    }
51
52    pub fn is_locked_all(&self) -> bool {
53        !self.to_lock.is_empty()
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::Locker;
60    use crate::vault::MarketId;
61    use near_sdk::{test_utils::VMContextBuilder, testing_env};
62
63    #[test]
64    fn lock_unlock_and_clear_track_state() {
65        testing_env!(VMContextBuilder::new().build());
66
67        let first = MarketId(1);
68        let second = MarketId(2);
69        let mut locker = Locker::default();
70
71        locker.lock(first);
72        locker.lock(second);
73        assert!(locker.is_locked(first));
74        assert!(locker.is_locked(second));
75        assert!(locker.is_locked_all());
76
77        locker.unlock(first);
78        assert!(!locker.is_locked(first));
79        assert!(locker.is_locked(second));
80
81        locker.clear();
82        assert!(!locker.is_locked_all());
83        assert!(!locker.is_locked(second));
84    }
85
86    #[test]
87    fn unlock_only_emits_when_state_changes() {
88        testing_env!(VMContextBuilder::new().build());
89
90        let market = MarketId(7);
91        let mut locker = Locker::default();
92
93        locker.unlock(market);
94        assert!(near_sdk::test_utils::get_logs().is_empty());
95
96        locker.lock(market);
97        locker.unlock(market);
98
99        let logs = near_sdk::test_utils::get_logs().join("\n");
100        assert!(logs.contains("\"event\":\"lock_change\""));
101        assert!(logs.contains("\"is_locked\":false"));
102    }
103}