From c3409388c83a4e33cf3fb1f38efe527c131c7ef1 Mon Sep 17 00:00:00 2001
From: Spablob <99089658+Spablob@users.noreply.github.com>
Date: Thu, 28 Nov 2024 10:35:36 +0000
Subject: [PATCH] update arbitration policy fix (#317)

---
 contracts/modules/dispute/DisputeModule.sol   |  4 ++++
 .../modules/dispute/DisputeModule.t.sol       | 21 +++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/contracts/modules/dispute/DisputeModule.sol b/contracts/modules/dispute/DisputeModule.sol
index 6674fe18..6826dfb4 100644
--- a/contracts/modules/dispute/DisputeModule.sol
+++ b/contracts/modules/dispute/DisputeModule.sol
@@ -173,6 +173,10 @@ contract DisputeModule is
         if (!$.isWhitelistedArbitrationPolicy[nextArbitrationPolicy])
             revert Errors.DisputeModule__NotWhitelistedArbitrationPolicy();
 
+        // if applicable updates the active arbitration policy to the queued arbitration policy
+        // before setting the new next arbitration policy
+        _updateActiveArbitrationPolicy(ipId);
+
         $.nextArbitrationPolicies[ipId] = nextArbitrationPolicy;
 
         uint256 nextArbitrationUpdateTimestamp = block.timestamp + $.arbitrationPolicyCooldown;
diff --git a/test/foundry/modules/dispute/DisputeModule.t.sol b/test/foundry/modules/dispute/DisputeModule.t.sol
index cb6042f4..0b7ab879 100644
--- a/test/foundry/modules/dispute/DisputeModule.t.sol
+++ b/test/foundry/modules/dispute/DisputeModule.t.sol
@@ -249,6 +249,27 @@ contract DisputeModuleTest is BaseTest {
         assertEq(disputeModule.nextArbitrationUpdateTimestamps(ipAddr), block.timestamp + 7 days);
     }
 
+    function test_DisputeModule_setArbitrationPolicy_WithQueuedPolicy() public {
+        vm.startPrank(u.admin);
+        disputeModule.whitelistArbitrationPolicy(address(mockArbitrationPolicy2), true);
+        vm.stopPrank();
+
+        vm.startPrank(ipAddr);
+
+        vm.expectEmit(true, true, true, true, address(disputeModule));
+        emit IDisputeModule.ArbitrationPolicySet(ipAddr, address(mockArbitrationPolicy2), block.timestamp + 7 days);
+
+        disputeModule.setArbitrationPolicy(ipAddr, address(mockArbitrationPolicy2));
+
+        vm.warp(block.timestamp + disputeModule.arbitrationPolicyCooldown() + 1);
+
+        disputeModule.setArbitrationPolicy(ipAddr, address(mockArbitrationPolicy));
+
+        assertEq(disputeModule.arbitrationPolicies(ipAddr), address(mockArbitrationPolicy2));
+        assertEq(disputeModule.nextArbitrationPolicies(ipAddr), address(mockArbitrationPolicy));
+        assertEq(disputeModule.nextArbitrationUpdateTimestamps(ipAddr), block.timestamp + 7 days);
+    }
+
     function test_DisputeModule_raiseDispute_revert_NotRegisteredIpId() public {
         vm.expectRevert(Errors.DisputeModule__NotRegisteredIpId.selector);
         disputeModule.raiseDispute(address(1), disputeEvidenceHashExample, "IMPROPER_REGISTRATION", "");