From b6930b3be7a312fab6892ef1c6810e24a74c920b Mon Sep 17 00:00:00 2001
From: Andy <andysim3d@gmail.com>
Date: Wed, 21 Aug 2024 12:45:16 -0400
Subject: [PATCH] fix(STO-020): fix unit test and restructure code logic

---
 crates/sim/src/simulation/simulator.rs | 63 ++++++++++++++++++++------
 1 file changed, 50 insertions(+), 13 deletions(-)

diff --git a/crates/sim/src/simulation/simulator.rs b/crates/sim/src/simulation/simulator.rs
index 55ae206a..6de30206 100644
--- a/crates/sim/src/simulation/simulator.rs
+++ b/crates/sim/src/simulation/simulator.rs
@@ -282,27 +282,33 @@ where
                             slot,
                         ) => {
                             let needs_stake_entity = needs_stake.and_then(|t| entity_infos.get(t));
+                            let mut entity_stacked_violation = true;
+                            let mut factory_stacked_violation = true;
+
+                            if needs_stake.is_none(){
+                                entity_stacked_violation = false;
+                            }
+
+                            if let Some(factory) = entity_infos.get(EntityType::Factory) {
+                                if factory.is_staked {
+                                    tracing::debug!("Associated storage accessed by staked entity during deploy, and factory is staked");
+                                    factory_stacked_violation = false;
+                                }
+                            }
                             if let Some(needs_stake_entity_info) = needs_stake_entity {
                                 if needs_stake_entity_info.is_staked {
                                     tracing::debug!("Associated storage accessed by staked entity during deploy, and entity is staked");
-                                    continue;
+                                    entity_stacked_violation = false;
                                 }
+                            }
+
+                            if entity_stacked_violation || factory_stacked_violation {
+                                // [STO-022]
                                 violations.push(SimulationViolation::AssociatedStorageDuringDeploy(
                                     needs_stake_entity.map(|ei| ei.entity),
                                     StorageSlot { address, slot },
                                 ))
                             }
-                            if let Some(factory) = entity_infos.get(EntityType::Factory) {
-                                if factory.is_staked {
-                                    tracing::debug!("Associated storage accessed by staked entity during deploy, and factory is staked");
-                                    continue;
-                                }
-                            }
-                            // [STO-022]
-                            violations.push(SimulationViolation::AssociatedStorageDuringDeploy(
-                                needs_stake_entity.map(|ei| ei.entity),
-                                StorageSlot { address, slot },
-                            ))
                         }
                         StorageRestriction::Banned(slot) => {
                             // [STO-*]
@@ -1178,10 +1184,41 @@ mod tests {
             )]
         );
 
-        // staked causes no errors
+        // staked only factory causes errors
+        context.entity_infos.factory.as_mut().unwrap().is_staked = true;
+        context.entity_infos.paymaster.as_mut().unwrap().is_staked = false;
+        let res = simulator.gather_context_violations(&mut context);
+        assert_eq!(
+            res.unwrap(),
+            vec![SimulationViolation::AssociatedStorageDuringDeploy(
+                Some(Entity::paymaster(paymaster_address)),
+                StorageSlot {
+                    address: external_access_address,
+                    slot: sender_address.as_bytes().into()
+                }
+            )]
+        );
+
+        // staked only paymaster causes errors
+        context.entity_infos.factory.as_mut().unwrap().is_staked = false;
+        context.entity_infos.paymaster.as_mut().unwrap().is_staked = true;
+        let res = simulator.gather_context_violations(&mut context);
+        assert_eq!(
+            res.unwrap(),
+            vec![SimulationViolation::AssociatedStorageDuringDeploy(
+                Some(Entity::paymaster(paymaster_address)),
+                StorageSlot {
+                    address: external_access_address,
+                    slot: sender_address.as_bytes().into()
+                }
+            )]
+        );
+        // staked both paymaster and factory cause no errors
         context.entity_infos.factory.as_mut().unwrap().is_staked = true;
+        context.entity_infos.paymaster.as_mut().unwrap().is_staked = true;
         let res = simulator.gather_context_violations(&mut context);
         assert!(res.unwrap().is_empty());
+
     }
 
     #[tokio::test]