feat: add support for frozen on asset holdings, including a new ledge…

…r function `update_asset_holdings` for setting asset holding balances and frozen states

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.

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
 

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:

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}")
 

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
 

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: