test(pulse): add tests for price updates removal and max price IDs validation (#2676)

* feat: add tests for price updates removal and max price IDs validation

Co-Authored-By: Tejas Badadare <[email protected]>

* style: fix formatting in PulseScheduler.t.sol

Co-Authored-By: Tejas Badadare <[email protected]>

* test: address PR comments for price updates removal test

Co-Authored-By: Tejas Badadare <[email protected]>

* test: fix refs to mockParsePriceFeedUpdatesWithSlots

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Tejas Badadare <[email protected]>
Co-authored-by: Tejas Badadare <[email protected]>

Author: devin-ai-integration[bot]

target_chains/ethereum/contracts/forge-test/PulseScheduler.t.sol

@@ -2231,4 +2231,128 @@ contract SchedulerTest is Test, SchedulerEvents, PulseSchedulerTestUtils {
 
     // Required to receive ETH when withdrawing funds
     receive() external payable {}
+
+    function testUpdateSubscriptionRemovesPriceUpdatesForRemovedPriceIds()
+        public
+    {
+        // 1. Setup: Add subscription with 3 price feeds, update prices
+        uint8 numInitialFeeds = 3;
+        uint256 subscriptionId = addTestSubscriptionWithFeeds(
+            scheduler,
+            numInitialFeeds,
+            address(reader)
+        );
+        scheduler.addFunds{value: 1 ether}(subscriptionId);
+
+        // Get initial price IDs and create mock price feeds
+        bytes32[] memory initialPriceIds = createPriceIds(numInitialFeeds);
+        uint64 publishTime = SafeCast.toUint64(block.timestamp);
+
+        // Setup and perform initial price update
+        (
+            PythStructs.PriceFeed[] memory priceFeeds,
+            uint64[] memory slots
+        ) = createMockPriceFeedsWithSlots(publishTime, numInitialFeeds);
+        mockParsePriceFeedUpdatesWithSlotsStrict(pyth, priceFeeds, slots);
+
+        vm.prank(pusher);
+        scheduler.updatePriceFeeds(
+            subscriptionId,
+            createMockUpdateData(priceFeeds)
+        );
+
+        // Store the removed price ID for later use
+        bytes32 removedPriceId = initialPriceIds[numInitialFeeds - 1];
+
+        // 2. Action: Update subscription to remove the last price feed
+        (SchedulerState.SubscriptionParams memory params, ) = scheduler
+            .getSubscription(subscriptionId);
+
+        // Create new price IDs array without the last ID
+        bytes32[] memory newPriceIds = new bytes32[](numInitialFeeds - 1);
+        for (uint i = 0; i < newPriceIds.length; i++) {
+            newPriceIds[i] = initialPriceIds[i];
+        }
+
+        params.priceIds = newPriceIds;
+
+        vm.expectEmit();
+        emit SubscriptionUpdated(subscriptionId);
+        scheduler.updateSubscription(subscriptionId, params);
+
+        // 3. Verification:
+        // - Verify that the removed price ID is no longer part of the subscription's price IDs
+        (SchedulerState.SubscriptionParams memory updatedParams, ) = scheduler
+            .getSubscription(subscriptionId);
+        assertEq(
+            updatedParams.priceIds.length,
+            numInitialFeeds - 1,
+            "Subscription should have one less price ID"
+        );
+
+        bool removedPriceIdFound = false;
+        for (uint i = 0; i < updatedParams.priceIds.length; i++) {
+            if (updatedParams.priceIds[i] == removedPriceId) {
+                removedPriceIdFound = true;
+                break;
+            }
+        }
+        assertFalse(
+            removedPriceIdFound,
+            "Removed price ID should not be in the subscription's price IDs"
+        );
+
+        // - Querying all feeds should return only the remaining feeds
+        PythStructs.Price[] memory allPricesAfterUpdate = scheduler
+            .getPricesUnsafe(subscriptionId, new bytes32[](0));
+        assertEq(
+            allPricesAfterUpdate.length,
+            newPriceIds.length,
+            "Querying all should only return remaining feeds"
+        );
+
+        // - Verify that trying to get the price of the removed feed directly reverts
+        bytes32[] memory removedIdArray = new bytes32[](1);
+        removedIdArray[0] = removedPriceId;
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                InvalidPriceId.selector,
+                removedPriceId,
+                bytes32(0)
+            )
+        );
+        scheduler.getPricesUnsafe(subscriptionId, removedIdArray);
+    }
+
+    function testUpdateSubscriptionRevertsWithTooManyPriceIds() public {
+        // 1. Setup: Create a subscription with a valid number of price IDs
+        uint8 initialNumFeeds = 2;
+        uint256 subscriptionId = addTestSubscriptionWithFeeds(
+            scheduler,
+            initialNumFeeds,
+            address(reader)
+        );
+
+        // 2. Prepare params with too many price IDs (MAX_PRICE_IDS_PER_SUBSCRIPTION + 1)
+        (SchedulerState.SubscriptionParams memory currentParams, ) = scheduler
+            .getSubscription(subscriptionId);
+
+        uint16 tooManyFeeds = uint16(
+            scheduler.MAX_PRICE_IDS_PER_SUBSCRIPTION()
+        ) + 1;
+        bytes32[] memory tooManyPriceIds = createPriceIds(tooManyFeeds);
+
+        SchedulerState.SubscriptionParams memory newParams = currentParams;
+        newParams.priceIds = tooManyPriceIds;
+
+        // 3. Expect revert when trying to update with too many price IDs
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                TooManyPriceIds.selector,
+                tooManyFeeds,
+                scheduler.MAX_PRICE_IDS_PER_SUBSCRIPTION()
+            )
+        );
+        scheduler.updateSubscription(subscriptionId, newParams);
+    }
 }