rialo_feature_management_interface/
state.rs

1// Copyright (c) Subzero Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Feature state management
5
6extern crate alloc;
7
8use alloc::{
9    collections::BTreeMap,
10    string::{String, ToString},
11    vec::Vec,
12};
13
14use borsh::{BorshDeserialize, BorshSerialize};
15use rialo_s_pubkey::Pubkey;
16
17/// The program's global state
18///
19/// This is the top-level structure stored in the storage account.
20#[derive(Clone, Debug, PartialEq, BorshSerialize, BorshDeserialize)]
21pub struct FeaturesState {
22    /// Authority pubkey
23    authority: Pubkey,
24    /// The features map
25    pub features_map: BTreeMap<String, (u64, u64)>, // feature name -> (start_time, end_time)
26}
27
28#[cfg(test)]
29pub const DETERMINISTIC_TEST_KEYPAIR: &str =
30    "57Vqb7tHij5NhQnTgrgYXA19pC8ZVHQoCHpapSiQ8LJaeUvTcSBzKoB6CazhR6VtxmyVAbWnoeDSzD1Vm672NaKp";
31
32impl FeaturesState {
33    /// Create a new state with the given authority
34    pub fn new(authority: Pubkey) -> Self {
35        Self {
36            authority,
37            features_map: BTreeMap::new(),
38        }
39    }
40
41    #[cfg(test)]
42    pub fn new_for_test() -> Self {
43        use rialo_s_keypair::Keypair;
44        use rialo_s_signer::Signer;
45
46        let deterministic_test_pubkey: Pubkey =
47            Keypair::from_base58_string(DETERMINISTIC_TEST_KEYPAIR)
48                .try_pubkey()
49                .expect("Failed to get pubkey from deterministic test keypair");
50        Self {
51            authority: deterministic_test_pubkey,
52            features_map: BTreeMap::new(),
53        }
54    }
55
56    pub fn get_authority(&self) -> &rialo_s_pubkey::Pubkey {
57        &self.authority
58    }
59
60    pub fn set_authority(&mut self, new_authority: rialo_s_pubkey::Pubkey) {
61        self.authority = new_authority;
62    }
63
64    /// Serialize the state
65    pub fn serialize(&self) -> Result<Vec<u8>, borsh::io::Error> {
66        borsh::to_vec(self)
67    }
68
69    /// Deserialize the state
70    pub fn deserialize(data: &[u8]) -> Result<Self, borsh::io::Error> {
71        borsh::from_slice(data)
72    }
73
74    /// Check if the feature is currently active
75    pub fn is_active(&self, feature_name: &str, current_time: u64) -> bool {
76        if let Some((start_time, end_time)) = self.features_map.get(feature_name) {
77            current_time >= *start_time && current_time < *end_time
78        } else {
79            false
80        }
81    }
82
83    /// Add or update a feature
84    pub fn upsert(&mut self, name: String, start_time: u64, end_time: u64) -> Result<(), String> {
85        if start_time >= end_time {
86            return Err("Invalid time range: start_time must be less than end_time".to_string());
87        }
88
89        // Check if adding a new feature would exceed the maximum count
90        if !self.features_map.contains_key(&name)
91            && self.features_map.len() >= crate::MAX_FEATURE_COUNT
92        {
93            return Err(alloc::format!(
94                "Maximum feature count ({}) exceeded",
95                crate::MAX_FEATURE_COUNT
96            ));
97        }
98
99        self.features_map.insert(name, (start_time, end_time));
100        Ok(())
101    }
102
103    /// Get a feature's time range in (start, end) format
104    pub fn get(&self, name: &str) -> Option<&(u64, u64)> {
105        self.features_map.get(name)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn test_new_state() {
115        let authority = Pubkey::new_unique();
116        let state = FeaturesState::new(authority);
117
118        assert_eq!(state.get_authority(), &authority);
119        assert!(state.features_map.is_empty());
120    }
121
122    #[test]
123    fn test_upsert_feature() {
124        let mut state = FeaturesState::new_for_test();
125
126        let result = state.upsert("feature1".to_string(), 100, 200);
127        assert!(result.is_ok());
128        assert_eq!(state.get("feature1"), Some(&(100, 200)));
129    }
130
131    #[test]
132    fn test_upsert_invalid_time_range() {
133        let mut state = FeaturesState::new_for_test();
134
135        // start_time >= end_time should fail
136        let result = state.upsert("feature1".to_string(), 200, 100);
137        assert!(result.is_err());
138
139        let result = state.upsert("feature2".to_string(), 100, 100);
140        assert!(result.is_err());
141    }
142
143    #[test]
144    fn test_is_active() {
145        let mut state = FeaturesState::new_for_test();
146        state.upsert("feature1".to_string(), 100, 200).unwrap();
147
148        // Before start time
149        assert!(!state.is_active("feature1", 50));
150
151        // At start time
152        assert!(state.is_active("feature1", 100));
153
154        // During active period
155        assert!(state.is_active("feature1", 150));
156
157        // At end time (should be inactive)
158        assert!(!state.is_active("feature1", 200));
159
160        // After end time
161        assert!(!state.is_active("feature1", 250));
162
163        // Non-existent feature
164        assert!(!state.is_active("nonexistent", 150));
165    }
166
167    #[test]
168    fn test_serialize_deserialize() {
169        let mut state = FeaturesState::new_for_test();
170        state.upsert("feature1".to_string(), 100, 200).unwrap();
171        state.upsert("feature2".to_string(), 300, 400).unwrap();
172
173        let serialized = state.serialize().expect("Serialization failed");
174        let deserialized = FeaturesState::deserialize(&serialized).expect("Deserialization failed");
175
176        assert_eq!(state, deserialized);
177    }
178
179    #[test]
180    fn test_set_authority() {
181        let mut state = FeaturesState::new_for_test();
182        let new_authority = Pubkey::new_unique();
183
184        state.set_authority(new_authority);
185        assert_eq!(state.get_authority(), &new_authority);
186    }
187}
rialo_feature_management_interface/state.rs

rialo_feature_management_interface/
state.rs
