From a1e3f263f157c7ff2b848f03cc238c46a7fd6f7f Mon Sep 17 00:00:00 2001
From: SHAKOTN <shakotn@protonmail.com>
Date: Wed, 12 Jul 2023 13:11:06 +0300
Subject: [PATCH] chore: add some basic tests

---
 bal_addresses/addresses.py | 25 ++++++++------
 tests/test_addresses.py    | 69 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 83 insertions(+), 11 deletions(-)

diff --git a/bal_addresses/addresses.py b/bal_addresses/addresses.py
index 07e2243c..6f0051a2 100644
--- a/bal_addresses/addresses.py
+++ b/bal_addresses/addresses.py
@@ -1,5 +1,6 @@
 import json
 from typing import Dict
+from typing import Optional
 
 import requests
 from dotmap import DotMap
@@ -57,23 +58,27 @@ def __init__(self, chain, jsonfile=False):
         self._deployments = None
 
     @property
-    def deployments(self) -> Munch:
+    def deployments(self) -> Optional[Munch]:
         """
         Get the deployments for all chains in a form of a Munch object
         """
         if self._deployments is not None:
             return self._deployments
-        self._deployments = Munch()
-        for chain_name in self.chains['CHAIN_IDS_BY_NAME'].keys():
-            chain_deployments = requests.get(
-                f"{GITHUB_DEPLOYMENTS_RAW}/addresses/{chain_name}.json"
-            )
-            if chain_deployments.ok:
-                # Remove date from key
-                processed_deployment = self._process_deployment(chain_deployments.json())
-                setattr(self._deployments, chain_name, Munch.fromDict(processed_deployment))
+        else:
+            self.populate_deployments()
+
         return self._deployments
 
+    def populate_deployments(self) -> None:
+        chain_deployments = requests.get(
+            f"{GITHUB_DEPLOYMENTS_RAW}/addresses/{self.chain}.json"
+        )
+        if chain_deployments.ok:
+            self._deployments = Munch()
+            # Remove date from key
+            processed_deployment = self._process_deployment(chain_deployments.json())
+            self._deployments = Munch.fromDict(processed_deployment)
+
     def _process_deployment(self, deployment: Dict) -> Dict:
         """
         Process deployment to remove date from key and replace - with _
diff --git a/tests/test_addresses.py b/tests/test_addresses.py
index 134407e8..b16dde31 100644
--- a/tests/test_addresses.py
+++ b/tests/test_addresses.py
@@ -1,5 +1,72 @@
+import pytest
+import responses
+
 from bal_addresses import AddrBook
 
 
-def test_addresses():
+@responses.activate
+def test_deployments_populated():
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/BalancerMaxis"
+        "/bal_addresses/main/outputs/deployments.json",
+        json={
+            "BFactory": "0x9424B1412450D0f8Fc2255FAf6046b98213B76Bd",
+        }
+    )
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/balancer"
+        "/balancer-deployments/master/addresses/mainnet.json",
+        json={
+            "20210418-vault": {
+                "contracts": [
+                    {
+                        "name": "Vault",
+                        "address": "0xBA12222222228d8Ba445958a75a0704d566BF2C8"
+                    },
+                    {
+                        "name": "BalancerHelpers",
+                        "address": "0x5aDDCCa35b7A0D07C74063c48700C8590E87864E"
+                    },
+                    {
+                        "name": "ProtocolFeesCollector",
+                        "address": "0xce88686553686DA562CE7Cea497CE749DA109f9F"
+                    }
+                ],
+                "status": "ACTIVE"
+            }
+        }
+    )
+    a = AddrBook("mainnet")
+
+    a.populate_deployments()
+    assert a.deployments.vault.status == "ACTIVE"
+    assert a.deployments.vault.contracts[0].name == "Vault"
+    assert a.deployments.vault.contracts[1].name == "BalancerHelpers"
+    # Make sure that when we try to access a non-existing attribute, we get an error
+    with pytest.raises(AttributeError):
+        assert a.deployments.vault.non_existing_attribute
+
+
+@responses.activate
+def test_deployments_not_populated():
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/BalancerMaxis"
+        "/bal_addresses/main/outputs/deployments.json",
+        json={
+            "BFactory": "0x9424B1412450D0f8Fc2255FAf6046b98213B76Bd",
+        }
+    )
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/balancer"
+        "/balancer-deployments/master/addresses/mainnet.json",
+        json={},
+        status=404
+    )
     a = AddrBook("mainnet")
+    assert a.deployments is None
+    with pytest.raises(AttributeError):
+        assert a.deployments.vault.non_existing_attribute