diff --git a/blockapi/test/v2/api/test_terra.py b/blockapi/test/v2/api/test_terra.py index 1ab58e4..b3debf1 100644 --- a/blockapi/test/v2/api/test_terra.py +++ b/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}', diff --git a/blockapi/v2/api/terra.py b/blockapi/v2/api/terra.py index a3b313b..8794be6 100644 --- a/blockapi/v2/api/terra.py +++ b/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)