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