fix(STO-020): fix unit test and restructure code logic

alchemyplatform · Aug 22, 2024 · e2b7b27 · e2b7b27
1 parent 2e56905
commit e2b7b27
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/crates/sim/src/simulation/simulator.rs b/crates/sim/src/simulation/simulator.rs
@@ -282,22 +282,23 @@ where
                             slot,
                         ) => {
                             let needs_stake_entity = needs_stake.and_then(|t| entity_infos.get(t));
+
+                            if needs_stake.is_none() {
+                                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;
+                                    }
+                                }
+                            }
+
                             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;
                                 }
-                                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),
@@ -1178,8 +1179,7 @@ mod tests {
             )]
         );
 
-        // staked causes 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());
     }