From 1364db3f9453e9c30576e00002c70bea87145240 Mon Sep 17 00:00:00 2001
From: kristof <kristof@incubaid.com>
Date: Thu, 3 Apr 2025 13:58:44 +0200
Subject: [PATCH] ...

---
 herodb/readme.md                            |   5 +
 herodb/src/core/db.rs                       | 342 -------------
 herodb/src/zaz/mod.rs                       |   3 +
 herodb/src/zaz/tests/db_integration_test.rs | 537 ++++----------------
 herodb/src/zaz/tests/mod.rs                 |   6 +
 herodb/src/zaz/tests/model_db_test.rs       |  13 +-
 6 files changed, 122 insertions(+), 784 deletions(-)
 create mode 100644 herodb/readme.md
 create mode 100644 herodb/src/zaz/tests/mod.rs

diff --git a/herodb/readme.md b/herodb/readme.md
new file mode 100644
index 0000000..cc7c647
--- /dev/null
+++ b/herodb/readme.md
@@ -0,0 +1,5 @@
+
+
+```bash
+cargo test zaz::tests -- --test-threads=1
+```
\ No newline at end of file
diff --git a/herodb/src/core/db.rs b/herodb/src/core/db.rs
index e7d1d2e..4c7f70c 100644
--- a/herodb/src/core/db.rs
+++ b/herodb/src/core/db.rs
@@ -424,345 +424,3 @@ impl DB {
         }
     }
 }
-
-// Test module with mocked models
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use chrono::Utc;
-    use tempfile::tempdir;
-    use serde::{Deserialize, Serialize};
-    
-    // Test model
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    struct TestUser {
-        id: u32,
-        name: String,
-        email: String,
-    }
-    
-    impl TestUser {
-        fn new(id: u32, name: String, email: String) -> Self {
-            Self { id, name, email }
-        }
-    }
-    
-    impl Storable for TestUser {}
-    
-    impl SledModel for TestUser {
-        fn get_id(&self) -> String {
-            self.id.to_string()
-        }
-        
-        fn db_prefix() -> &'static str {
-            "test_user"
-        }
-    }
-    
-    #[test]
-    fn test_db_builder() {
-        // Create a temporary directory for the test
-        let dir = tempdir().expect("Failed to create temp dir");
-        
-        // Create a DB with the builder
-        let db = DBBuilder::new(dir.path())
-            .register_model::<TestUser>()
-            .build()
-            .expect("Failed to build DB");
-        
-        // Create a test user
-        let user = TestUser::new(1, "Test User".to_string(), "test@example.com".to_string());
-        
-        // Set the user
-        db.set(&user).expect("Failed to set user");
-        
-        // Get the user
-        let retrieved: TestUser = db.get(&user.id.to_string()).expect("Failed to get user");
-        
-        // Check that it matches
-        assert_eq!(user, retrieved);
-    }
-    
-    #[test]
-    fn test_dynamic_registration() {
-        // Create a temporary directory for the test
-        let dir = tempdir().expect("Failed to create temp dir");
-        
-        // Create an empty DB
-        let mut db = DB::new(dir.path()).expect("Failed to create DB");
-        
-        // Register the TestUser model
-        db.register::<TestUser>().expect("Failed to register TestUser");
-        
-        // Create a test user
-        let user = TestUser::new(2, "Dynamic User".to_string(), "dynamic@example.com".to_string());
-        
-        // Set the user
-        db.set(&user).expect("Failed to set user");
-        
-        // Get the user
-        let retrieved: TestUser = db.get(&user.id.to_string()).expect("Failed to get user");
-        
-        // Check that it matches
-        assert_eq!(user, retrieved);
-    }
-}
-
-// The as_type function is no longer needed with our type-map based implementation
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use chrono::Utc;
-    use tempfile::tempdir;
-    
-    #[test]
-    fn test_read_your_writes() {
-        // Create a temporary directory for the test
-        let dir = tempdir().expect("Failed to create temp dir");
-        let db = DB::new(dir.path()).expect("Failed to create DB");
-        
-        // Create a user
-        let user = User::new(
-            10,
-            "Original User".to_string(),
-            "original@example.com".to_string(),
-            "password".to_string(),
-            "Original Corp".to_string(),
-            "User".to_string(),
-        );
-        
-        // Insert the user directly (no transaction)
-        db.set(&user).expect("Failed to insert user");
-        
-        // Begin a transaction
-        db.begin_transaction().expect("Failed to begin transaction");
-        
-        // Verify we can read the original user
-        let original = db.get::<User>(&user.id.to_string()).expect("Failed to get original user");
-        assert_eq!(original.name, "Original User");
-        
-        // Create a modified user with the same ID
-        let modified_user = User::new(
-            10,
-            "Modified User".to_string(),
-            "modified@example.com".to_string(),
-            "new_password".to_string(),
-            "Modified Corp".to_string(),
-            "Admin".to_string(),
-        );
-        
-        // Update the user in the transaction
-        db.set(&modified_user).expect("Failed to update user in transaction");
-        
-        // Verify we can read our own writes within the transaction
-        let in_transaction = db.get::<User>(&user.id.to_string()).expect("Failed to get user from transaction");
-        assert_eq!(in_transaction.name, "Modified User");
-        assert_eq!(in_transaction.email, "modified@example.com");
-        
-        // Create a new user that only exists in the transaction
-        let new_user = User::new(
-            20,
-            "Transaction Only User".to_string(),
-            "tx@example.com".to_string(),
-            "password".to_string(),
-            "TX Corp".to_string(),
-            "Admin".to_string(),
-        );
-        
-        // Add the new user in the transaction
-        db.set(&new_user).expect("Failed to add new user in transaction");
-        
-        // Verify we can read the new user within the transaction
-        let new_in_tx = db.get::<User>(&new_user.id.to_string()).expect("Failed to get new user from transaction");
-        assert_eq!(new_in_tx.name, "Transaction Only User");
-        
-        // Delete a user in the transaction and verify it appears deleted within the transaction
-        db.delete::<User>(&user.id.to_string()).expect("Failed to delete user in transaction");
-        match db.get::<User>(&user.id.to_string()) {
-            Err(SledDBError::NotFound(_)) => (), // Expected result
-            Ok(_) => panic!("User should appear deleted within transaction"),
-            Err(e) => panic!("Unexpected error: {}", e),
-        }
-        
-        // Rollback the transaction
-        db.rollback_transaction().expect("Failed to rollback transaction");
-        
-        // Verify the original user is still available and unchanged after rollback
-        let after_rollback = db.get::<User>(&user.id.to_string()).expect("Failed to get user after rollback");
-        assert_eq!(after_rollback.name, "Original User");
-        
-        // Verify the transaction-only user doesn't exist after rollback
-        assert!(db.get::<User>(&new_user.id.to_string()).is_err());
-    }
-
-    #[test]
-    fn test_transactions() {
-        // Create a temporary directory for the test
-        let dir = tempdir().expect("Failed to create temp dir");
-        let db = DB::new(dir.path()).expect("Failed to create DB");
-        
-        // Create a sample user and company for testing
-        let user = User::new(
-            1,
-            "Transaction Test User".to_string(),
-            "tx@example.com".to_string(),
-            "password".to_string(),
-            "Test Corp".to_string(),
-            "Admin".to_string(),
-        );
-        
-        let incorporation_date = Utc::now();
-        let company = Company::new(
-            1,
-            "Transaction Test Corp".to_string(),
-            "TX123".to_string(),
-            incorporation_date,
-            "12-31".to_string(),
-            "tx@corp.com".to_string(),
-            "123-456-7890".to_string(),
-            "www.testcorp.com".to_string(),
-            "123 Test St".to_string(),
-            BusinessType::Global,
-            "Tech".to_string(),
-            "A test company for transactions".to_string(),
-            CompanyStatus::Active,
-        );
-        
-        // Test successful transaction (multiple operations committed at once)
-        {
-            // Start a transaction
-            db.begin_transaction().expect("Failed to begin transaction");
-            assert!(db.has_active_transaction());
-            
-            // Perform multiple operations within the transaction
-            db.set(&user).expect("Failed to add user to transaction");
-            db.set(&company).expect("Failed to add company to transaction");
-            
-            // Commit the transaction
-            db.commit_transaction().expect("Failed to commit transaction");
-            assert!(!db.has_active_transaction());
-            
-            // Verify both operations were applied
-            let retrieved_user: User = db.get(&user.id.to_string()).expect("Failed to get user after commit");
-            let retrieved_company: Company = db.get(&company.id.to_string()).expect("Failed to get company after commit");
-            
-            assert_eq!(user.name, retrieved_user.name);
-            assert_eq!(company.name, retrieved_company.name);
-        }
-        
-        // Test transaction rollback
-        {
-            // Create another user that should not be persisted
-            let temp_user = User::new(
-                2,
-                "Temporary User".to_string(),
-                "temp@example.com".to_string(),
-                "password".to_string(),
-                "Temp Corp".to_string(),
-                "Temp".to_string(),
-            );
-            
-            // Start a transaction
-            db.begin_transaction().expect("Failed to begin transaction");
-            
-            // Add the temporary user
-            db.set(&temp_user).expect("Failed to add temporary user to transaction");
-            
-            // Perform a delete operation in the transaction
-            db.delete::<Company>(&company.id.to_string()).expect("Failed to delete company in transaction");
-            
-            // Rollback the transaction - should discard all operations
-            db.rollback_transaction().expect("Failed to rollback transaction");
-            assert!(!db.has_active_transaction());
-            
-            // Verify the temporary user was not added
-            match db.get::<User>(&temp_user.id.to_string()) {
-                Err(SledDBError::NotFound(_)) => (), // Expected outcome
-                Ok(_) => panic!("Temporary user should not exist after rollback"),
-                Err(e) => panic!("Unexpected error: {}", e),
-            }
-            
-            // Verify the company was not deleted
-            let company_still_exists = db.get::<Company>(&company.id.to_string()).is_ok();
-            assert!(company_still_exists, "Company should still exist after transaction rollback");
-        }
-    }
-    
-    #[test]
-    fn test_generic_db_operations() {
-        // Create a temporary directory for the test
-        let dir = tempdir().expect("Failed to create temp dir");
-        let db = DB::new(dir.path()).expect("Failed to create DB");
-        
-        // Test simple transaction functionality
-        assert!(!db.has_active_transaction());
-        db.begin_transaction().expect("Failed to begin transaction");
-        assert!(db.has_active_transaction());
-        db.rollback_transaction().expect("Failed to rollback transaction");
-        assert!(!db.has_active_transaction());
-        
-        // Create a sample user
-        let user = User::new(
-            1,
-            "Test User".to_string(),
-            "test@example.com".to_string(),
-            "password".to_string(),
-            "Test Corp".to_string(),
-            "Admin".to_string(),
-        );
-        
-        // Insert the user
-        db.set(&user).expect("Failed to insert user");
-        
-        // Get the user
-        let retrieved_user: User = db.get(&user.id.to_string()).expect("Failed to get user");
-        assert_eq!(user.name, retrieved_user.name);
-        
-        // Create a sample company
-        let incorporation_date = Utc::now();
-        let company = Company::new(
-            1,
-            "Test Corp".to_string(),
-            "REG123".to_string(),
-            incorporation_date,
-            "12-31".to_string(),
-            "test@corp.com".to_string(),
-            "123-456-7890".to_string(),
-            "www.testcorp.com".to_string(),
-            "123 Test St".to_string(),
-            BusinessType::Global,
-            "Tech".to_string(),
-            "A test company".to_string(),
-            CompanyStatus::Active,
-        );
-        
-        // Insert the company
-        db.set(&company).expect("Failed to insert company");
-        
-        // Get the company
-        let retrieved_company: Company = db.get(&company.id.to_string())
-            .expect("Failed to get company");
-        assert_eq!(company.name, retrieved_company.name);
-        
-        // List all companies
-        let companies: Vec<Company> = db.list().expect("Failed to list companies");
-        assert_eq!(companies.len(), 1);
-        assert_eq!(companies[0].name, company.name);
-        
-        // List all users
-        let users: Vec<User> = db.list().expect("Failed to list users");
-        assert_eq!(users.len(), 1);
-        assert_eq!(users[0].name, user.name);
-        
-        // Delete the company
-        db.delete::<Company>(&company.id.to_string())
-            .expect("Failed to delete company");
-        
-        // Try to get the deleted company (should fail)
-        match db.get::<Company>(&company.id.to_string()) {
-            Err(SledDBError::NotFound(_)) => (),
-            _ => panic!("Expected NotFound error"),
-        }
-    }
-}
diff --git a/herodb/src/zaz/mod.rs b/herodb/src/zaz/mod.rs
index 0916b49..9557194 100644
--- a/herodb/src/zaz/mod.rs
+++ b/herodb/src/zaz/mod.rs
@@ -14,3 +14,6 @@ pub mod examples;
 // Expose the cmd module
 pub mod cmd;
 
+// Include tests module when running tests
+#[cfg(test)]
+pub mod tests;
diff --git a/herodb/src/zaz/tests/db_integration_test.rs b/herodb/src/zaz/tests/db_integration_test.rs
index 7bc2e19..ef4c3e6 100644
--- a/herodb/src/zaz/tests/db_integration_test.rs
+++ b/herodb/src/zaz/tests/db_integration_test.rs
@@ -1,7 +1,8 @@
 //! Integration tests for the zaz database module
 
-use sled;
-use bincode;
+use crate::core::{DB, DBBuilder, SledDBResult, Storable, SledModel, SledDB};
+use crate::zaz::models::user::User;
+use crate::zaz::models::company::{Company, BusinessType, CompanyStatus};
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
 use std::fs;
@@ -19,113 +20,6 @@ fn test_basic_database_operations() {
 }
 
 fn run_comprehensive_test() -> Result<(), Box<dyn std::error::Error>> {
-    // User model
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    struct User {
-        id: u32,
-        name: String,
-        email: String,
-        password: String,
-        company: String,
-        role: String,
-        created_at: DateTime<Utc>,
-        updated_at: DateTime<Utc>,
-    }
-    
-    impl User {
-        fn new(
-            id: u32,
-            name: String,
-            email: String,
-            password: String,
-            company: String,
-            role: String,
-        ) -> Self {
-            let now = Utc::now();
-            Self {
-                id,
-                name,
-                email,
-                password,
-                company,
-                role,
-                created_at: now,
-                updated_at: now,
-            }
-        }
-    }
-    
-    // Company model
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    enum BusinessType {
-        Local,
-        National,
-        Global,
-    }
-    
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    enum CompanyStatus {
-        Active,
-        Inactive,
-        Pending,
-    }
-    
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    struct Company {
-        id: u32,
-        name: String,
-        registration_number: String,
-        registration_date: DateTime<Utc>,
-        fiscal_year_end: String,
-        email: String,
-        phone: String,
-        website: String,
-        address: String,
-        business_type: BusinessType,
-        industry: String,
-        description: String,
-        status: CompanyStatus,
-        created_at: DateTime<Utc>,
-        updated_at: DateTime<Utc>,
-    }
-    
-    impl Company {
-        fn new(
-            id: u32,
-            name: String,
-            registration_number: String,
-            registration_date: DateTime<Utc>,
-            fiscal_year_end: String,
-            email: String,
-            phone: String,
-            website: String,
-            address: String,
-            business_type: BusinessType,
-            industry: String,
-            description: String,
-            status: CompanyStatus,
-        ) -> Self {
-            let now = Utc::now();
-            Self {
-                id,
-                name,
-                registration_number,
-                registration_date,
-                fiscal_year_end,
-                email,
-                phone,
-                website,
-                address,
-                business_type,
-                industry,
-                description,
-                status,
-                created_at: now,
-                updated_at: now,
-            }
-        }
-    }
-    
     // Create a temporary directory for testing
     let temp_dir = tempdir()?;
     println!("Using temporary directory: {:?}", temp_dir.path());
@@ -147,44 +41,8 @@ fn run_comprehensive_test() -> Result<(), Box<dyn std::error::Error>> {
 }
 
 fn test_user_operations(base_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
-    // User model (duplicate for scope)
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    struct User {
-        id: u32,
-        name: String,
-        email: String,
-        password: String,
-        company: String,
-        role: String,
-        created_at: DateTime<Utc>,
-        updated_at: DateTime<Utc>,
-    }
-    
-    impl User {
-        fn new(
-            id: u32,
-            name: String,
-            email: String,
-            password: String,
-            company: String,
-            role: String,
-        ) -> Self {
-            let now = Utc::now();
-            Self {
-                id,
-                name,
-                email,
-                password,
-                company,
-                role,
-                created_at: now,
-                updated_at: now,
-            }
-        }
-    }
-    
     // Open the user database
-    let db = sled::open(base_path.join("users"))?;
+    let db = SledDB::<User>::open(base_path.join("users"))?;
     println!("Opened user database at: {:?}", base_path.join("users"));
     
     // Create a test user
@@ -198,21 +56,14 @@ fn test_user_operations(base_path: &Path) -> Result<(), Box<dyn std::error::Erro
     );
     
     // Insert the user
-    let user_id = user.id.to_string();
-    let user_bytes = bincode::serialize(&user)?;
-    db.insert(user_id.as_bytes(), user_bytes)?;
-    db.flush()?;
+    db.insert(&user)?;
     println!("Inserted user: {}", user.name);
     
     // Retrieve the user
-    if let Some(data) = db.get(user_id.as_bytes())? {
-        let retrieved_user: User = bincode::deserialize(&data)?;
-        println!("Retrieved user: {}", retrieved_user.name);
-        assert_eq!(user.name, retrieved_user.name);
-        assert_eq!(user.email, retrieved_user.email);
-    } else {
-        return Err("Failed to retrieve user".into());
-    }
+    let retrieved_user = db.get(&user.id.to_string())?;
+    println!("Retrieved user: {}", retrieved_user.name);
+    assert_eq!(user.name, retrieved_user.name);
+    assert_eq!(user.email, retrieved_user.email);
     
     // Update the user
     let updated_user = User::new(
@@ -224,108 +75,30 @@ fn test_user_operations(base_path: &Path) -> Result<(), Box<dyn std::error::Erro
         "SuperAdmin".to_string(),
     );
     
-    let updated_bytes = bincode::serialize(&updated_user)?;
-    db.insert(user_id.as_bytes(), updated_bytes)?;
-    db.flush()?;
+    db.insert(&updated_user)?;
     println!("Updated user: {}", updated_user.name);
     
     // Retrieve the updated user
-    if let Some(data) = db.get(user_id.as_bytes())? {
-        let retrieved_user: User = bincode::deserialize(&data)?;
-        println!("Retrieved updated user: {}", retrieved_user.name);
-        assert_eq!(updated_user.name, retrieved_user.name);
-        assert_eq!(updated_user.email, retrieved_user.email);
-    } else {
-        return Err("Failed to retrieve updated user".into());
-    }
+    let retrieved_user = db.get(&user.id.to_string())?;
+    println!("Retrieved updated user: {}", retrieved_user.name);
+    assert_eq!(updated_user.name, retrieved_user.name);
+    assert_eq!(updated_user.email, retrieved_user.email);
     
     // Delete the user
-    db.remove(user_id.as_bytes())?;
-    db.flush()?;
+    db.delete(&user.id.to_string())?;
     println!("Deleted user: {}", user.name);
     
     // Try to retrieve the deleted user (should fail)
-    let result = db.get(user_id.as_bytes())?;
-    assert!(result.is_none(), "User should be deleted");
+    let result = db.get(&user.id.to_string());
+    assert!(result.is_err(), "User should be deleted");
     println!("Verified user was deleted");
     
     Ok(())
 }
 
 fn test_company_operations(base_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
-    // Company model (duplicate for scope)
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    enum BusinessType {
-        Local,
-        National,
-        Global,
-    }
-    
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    enum CompanyStatus {
-        Active,
-        Inactive,
-        Pending,
-    }
-    
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    struct Company {
-        id: u32,
-        name: String,
-        registration_number: String,
-        registration_date: DateTime<Utc>,
-        fiscal_year_end: String,
-        email: String,
-        phone: String,
-        website: String,
-        address: String,
-        business_type: BusinessType,
-        industry: String,
-        description: String,
-        status: CompanyStatus,
-        created_at: DateTime<Utc>,
-        updated_at: DateTime<Utc>,
-    }
-    
-    impl Company {
-        fn new(
-            id: u32,
-            name: String,
-            registration_number: String,
-            registration_date: DateTime<Utc>,
-            fiscal_year_end: String,
-            email: String,
-            phone: String,
-            website: String,
-            address: String,
-            business_type: BusinessType,
-            industry: String,
-            description: String,
-            status: CompanyStatus,
-        ) -> Self {
-            let now = Utc::now();
-            Self {
-                id,
-                name,
-                registration_number,
-                registration_date,
-                fiscal_year_end,
-                email,
-                phone,
-                website,
-                address,
-                business_type,
-                industry,
-                description,
-                status,
-                created_at: now,
-                updated_at: now,
-            }
-        }
-    }
-    
     // Open the company database
-    let db = sled::open(base_path.join("companies"))?;
+    let db = SledDB::<Company>::open(base_path.join("companies"))?;
     println!("Opened company database at: {:?}", base_path.join("companies"));
     
     // Create a test company
@@ -346,43 +119,25 @@ fn test_company_operations(base_path: &Path) -> Result<(), Box<dyn std::error::E
     );
     
     // Insert the company
-    let company_id = company.id.to_string();
-    let company_bytes = bincode::serialize(&company)?;
-    db.insert(company_id.as_bytes(), company_bytes)?;
-    db.flush()?;
+    db.insert(&company)?;
     println!("Inserted company: {}", company.name);
     
     // Retrieve the company
-    if let Some(data) = db.get(company_id.as_bytes())? {
-        let retrieved_company: Company = bincode::deserialize(&data)?;
-        println!("Retrieved company: {}", retrieved_company.name);
-        assert_eq!(company.name, retrieved_company.name);
-    } else {
-        return Err("Failed to retrieve company".into());
-    }
+    let retrieved_company = db.get(&company.id.to_string())?;
+    println!("Retrieved company: {}", retrieved_company.name);
+    assert_eq!(company.name, retrieved_company.name);
     
     // List all companies
-    let mut companies = Vec::new();
-    for item in db.iter() {
-        let (_key, value) = item?;
-        let company: Company = bincode::deserialize(&value)?;
-        companies.push(company);
-    }
+    let companies = db.list()?;
     println!("Found {} companies", companies.len());
     assert_eq!(companies.len(), 1);
     
     // Delete the company
-    db.remove(company_id.as_bytes())?;
-    db.flush()?;
+    db.delete(&company.id.to_string())?;
     println!("Deleted company: {}", company.name);
     
     // List companies again (should be empty)
-    let mut companies = Vec::new();
-    for item in db.iter() {
-        let (_key, value) = item?;
-        let company: Company = bincode::deserialize(&value)?;
-        companies.push(company);
-    }
+    let companies = db.list()?;
     assert_eq!(companies.len(), 0);
     println!("Verified company was deleted");
     
@@ -390,45 +145,11 @@ fn test_company_operations(base_path: &Path) -> Result<(), Box<dyn std::error::E
 }
 
 fn test_transaction_simulation(base_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
-    // User model (duplicate for scope)
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-    struct User {
-        id: u32,
-        name: String,
-        email: String,
-        password: String,
-        company: String,
-        role: String,
-        created_at: DateTime<Utc>,
-        updated_at: DateTime<Utc>,
-    }
+    // Create a DB instance with User model registered
+    let mut db = DB::new(base_path.join("transaction"))?;
+    db.register::<User>()?;
     
-    impl User {
-        fn new(
-            id: u32,
-            name: String,
-            email: String,
-            password: String,
-            company: String,
-            role: String,
-        ) -> Self {
-            let now = Utc::now();
-            Self {
-                id,
-                name,
-                email,
-                password,
-                company,
-                role,
-                created_at: now,
-                updated_at: now,
-            }
-        }
-    }
-    
-    // Open the user database
-    let db = sled::open(base_path.join("tx_users"))?;
-    println!("Opened transaction test database at: {:?}", base_path.join("tx_users"));
+    println!("Created DB with User model registered at: {:?}", base_path.join("transaction"));
     
     // Add a user outside of transaction
     let user = User::new(
@@ -440,28 +161,15 @@ fn test_transaction_simulation(base_path: &Path) -> Result<(), Box<dyn std::erro
         "User".to_string(),
     );
     
-    let user_id = user.id.to_string();
-    let user_bytes = bincode::serialize(&user)?;
-    db.insert(user_id.as_bytes(), user_bytes)?;
-    db.flush()?;
+    // Add the user without a transaction
+    db.set(&user)?;
     println!("Added initial user: {}", user.name);
     
-    // Since sled doesn't have explicit transaction support like the DB mock in the original code,
-    // we'll simulate transaction behavior by:
-    // 1. Making changes in memory
-    // 2. Only writing to the database when we "commit"
-    println!("Simulating transaction operations...");
+    // Begin a transaction
+    db.begin_transaction()?;
+    println!("Transaction started");
     
-    // Create in-memory copy of our data (transaction workspace)
-    let mut tx_workspace = std::collections::HashMap::new();
-    
-    // Retrieve initial state from db
-    if let Some(data) = db.get(user_id.as_bytes())? {
-        let retrieved_user: User = bincode::deserialize(&data)?;
-        tx_workspace.insert(user_id.clone(), retrieved_user);
-    }
-    
-    // Update user in transaction workspace
+    // Update user in transaction
     let updated_user = User::new(
         200,
         "Updated in TX".to_string(),
@@ -470,10 +178,10 @@ fn test_transaction_simulation(base_path: &Path) -> Result<(), Box<dyn std::erro
         "New Corp".to_string(),
         "Admin".to_string(),
     );
-    tx_workspace.insert(user_id.clone(), updated_user.clone());
-    println!("Updated user in transaction workspace");
+    db.set(&updated_user)?;
+    println!("Updated user in transaction");
     
-    // Add new user in transaction workspace
+    // Add new user in transaction
     let new_user = User::new(
         201,
         "New in TX".to_string(),
@@ -482,85 +190,59 @@ fn test_transaction_simulation(base_path: &Path) -> Result<(), Box<dyn std::erro
         "New Corp".to_string(),
         "User".to_string(),
     );
-    let new_user_id = new_user.id.to_string();
-    tx_workspace.insert(new_user_id.clone(), new_user.clone());
-    println!("Added new user in transaction workspace");
+    db.set(&new_user)?;
+    println!("Added new user in transaction");
     
-    // Verify the transaction workspace state
-    let tx_user = tx_workspace.get(&user_id).unwrap();
+    // Verify transaction changes are visible within transaction
+    let tx_user: User = db.get(&user.id.to_string())?;
     assert_eq!(tx_user.name, "Updated in TX");
-    println!("Verified transaction changes are visible within workspace");
+    println!("Verified transaction changes are visible");
     
-    // Simulate a rollback by discarding our workspace without writing to db
-    println!("Rolled back transaction (discarded workspace without writing to db)");
+    // Commit the transaction
+    db.commit_transaction()?;
+    println!("Transaction committed");
     
-    // Verify original user is unchanged in the database
-    if let Some(data) = db.get(user_id.as_bytes())? {
-        let original: User = bincode::deserialize(&data)?;
-        assert_eq!(original.name, "Transaction Test");
-        println!("Verified original user is unchanged after rollback");
-    } else {
-        return Err("Failed to retrieve user after rollback".into());
-    }
+    // Verify changes persisted after commit
+    let committed_user: User = db.get(&user.id.to_string())?;
+    assert_eq!(committed_user.name, "Updated in TX");
+    println!("Verified changes persisted after commit");
     
-    // Verify new user was not added to the database
-    let result = db.get(new_user_id.as_bytes())?;
-    assert!(result.is_none());
-    println!("Verified new user was not added after rollback");
+    // Test transaction rollback
     
-    // Test commit transaction
-    println!("Simulating a new transaction...");
+    // Begin another transaction
+    db.begin_transaction()?;
+    println!("New transaction started");
     
-    // Create new transaction workspace
-    let mut tx_workspace = std::collections::HashMap::new();
-    
-    // Retrieve current state from db
-    if let Some(data) = db.get(user_id.as_bytes())? {
-        let retrieved_user: User = bincode::deserialize(&data)?;
-        tx_workspace.insert(user_id.clone(), retrieved_user);
-    }
-    
-    // Update user in new transaction workspace
-    let committed_user = User::new(
+    // Make changes that will be rolled back
+    let rollback_user = User::new(
         200,
-        "Committed Update".to_string(),
-        "commit@example.com".to_string(),
-        "commit_pass".to_string(),
-        "Commit Corp".to_string(),
-        "Manager".to_string(),
+        "Will Be Rolled Back".to_string(),
+        "rollback@example.com".to_string(),
+        "temppass".to_string(),
+        "Temp Corp".to_string(),
+        "TempAdmin".to_string(),
     );
-    tx_workspace.insert(user_id.clone(), committed_user.clone());
-    println!("Updated user in new transaction");
+    db.set(&rollback_user)?;
+    println!("Updated user in transaction that will be rolled back");
     
-    // Commit the transaction by writing the workspace changes to the database
-    println!("Committing transaction by writing changes to database");
-    for (key, user) in tx_workspace {
-        let user_bytes = bincode::serialize(&user)?;
-        db.insert(key.as_bytes(), user_bytes)?;
-    }
-    db.flush()?;
+    // Rollback the transaction
+    db.rollback_transaction()?;
+    println!("Transaction rolled back");
     
-    // Verify changes persisted to the database
-    if let Some(data) = db.get(user_id.as_bytes())? {
-        let final_user: User = bincode::deserialize(&data)?;
-        assert_eq!(final_user.name, "Committed Update");
-        println!("Verified changes persisted after commit");
-    } else {
-        return Err("Failed to retrieve user after commit".into());
-    }
+    // Verify original data is intact
+    let after_rollback: User = db.get(&user.id.to_string())?;
+    assert_eq!(after_rollback.name, "Updated in TX");
+    println!("Verified data is intact after rollback: {}", after_rollback.name);
     
     Ok(())
 }
 
-/// Test the basic CRUD functionality with a single model
 #[test]
 fn test_simple_db() {
-    // Create a temporary directory for testing
-    let temp_dir = tempdir().expect("Failed to create temp directory");
-    println!("Created temporary directory at: {:?}", temp_dir.path());
-    
-    // Create a test user
-    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+    // A simpler test that uses a basic DB setup
+
+    // Test model
+    #[derive(Debug, Clone, Serialize, Deserialize)]
     struct User {
         id: u32,
         name: String,
@@ -573,56 +255,39 @@ fn test_simple_db() {
         }
     }
     
-    // Open a sled database in the temporary directory
-    let db = sled::open(temp_dir.path().join("simple_users")).expect("Failed to open database");
-    println!("Opened database at: {:?}", temp_dir.path().join("simple_users"));
+    impl Storable for User {}
     
-    // CREATE: Create a user
-    let user = User::new(1, "Simple User".to_string(), "simple@example.com".to_string());
-    let user_key = user.id.to_string();
-    let user_value = bincode::serialize(&user).expect("Failed to serialize user");
-    db.insert(user_key.as_bytes(), user_value).expect("Failed to insert user");
-    db.flush().expect("Failed to flush database");
-    println!("Created user: {} ({})", user.name, user.email);
-    
-    // READ: Retrieve the user
-    let result = db.get(user_key.as_bytes()).expect("Failed to query database");
-    assert!(result.is_some(), "User should exist");
-    if let Some(data) = result {
-        let retrieved_user: User = bincode::deserialize(&data).expect("Failed to deserialize user");
-        println!("Retrieved user: {} ({})", retrieved_user.name, retrieved_user.email);
-        assert_eq!(user, retrieved_user, "Retrieved user should match original");
+    impl SledModel for User {
+        fn get_id(&self) -> String {
+            self.id.to_string()
+        }
+        
+        fn db_prefix() -> &'static str {
+            "test_simple_user"
+        }
     }
     
-    // UPDATE: Update the user
-    let updated_user = User::new(1, "Updated User".to_string(), "updated@example.com".to_string());
-    let updated_value = bincode::serialize(&updated_user).expect("Failed to serialize updated user");
-    db.insert(user_key.as_bytes(), updated_value).expect("Failed to update user");
-    db.flush().expect("Failed to flush database");
-    println!("Updated user: {} ({})", updated_user.name, updated_user.email);
+    // Create a temporary directory for the test
+    let dir = tempdir().expect("Failed to create temp dir");
     
-    // Verify update
-    let result = db.get(user_key.as_bytes()).expect("Failed to query database");
-    assert!(result.is_some(), "Updated user should exist");
-    if let Some(data) = result {
-        let retrieved_user: User = bincode::deserialize(&data).expect("Failed to deserialize user");
-        println!("Retrieved updated user: {} ({})", retrieved_user.name, retrieved_user.email);
-        assert_eq!(updated_user, retrieved_user, "Retrieved user should match updated version");
-    }
+    // Create a DB with the builder
+    let db = DBBuilder::new(dir.path())
+        .register_model::<User>()
+        .build()
+        .expect("Failed to build DB");
     
-    // DELETE: Delete the user
-    db.remove(user_key.as_bytes()).expect("Failed to delete user");
-    db.flush().expect("Failed to flush database");
-    println!("Deleted user");
+    // Create a test user
+    let user = User::new(1, "Simple Test User".to_string(), "simple@example.com".to_string());
     
-    // Verify deletion
-    let result = db.get(user_key.as_bytes()).expect("Failed to query database");
-    assert!(result.is_none(), "User should be deleted");
-    println!("Verified user deletion");
+    // Set the user
+    db.set(&user).expect("Failed to set user");
     
-    // Clean up
-    drop(db);
-    temp_dir.close().expect("Failed to cleanup temporary directory");
+    // Get the user
+    let retrieved: User = db.get(&user.id.to_string()).expect("Failed to get user");
     
-    println!("Simple DB test completed successfully!");
+    // Check that it matches
+    assert_eq!(user.name, retrieved.name);
+    assert_eq!(user.email, retrieved.email);
+    
+    println!("Simple DB test passed!");
 }
diff --git a/herodb/src/zaz/tests/mod.rs b/herodb/src/zaz/tests/mod.rs
new file mode 100644
index 0000000..b8a7d02
--- /dev/null
+++ b/herodb/src/zaz/tests/mod.rs
@@ -0,0 +1,6 @@
+//! Tests for the Zaz module
+
+// Re-export the test modules
+pub mod model_db_test;
+pub mod db_integration_test;
+pub mod transaction_test;
\ No newline at end of file
diff --git a/herodb/src/zaz/tests/model_db_test.rs b/herodb/src/zaz/tests/model_db_test.rs
index 19e53e4..2ca13ba 100644
--- a/herodb/src/zaz/tests/model_db_test.rs
+++ b/herodb/src/zaz/tests/model_db_test.rs
@@ -2,7 +2,9 @@
 
 use crate::core::{DB, DBBuilder, SledDBResult, Storable, SledModel};
 use crate::zaz::factory::create_zaz_db;
-use crate::zaz::models::*;
+use crate::zaz::models::user::User;
+use crate::zaz::models::company::{Company, BusinessType, CompanyStatus};
+use crate::zaz::models::product::{Product, ProductStatus, ProductType, Currency};
 use chrono::Utc;
 use std::path::Path;
 use tempfile::tempdir;
@@ -103,14 +105,13 @@ fn test_db_builder() {
     let product = Product::new(
         1,
         "Unregistered Product".to_string(),
-        "PROD-123".to_string(),
         "A test product".to_string(),
+        Currency::new(100.0, "USD".to_string()),
+        ProductType::Product,
         "Test".to_string(),
-        ProductType::Standard,
         ProductStatus::Available,
-        Currency::USD,
-        100.0,
-        vec![],
+        10, // max_amount
+        30, // validity_days
     );
     
     // This should fail because Product was not registered