feat: zksync test fixes and get_account func added (#202)

* feat: zksync test fixes and get_account func added

* lint: fix

moccasin/moccasin_account.py

@@ -106,6 +106,12 @@ def unlock(
             self._init_key(decrypted_key)
         return cast(HexBytes, self.private_key)
 
+    def get_balance(self) -> int:
+        # This might be dumb? Idk
+        import boa
+
+        return boa.env.get_balance(self.address)
+
     @classmethod
     def from_boa_address(cls, address: Address) -> "MoccasinAccount":
         return cls()

tests/conftest.py

@@ -254,7 +254,7 @@ def anvil(anvil_process, anvil_keystore):
     yield
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def anvil_two_no_state():
     with AnvilProcess(args=["-p", "8546"], port=8546):
         yield

tests/unit/test_unit_moccasin_account.py

@@ -0,0 +1,21 @@
+import boa
+from boa import Env
+
+from moccasin.moccasin_account import MoccasinAccount
+from tests.constants import ANVIL1_PRIVATE_KEY
+
+STARTING_ANVIL1_BALANCE = 10000000000000000000000
+
+
+def test_get_balance_env(anvil_two_no_state, complex_project_config):
+    network = "anvil-fork"  # Not a real fork
+    complex_project_config.set_active_network(network)
+    mox_account = MoccasinAccount(private_key=ANVIL1_PRIVATE_KEY)
+    assert mox_account.get_balance() >= STARTING_ANVIL1_BALANCE
+    # We should do this more places to make tests more isolated.
+    boa.set_env(Env())
+
+
+def test_get_balance_no_env():
+    mox_account = MoccasinAccount(private_key=ANVIL1_PRIVATE_KEY)
+    assert mox_account.get_balance() == 0

tests/utils/anvil.py

@@ -56,12 +56,21 @@ def terminate(self):
         if self.process:
             self.process.terminate()
             try:
-                self.process.wait(timeout=0.5)
+                self.process.wait(timeout=0.75)
             except subprocess.TimeoutExpired:
                 self.process.kill()
-            self.process = None
+            max_retries = 5
+            for _ in range(max_retries):
+                if not self.is_running():
+                    break
+                self.process.wait(timeout=0.5)
+            if self.is_running():
+                raise RuntimeError("Anvil process failed to terminate")
         atexit.unregister(self.terminate)
 
+    def is_running(self):
+        return self.process is not None and self.process.poll() is None
+
     @property
     def pid(self):
         return self.process.pid if self.process else None

tests/zksync/conftest.py

@@ -8,7 +8,7 @@
 from boa_zksync import set_zksync_test_env
 
 from moccasin.commands.wallet import save_to_keystores
-from moccasin.config import Config, initialize_global_config
+from moccasin.config import Config, get_or_initialize_config
 from moccasin.constants.vars import DEPENDENCIES_FOLDER
 from tests.constants import ZKSYNC_PROJECT_PATH
 
@@ -53,7 +53,7 @@ def zk_temp_path() -> Generator[Path, None, None]:
 
 @pytest.fixture(scope="module")
 def zksync_project_config(zk_temp_path) -> Config:
-    return initialize_global_config(zk_temp_path)
+    return get_or_initialize_config(zk_temp_path)
 
 
 @pytest.fixture(scope="module")

tests/zksync/test_unit_zksync.py

@@ -7,13 +7,14 @@
 
 from moccasin.commands.compile import compile_
 from moccasin.commands.run import run_script
-from tests.constants import ZKSYNC_PROJECT_PATH
 
 
-def test_compile_zksync_pyevm(zksync_cleanup_out_folder, zksync_out_folder, mox_path):
+def test_compile_zksync_pyevm(
+    zksync_cleanup_out_folder, zk_temp_path, zksync_out_folder, mox_path
+):
     current_dir = Path.cwd()
     try:
-        os.chdir(current_dir.joinpath(ZKSYNC_PROJECT_PATH))
+        os.chdir(current_dir.joinpath(zk_temp_path))
         result = subprocess.run(
             [mox_path, "build", "Difficulty.vy", "--network", "pyevm"],
             check=True,
@@ -26,35 +27,35 @@ def test_compile_zksync_pyevm(zksync_cleanup_out_folder, zksync_out_folder, mox_
     assert result.returncode == 0
     # read the Difficulty.json in zksync_out_folder
     with open(
-        ZKSYNC_PROJECT_PATH.joinpath(zksync_out_folder).joinpath("Difficulty.json"), "r"
+        zk_temp_path.joinpath(zksync_out_folder).joinpath("Difficulty.json"), "r"
     ) as f:
         data = json.load(f)
         assert data["vm"] == "evm"
 
 
 def test_compile_zksync_one(
-    zksync_cleanup_out_folder, zksync_out_folder, zksync_test_env
+    zksync_cleanup_out_folder, zk_temp_path, zksync_out_folder, zksync_test_env
 ):
     compile_(
-        ZKSYNC_PROJECT_PATH.joinpath("src/Difficulty.vy"),
-        ZKSYNC_PROJECT_PATH.joinpath(zksync_out_folder),
+        zk_temp_path.joinpath("src/Difficulty.vy"),
+        zk_temp_path.joinpath(zksync_out_folder),
         is_zksync=True,
         write_data=True,
     )
     with open(
-        ZKSYNC_PROJECT_PATH.joinpath(zksync_out_folder).joinpath("Difficulty.json"), "r"
+        zk_temp_path.joinpath(zksync_out_folder).joinpath("Difficulty.json"), "r"
     ) as f:
         data = json.load(f)
         assert data["vm"] == "eravm"
 
 
 def test_compile_zksync_bad(
-    zksync_cleanup_out_folder, zksync_out_folder, zksync_test_env
+    zksync_cleanup_out_folder, zk_temp_path, zksync_out_folder, zksync_test_env
 ):
     with pytest.raises(AssertionError) as excinfo:
         compile_(
-            ZKSYNC_PROJECT_PATH.joinpath("src/SelfDestruct.vy"),
-            ZKSYNC_PROJECT_PATH.joinpath(zksync_out_folder),
+            zk_temp_path.joinpath("src/SelfDestruct.vy"),
+            zk_temp_path.joinpath(zksync_out_folder),
             is_zksync=True,
             write_data=True,
         )
@@ -64,7 +65,7 @@ def test_compile_zksync_bad(
     assert "The `SELFDESTRUCT` instruction is not supported" in error_message
 
 
-def test_run_zksync_good(zksync_cleanup_out_folder, zksync_test_env):
-    difficulty_contract = run_script(ZKSYNC_PROJECT_PATH.joinpath("script/deploy.py"))
+def test_run_zksync_good(zksync_cleanup_out_folder, zk_temp_path, zksync_test_env):
+    difficulty_contract = run_script(zk_temp_path.joinpath("script/deploy.py"))
     difficulty = difficulty_contract.get_difficulty()
     assert difficulty == 2500000000000000