tests for upgrade_agnet

Author: remollemo

packages/utils/src/components/roles/event_test_utils.cairo

@@ -197,6 +197,50 @@ pub(crate) fn assert_upgrade_governor_added_event(
     }
 }
 
+pub(crate) fn assert_upgrade_agent_added_event(
+    spied_event: @(ContractAddress, Event),
+    added_account: ContractAddress,
+    added_by: ContractAddress,
+) {
+    let expected_event = @MockContract::Event::RolesEvent(
+        RolesEvent::UpgradeAgentAdded(
+            RolesInterface::UpgradeAgentAdded { added_account, added_by },
+        ),
+    );
+    let (expected_emitted_by, raw_event) = spied_event;
+    let wrapped_spied_event = Events { events: array![(*expected_emitted_by, raw_event.clone())] };
+    let emitted = is_emitted(self: @wrapped_spied_event, :expected_emitted_by, :expected_event);
+    if !emitted {
+        let details = format!(
+            "UpgradeAgentAdded{{added_account: {:?}, added_by: {:?}}}", added_account, added_by,
+        );
+        panic_with_event_details(:expected_emitted_by, :details);
+    }
+}
+
+pub(crate) fn assert_upgrade_agent_removed_event(
+    spied_event: @(ContractAddress, Event),
+    removed_account: ContractAddress,
+    removed_by: ContractAddress,
+) {
+    let expected_event = @MockContract::Event::RolesEvent(
+        RolesEvent::UpgradeAgentRemoved(
+            RolesInterface::UpgradeAgentRemoved { removed_account, removed_by },
+        ),
+    );
+    let (expected_emitted_by, raw_event) = spied_event;
+    let wrapped_spied_event = Events { events: array![(*expected_emitted_by, raw_event.clone())] };
+    let emitted = is_emitted(self: @wrapped_spied_event, :expected_emitted_by, :expected_event);
+    if !emitted {
+        let details = format!(
+            "UpgradeAgentRemoved{{removed_account: {:?}, removed_by: {:?}}}",
+            removed_account,
+            removed_by,
+        );
+        panic_with_event_details(:expected_emitted_by, :details);
+    }
+}
+
 pub(crate) fn assert_upgrade_governor_removed_event(
     spied_event: @(ContractAddress, Event),
     removed_account: ContractAddress,

packages/utils/src/components/roles/test.cairo

@@ -210,6 +210,100 @@ fn test_remove_upgrade_governor() {
     assert!(!roles_dispatcher.is_upgrade_governor(account: upgrade_governor));
 }
 
+#[feature("safe_dispatcher")]
+#[test]
+fn test_register_upgrade_agent() {
+    // Deploy mock contract.
+    let contract_address = test_utils::deploy_mock_contract();
+    let roles_dispatcher = IRolesDispatcher { contract_address };
+    let roles_safe_dispatcher = IRolesSafeDispatcher { contract_address };
+    let app_role_admin = constants::APP_ROLE_ADMIN;
+    let upgrade_governor = constants::UPGRADE_GOVERNOR;
+    let wrong_admin = constants::WRONG_ADMIN;
+
+    // Try to add zero address as upgrade governor.
+    cheat_caller_address_once(:contract_address, caller_address: app_role_admin);
+    let result = roles_safe_dispatcher.register_upgrade_governor(account: Zero::zero());
+    assert_panic_with_error(:result, expected_error: AccessErrors::ZERO_ADDRESS.describe());
+
+    // Try to add upgrade governor with unqualified caller.
+    cheat_caller_address_once(:contract_address, caller_address: wrong_admin);
+    let result = roles_safe_dispatcher.register_upgrade_governor(account: upgrade_governor);
+    assert_panic_with_felt_error(:result, expected_error: OZAccessErrors::MISSING_ROLE);
+
+    // Register upgrade governor and perform the corresponding checks.
+    let mut spy = snforge_std::spy_events();
+    assert!(!roles_dispatcher.is_upgrade_governor(account: upgrade_governor));
+    cheat_caller_address_once(:contract_address, caller_address: app_role_admin);
+    roles_dispatcher.register_upgrade_governor(account: upgrade_governor);
+
+    let events = spy.get_events().emitted_by(:contract_address).events;
+    // We only check events[1] because events[0] is an event emitted by OZ AccessControl.
+    assert_number_of_events(
+        actual: events.len(), expected: 2, message: "test_register_upgrade_governor first",
+    );
+    event_test_utils::assert_upgrade_governor_added_event(
+        events[1], added_account: upgrade_governor, added_by: app_role_admin,
+    );
+    assert!(roles_dispatcher.is_upgrade_governor(account: upgrade_governor));
+
+    // Register upgrade governor that is already registered (should not emit events).
+    let mut spy = snforge_std::spy_events();
+    cheat_caller_address_once(:contract_address, caller_address: app_role_admin);
+    roles_dispatcher.register_upgrade_governor(account: upgrade_governor);
+    let events = spy.get_events().emitted_by(:contract_address).events;
+    assert_number_of_events(
+        actual: events.len(), expected: 0, message: "test_register_upgrade_governor second",
+    );
+}
+
+
+#[feature("safe_dispatcher")]
+#[test]
+fn test_remove_upgrade_agent() {
+    // Deploy mock contract.
+    let contract_address = test_utils::deploy_mock_contract();
+    let roles_dispatcher = IRolesDispatcher { contract_address };
+    let roles_safe_dispatcher = IRolesSafeDispatcher { contract_address };
+    let app_role_admin = constants::APP_ROLE_ADMIN;
+    let upgrade_agent = constants::UPGRADE_AGENT;
+    let wrong_admin = constants::WRONG_ADMIN;
+
+    // Remove upgrade governor that was not registered (should not emit events).
+    assert!(!roles_dispatcher.is_upgrade_agent(account: upgrade_agent));
+    let mut spy = snforge_std::spy_events();
+    cheat_caller_address_once(:contract_address, caller_address: app_role_admin);
+    roles_dispatcher.remove_upgrade_agent(account: upgrade_agent);
+    let events = spy.get_events().emitted_by(:contract_address).events;
+    assert_number_of_events(
+        actual: events.len(), expected: 0, message: "test_remove_upgrade_agent first",
+    );
+
+    // Register upgrade governor.
+    cheat_caller_address_once(:contract_address, caller_address: app_role_admin);
+    roles_dispatcher.register_upgrade_agent(account: upgrade_agent);
+
+    // Try to remove upgrade governor with unqualified caller.
+    cheat_caller_address_once(:contract_address, caller_address: wrong_admin);
+    let result = roles_safe_dispatcher.remove_upgrade_agent(account: upgrade_agent);
+    assert_panic_with_felt_error(:result, expected_error: OZAccessErrors::MISSING_ROLE);
+
+    // Remove upgrade governor and perform the corresponding checks.
+    assert!(roles_dispatcher.is_upgrade_agent(account: upgrade_agent));
+    let mut spy = snforge_std::spy_events();
+    cheat_caller_address_once(:contract_address, caller_address: app_role_admin);
+    roles_dispatcher.remove_upgrade_agent(account: upgrade_agent);
+    let events = spy.get_events().emitted_by(:contract_address).events;
+    // We only check events[1] because events[0] is an event emitted by OZ AccessControl.
+    assert_number_of_events(
+        actual: events.len(), expected: 2, message: "test_remove_upgrade_agent second",
+    );
+    event_test_utils::assert_upgrade_agent_removed_event(
+        events[1], removed_account: upgrade_agent, removed_by: app_role_admin,
+    );
+    assert!(!roles_dispatcher.is_upgrade_agent(account: upgrade_agent));
+}
+
 
 #[test]
 #[feature("safe_dispatcher")]