refactor: expand accepted types for account value generator

Author: daniel-makerx

docs/testing-guide/avm-types.md

@@ -103,7 +103,7 @@ asset = context.ledger.get_asset(asset_id=random_asset.id)
 
 # Update an asset
 context.ledger.update_asset(
-    asset=random_asset,
+    random_asset,
     name=...,  # Optional: New asset name
     total=...,  # Optional: New total supply
     decimals=...,  # Optional: Number of decimals
@@ -130,7 +130,7 @@ account = algopy.Account(raw_address) # zero address by default
 # Generate a random account
 random_account = context.any.account(
     address=str(raw_address),  # Optional: Specify a custom address, defaults to a random address
-    opted_asset_balances={},  # Optional: Specify opted asset balances as dict of algopy.UInt64 as key and algopy.UInt64 as value
+    opted_asset_balances={},  # Optional: Specify opted asset balances as dict of assets to balance
     opted_apps=[],  # Optional: Specify opted apps as sequence of algopy.Application objects
     balance=...,  # Optional: Specify an initial balance
     min_balance=...,  # Optional: Specify a minimum balance
@@ -145,15 +145,15 @@ random_account = context.any.account(
 # Generate a random account that is opted into a specific asset
 mock_asset = context.any.asset()
 mock_account = context.any.account(
-    opted_asset_balances={mock_asset.id: algopy.UInt64(123)}
+    opted_asset_balances={mock_asset: 123}
 )
 
 # Get an account by address
 account = context.ledger.get_account(str(mock_account))
 
 # Update an account
 context.ledger.update_account(
-    str(mock_account),
+    mock_account,
     balance=...,  # Optional: New balance
     min_balance=...,  # Optional: New minimum balance
     auth_address=context.any.account(),  # Optional: New auth address
@@ -196,7 +196,7 @@ app = context.ledger.get_app(app_id=random_app.id)
 
 # Update an application
 context.ledger.update_app(
-    app=random_app,
+    random_app,
     approval_program=...,  # Optional: New approval program
     clear_state_program=...,  # Optional: New clear state program
     global_num_uint=...,  # Optional: New number of global uint values

src/_algopy_testing/context_helpers/ledger_context.py

@@ -107,10 +107,11 @@ def update_asset_holdings(
 
         address = _get_address(account)
         account_data = self._account_data[address]
-        asset = self.get_asset(_get_asset_id(asset))
+        asset_id = _get_asset_id(asset)
+        asset = self.get_asset(asset_id)
 
         holdings = account_data.opted_assets.setdefault(
-            asset.id,
+            asset_id,
             AssetHolding(balance=UInt64(), frozen=asset.default_frozen),
         )
         if balance is not None:

src/_algopy_testing/models/account.py

@@ -33,22 +33,21 @@ class AccountFields(typing.TypedDict, total=False):
 
 
 def get_empty_account() -> AccountContextData:
-    zero = UInt64()
     return AccountContextData(
         fields={
-            "balance": zero,
+            "balance": UInt64(),
             "min_balance": UInt64(DEFAULT_ACCOUNT_MIN_BALANCE),
             "auth_address": Account(),
-            "total_num_uint": zero,
-            "total_num_byte_slice": zero,
-            "total_extra_app_pages": zero,
-            "total_apps_created": zero,
-            "total_apps_opted_in": zero,
-            "total_assets_created": zero,
-            "total_assets": zero,
-            "total_boxes": zero,
-            "total_box_bytes": zero,
-        }
+            "total_num_uint": UInt64(),
+            "total_num_byte_slice": UInt64(),
+            "total_extra_app_pages": UInt64(),
+            "total_apps_created": UInt64(),
+            "total_apps_opted_in": UInt64(),
+            "total_assets_created": UInt64(),
+            "total_assets": UInt64(),
+            "total_boxes": UInt64(),
+            "total_box_bytes": UInt64(),
+        },
     )
 
 
@@ -64,12 +63,12 @@ class AccountContextData:
 
     Attributes:
         opted_assets (dict[int, AssetHolding]): Mapping of asset IDs to holdings.
-        opted_apps (dict[int, algopy.UInt64]): Mapping of application IDs to instances.
+        opted_apps (dict[int, algopy.Application]): Mapping of application IDs to instances.
         fields (AccountFields): Additional account fields.
     """
 
-    opted_assets: dict[algopy.UInt64, AssetHolding] = dataclasses.field(default_factory=dict)
-    opted_apps: dict[algopy.UInt64, algopy.Application] = dataclasses.field(default_factory=dict)
+    opted_assets: dict[int, AssetHolding] = dataclasses.field(default_factory=dict)
+    opted_apps: dict[int, algopy.Application] = dataclasses.field(default_factory=dict)
     fields: AccountFields = dataclasses.field(default_factory=AccountFields)  # type: ignore[arg-type]
 
 

src/_algopy_testing/models/asset.py

@@ -45,14 +45,14 @@ def balance(self, account: algopy.Account) -> algopy.UInt64:
 
         account_data = lazy_context.get_account_data(account.public_key)
 
-        if int(self.id) not in account_data.opted_assets:
+        if self.int_ 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].balance
+        return account_data.opted_assets[self.int_].balance
 
     def frozen(self, account: algopy.Account) -> bool:
         from _algopy_testing.context_helpers import lazy_context
@@ -64,7 +64,7 @@ def frozen(self, account: algopy.Account) -> bool:
                 "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
+        return account_data.opted_assets[self.int_].frozen
 
     @property
     def fields(self) -> AssetFields:

src/_algopy_testing/op/misc.py

@@ -171,8 +171,9 @@ def _get_asset_holding(
             return UInt64(0), False
 
         account_data = lazy_context.get_account_data(account.public_key)
-        holding = account_data.opted_assets.get(asset.id)
-        if holding is None:
+        try:
+            holding = account_data.opted_assets[asset.id.value]
+        except KeyError:
             return UInt64(0), False
 
         if field == "balance":

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, AssetHolding
+from _algopy_testing.models.account import AccountFields
 from _algopy_testing.models.application import ApplicationContextData, ApplicationFields
 from _algopy_testing.models.asset import AssetFields
 from _algopy_testing.utils import generate_random_int
@@ -70,10 +70,13 @@ def string(self, length: int = MAX_BYTES_SIZE) -> algopy.String:
     def account(
         self,
         address: str | None = None,
-        opted_asset_balances: dict[algopy.UInt64, algopy.UInt64] | None = None,
-        opted_apps: typing.Sequence[algopy.Application] = (),
+        opted_asset_balances: (
+            dict[algopy.Asset | algopy.UInt64 | int, algopy.UInt64 | int] | None
+        ) = None,
+        opted_apps: typing.Sequence[algopy.Application | algopy.UInt64 | int] = (),
         **account_fields: typing.Unpack[AccountFields],
     ) -> algopy.Account:
+        """Initialize a new account with specified fields and balances."""
         import algopy
 
         if address is not None and address in lazy_context.ledger._account_data:
@@ -86,26 +89,26 @@ def account(
             if key not in AccountFields.__annotations__:
                 raise AttributeError(f"Invalid field '{key}' for Account")
 
+        ledger = lazy_context.ledger
         new_account_address = address or algosdk.account.generate_account()[1]
         new_account = algopy.Account(new_account_address)
         # defaultdict of account_data ensures we get a new initialized account
         account_data = lazy_context.get_account_data(new_account_address)
         # update so defaults are preserved
         account_data.fields.update(account_fields)
-        # can set these since it is a new account
-        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"]
+            ledger.update_asset_holdings(
+                new_account_address,
+                asset_id,
+                balance=balance,
             )
-        account_data.opted_apps = {app.id: app for app in opted_apps}
+        account_data.opted_apps = {_get_app_id(app): ledger.get_app(app) for app in opted_apps}
         return new_account
 
     def asset(
         self, asset_id: int | None = None, **asset_fields: typing.Unpack[AssetFields]
     ) -> algopy.Asset:
-        r"""Generate and add a new asset with a unique ID."""
+        """Generate and add a new asset with a unique ID."""
         import algopy
 
         if asset_id and asset_id in lazy_context.ledger._asset_data:
@@ -193,3 +196,10 @@ def bytes(self, length: int | None = None) -> algopy.Bytes:
         """
         length = length or MAX_BYTES_SIZE
         return _algopy_testing.Bytes(secrets.token_bytes(length))
+
+
+def _get_app_id(app: algopy.Application | algopy.UInt64 | int) -> int:
+    from _algopy_testing.models import Application
+
+    app_id = app.id if isinstance(app, Application) else app
+    return int(app_id)