feat: TerraApi: merging balances by str(coingecko_id), asset_type

blockapi/test/v2/api/test_terra.py

@@ -3,7 +3,7 @@
 import pytest
 
 from blockapi.v2.api.terra import TerraApi
-from blockapi.v2.models import AssetType, Blockchain
+from blockapi.v2.models import AssetType, BalanceItem, Blockchain, Coin, CoinInfo
 
 
 @pytest.fixture()
@@ -257,6 +257,52 @@ def test_get_balance_with_ibc_tokens(
     assert 'ibc' in rowan.coin.standards
 
 
+def _make_balance(address, balance_raw, coingecko_id):
+    return BalanceItem(
+        balance=Decimal(balance_raw) / Decimal(10**6),
+        balance_raw=Decimal(balance_raw),
+        raw={'denom': address, 'amount': balance_raw},
+        coin=Coin(
+            symbol='X',
+            name='X',
+            decimals=6,
+            blockchain=Blockchain.TERRA,
+            address=address,
+            standards=[],
+            info=CoinInfo(coingecko_id=coingecko_id),
+        ),
+        asset_type=AssetType.AVAILABLE,
+    )
+
+
+def test_merge_duplicate_balances():
+    """Two balances share a coingecko_id and merge; two distinct ones do not."""
+    balances = [
+        _make_balance('ibc/AAAA', '100', 'secret'),
+        _make_balance('ibc/BBBB', '200', 'secret'),
+        _make_balance('uluna', '1000000', 'terra-luna'),
+        _make_balance('ibc/CCCC', '50', None),
+    ]
+
+    merged = TerraApi._merge_duplicate_balances(balances)
+
+    assert len(merged) == 3
+
+    secret = next(b for b in merged if b.coin.info.coingecko_id == 'secret')
+    assert secret.balance_raw == Decimal('300')
+    assert secret.balance == Decimal('0.0003')
+    assert secret.raw['merged_count'] == 2
+    assert len(secret.raw['merged']) == 2
+
+    lunc = next(b for b in merged if b.coin.info.coingecko_id == 'terra-luna')
+    assert lunc.balance_raw == Decimal('1000000')
+    assert 'merged' not in lunc.raw
+
+    no_cg = next(b for b in merged if b.coin.info.coingecko_id is None)
+    assert no_cg.balance_raw == Decimal('50')
+    assert 'merged' not in no_cg.raw
+
+
 def test_unbonding_included_in_staked(terra_api, requests_mock):
     requests_mock.get(
         f'{BASE_URL}/cosmos/bank/v1beta1/balances/{ADDRESS}',

blockapi/v2/api/terra.py

@@ -1,10 +1,13 @@
 import logging
+from decimal import Decimal
 from functools import lru_cache
 
+import attr
+
 from blockapi.v2.api.cosmos import CosmosApiBase
 from blockapi.v2.base import ApiException, ApiOptions
 from blockapi.v2.coins import COIN_TERRA
-from blockapi.v2.models import Blockchain, Coin
+from blockapi.v2.models import BalanceItem, Blockchain, Coin
 
 logger = logging.getLogger(__name__)
 
@@ -28,6 +31,62 @@ class TerraApi(CosmosApiBase):
         'get_ibc_denom_trace': '/ibc/apps/transfer/v1/denom_traces/{hash}',
     }
 
+    def get_balance(
+        self, address: str, merge_duplicates: bool = True
+    ) -> list[BalanceItem]:
+        balances = super().get_balance(address)
+        if merge_duplicates:
+            balances = self._merge_duplicate_balances(balances)
+        return balances
+
+    @staticmethod
+    def _merge_duplicate_balances(
+        balances: list[BalanceItem],
+    ) -> list[BalanceItem]:
+        """
+        Merge balances that map to the same upstream coin (same coingecko_id
+        + asset_type) into a single BalanceItem. Distinct IBC denoms that
+        resolve to the same currency (e.g. SCRT via two channels) collapse
+        into one entry. Items without a coingecko_id are left untouched to
+        avoid collapsing unrelated unknown tokens.
+
+        Merged items carry `raw = {'merged': [...originals], 'merged_count': N}`.
+        """
+        groups: dict[tuple, list[BalanceItem]] = {}
+        order: list[tuple] = []
+
+        for idx, b in enumerate(balances):
+            coingecko_id = b.coin.info.coingecko_id if b.coin and b.coin.info else None
+            if coingecko_id:
+                key = (str(coingecko_id), b.asset_type)
+            else:
+                key = ('__unique__', idx)
+            if key not in groups:
+                groups[key] = []
+                order.append(key)
+            groups[key].append(b)
+
+        merged: list[BalanceItem] = []
+        for key in order:
+            group = groups[key]
+            if len(group) == 1:
+                merged.append(group[0])
+                continue
+            head = group[0]
+            merged.append(
+                attr.evolve(
+                    head,
+                    balance=sum((b.balance for b in group), Decimal(0)),
+                    balance_raw=sum((b.balance_raw for b in group), Decimal(0)),
+                    raw={
+                        'merged': [b.raw for b in group],
+                        'merged_count': len(group),
+                    },
+                )
+            )
+
+        return merged
+
     def create_default_coin(self, denom: str) -> Coin:
         if denom.startswith('ibc/'):
             return self._resolve_ibc_denom(denom)