@@ -210,6 +210,100 @@ fn test_remove_upgrade_governor() {
210210 assert! (! roles_dispatcher . is_upgrade_governor (account : upgrade_governor ));
211211}
212212
213+ #[feature(" safe_dispatcher" )]
214+ #[test]
215+ fn test_register_upgrade_agent () {
216+ // Deploy mock contract.
217+ let contract_address = test_utils :: deploy_mock_contract ();
218+ let roles_dispatcher = IRolesDispatcher { contract_address };
219+ let roles_safe_dispatcher = IRolesSafeDispatcher { contract_address };
220+ let app_role_admin = constants :: APP_ROLE_ADMIN ;
221+ let upgrade_agent = constants :: UPGRADE_AGENT ;
222+ let wrong_admin = constants :: WRONG_ADMIN ;
223+
224+ // Try to add zero address as upgrade governor.
225+ cheat_caller_address_once (: contract_address , caller_address : app_role_admin );
226+ let result = roles_safe_dispatcher . register_upgrade_agent (account : Zero :: zero ());
227+ assert_panic_with_error (: result , expected_error : AccessErrors :: ZERO_ADDRESS . describe ());
228+
229+ // Try to add upgrade governor with unqualified caller.
230+ cheat_caller_address_once (: contract_address , caller_address : wrong_admin );
231+ let result = roles_safe_dispatcher . register_upgrade_agent (account : upgrade_agent );
232+ assert_panic_with_felt_error (: result , expected_error : OZAccessErrors :: MISSING_ROLE );
233+
234+ // Register upgrade governor and perform the corresponding checks.
235+ let mut spy = snforge_std :: spy_events ();
236+ assert! (! roles_dispatcher . is_upgrade_agent (account : upgrade_agent ));
237+ cheat_caller_address_once (: contract_address , caller_address : app_role_admin );
238+ roles_dispatcher . register_upgrade_agent (account : upgrade_agent );
239+
240+ let events = spy . get_events (). emitted_by (: contract_address ). events;
241+ // We only check events[1] because events[0] is an event emitted by OZ AccessControl.
242+ assert_number_of_events (
243+ actual : events . len (), expected : 2 , message : " test_register_upgrade_agent first" ,
244+ );
245+ event_test_utils :: assert_upgrade_agent_added_event (
246+ events [1 ], added_account : upgrade_agent , added_by : app_role_admin ,
247+ );
248+ assert! (roles_dispatcher . is_upgrade_agent (account : upgrade_agent ));
249+
250+ // Register upgrade governor that is already registered (should not emit events).
251+ let mut spy = snforge_std :: spy_events ();
252+ cheat_caller_address_once (: contract_address , caller_address : app_role_admin );
253+ roles_dispatcher . register_upgrade_agent (account : upgrade_agent );
254+ let events = spy . get_events (). emitted_by (: contract_address ). events;
255+ assert_number_of_events (
256+ actual : events . len (), expected : 0 , message : " test_register_upgrade_agent second" ,
257+ );
258+ }
259+
260+
261+ #[feature(" safe_dispatcher" )]
262+ #[test]
263+ fn test_remove_upgrade_agent () {
264+ // Deploy mock contract.
265+ let contract_address = test_utils :: deploy_mock_contract ();
266+ let roles_dispatcher = IRolesDispatcher { contract_address };
267+ let roles_safe_dispatcher = IRolesSafeDispatcher { contract_address };
268+ let app_role_admin = constants :: APP_ROLE_ADMIN ;
269+ let upgrade_agent = constants :: UPGRADE_AGENT ;
270+ let wrong_admin = constants :: WRONG_ADMIN ;
271+
272+ // Remove upgrade governor that was not registered (should not emit events).
273+ assert! (! roles_dispatcher . is_upgrade_agent (account : upgrade_agent ));
274+ let mut spy = snforge_std :: spy_events ();
275+ cheat_caller_address_once (: contract_address , caller_address : app_role_admin );
276+ roles_dispatcher . remove_upgrade_agent (account : upgrade_agent );
277+ let events = spy . get_events (). emitted_by (: contract_address ). events;
278+ assert_number_of_events (
279+ actual : events . len (), expected : 0 , message : " test_remove_upgrade_agent first" ,
280+ );
281+
282+ // Register upgrade governor.
283+ cheat_caller_address_once (: contract_address , caller_address : app_role_admin );
284+ roles_dispatcher . register_upgrade_agent (account : upgrade_agent );
285+
286+ // Try to remove upgrade governor with unqualified caller.
287+ cheat_caller_address_once (: contract_address , caller_address : wrong_admin );
288+ let result = roles_safe_dispatcher . remove_upgrade_agent (account : upgrade_agent );
289+ assert_panic_with_felt_error (: result , expected_error : OZAccessErrors :: MISSING_ROLE );
290+
291+ // Remove upgrade governor and perform the corresponding checks.
292+ assert! (roles_dispatcher . is_upgrade_agent (account : upgrade_agent ));
293+ let mut spy = snforge_std :: spy_events ();
294+ cheat_caller_address_once (: contract_address , caller_address : app_role_admin );
295+ roles_dispatcher . remove_upgrade_agent (account : upgrade_agent );
296+ let events = spy . get_events (). emitted_by (: contract_address ). events;
297+ // We only check events[1] because events[0] is an event emitted by OZ AccessControl.
298+ assert_number_of_events (
299+ actual : events . len (), expected : 2 , message : " test_remove_upgrade_agent second" ,
300+ );
301+ event_test_utils :: assert_upgrade_agent_removed_event (
302+ events [1 ], removed_account : upgrade_agent , removed_by : app_role_admin ,
303+ );
304+ assert! (! roles_dispatcher . is_upgrade_agent (account : upgrade_agent ));
305+ }
306+
213307
214308#[test]
215309#[feature(" safe_dispatcher" )]
0 commit comments