diff --git a/fin_depo/data.py b/fin_depo/data.py index d0aed92..ed291ad 100644 --- a/fin_depo/data.py +++ b/fin_depo/data.py @@ -59,9 +59,9 @@ class DepoGroup(Depo): nested: list[Depo] def assets(self) -> Iterable[Asset]: - assets: list[Asset] = [] + assets: set[Asset] = set() for nested_depo in self.nested: - assets.extend(nested_depo.assets()) + assets.update(nested_depo.assets()) return assets def get_amount_of_asset(self, asset: Asset) -> Decimal: diff --git a/fin_depo/defi_kucoin.py b/fin_depo/defi_kucoin.py index d4bfb80..69354e8 100644 --- a/fin_depo/defi_kucoin.py +++ b/fin_depo/defi_kucoin.py @@ -38,8 +38,14 @@ class KucoinDepoFetcher(DepoFetcher): ) def get_depo(self) -> DepoGroup: + # We would ideally get timestamp from request, + # but this is fine for now. now = datetime.datetime.now(tz=datetime.UTC) + + # Assets are spread across account types, but we would like them + # clustered into different depos. assets_by_account_type: dict[str, dict[fin_defs.Asset, Decimal]] = {} + for account_data in self.client.get_accounts(): asset = fin_defs.WELL_KNOWN_SYMBOLS[account_data['currency']] balance = Decimal(account_data['balance']) @@ -51,6 +57,7 @@ class KucoinDepoFetcher(DepoFetcher): assets_for_account_type.get(asset, Decimal(0)) + balance ) del account_data, asset, balance, assets_for_account_type + return DepoGroup( 'Kucoin', now, diff --git a/test/test_data.py b/test/test_data.py index a3e981b..bd9f256 100644 --- a/test/test_data.py +++ b/test/test_data.py @@ -1,2 +1,17 @@ -def test_nothing(): - pass +import datetime +from decimal import Decimal + +import fin_defs + +from fin_depo.data import DepoGroup, DepoSingle + + +def test_duplicate_assets(): + now = datetime.datetime.now(tz=datetime.UTC) + single1 = DepoSingle('test1', now, {fin_defs.DKK: Decimal(1000)}) + single2 = DepoSingle('test2', now, {fin_defs.DKK: Decimal(1000)}) + group = DepoGroup('test3', now, [single1, single2]) + + assert len(group.nested) == 2 + + assert len(list(group.assets())) == 1 diff --git a/test/test_nordnet.py b/test/test_nordnet.py index 55c27dc..2eca71c 100644 --- a/test/test_nordnet.py +++ b/test/test_nordnet.py @@ -1,7 +1,7 @@ import pytest import requests -from fin_depo import investbank_nordnet +from fin_depo import data, investbank_nordnet from . import secrets @@ -11,6 +11,13 @@ needs_secrets = pytest.mark.skipif( ) +def print_pretty(depo: data.Depo): + for asset in depo.assets(): + amount = depo.get_amount_of_asset(asset) + print(f'{str(asset):15} : {amount:10}') + del asset, amount + + @needs_secrets def test_get_depo(): session = requests.Session() @@ -21,4 +28,5 @@ def test_get_depo(): ) depo = fetcher.get_depo() + print_pretty(depo) assert depo is not None