models for new queries, tests

bal_tools/models.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass, fields
 from enum import Enum
 from pydantic import BaseModel, field_validator, model_validator, Field
+import json_fix
 
 
 class GqlChain(Enum):
@@ -43,6 +44,8 @@ def __iter__(self):
     def __len__(self):
         return len(self.pools)
 
+    def __json__(self):
+        return self.pools
 
 
 @dataclass
@@ -154,3 +157,21 @@ def validate_model(cls, values):
             if field in values:
                 values[field] = cls.str_to_decimal(values[field])
         return values
+
+
+class PoolData(BaseModel):
+    address: str
+    symbol: str
+    dynamicData: dict
+
+
+class GaugePoolData(BaseModel):
+    staking: Optional[dict]
+    chain: str
+    symbol: str
+
+
+@dataclass
+class GaugeData:
+    address: 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, CorePools, PoolId, Symbol
 
 GITHUB_RAW_OUTPUTS = (
     "https://raw.githubusercontent.com/BalancerMaxis/bal_addresses/main/outputs"
@@ -116,31 +116,39 @@ 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_gauges(self, include_other_gauges=True) -> List[GaugeData]:
         """
         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"})
+        for pool in data["poolGetPools"]:
+            gauge_pool = GaugePoolData(**flatten_nested_dict(pool))
+            if gauge_pool.staking is not None and gauge_pool.staking.get('gauge') is not None:
+                gauge = gauge_pool.staking['gauge']
+                all_gauges.append(GaugeData(
+                    address=gauge['gaugeAddress'],
+                    symbol=f"{gauge_pool.symbol}-gauge"
+                ))
+                if include_other_gauges:
+                    for other_gauge in gauge.get('otherGauges', []):
+                        all_gauges.append(GaugeData(
+                            address=other_gauge['id'],
+                            symbol=f"{gauge_pool.symbol}-gauge"
+                        ))
         return all_gauges
 
-    def query_all_pools(self) -> 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()})
         all_pools = []
         for pool in data["poolGetPools"]:
-            if pool['dynamicData']['swapEnabled']:
-                all_pools.append({"address": pool['address'], "symbol": pool['symbol']})
+            pool_data = PoolData(**flatten_nested_dict(pool))
+            if pool_data.dynamicData['swapEnabled']:
+                all_pools.append(pool_data)
         return all_pools
 
     def get_last_join_exit(self, pool_id: int) -> int:
@@ -162,7 +170,6 @@ def get_pool_tvl(self, pool_id: str) -> float:
         """
         Returns the TVL of a pool as per the API V3 subgraph
         """
-        print(pool_id)
         try:
             data = self.subgraph.fetch_graphql_data(
                 "apiv3", "get_pool_tvl", {"chain": self.chain.upper(), "poolId": pool_id}

bal_tools/requirements.txt

@@ -7,5 +7,6 @@ munch==4.0.0
 gql[requests]
 python-dotenv
 pydantic==2.7.4
+json-fix
 
 git+https://github.com/BalancerMaxis/[email protected]

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, GaugeData
 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], GaugeData)