feat: add ChainSimulatorSetStateStep for chain simulator state patching

docs/source/dev_documentation/changelog.md

@@ -4,6 +4,7 @@
 
 ### Added
 
+- `ChainSimulatorSetStateStep` to set specific hex-encoded storage key-value pairs for an address on the chain simulator
 - Explicit test for variadic counted values
 - Dynamic batch size recovery: after a timeout reduces the batch size, subsequent successful requests with smaller payloads automatically double the batch size back toward the original (for both storage fetch and push operations)
 - Tests verifying batch size resets between accounts

docs/source/user_documentation/steps.md

@@ -676,6 +676,20 @@ targets:
   - erd1y3296u7m2v5653pddey3p7l5zacqmsgqc7vsu3w74p9jm2qp3tqqz950yl  # or direct bech32
 ```
 
+(chain_simulator_set_state_target)=
+### Chain Simulator Set State Step
+
+Exclusive to the chain simulator.
+This step allows you to set specific storage key-value pairs for an address on the chain simulator. Keys and values must be hex-encoded. This is useful after cloning contracts from mainnet when some storage values (e.g. round numbers, epoch values) are incompatible with the simulator's timeline.
+
+```yaml
+type: ChainSimulatorSetState
+address: "%my_contract.address"
+keys:
+  "736166655f70726963655f63757272656e745f696e646578": "00000000"
+  "6d795f6f746865725f6b6579": "01"
+```
+
 (account_clone_target)=
 ### Account Clone Step
 

mxops/common/providers.py

@@ -423,6 +423,16 @@ def set_state(self, states: list[dict]) -> GenericResponse:
         url = "simulator/set-state"
         return self.do_post_generic(url, states)
 
+    def set_address_state(self, address: str, keys: dict[str, str]) -> GenericResponse:
+        """Set specific storage key-value pairs for an address on the chain simulator.
+
+        :param address: bech32 address to set state for
+        :param keys: hex-encoded key-value pairs to set
+        :return: response from the chain simulator
+        """
+        url = f"simulator/address/{address}/set-state"
+        return self.do_post_generic(url, keys)
+
     def set_state_overwrite(self, states: list[dict]) -> GenericResponse:
         url = "simulator/set-state-overwrite"
         return self.do_post_generic(url, states)

mxops/execution/steps/__init__.py

@@ -17,6 +17,7 @@
 from mxops.execution.steps.setup import (
     AccountCloneStep,
     ChainSimulatorFaucetStep,
+    ChainSimulatorSetStateStep,
     GenerateWalletsStep,
     R3D4FaucetStep,
 )
@@ -45,6 +46,7 @@
     "AccountCloneStep",
     "AssertStep",
     "ChainSimulatorFaucetStep",
+    "ChainSimulatorSetStateStep",
     "ContractCallStep",
     "ContractDeployStep",
     "ContractQueryStep",

mxops/execution/steps/setup.py

@@ -37,7 +37,7 @@
 from mxops.enums import LogGroupEnum, NetworkEnum, parse_network_enum
 from mxops.execution import utils
 from mxops.execution.account import AccountsManager
-from mxops.smart_values import SmartInt, SmartPath, SmartValue
+from mxops.smart_values import SmartDict, SmartInt, SmartPath, SmartValue
 from mxops.smart_values.mx_sdk import SmartAddress, SmartAddresses
 from mxops.smart_values.native import SmartBool, SmartDatetime, SmartStr
 from mxops.execution.steps.base import Step
@@ -258,6 +258,60 @@ def _execute(self):
             egld_transfer_step.execute()
 
 
+@dataclass
+class ChainSimulatorSetStateStep(Step):
+    """
+    Represents a step to set specific storage key-value pairs for an address
+    on the chain simulator using hex-encoded keys and values.
+    """
+
+    address: SmartAddress
+    keys: SmartDict
+    ALLOWED_NETWORKS: ClassVar[set] = (NetworkEnum.CHAIN_SIMULATOR,)
+
+    @staticmethod
+    def _validate_hex_keys(keys: dict[str, str]):
+        """Validate that all keys and values are valid hex strings."""
+        if not keys:
+            raise errors.InvalidSceneDefinition(
+                "ChainSimulatorSetState requires at least one key-value pair"
+            )
+        for k, v in keys.items():
+            try:
+                bytes.fromhex(k)
+            except ValueError as exc:
+                raise errors.InvalidSceneDefinition(
+                    f"ChainSimulatorSetState key must be hex-encoded, got: '{k}'"
+                ) from exc
+            try:
+                bytes.fromhex(v)
+            except ValueError as exc:
+                raise errors.InvalidSceneDefinition(
+                    f"ChainSimulatorSetState value must be hex-encoded, got: '{v}'"
+                ) from exc
+
+    def _execute(self):
+        """
+        Post hex-encoded key-value pairs to the chain simulator's
+        address-specific set-state endpoint.
+        """
+        logger = ScenarioData.get_scenario_logger(LogGroupEnum.EXEC)
+        scenario_data = ScenarioData.get()
+        if scenario_data.network not in self.ALLOWED_NETWORKS:
+            raise errors.WrongNetworkForStep(
+                scenario_data.network, self.ALLOWED_NETWORKS
+            )
+        proxy = MyProxyNetworkProvider()
+        bech32 = self.address.get_evaluated_value().to_bech32()
+        keys = self.keys.get_evaluated_value()
+        self._validate_hex_keys(keys)
+        logger.info(
+            f"Setting {len(keys)} storage key(s) for {bech32} on chain simulator"
+        )
+        response = proxy.set_address_state(bech32, keys)
+        logger.debug(f"set-state response: {response.to_dictionary()}")
+
+
 @dataclass
 class AccountCloneStep(Step):
     """

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mxops"
-version = "3.1.1-dev2"
+version = "3.1.1-dev3"
 authors = [
   {name="Etienne Wallet"},
 ]

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 3.1.1-dev2
+current_version = 3.1.1-dev3
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(-(?P<release>[^-0-9]+)(?P<build>\d+))?
 serialize = 
 	{major}.{minor}.{patch}-{release}{build}

tests/data/all_steps.yaml

@@ -69,6 +69,11 @@ steps:
       - bob
     amount: "={10**18}"  # amount to send to each wallet
 
+  - type: ChainSimulatorSetState
+    address: "%my_contract.address"
+    keys:
+      "736166655f70726963655f63757272656e745f696e646578": "00000000"
+
   - type: GenerateWallets
     save_folder: path/to/wallets/folder  # folder where to save the generated wallets
     wallets: # wallets to generate, can also just supply a number of wallets

tests/test_setup_steps.py

@@ -1,9 +1,14 @@
 import os
 from pathlib import Path
 import shutil
+from unittest.mock import patch
 
 from multiversx_sdk import Account
-from mxops.execution.steps import GenerateWalletsStep
+import pytest
+
+from mxops import errors
+from mxops.enums import NetworkEnum
+from mxops.execution.steps import ChainSimulatorSetStateStep, GenerateWalletsStep
 
 
 def test_generate_n_wallet_step():
@@ -100,3 +105,56 @@ def test_generate_n_keystore_wallet_step():
     # Cleanup
     shutil.rmtree("./tests/data/TEMP_UNIT_TEST")
     del os.environ["TEST_WALLET_PASSWORD"]
+
+
+MOCK_BECH32 = "erd1qqqqqqqqqqqqqpgqdmq43snzxutandvqefxgj89r6fh528v9dwnswvgq9t"
+MOCK_KEYS = {"736166655f70726963655f63757272656e745f696e646578": "00000000"}
+
+
+@pytest.fixture
+def chain_simulator_scenario(scenario_data, chain_simulator_network):
+    """Set both Config and ScenarioData to chain simulator network."""
+    original = scenario_data.network
+    scenario_data.network = NetworkEnum.CHAIN_SIMULATOR
+    yield scenario_data
+    scenario_data.network = original
+
+
+def test_chain_simulator_set_state_wrong_network(scenario_data):
+    """Network is LOCAL by default, step should reject it."""
+    step = ChainSimulatorSetStateStep(address=MOCK_BECH32, keys=MOCK_KEYS)
+    with pytest.raises(errors.WrongNetworkForStep):
+        step.execute()
+
+
+def test_chain_simulator_set_state_success(chain_simulator_scenario):
+    step = ChainSimulatorSetStateStep(address=MOCK_BECH32, keys=MOCK_KEYS)
+    with patch(
+        "mxops.execution.steps.setup.MyProxyNetworkProvider.set_address_state"
+    ) as mock_set:
+        step.execute()
+        mock_set.assert_called_once_with(MOCK_BECH32, MOCK_KEYS)
+
+
+def test_chain_simulator_set_state_empty_keys(chain_simulator_scenario):
+    step = ChainSimulatorSetStateStep(address=MOCK_BECH32, keys={})
+    with pytest.raises(errors.InvalidSceneDefinition):
+        step.execute()
+
+
+def test_chain_simulator_set_state_invalid_hex_key(chain_simulator_scenario):
+    step = ChainSimulatorSetStateStep(
+        address=MOCK_BECH32, keys={"not_hex!": "00"}
+    )
+    with pytest.raises(errors.InvalidSceneDefinition, match="key must be hex-encoded"):
+        step.execute()
+
+
+def test_chain_simulator_set_state_invalid_hex_value(chain_simulator_scenario):
+    step = ChainSimulatorSetStateStep(
+        address=MOCK_BECH32, keys={"00": "not_hex!"}
+    )
+    with pytest.raises(
+        errors.InvalidSceneDefinition, match="value must be hex-encoded"
+    ):
+        step.execute()