Fix get_state_update when forking (#525)

starknet_devnet/blocks.py

@@ -171,11 +171,14 @@ async def get_state_update(
         # now either an int or "latest"
         if block_number != LATEST_BLOCK_ID:
             self.__assert_block_number_in_range(block_number)
-            numeric_hash = self.__num2hash[block_number]
-            if numeric_hash in self.__state_updates:
-                return self.__state_updates[numeric_hash]
 
-            return await self.origin.get_state_update(block_hash=numeric_hash)
+            # if block_number not in mapping, poll origin
+            if block_number in self.__num2hash:
+                numeric_hash = self.__num2hash[block_number]
+                if numeric_hash in self.__state_updates:
+                    return self.__state_updates[numeric_hash]
+
+            return await self.origin.get_state_update(block_number=block_number)
 
         # now it's the latest
         return (

starknet_devnet/origin.py

@@ -8,6 +8,7 @@
     FeederGatewayClient,
 )
 from starkware.starknet.services.api.feeder_gateway.response_objects import (
+    BlockStateUpdate,
     FinalityStatus,
     StarknetBlock,
     TransactionInfo,
@@ -61,7 +62,7 @@ def get_number_of_blocks(self):
 
     async def get_state_update(
         self, block_hash: str = None, block_number: int = None
-    ) -> dict:
+    ) -> BlockStateUpdate:
         """
         Returns the state update for provided block hash or block number.
         If none are provided return the last state update
@@ -136,7 +137,7 @@ def get_number_of_blocks(self):
 
     async def get_state_update(
         self, block_hash: str = None, block_number: int = None
-    ) -> dict:
+    ) -> BlockStateUpdate:
         if block_hash:
             error_message = (
                 f"No state updates saved for the provided block hash {block_hash}"
@@ -225,13 +226,14 @@ def get_number_of_blocks(self):
 
     async def get_state_update(
         self, block_hash: str = None, block_number: int = None
-    ) -> dict:
+    ) -> BlockStateUpdate:
         try:
             with suppress_feeder_gateway_client_logger:
-                return await self.__feeder_gateway_client.get_state_update(
+                state_update_dict = await self.__feeder_gateway_client.get_state_update(
                     block_hash=block_hash,
                     block_number=block_number,
                 )
+                return BlockStateUpdate.load(state_update_dict)
         except BadRequest as bad_request:
             if is_originally_starknet_exception(bad_request):
                 raise StarknetDevnetException(

test/test_fork_feeder_gateway.py

@@ -429,3 +429,15 @@ def test_block_responses_by_hash():
     state_update_by_hash = get_block_traces({"blockHash": latest_block_hash})
     state_update_by_number = get_block_traces({"blockNumber": "latest"})
     assert state_update_by_hash == state_update_by_number
+
+
+@devnet_in_background(*TESTNET_FORK_PARAMS)
+def test_get_state_update():
+    """Test if get_state_update works on blocks not locally present"""
+
+    state_update = get_state_update(block_number=0)
+
+    # make sure we aren't receiving the block that devnet mines on startup
+    assert state_update["block_hash"].startswith("0x")
+    assert state_update["old_root"] == "0x0"
+    assert state_update["new_root"] != "0x0"