envelope_cli/storage/
budget.rs

1//! Budget allocation repository for JSON storage
2//!
3//! Manages loading and saving budget allocations (shares budget.json with categories)
4
5use std::collections::HashMap;
6use std::path::PathBuf;
7use std::sync::RwLock;
8
9use crate::error::EnvelopeError;
10use crate::models::{BudgetAllocation, BudgetPeriod, CategoryId};
11
12use super::file_io::{read_json, write_json_atomic};
13
14/// Serializable budget data (extends CategoryData)
15#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
16struct BudgetData {
17    #[serde(default)]
18    allocations: Vec<BudgetAllocation>,
19}
20
21/// Composite key for budget allocations
22#[derive(Debug, Clone, PartialEq, Eq, Hash)]
23pub struct AllocationKey {
24    pub category_id: CategoryId,
25    pub period: BudgetPeriod,
26}
27
28impl AllocationKey {
29    pub fn new(category_id: CategoryId, period: BudgetPeriod) -> Self {
30        Self {
31            category_id,
32            period,
33        }
34    }
35}
36
37/// Repository for budget allocation persistence
38pub struct BudgetRepository {
39    path: PathBuf,
40    allocations: RwLock<HashMap<AllocationKey, BudgetAllocation>>,
41}
42
43impl BudgetRepository {
44    /// Create a new budget repository
45    pub fn new(path: PathBuf) -> Self {
46        Self {
47            path,
48            allocations: RwLock::new(HashMap::new()),
49        }
50    }
51
52    /// Load allocations from disk
53    pub fn load(&self) -> Result<(), EnvelopeError> {
54        let file_data: BudgetData = read_json(&self.path)?;
55
56        let mut allocations = self
57            .allocations
58            .write()
59            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire write lock: {}", e)))?;
60
61        allocations.clear();
62        for alloc in file_data.allocations {
63            let key = AllocationKey::new(alloc.category_id, alloc.period.clone());
64            allocations.insert(key, alloc);
65        }
66
67        Ok(())
68    }
69
70    /// Save allocations to disk
71    pub fn save(&self) -> Result<(), EnvelopeError> {
72        let allocations = self
73            .allocations
74            .read()
75            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire read lock: {}", e)))?;
76
77        let mut alloc_list: Vec<_> = allocations.values().cloned().collect();
78        alloc_list.sort_by(|a, b| a.period.cmp(&b.period));
79
80        let file_data = BudgetData {
81            allocations: alloc_list,
82        };
83
84        write_json_atomic(&self.path, &file_data)
85    }
86
87    /// Get an allocation for a category and period
88    pub fn get(
89        &self,
90        category_id: CategoryId,
91        period: &BudgetPeriod,
92    ) -> Result<Option<BudgetAllocation>, EnvelopeError> {
93        let allocations = self
94            .allocations
95            .read()
96            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire read lock: {}", e)))?;
97
98        let key = AllocationKey::new(category_id, period.clone());
99        Ok(allocations.get(&key).cloned())
100    }
101
102    /// Get or create an allocation (returns default if not found)
103    pub fn get_or_default(
104        &self,
105        category_id: CategoryId,
106        period: &BudgetPeriod,
107    ) -> Result<BudgetAllocation, EnvelopeError> {
108        if let Some(alloc) = self.get(category_id, period)? {
109            Ok(alloc)
110        } else {
111            Ok(BudgetAllocation::new(category_id, period.clone()))
112        }
113    }
114
115    /// Get all allocations for a period
116    pub fn get_for_period(
117        &self,
118        period: &BudgetPeriod,
119    ) -> Result<Vec<BudgetAllocation>, EnvelopeError> {
120        let allocations = self
121            .allocations
122            .read()
123            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire read lock: {}", e)))?;
124
125        Ok(allocations
126            .values()
127            .filter(|a| &a.period == period)
128            .cloned()
129            .collect())
130    }
131
132    /// Get all allocations for a category
133    pub fn get_for_category(
134        &self,
135        category_id: CategoryId,
136    ) -> Result<Vec<BudgetAllocation>, EnvelopeError> {
137        let allocations = self
138            .allocations
139            .read()
140            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire read lock: {}", e)))?;
141
142        let mut list: Vec<_> = allocations
143            .values()
144            .filter(|a| a.category_id == category_id)
145            .cloned()
146            .collect();
147        list.sort_by(|a, b| a.period.cmp(&b.period));
148        Ok(list)
149    }
150
151    /// Insert or update an allocation
152    pub fn upsert(&self, allocation: BudgetAllocation) -> Result<(), EnvelopeError> {
153        let mut allocations = self
154            .allocations
155            .write()
156            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire write lock: {}", e)))?;
157
158        let key = AllocationKey::new(allocation.category_id, allocation.period.clone());
159        allocations.insert(key, allocation);
160        Ok(())
161    }
162
163    /// Delete an allocation
164    pub fn delete(
165        &self,
166        category_id: CategoryId,
167        period: &BudgetPeriod,
168    ) -> Result<bool, EnvelopeError> {
169        let mut allocations = self
170            .allocations
171            .write()
172            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire write lock: {}", e)))?;
173
174        let key = AllocationKey::new(category_id, period.clone());
175        Ok(allocations.remove(&key).is_some())
176    }
177
178    /// Delete all allocations for a category
179    pub fn delete_for_category(&self, category_id: CategoryId) -> Result<usize, EnvelopeError> {
180        let mut allocations = self
181            .allocations
182            .write()
183            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire write lock: {}", e)))?;
184
185        let initial_count = allocations.len();
186        allocations.retain(|k, _| k.category_id != category_id);
187        Ok(initial_count - allocations.len())
188    }
189
190    /// Count allocations
191    pub fn count(&self) -> Result<usize, EnvelopeError> {
192        let allocations = self
193            .allocations
194            .read()
195            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire read lock: {}", e)))?;
196        Ok(allocations.len())
197    }
198
199    /// Get all allocations
200    pub fn get_all(&self) -> Result<Vec<BudgetAllocation>, EnvelopeError> {
201        let allocations = self
202            .allocations
203            .read()
204            .map_err(|e| EnvelopeError::Storage(format!("Failed to acquire read lock: {}", e)))?;
205
206        let mut list: Vec<_> = allocations.values().cloned().collect();
207        list.sort_by(|a, b| a.period.cmp(&b.period));
208        Ok(list)
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215    use crate::models::Money;
216    use tempfile::TempDir;
217
218    fn create_test_repo() -> (TempDir, BudgetRepository) {
219        let temp_dir = TempDir::new().unwrap();
220        let path = temp_dir.path().join("budget.json");
221        let repo = BudgetRepository::new(path);
222        (temp_dir, repo)
223    }
224
225    fn test_period() -> BudgetPeriod {
226        BudgetPeriod::monthly(2025, 1)
227    }
228
229    #[test]
230    fn test_empty_load() {
231        let (_temp_dir, repo) = create_test_repo();
232        repo.load().unwrap();
233        assert_eq!(repo.count().unwrap(), 0);
234    }
235
236    #[test]
237    fn test_upsert_and_get() {
238        let (_temp_dir, repo) = create_test_repo();
239        repo.load().unwrap();
240
241        let category_id = CategoryId::new();
242        let period = test_period();
243
244        let alloc =
245            BudgetAllocation::with_budget(category_id, period.clone(), Money::from_cents(50000));
246
247        repo.upsert(alloc).unwrap();
248
249        let retrieved = repo.get(category_id, &period).unwrap().unwrap();
250        assert_eq!(retrieved.budgeted.cents(), 50000);
251    }
252
253    #[test]
254    fn test_get_or_default() {
255        let (_temp_dir, repo) = create_test_repo();
256        repo.load().unwrap();
257
258        let category_id = CategoryId::new();
259        let period = test_period();
260
261        // Should return default (zero) if not found
262        let alloc = repo.get_or_default(category_id, &period).unwrap();
263        assert_eq!(alloc.budgeted.cents(), 0);
264
265        // Now insert
266        let alloc2 =
267            BudgetAllocation::with_budget(category_id, period.clone(), Money::from_cents(100));
268        repo.upsert(alloc2).unwrap();
269
270        // Should return the inserted value
271        let alloc3 = repo.get_or_default(category_id, &period).unwrap();
272        assert_eq!(alloc3.budgeted.cents(), 100);
273    }
274
275    #[test]
276    fn test_get_for_period() {
277        let (_temp_dir, repo) = create_test_repo();
278        repo.load().unwrap();
279
280        let cat1 = CategoryId::new();
281        let cat2 = CategoryId::new();
282        let jan = BudgetPeriod::monthly(2025, 1);
283        let feb = BudgetPeriod::monthly(2025, 2);
284
285        repo.upsert(BudgetAllocation::with_budget(
286            cat1,
287            jan.clone(),
288            Money::from_cents(100),
289        ))
290        .unwrap();
291        repo.upsert(BudgetAllocation::with_budget(
292            cat2,
293            jan.clone(),
294            Money::from_cents(200),
295        ))
296        .unwrap();
297        repo.upsert(BudgetAllocation::with_budget(
298            cat1,
299            feb.clone(),
300            Money::from_cents(300),
301        ))
302        .unwrap();
303
304        let jan_allocs = repo.get_for_period(&jan).unwrap();
305        assert_eq!(jan_allocs.len(), 2);
306
307        let feb_allocs = repo.get_for_period(&feb).unwrap();
308        assert_eq!(feb_allocs.len(), 1);
309    }
310
311    #[test]
312    fn test_save_and_reload() {
313        let (temp_dir, repo) = create_test_repo();
314        repo.load().unwrap();
315
316        let category_id = CategoryId::new();
317        let period = test_period();
318        let alloc =
319            BudgetAllocation::with_budget(category_id, period.clone(), Money::from_cents(50000));
320
321        repo.upsert(alloc).unwrap();
322        repo.save().unwrap();
323
324        // Create new repo and load
325        let path = temp_dir.path().join("budget.json");
326        let repo2 = BudgetRepository::new(path);
327        repo2.load().unwrap();
328
329        let retrieved = repo2.get(category_id, &period).unwrap().unwrap();
330        assert_eq!(retrieved.budgeted.cents(), 50000);
331    }
332
333    #[test]
334    fn test_delete() {
335        let (_temp_dir, repo) = create_test_repo();
336        repo.load().unwrap();
337
338        let category_id = CategoryId::new();
339        let period = test_period();
340        let alloc =
341            BudgetAllocation::with_budget(category_id, period.clone(), Money::from_cents(100));
342
343        repo.upsert(alloc).unwrap();
344        assert_eq!(repo.count().unwrap(), 1);
345
346        repo.delete(category_id, &period).unwrap();
347        assert_eq!(repo.count().unwrap(), 0);
348    }
349
350    #[test]
351    fn test_delete_for_category() {
352        let (_temp_dir, repo) = create_test_repo();
353        repo.load().unwrap();
354
355        let cat1 = CategoryId::new();
356        let cat2 = CategoryId::new();
357        let jan = BudgetPeriod::monthly(2025, 1);
358        let feb = BudgetPeriod::monthly(2025, 2);
359
360        repo.upsert(BudgetAllocation::new(cat1, jan.clone()))
361            .unwrap();
362        repo.upsert(BudgetAllocation::new(cat1, feb.clone()))
363            .unwrap();
364        repo.upsert(BudgetAllocation::new(cat2, jan.clone()))
365            .unwrap();
366
367        assert_eq!(repo.count().unwrap(), 3);
368
369        let deleted = repo.delete_for_category(cat1).unwrap();
370        assert_eq!(deleted, 2);
371        assert_eq!(repo.count().unwrap(), 1);
372    }
373}
envelope_cli/storage/budget.rs

envelope_cli/storage/
budget.rs
