Merge remote-tracking branch 'origin/add_aura_balances_query' into ad…

…d_aura_balances_query

bal_addresses/ecosystem.py

@@ -5,7 +5,9 @@
 from .subgraph import Subgraph
 from .utils import to_checksum_address
 
-AURA_L2_DEFAULT_GAUGE_STAKER = to_checksum_address("0xC181Edc719480bd089b94647c2Dc504e2700a2B0")
+AURA_L2_DEFAULT_GAUGE_STAKER = to_checksum_address(
+    "0xC181Edc719480bd089b94647c2Dc504e2700a2B0"
+)
 
 
 class KeyAsDefaultDict(defaultdict):
@@ -16,7 +18,9 @@ def __missing__(self, key):
 class Aura:
     ## Aura seems to stake from the same address on all chains except mainnet
     AURA_GAUGE_STAKER_BY_CHAIN = defaultdict(lambda: AURA_L2_DEFAULT_GAUGE_STAKER)
-    AURA_GAUGE_STAKER_BY_CHAIN["mainnet"] = to_checksum_address("0xaF52695E1bB01A16D33D7194C28C42b10e0Dbec2")
+    AURA_GAUGE_STAKER_BY_CHAIN["mainnet"] = to_checksum_address(
+        "0xaF52695E1bB01A16D33D7194C28C42b10e0Dbec2"
+    )
 
     def __init__(self, chain):
         self.chain = chain
@@ -32,21 +36,24 @@ def get_aura_gauge_mappings(self) -> Dict[str, int]:
         Get a dict with gauge_address as key and aura PID as value for the running chain.
         """
         data = self.subgraph.fetch_graphql_data("aura", "get_aura_gauge_mappings")
-        #print(json.dumps(data, indent=1))
+        # print(json.dumps(data, indent=1))
         aura_pid_by_gauge = {}
         for result_item in data["gauges"]:
             gauge_address = to_checksum_address(result_item["pool"]["gauge"]["id"])
             pid = result_item["pool"]["id"]
             # Seems like pid can be a string or a list
             if isinstance(pid, list):
                 if len(pid > 1):
-                    raise MultipleMatchesError(f"Gauge: {gauge_address} is returning multiple aura PIDs: {pid}")
+                    raise MultipleMatchesError(
+                        f"Gauge: {gauge_address} is returning multiple aura PIDs: {pid}"
+                    )
                 else:
                     pid = [pid][0]
 
             if gauge_address in aura_pid_by_gauge:
                 raise MultipleMatchesError(
-                    f"Gauge with address{gauge_address} already found with PID {aura_pid_by_gauge[gauge_address]} when trying to insert new PID {pid}")
+                    f"Gauge with address{gauge_address} already found with PID {aura_pid_by_gauge[gauge_address]} when trying to insert new PID {pid}"
+                )
             aura_pid_by_gauge[gauge_address] = pid
         return aura_pid_by_gauge
 
@@ -65,14 +72,16 @@ def get_aura_pool_shares(self, gauge_address, block) -> Dict[str, int]:
         gauge_address = to_checksum_address(gauge_address)
         aura_pid = self.aura_pids_by_address.get(gauge_address)
         variables = {"poolId": aura_pid, "block": int(block)}
-        data = self.subgraph.fetch_graphql_data("aura", "get_aura_user_pool_balances", variables)
+        data = self.subgraph.fetch_graphql_data(
+            "aura", "get_aura_user_pool_balances", variables
+        )
         results = {}
         # Parse the data if the query was successful
-        if data and  'leaderboard' in data and data['leaderboard']['accounts']:
-            for account in data['leaderboard']['accounts']:
+        if data and "leaderboard" in data and data["leaderboard"]["accounts"]:
+            for account in data["leaderboard"]["accounts"]:
                 ## Aura amounts are WEI denominated and others are float.  Transform
-                amount = float(int(account['staked']) / 1e18)
-                user_address = to_checksum_address(account['account']['id'])
+                amount = float(int(account["staked"]) / 1e18)
+                user_address = to_checksum_address(account["account"]["id"])
                 results[user_address] = amount
         # TODO better handle pagination with the query and this function/pull multiple pages if required
         assert len(results) < 1000, "Pagination limit hit on Aura query"
@@ -97,5 +106,6 @@ def get_aura_pid_from_gauge(self, deposit_gauge_address: str) -> int:
         else:
             if len(result) != 1:
                 raise UnexpectedListLengthError(
-                    f"Got a list result with something other than 1 member when compiling aura PID mapping: {result}")
+                    f"Got a list result with something other than 1 member when compiling aura PID mapping: {result}"
+                )
             return result[0]

bal_addresses/pools_gauges.py

@@ -55,7 +55,9 @@ def get_gauge_deposit_shares(
     ) -> Dict[str, int]:
         gauge_address = to_checksum_address(gauge_address)
         variables = {"gaugeAddress": gauge_address, "block": int(block)}
-        data = self.subgraph.fetch_graphql_data("gauges", "fetch_gauge_shares", variables)
+        data = self.subgraph.fetch_graphql_data(
+            "gauges", "fetch_gauge_shares", variables
+        )
         results = {}
         if "gaugeShares" in data:
             for share in data["gaugeShares"]:

tests/conftest.py

@@ -5,6 +5,7 @@
 from bal_addresses.subgraph import Subgraph
 from bal_addresses.ecosystem import Aura
 
+
 @pytest.fixture(scope="module", params=list(AddrBook.chains["CHAIN_IDS_BY_NAME"]))
 def chain(request):
     chain = request.param
@@ -20,6 +21,7 @@ def bal_pools_gauges(chain):
 def subgraph(chain):
     return Subgraph(chain)
 
+
 @pytest.fixture(scope="module")
 def aura(chain):
-    return Aura(chain)
+    return Aura(chain)