From d777ca0a318a8ade7a20363c9ce77fe8a8bf5d68 Mon Sep 17 00:00:00 2001 From: Daniel McGregor Date: Tue, 27 Aug 2024 09:23:26 +0800 Subject: [PATCH] feat: add support for frozen on asset holdings, including a new ledger function `update_asset_holdings` for setting asset holding balances and frozen states --- .../context_helpers/ledger_context.py | 33 +++++++++++++++---- src/_algopy_testing/models/account.py | 14 +++++--- src/_algopy_testing/models/asset.py | 25 +++++++------- src/_algopy_testing/op/misc.py | 13 +++----- src/_algopy_testing/value_generators/avm.py | 9 +++-- tests/models/test_asset.py | 23 ++++++------- 6 files changed, 73 insertions(+), 44 deletions(-) diff --git a/src/_algopy_testing/context_helpers/ledger_context.py b/src/_algopy_testing/context_helpers/ledger_context.py index 8663248..e95e522 100644 --- a/src/_algopy_testing/context_helpers/ledger_context.py +++ b/src/_algopy_testing/context_helpers/ledger_context.py @@ -79,23 +79,44 @@ def account_is_funded(self, account: algopy.Account | str) -> bool: def update_account( self, account: algopy.Account | str, - opted_asset_balances: dict[int, algopy.UInt64] | None = None, **account_fields: typing.Unpack[AccountFields], ) -> None: """Update account fields. Args: - address (str): The account address. - opted_asset_balances (dict[int, algopy.UInt64] | None): The opted asset balances . + account: The account. **account_fields: The fields to update. """ address = _get_address(account) assert_address_is_valid(address) self._account_data[address].fields.update(account_fields) - if opted_asset_balances is not None: - for asset_id, balance in opted_asset_balances.items(): - self._account_data[address].opted_asset_balances[UInt64(asset_id)] = balance + def update_asset_holdings( + self, + account: algopy.Account | str, + asset: algopy.Asset | algopy.UInt64 | int, + *, + balance: algopy.UInt64 | int | None = None, + frozen: bool | None = None, + ) -> None: + """Update asset holdings for account, only specified values will be updated. + + Account will also be opted-in to asset + """ + from _algopy_testing.models.account import AssetHolding + + address = _get_address(account) + account_data = self._account_data[address] + asset = self.get_asset(_get_asset_id(asset)) + + holdings = account_data.opted_assets.setdefault( + asset.id, + AssetHolding(balance=UInt64(), frozen=asset.default_frozen), + ) + if balance is not None: + holdings.balance = UInt64(int(balance)) + if frozen is not None: + holdings.frozen = frozen def get_asset(self, asset_id: algopy.UInt64 | int) -> algopy.Asset: """Get an asset by ID. diff --git a/src/_algopy_testing/models/account.py b/src/_algopy_testing/models/account.py index a77041b..f01b286 100644 --- a/src/_algopy_testing/models/account.py +++ b/src/_algopy_testing/models/account.py @@ -52,19 +52,23 @@ def get_empty_account() -> AccountContextData: ) +@dataclasses.dataclass +class AssetHolding: + balance: algopy.UInt64 + frozen: bool + + @dataclasses.dataclass class AccountContextData: """Stores account-related information. Attributes: - opted_asset_balances (dict[int, algopy.UInt64]): Mapping of asset IDs to balances. + opted_assets (dict[int, AssetHolding]): Mapping of asset IDs to holdings. opted_apps (dict[int, algopy.UInt64]): Mapping of application IDs to instances. fields (AccountFields): Additional account fields. """ - opted_asset_balances: dict[algopy.UInt64, algopy.UInt64] = dataclasses.field( - default_factory=dict - ) + opted_assets: dict[algopy.UInt64, AssetHolding] = dataclasses.field(default_factory=dict) opted_apps: dict[algopy.UInt64, algopy.Application] = dataclasses.field(default_factory=dict) fields: AccountFields = dataclasses.field(default_factory=AccountFields) # type: ignore[arg-type] @@ -96,7 +100,7 @@ def is_opted_in(self, asset_or_app: algopy.Asset | algopy.Application, /) -> boo from _algopy_testing.models import Application, Asset if isinstance(asset_or_app, Asset): - return asset_or_app.id in self.data.opted_asset_balances + return asset_or_app.id in self.data.opted_assets elif isinstance(asset_or_app, Application): return asset_or_app.id in self.data.opted_apps diff --git a/src/_algopy_testing/models/asset.py b/src/_algopy_testing/models/asset.py index 276c1eb..7648f0d 100644 --- a/src/_algopy_testing/models/asset.py +++ b/src/_algopy_testing/models/asset.py @@ -46,25 +46,26 @@ def balance(self, account: algopy.Account) -> algopy.UInt64: account_data = lazy_context.get_account_data(account.public_key) - if not account_data: - raise ValueError("Account not found in testing context!") - - if int(self.id) not in account_data.opted_asset_balances: + if int(self.id) not in account_data.opted_assets: raise ValueError( "The asset is not opted into the account! " "Use `ctx.any.account(opted_asset_balances={{ASSET_ID: VALUE}})` " "to set emulated opted asset into the account." ) - return account_data.opted_asset_balances[self.id] + return account_data.opted_assets[self.id].balance + + def frozen(self, account: algopy.Account) -> bool: + from _algopy_testing.context_helpers import lazy_context - def frozen(self, _account: algopy.Account) -> bool: - # TODO: 1.0 expand data structure on AccountContextData.opted_asset_balances - # to support frozen attribute - raise NotImplementedError( - "The 'frozen' method is being executed in a python testing context. " - "Please mock this method using your python testing framework of choice." - ) + account_data = lazy_context.get_account_data(account.public_key) + if int(self.id) not in account_data.opted_assets: + raise ValueError( + "The asset is not opted into the account! " + "Use `ctx.any.account(opted_asset_balances={{ASSET_ID: VALUE}})` " + "to set emulated opted asset into the account." + ) + return account_data.opted_assets[self.id].frozen @property def fields(self) -> AssetFields: diff --git a/src/_algopy_testing/op/misc.py b/src/_algopy_testing/op/misc.py index c3ba2dd..440fe6c 100644 --- a/src/_algopy_testing/op/misc.py +++ b/src/_algopy_testing/op/misc.py @@ -91,6 +91,7 @@ def balance(a: algopy.Account | algopy.UInt64 | int, /) -> algopy.UInt64: account = _get_account(a) return account.balance + def min_balance(a: algopy.Account | algopy.UInt64 | int, /) -> algopy.UInt64: account = _get_account(a) return account.min_balance @@ -170,18 +171,14 @@ def _get_asset_holding( return UInt64(0), False account_data = lazy_context.get_account_data(account.public_key) - asset_balance = account_data.opted_asset_balances.get(asset.id) - if asset_balance is None: + holding = account_data.opted_assets.get(asset.id) + if holding is None: return UInt64(0), False if field == "balance": - return asset_balance, True + return holding.balance, True elif field == "frozen": - try: - asset_data = lazy_context.ledger.get_asset(asset.id) - except KeyError: - return UInt64(0), False - return asset_data.default_frozen, True + return holding.frozen, True else: raise ValueError(f"Invalid asset holding field: {field}") diff --git a/src/_algopy_testing/value_generators/avm.py b/src/_algopy_testing/value_generators/avm.py index ef083ff..a876c67 100644 --- a/src/_algopy_testing/value_generators/avm.py +++ b/src/_algopy_testing/value_generators/avm.py @@ -15,7 +15,7 @@ MAX_UINT512, ) from _algopy_testing.context_helpers import lazy_context -from _algopy_testing.models.account import AccountFields +from _algopy_testing.models.account import AccountFields, AssetHolding from _algopy_testing.models.application import ApplicationContextData, ApplicationFields from _algopy_testing.models.asset import AssetFields from _algopy_testing.utils import generate_random_int @@ -102,7 +102,12 @@ def account( # update so defaults are preserved account_data.fields.update(account_fields) # can set these since it is a new account - account_data.opted_asset_balances = opted_asset_balances or {} + account_data.opted_assets = {} + for asset_id, balance in (opted_asset_balances or {}).items(): + asset = lazy_context.get_asset_data(asset_id) + account_data.opted_assets[asset_id] = AssetHolding( + balance=balance, frozen=asset["default_frozen"] + ) account_data.opted_apps = {app.id: app for app in opted_apps} return new_account diff --git a/tests/models/test_asset.py b/tests/models/test_asset.py index 2f8616c..b1eabfb 100644 --- a/tests/models/test_asset.py +++ b/tests/models/test_asset.py @@ -37,9 +37,7 @@ def test_asset_from_int() -> None: def test_asset_balance(context: AlgopyTestContext) -> None: account = context.any.account() asset = context.any.asset() - context.ledger.update_account( - account.public_key, opted_asset_balances={asset.id.value: UInt64(1000)} - ) + context.ledger.update_asset_holdings(account, asset, balance=1000) assert asset.balance(account) == UInt64(1000) @@ -52,15 +50,18 @@ def test_asset_balance_not_opted_in(context: AlgopyTestContext) -> None: asset.balance(account) -def test_asset_frozen() -> None: - asset = Asset(1) - account = Account() +@pytest.mark.parametrize( + "default_frozen", + [ + True, + False, + ], +) +def test_asset_frozen(context: AlgopyTestContext, *, default_frozen: bool) -> None: + asset = context.any.asset(default_frozen=default_frozen) + account = context.any.account(opted_asset_balances={asset.id: UInt64()}) - with pytest.raises( - NotImplementedError, - match="The 'frozen' method is being executed in a python testing context", - ): - asset.frozen(account) + assert asset.frozen(account) == default_frozen def test_asset_attributes(context: AlgopyTestContext) -> None: