From a35ba6d97a42e5b978ac241ae59fac7bfae7bee8 Mon Sep 17 00:00:00 2001
From: antazoey <jules@apeworx.io>
Date: Thu, 6 Feb 2025 19:37:45 -0600
Subject: [PATCH] fix: `KeyError` on deployments-cache (#2496)

---
 src/ape/managers/_deploymentscache.py     |  5 +-
 src/ape/utils/basemodel.py                |  1 +
 tests/functional/test_deploymentscache.py | 67 +++++++++++++++++++++++
 3 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 tests/functional/test_deploymentscache.py

diff --git a/src/ape/managers/_deploymentscache.py b/src/ape/managers/_deploymentscache.py
index 9a496f83b0..12f21078c1 100644
--- a/src/ape/managers/_deploymentscache.py
+++ b/src/ape/managers/_deploymentscache.py
@@ -83,6 +83,9 @@ def __setitem__(self, contract_name, deployments: list[Deployment]):
     def __delitem__(self, contract_name: str):
         self.remove_deployments(contract_name)
 
+    def __contains__(self, contract_name: str):
+        return bool(self.get_deployments(contract_name))
+
     def get_deployments(
         self,
         contract_name: str,
@@ -172,7 +175,7 @@ def _set_deployments(
         self._all_deployments.ecosystems[ecosystem_name][network_name][contract_name] = deployments
 
         # For live networks, cache the deployments to a file as well.
-        if self._is_live_network:
+        if self._is_live_network and ecosystem_name in self._deployments:
             self._deployments[ecosystem_name].model_dump_file()
 
     def remove_deployments(self, contract_name: str):
diff --git a/src/ape/utils/basemodel.py b/src/ape/utils/basemodel.py
index be4ecee9e4..e96cbeda44 100644
--- a/src/ape/utils/basemodel.py
+++ b/src/ape/utils/basemodel.py
@@ -640,6 +640,7 @@ def model_dump_file(self, path: Optional[Path] = None, **kwargs):
         path = self._get_path(path=path)
         json_str = self.model_dump_json(**kwargs)
         path.unlink(missing_ok=True)
+        path.parent.mkdir(parents=True, exist_ok=True)
         path.write_text(json_str)
 
     @classmethod
diff --git a/tests/functional/test_deploymentscache.py b/tests/functional/test_deploymentscache.py
new file mode 100644
index 0000000000..f59946d9d8
--- /dev/null
+++ b/tests/functional/test_deploymentscache.py
@@ -0,0 +1,67 @@
+import json
+
+import pytest
+
+from ape.managers._deploymentscache import DeploymentDiskCache
+
+
+class TestDeploymentDiskCache:
+    CONTRACT_NAME = "DeploymentTestContractName"
+
+    @pytest.fixture(scope="class")
+    def contract_name(self):
+        return self.CONTRACT_NAME
+
+    @pytest.fixture
+    def cache(self):
+        return DeploymentDiskCache()
+
+    def test_cache_deployment(self, zero_address, cache, contract_name):
+        cache.cache_deployment(zero_address, contract_name)
+        assert contract_name in cache
+        assert cache[contract_name][-1].address == zero_address
+
+    def test_cache_deployment_live_network(
+        self, zero_address, cache, contract_name, mock_sepolia, eth_tester_provider
+    ):
+        local = eth_tester_provider.network
+        ecosystem_name = mock_sepolia.ecosystem.name
+
+        eth_tester_provider.network = mock_sepolia
+        cache.cache_deployment(zero_address, contract_name)
+        eth_tester_provider.network = local
+
+        assert contract_name in cache
+        assert cache[contract_name][-1].address == zero_address
+        # Show it is also cached on disk.
+        disk_data = json.loads(cache.cachefile.read_text())
+        assert (
+            disk_data["ecosystems"][ecosystem_name][mock_sepolia.name][contract_name][0]["address"]
+            == zero_address
+        )
+
+    def test_cache_deployment_live_network_new_ecosystem(
+        self, zero_address, cache, contract_name, mock_sepolia, eth_tester_provider
+    ):
+        """
+        Tests the case when caching a deployment in a new ecosystem.
+        """
+        ecosystem_name = mock_sepolia.ecosystem.name
+        local = eth_tester_provider.network
+        eth_tester_provider.network = mock_sepolia
+        # Make the ecosystem key not exist.
+        deployments = cache._deployments.pop(ecosystem_name, None)
+        cache.cache_deployment(zero_address, contract_name)
+        eth_tester_provider.network = local
+        if deployments is not None:
+            cache._deployments[ecosystem_name] = deployments
+        cache.cachefile.unlink(missing_ok=True)
+
+        # In memory cached still work.
+        assert contract_name in cache
+        assert cache[contract_name][-1].address == zero_address
+
+        # Show it did NOT cache to disk.
+        if cache.cachefile.is_file():
+            disk_data = json.loads(cache.cachefile.read_text())
+            assert contract_name not in disk_data["ecosystems"][ecosystem_name]["sepolia"]