Slava Barinov 2 лет назад
Родитель
Сommit
b206fbc56b
1 измененных файлов с 73 добавлено и 4 удалено
  1. 73 4
      src/user.rs

+ 73 - 4
src/user.rs

@@ -1,10 +1,12 @@
 use crate::categories::CatStats;
+use derive_more::From;
 use pickledb::{PickleDb, PickleDbDumpPolicy, SerializationMethod};
 use radix_trie::Trie;
 use shellexpand::tilde;
 use std::collections::HashSet;
 use std::path::PathBuf;
 use std::time::Duration;
+use thiserror::Error;
 
 /// Configuration for single user
 pub struct User {
@@ -24,9 +26,15 @@ pub const DEFAULT_DB_PATH: &str = "~/.config/receqif/";
 #[cfg(feature = "docker")]
 pub const DEFAULT_DB_PATH: &str = "/etc/receqif/";
 
+#[derive(Debug, Error, From)]
+pub enum UserError {
+    #[error("Database error: {0}")]
+    DbError(#[source] pickledb::error::Error),
+}
+
 impl Drop for User {
     fn drop(&mut self) {
-        self.save_data();
+        self.save_data().unwrap();
     }
 }
 
@@ -76,9 +84,70 @@ impl User {
         self.accounts = HashSet::from_iter(acc);
     }
 
-    pub fn save_data(&mut self) {
+    pub fn save_data(&mut self) -> Result<(), UserError> {
         log::debug!("Saving user data");
-        self.db.set("catmap", &self.catmap).unwrap();
-        self.db.dump().unwrap();
+        self.db
+            .set("catmap", &self.catmap)
+            .map_err(UserError::DbError)?;
+
+        self.db
+            .set("accounts", &self.accounts)
+            .map_err(UserError::DbError)?;
+
+        self.db.dump().map_err(UserError::DbError)?;
+
+        Ok(())
+    }
+
+    pub fn new_account(&mut self, acc: String) {
+        self.accounts.insert(acc);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs::create_dir_all;
+
+    const TEST_DB_DIR: &str = "/tmp/receqif_test/";
+
+    fn setup(db_suffix: &str) -> Result<User, Box<dyn std::error::Error>> {
+        create_dir_all(TEST_DB_DIR)?;
+
+        let temp_db_path = format!("{}test_user_{}.db", TEST_DB_DIR, db_suffix);
+        let user = User::new(123, &Some(temp_db_path));
+        Ok(user)
+    }
+
+    #[test]
+    fn test_user_initialization() {
+        let user = setup("init").expect("Failed to initialize user");
+
+        assert!(user.accounts.is_empty());
+        assert_eq!(user.catmap, Trie::new());
+    }
+
+    #[test]
+    fn test_adding_account() {
+        let mut user = setup("add_acc").expect("Failed to set up user for adding account");
+
+        user.new_account("account".to_string());
+        assert!(user.accounts.contains("account"));
+    }
+
+    #[test]
+    fn test_saving_data() {
+        let mut user = setup("save_data").expect("Failed to set up user for saving data");
+
+        user.new_account("account".to_string());
+        user.save_data().expect("Failed to save data");
+
+        // Recreate the user object to verify data persistence
+        let reloaded_user = setup("save_data").expect("Failed to reload user for verifying data");
+
+        assert!(
+            reloaded_user.accounts.contains("account"),
+            "Account not found in reloaded user"
+        );
     }
 }