[CHIA-3707] Port create_new_wallet to @marshal

chia/_tests/cmds/cmd_test_utils.py

@@ -13,7 +13,7 @@
 
 import chia.cmds.wallet_funcs
 from chia._tests.cmds.testing_classes import create_test_block_record
-from chia._tests.cmds.wallet.test_consts import get_bytes32
+from chia._tests.cmds.wallet.test_consts import STD_TX, STD_UTX, get_bytes32
 from chia.cmds.chia import cli as chia_cli
 from chia.cmds.cmds_util import _T_RpcClient, node_config_section_names
 from chia.consensus.default_constants import DEFAULT_CONSTANTS
@@ -25,17 +25,21 @@
 from chia.types.signing_mode import SigningMode
 from chia.util.bech32m import encode_puzzle_hash
 from chia.util.config import load_config
-from chia.wallet.conditions import ConditionValidTimes
+from chia.wallet.conditions import Condition, ConditionValidTimes
 from chia.wallet.nft_wallet.nft_info import NFTInfo
 from chia.wallet.nft_wallet.nft_wallet import NFTWallet
 from chia.wallet.transaction_record import TransactionRecord
 from chia.wallet.util.transaction_type import TransactionType
+from chia.wallet.util.tx_config import TXConfig
 from chia.wallet.util.wallet_types import WalletType
 from chia.wallet.wallet_request_types import (
     CATAssetIDToName,
     CATAssetIDToNameResponse,
     CATGetName,
     CATGetNameResponse,
+    CreateNewWallet,
+    CreateNewWalletResponse,
+    CreateNewWalletType,
     GetSyncStatusResponse,
     GetTransaction,
     GetTransactionResponse,
@@ -233,6 +237,24 @@ async def nft_calculate_royalties(
             )
         )
 
+    async def create_new_wallet(
+        self,
+        request: CreateNewWallet,
+        tx_config: TXConfig,
+        extra_conditions: tuple[Condition, ...] = tuple(),
+        timelock_info: ConditionValidTimes = ConditionValidTimes(),
+    ) -> CreateNewWalletResponse:
+        self.add_to_log("create_new_wallet", (request, tx_config, extra_conditions, timelock_info))
+        return CreateNewWalletResponse(
+            [STD_UTX],
+            [STD_TX],
+            type=(
+                WalletType.NFT if request.wallet_type == CreateNewWalletType.NFT_WALLET else WalletType.DECENTRALIZED_ID
+            ).name,
+            wallet_id=uint32(4 if request.wallet_type == CreateNewWalletType.NFT_WALLET else 3),
+            my_did="did:chia:1qgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpq4msw0c",
+        )
+
 
 @dataclass
 class TestFullNodeRpcClient(TestRpcClient):

chia/_tests/cmds/wallet/test_did.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Optional, Union
 
 import pytest
 from chia_rs import G2Element
@@ -18,6 +17,8 @@
 from chia.wallet.util.curry_and_treehash import NIL_TREEHASH
 from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG, TXConfig
 from chia.wallet.wallet_request_types import (
+    CreateNewWallet,
+    CreateNewWalletType,
     DIDFindLostDID,
     DIDFindLostDIDResponse,
     DIDGetDID,
@@ -30,6 +31,7 @@
     DIDSetWalletNameResponse,
     DIDTransferDID,
     DIDTransferDIDResponse,
+    DIDType,
     DIDUpdateMetadata,
     DIDUpdateMetadataResponse,
 )
@@ -55,27 +57,7 @@ def test_did_recovery_is_nil(program: Program, result: bool) -> None:
 def test_did_create(capsys: object, get_test_cli_clients: tuple[TestRpcClients, Path]) -> None:
     test_rpc_clients, root_dir = get_test_cli_clients
 
-    # set RPC Client
-    class DidCreateRpcClient(TestWalletRpcClient):
-        async def create_new_did_wallet(
-            self,
-            amount: int,
-            tx_config: TXConfig,
-            fee: int = 0,
-            name: Optional[str] = "DID Wallet",
-            backup_ids: Optional[list[str]] = None,
-            required_num: int = 0,
-            push: bool = True,
-            timelock_info: ConditionValidTimes = ConditionValidTimes(),
-        ) -> dict[str, Union[str, int]]:
-            if backup_ids is None:
-                backup_ids = []
-            self.add_to_log(
-                "create_new_did_wallet", (amount, tx_config, fee, name, backup_ids, required_num, push, timelock_info)
-            )
-            return {"wallet_id": 3, "my_did": "did:chia:testdid123456"}
-
-    inst_rpc_client = DidCreateRpcClient()
+    inst_rpc_client = TestWalletRpcClient()
     test_rpc_clients.wallet_rpc_client = inst_rpc_client
     command_args = [
         "wallet",
@@ -93,12 +75,27 @@ async def create_new_did_wallet(
     # these are various things that should be in the output
     assert_list = [
         "Successfully created a DID wallet with name test and id 3 on key 123456",
-        "Successfully created a DID did:chia:testdid123456 in the newly created DID wallet",
+        (
+            "Successfully created a DID did:chia:1qgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpq4msw0c"
+            " in the newly created DID wallet"
+        ),
     ]
     run_cli_command_and_assert(capsys, root_dir, command_args, assert_list)
     expected_calls: logType = {
-        "create_new_did_wallet": [
-            (3, DEFAULT_TX_CONFIG, 100000000000, "test", [], 0, True, test_condition_valid_times)
+        "create_new_wallet": [
+            (
+                CreateNewWallet(
+                    wallet_type=CreateNewWalletType.DID_WALLET,
+                    did_type=DIDType.NEW,
+                    amount=uint64(3),
+                    wallet_name="test",
+                    fee=uint64(100_000_000_000),
+                    push=True,
+                ),
+                DEFAULT_TX_CONFIG,
+                tuple(),
+                test_condition_valid_times,
+            )
         ],
     }
     test_rpc_clients.wallet_rpc_client.check_log(expected_calls)

chia/_tests/cmds/wallet/test_nft.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, Optional
 
 from chia_rs import G2Element
 from chia_rs.sized_bytes import bytes32
@@ -15,6 +14,8 @@
 from chia.wallet.nft_wallet.nft_info import NFTInfo
 from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG, TXConfig
 from chia.wallet.wallet_request_types import (
+    CreateNewWallet,
+    CreateNewWalletType,
     NFTAddURI,
     NFTAddURIResponse,
     NFTGetNFTs,
@@ -38,21 +39,22 @@
 def test_nft_create(capsys: object, get_test_cli_clients: tuple[TestRpcClients, Path]) -> None:
     test_rpc_clients, root_dir = get_test_cli_clients
 
-    # set RPC Client
-    class NFTCreateRpcClient(TestWalletRpcClient):
-        async def create_new_nft_wallet(self, did_id: str, name: Optional[str] = None) -> dict[str, Any]:
-            self.add_to_log("create_new_nft_wallet", (did_id, name))
-            return {"wallet_id": 4}
-
-    inst_rpc_client = NFTCreateRpcClient()
-    did_id = encode_puzzle_hash(get_bytes32(2), "did:chia:")
+    inst_rpc_client = TestWalletRpcClient()
+    did_id = "did:chia:1qgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpq4msw0c"
     test_rpc_clients.wallet_rpc_client = inst_rpc_client
     command_args = ["wallet", "nft", "create", FINGERPRINT_ARG, "-ntest", "--did-id", did_id]
     # these are various things that should be in the output
     assert_list = [f"Successfully created an NFT wallet with id 4 on key {FINGERPRINT}"]
     run_cli_command_and_assert(capsys, root_dir, command_args, assert_list)
     expected_calls: logType = {
-        "create_new_nft_wallet": [(did_id, "test")],
+        "create_new_wallet": [
+            (
+                CreateNewWallet(wallet_type=CreateNewWalletType.NFT_WALLET, name="test", did_id=did_id, push=True),
+                DEFAULT_TX_CONFIG,
+                tuple(),
+                ConditionValidTimes(),
+            )
+        ],
     }
     test_rpc_clients.wallet_rpc_client.check_log(expected_calls)
 

chia/_tests/cmds/wallet/test_wallet.py

@@ -51,6 +51,8 @@
     CATSpend,
     CATSpendResponse,
     ClawbackPuzzleDecoratorOverride,
+    CreateNewWallet,
+    CreateNewWalletType,
     CreateOfferForIDs,
     CreateOfferForIDsResponse,
     DeleteUnconfirmedTransactions,
@@ -83,6 +85,7 @@
     TakeOffer,
     TakeOfferResponse,
     TransactionRecordWithMetadata,
+    WalletCreationMode,
     WalletInfoResponse,
 )
 from chia.wallet.wallet_spend_bundle import WalletSpendBundle
@@ -707,10 +710,6 @@ def test_add_token(capsys: object, get_test_cli_clients: tuple[TestRpcClients, P
 
     # set RPC Client
     class AddTokenRpcClient(TestWalletRpcClient):
-        async def create_wallet_for_existing_cat(self, asset_id: bytes) -> dict[str, int]:
-            self.add_to_log("create_wallet_for_existing_cat", (asset_id,))
-            return {"wallet_id": 3}
-
         async def set_cat_name(self, request: CATSetName) -> CATSetNameResponse:
             self.add_to_log("set_cat_name", (request.wallet_id, request.name))
             return CATSetNameResponse(wallet_id=request.wallet_id)
@@ -728,7 +727,19 @@ async def set_cat_name(self, request: CATSetName) -> CATSetNameResponse:
 
     expected_calls: logType = {
         "cat_asset_id_to_name": [(get_bytes32(1),), (bytes32([1, 2] * 16),)],
-        "create_wallet_for_existing_cat": [(bytes32([1, 2] * 16),)],
+        "create_new_wallet": [
+            (
+                CreateNewWallet(
+                    wallet_type=CreateNewWalletType.CAT_WALLET,
+                    mode=WalletCreationMode.EXISTING,
+                    asset_id=bytes32([1, 2] * 16),
+                    push=True,
+                ),
+                DEFAULT_TX_CONFIG,
+                tuple(),
+                ConditionValidTimes(),
+            )
+        ],
         "set_cat_name": [(2, "examplecat"), (3, "examplecat")],
     }
     test_rpc_clients.wallet_rpc_client.check_log(expected_calls)

chia/_tests/pools/test_pool_cmdline.py

@@ -992,9 +992,7 @@ async def test_plotnft_cli_misc(mocker: MockerFixture, consensus_mode: Consensus
         )
 
     # Test fall-through raise in create
-    mocker.patch.object(
-        test_rpc_client, "create_new_pool_wallet", create=True, side_effect=ValueError("Injected error")
-    )
+    mocker.patch.object(test_rpc_client, "create_new_wallet", create=True, side_effect=ValueError("Injected error"))
     with pytest.raises(CliRpcConnectionError, match="Error creating plot NFT: Injected error"):
         await create(
             wallet_info=WalletClientInfo(client=cast(WalletRpcClient, test_rpc_client), fingerprint=0, config=dict()),