From e56e15a9abc3d1ab19c7b1762f6a7bf58cf9e156 Mon Sep 17 00:00:00 2001
From: Tritium <tritium-vlk@telegram.works>
Date: Wed, 26 Jul 2023 23:38:08 +0200
Subject: [PATCH] make tests pass.

---
 Makefile                   |  2 +-
 bal_addresses/addresses.py | 25 +++++++++++++
 tests/test_addresses.py    | 77 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 102 insertions(+), 2 deletions(-)

diff --git a/Makefile b/Makefile
index 4dd849d0..088944a0 100644
--- a/Makefile
+++ b/Makefile
@@ -4,4 +4,4 @@ init:
 	pip install -r bal_addresses/requirements.txt
 	pip install -r bal_addresses/requirements-dev.txt
 ci:
-	pytest
+	pytest 
diff --git a/bal_addresses/addresses.py b/bal_addresses/addresses.py
index c0352ed9..ef692ccb 100644
--- a/bal_addresses/addresses.py
+++ b/bal_addresses/addresses.py
@@ -20,6 +20,9 @@
 GITHUB_RAW_OUTPUTS = (
     "https://raw.githubusercontent.com/BalancerMaxis/bal_addresses/main/outputs"
 )
+GITHUB_RAW_EXTRAS = (
+    "https://raw.githubusercontent.com/BalancerMaxis/bal_addresses/main/extras"
+)
 ZERO_ADDRESS = "0x0000000000000000000000000000000000000000"
 
 
@@ -39,6 +42,7 @@ class AddrBook:
     fx_description_by_name = DotMap(requests.get(
         "https://raw.githubusercontent.com/BalancerMaxis/bal_addresses/main/extras/func_desc_by_name.json"
     ).json())
+    chain_ids_by_name = chains.CHAIN_IDS_BY_NAME
 
     def __init__(self, chain, jsonfile=False):
         self.jsonfile = jsonfile
@@ -124,6 +128,27 @@ def _process_deployment(self, deployment: Dict) -> Dict:
             processed_deployment[deployment_identifier] = v
         return processed_deployment
 
+
+    def populate_extras(self) -> None:
+        chain_extras = requests.get(
+            f"{GITHUB_RAW_EXTRAS}/{self.chain}.json"
+        )
+        if chain_extras.ok:
+            self._extras = Munch.fromDict(chain_extras.json())
+
+
+    def populate_multisigs(self) -> None:
+        msigs = requests.get(
+            f"{GITHUB_RAW_EXTRAS}/multisigs.json"
+        ).json()
+        if msigs.get(self.chain):
+            self._multisigs = Munch.fromDict(msigs[self.chain])
+        else:
+            print(f"Warning: No multisigs for chain {self.chain}, multisigs must be added in extras/multisig.json")
+            self._multisigs = {}
+
+
+
     def search_unique(self, substr):
         results = [s for s in self.flatbook.keys() if substr in s]
         if len(results) > 1:
diff --git a/tests/test_addresses.py b/tests/test_addresses.py
index a97d6a97..e41d65f7 100644
--- a/tests/test_addresses.py
+++ b/tests/test_addresses.py
@@ -38,6 +38,28 @@ def test_deployments_populated():
             }
         }
     )
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/BalancerMaxis"
+        "/bal_addresses/main/extras/mainnet.json",
+        json={
+            "zero": {
+                "zero": "0x0000000000000000000000000000000000000000"
+            },
+            "balancer": {}
+
+            }
+    )
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/BalancerMaxis"
+        "/bal_addresses/main/extras/multisigs.json",
+        json={
+            "mainnet": {
+                "dao": "0x10A19e7eE7d7F8a52822f6817de8ea18204F2e4f",
+            }
+        }
+    )
     a = AddrBook("mainnet")
 
     a.populate_deployments()
@@ -50,7 +72,14 @@ def test_deployments_populated():
     # 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
-
+    a.populate_extras()
+    assert a.extras.zero.zero == "0x0000000000000000000000000000000000000000"
+    # Make sure that when we try to access a non-existing attribute, we get an error
+    with pytest.raises(AttributeError):
+        assert a.extras.balancer.non_existing_attribute
+    a.populate_multisigs()
+    print(a.multisigs)
+    assert a.multisigs.dao == "0x10A19e7eE7d7F8a52822f6817de8ea18204F2e4f"
 
 @responses.activate
 def test_deployments_invalid_format():
@@ -76,10 +105,36 @@ def test_deployments_invalid_format():
             }
         }
     )
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/BalancerMaxis"
+        "/bal_addresses/main/extras/mainnet.json",
+        json={
+            "vault": {
+                "contracts": {'name': 'Vault'},
+                "status": "ACTIVE"
+            }
+        }
+    )
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/BalancerMaxis"
+        "/bal_addresses/main/extras/multisigs.json",
+        json={
+            "mainnet": {
+                "contracts": {'name': 'Vault'},
+                "status": "ACTIVE"
+            }
+        }
+    )
     a = AddrBook("mainnet")
 
     a.populate_deployments()
     assert a.deployments.vault.contracts.name == "Vault"
+    a.populate_extras()
+    assert a.extras.vault.status == "ACTIVE"
+    a.populate_multisigs()
+    assert a.multisigs.status == "ACTIVE"
 
 
 @responses.activate
@@ -99,7 +154,27 @@ def test_deployments_not_populated():
         json={},
         status=404
     )
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/BalancerMaxis"
+        "/bal_addresses/main/extras/mainnet.json",
+        json={},
+        status=404
+    )
+    responses.add(
+        responses.GET,
+        "https://raw.githubusercontent.com/BalancerMaxis"
+        "/bal_addresses/main/extras/multisigs.json",
+        json={},
+        status=404
+    )
     a = AddrBook("mainnet")
     assert a.deployments is None
     with pytest.raises(AttributeError):
         assert a.deployments.vault.non_existing_attribute
+    assert a.extras is None
+    with pytest.raises(AttributeError):
+        assert a.extras.non_existing_attribute
+    assert a.multisigs is None
+    with pytest.raises(AttributeError):
+        assert a.multisigs.non_existing_attribute
\ No newline at end of file