declare_contract cheatcode (#354)

MaksymilianDemitraszek · web-flow · commit c9372d9a2b63 · 2022-06-09T18:17:01.000+02:00
* Added tests

* Added cheatcode

* Lint
diff --git a/protostar/commands/test/test_execution_environment.py b/protostar/commands/test/test_execution_environment.py
@@ -8,6 +8,7 @@
 from starkware.starknet.services.api.contract_class import ContractClass
 from starkware.starknet.testing.contract import StarknetContract
 from starkware.starkware_utils.error_handling import StarkException
+from starkware.starknet.testing.contract import DeclaredClass
 
 from protostar.commands.test.cheatcodes import (
     Cheatcode,
@@ -21,6 +22,8 @@
 )
 from protostar.commands.test.starkware.forkable_starknet import ForkableStarknet
 from protostar.commands.test.test_context import TestContext
+
+
 from protostar.commands.test.test_environment_exceptions import (
     CheatcodeException,
     ExpectedEventMissingException,
@@ -44,6 +47,15 @@ def contract_address(self):
         return self._starknet_contract.contract_address
 
 
+class ProtostarDeclaredClass:
+    def __init__(self, declared_class: DeclaredClass):
+        self._declared_class = declared_class
+
+    @property
+    def class_hash(self):
+        return self._declared_class.class_hash
+
+
 class TestExecutionEnvironment:
     def __init__(
         self,
@@ -98,6 +110,17 @@ def deploy_in_env(
         )
         return contract
 
+    def declare_in_env(self, contract_path: str):
+        contract = ProtostarDeclaredClass(
+            asyncio.run(
+                self.starknet.declare(
+                    source=contract_path,
+                    cairo_path=self._include_paths,
+                )
+            )
+        )
+        return contract
+
     async def invoke_setup_hook(self, fn_name: str) -> None:
         await self.invoke_test_case(fn_name)
 
@@ -291,6 +314,10 @@ def deploy_contract(
         ):
             return self.deploy_in_env(contract_path, constructor_calldata)
 
+        @register_cheatcode
+        def declare_contract(contract_path: str):
+            return self.declare_in_env(contract_path)
+
         cheatcodes: List[Cheatcode] = [
             ExpectRevertCheatcode(self),
             RollCheatcode(cheatable_syscall_handler),
diff --git a/tests/integration/cheatcodes/declare_contract/basic_contract.cairo b/tests/integration/cheatcodes/declare_contract/basic_contract.cairo
@@ -0,0 +1,22 @@
+%lang starknet
+
+from starkware.cairo.common.cairo_builtins import HashBuiltin
+
+@storage_var
+func balance() -> (res : felt):
+end
+
+@external
+func increase_balance{syscall_ptr : felt*, pedersen_ptr : HashBuiltin*, range_check_ptr}(
+        amount : felt):
+    let (res) = balance.read()
+    balance.write(res + amount)
+    return ()
+end
+
+@view
+func get_balance{syscall_ptr : felt*, pedersen_ptr : HashBuiltin*, range_check_ptr}() -> (
+        res : felt):
+    let (res) = balance.read()
+    return (res)
+end
diff --git a/tests/integration/cheatcodes/declare_contract/declare_contract_test.cairo b/tests/integration/cheatcodes/declare_contract/declare_contract_test.cairo
@@ -0,0 +1,60 @@
+%lang starknet
+
+from starkware.cairo.common.alloc import alloc
+from starkware.starknet.common.syscalls import deploy
+
+@contract_interface
+namespace BasicContract:
+    func increase_balance(amount : felt):
+    end
+
+    func get_balance() -> (res : felt):
+    end
+end
+
+@contract_interface
+namespace ProxyContract:
+    func deploy_contract_from_proxy(class_hash_d : felt) -> (address : felt):
+    end
+end
+
+
+@external
+func test_deploy_declared_contract{syscall_ptr : felt*, range_check_ptr}():
+    alloc_locals
+
+    local class_hash : felt
+    %{
+        ids.class_hash = declare_contract("./tests/integration/cheatcodes/declare_contract/basic_contract.cairo").class_hash
+    %}
+
+    let (local calldata: felt*) = alloc()
+    let (contract_address) = deploy(class_hash, 42, 0, calldata)
+
+    BasicContract.increase_balance(contract_address, 12)
+
+    let (balance) = BasicContract.get_balance(contract_address)
+    assert balance = 12
+    return ()
+end
+
+@external
+func test_deploy_declared_contract_in_proxy{syscall_ptr : felt*, range_check_ptr}():
+    alloc_locals
+
+    local proxy_address : felt
+    local class_hash : felt
+    %{
+        ids.proxy_address = deploy_contract("./tests/integration/cheatcodes/declare_contract/proxy_contract.cairo").contract_address
+        ids.class_hash = declare_contract("./tests/integration/cheatcodes/declare_contract/basic_contract.cairo").class_hash
+    %}
+
+    let (contract_address) = ProxyContract.deploy_contract_from_proxy(proxy_address, class_hash)
+
+    BasicContract.increase_balance(contract_address, 12)
+
+    let (balance) = BasicContract.get_balance(contract_address)
+    assert balance = 12
+    return ()
+end
+
diff --git a/tests/integration/cheatcodes/declare_contract/declare_contract_test.py b/tests/integration/cheatcodes/declare_contract/declare_contract_test.py
@@ -0,0 +1,23 @@
+from pathlib import Path
+
+import pytest
+
+from protostar.commands.test.test_command import TestCommand
+from tests.integration.conftest import assert_cairo_test_cases
+
+
+@pytest.mark.asyncio
+async def test_declare_contract(mocker):
+    testing_summary = await TestCommand(
+        project=mocker.MagicMock(),
+        protostar_directory=mocker.MagicMock(),
+    ).test(targets=[str(Path(__file__).parent / "declare_contract_test.cairo")])
+
+    assert_cairo_test_cases(
+        testing_summary,
+        expected_passed_test_cases_names=[
+            "test_deploy_declared_contract",
+            "test_deploy_declared_contract_in_proxy",
+        ],
+        expected_failed_test_cases_names=[],
+    )
diff --git a/tests/integration/cheatcodes/declare_contract/proxy_contract.cairo b/tests/integration/cheatcodes/declare_contract/proxy_contract.cairo
@@ -0,0 +1,26 @@
+# Declare this file as a StarkNet contract and set the required
+# builtins.
+%lang starknet
+
+from starkware.cairo.common.cairo_builtins import HashBuiltin
+from starkware.cairo.common.alloc import alloc
+from starkware.starknet.common.syscalls import deploy
+
+@contract_interface
+namespace BasicContract:
+    func increase_balance(amount : felt):
+    end
+
+    func get_balance() -> (res : felt):
+    end
+end
+
+
+@external
+func deploy_contract_from_proxy{syscall_ptr : felt*, pedersen_ptr : HashBuiltin*, range_check_ptr}(
+        class_hash_d : felt) -> (address : felt):
+    alloc_locals
+    let (local calldata: felt*) = alloc()
+    let (contract_address) = deploy(class_hash_d, 42, 0, calldata)
+    return (contract_address)
+end
