models for new queries, tests

bal_tools/models.py

@@ -154,3 +154,28 @@ def validate_model(cls, values):
             if field in values:
                 values[field] = cls.str_to_decimal(values[field])
         return values
+
+
+class DynamicData(BaseModel):
+    swapEnabled: bool
+
+
+class PoolData(BaseModel):
+    id: str
+    symbol: str
+    dynamicData: DynamicData
+
+
+class GaugeData(BaseModel):
+    gaugeAddress: Optional[str] = None
+    otherGauges: List[dict] = Field(default_factory=list)
+
+
+class StakingData(BaseModel):
+    gauge: Optional[GaugeData] = None
+
+
+class GaugePoolData(BaseModel):
+    staking: Optional[StakingData] = None
+    chain: str
+    symbol: str

bal_tools/pools_gauges.py

@@ -1,11 +1,11 @@
-from typing import Dict
+from typing import Dict, List
 import requests
-from .utils import to_checksum_address
-from .models import CorePools, PoolId, Symbol
+from .utils import to_checksum_address, flatten_nested_dict
 
 from gql.transport.exceptions import TransportQueryError
 from bal_tools.subgraph import Subgraph
 from bal_tools.errors import NoResultError
+from bal_tools.models import PoolData, GaugePoolData, GaugeData, StakingData, CorePools, PoolId, Symbol
 
 GITHUB_RAW_OUTPUTS = (
     "https://raw.githubusercontent.com/BalancerMaxis/bal_addresses/main/outputs"
@@ -116,33 +116,38 @@ def query_root_gauges(self, skip=0, step_size=100) -> list:
             result += self.query_root_gauges(skip + step_size, step_size)
         return result
 
-    def query_all_gauges(self, include_other_gauges=True) -> list:
+    def query_all_pools(self) -> List[PoolData]:
+        """
+        query all pools from the apiv3 subgraph
+        filters out disabled pools
+        """
+        data = self.subgraph.fetch_graphql_data("apiv3", "get_pools", {"chain": self.chain.upper()})
+        pools = [PoolData(**flatten_nested_dict(pool)) for pool in data["poolGetPools"]]
+        return [
+            pool for pool in pools
+            if pool.dynamicData.swapEnabled
+        ]
+
+    def query_all_gauges(self, include_other_gauges=True) -> List[GaugePoolData]:
         """
         query all gauges from the apiv3 subgraph
         """
         data = self.subgraph.fetch_graphql_data("apiv3", "get_gauges", {"chain": self.chain.upper()})
         all_gauges = []
         for gauge in data["poolGetPools"]:
-            if gauge['staking'] is not None:
-                if gauge['staking']['gauge'] is not None:  # this is an edge case for the 80bal20weth pool
-                    all_gauges.append({"id": gauge['staking']['gauge']['gaugeAddress'], "symbol": f"{gauge['symbol']}-gauge"})
-                    if include_other_gauges:
-                        for other_gauge in gauge['staking']['gauge']['otherGauges']:
-                            all_gauges.append({"id": other_gauge['id'], "symbol": f"{gauge['symbol']}-gauge"})
+            pool = GaugePoolData(**flatten_nested_dict(gauge))
+            if pool.staking is not None and pool.staking.gauge is not None:
+                gauge_pool = pool.model_copy(deep=True)
+                gauge_pool.symbol = f"{pool.symbol}-gauge"
+                all_gauges.append(gauge_pool)
+                if include_other_gauges and pool.staking.gauge.otherGauges:
+                    for other_gauge in pool.staking.gauge.otherGauges:
+                        other_pool = pool.model_copy(deep=True)
+                        other_pool.symbol = f"{pool.symbol}-gauge"
+                        other_pool.staking.gauge.gaugeAddress = other_gauge['id']
+                        all_gauges.append(other_pool)
         return all_gauges
 
-    def query_all_pools(self) -> list:
-        """
-        query all pools from the apiv3 subgraph
-        filters out disabled pools
-        """
-        data = self.subgraph.fetch_graphql_data("apiv3", "get_pools", {"chain": self.chain.upper()})
-        all_pools = []
-        for pool in data["poolGetPools"]:
-            if pool['dynamicData']['swapEnabled']:
-                all_pools.append({"address": pool['id'], "symbol": pool['symbol']})
-        return all_pools
-
     def get_last_join_exit(self, pool_id: int) -> int:
         """
         Returns a timestamp of the last join/exit for a given pool id

bal_tools/utils.py

@@ -27,8 +27,10 @@ def get_abi(contract_name: str) -> Union[Dict, List[Dict]]:
 
 
 def flatten_nested_dict(d):
-    for key, value in list(d.items()):
+    result = d.copy()
+    for key, value in list(result.items()):
         if isinstance(value, dict):
-            d.pop(key)
-            d.update(value)
-    return d
+            if key not in ['dynamicData', 'staking']:
+                result.pop(key)
+                result.update(value)
+    return result

tests/test_pools_gauges.py

@@ -1,5 +1,6 @@
 import pytest
 from gql.transport.exceptions import TransportQueryError
+from bal_tools.models import PoolData, GaugePoolData
 from bal_tools.models import CorePools
 import json
 
@@ -130,3 +131,22 @@ def test_get_preferential_gauge(bal_pools_gauges):
         pytest.skip(f"Skipping {bal_pools_gauges.chain}, no example preferential gauge")
 
     assert bal_pools_gauges.get_preferential_gauge(example[0]) == example[1]
+
+
+def test_query_all_pools(bal_pools_gauges):
+    """
+    test return data of v3 AllGauges query
+    """
+    response = bal_pools_gauges.query_all_pools()
+
+    if len(response) > 0:
+        assert isinstance(response[0], PoolData)
+
+def test_query_all_gauges(bal_pools_gauges):
+    """
+    test return data of v3 AllPools query
+    """
+    response = bal_pools_gauges.query_all_gauges()
+
+    if len(response) > 0:
+        assert isinstance(response[0], GaugePoolData)