From 17d6839a7ef4f3e311f69c9d0a34bf9c0122f707 Mon Sep 17 00:00:00 2001
From: John Peterson <john.peterson@coinbase.com>
Date: Thu, 16 Jan 2025 18:42:21 -0500
Subject: [PATCH] chore: Fix Mnemonic Seed Wallet import Specify Network ID

---
 CHANGELOG.md         |  6 +++++
 Makefile             |  2 +-
 cdp/wallet.py        |  5 ++--
 tests/test_wallet.py | 62 ++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cb2d1b3..9d35c6d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # CDP Python SDK Changelog
 
+## Unreleased
+
+### Fixed
+
+- Allow wallet mnemonic seed import to optionally accept `network_id` input argument.
+
 ## [0.14.0] - 2025-01-14
 
 ### Added
diff --git a/Makefile b/Makefile
index 19af360..9bb4246 100644
--- a/Makefile
+++ b/Makefile
@@ -12,7 +12,7 @@ lint-fix:
 
 .PHONY: test
 test:
-	poetry run pytest
+	poetry run pytest -m "not e2e"
 
 .PHONY: e2e
 e2e:
diff --git a/cdp/wallet.py b/cdp/wallet.py
index 542f7e3..e2219a8 100644
--- a/cdp/wallet.py
+++ b/cdp/wallet.py
@@ -260,13 +260,14 @@ def list(cls) -> Iterator["Wallet"]:
             page = response.next_page
 
     @classmethod
-    def import_wallet(cls, data: WalletData | MnemonicSeedPhrase) -> "Wallet":
+    def import_wallet(cls, data: WalletData | MnemonicSeedPhrase, network_id: str = "base-sepolia") -> "Wallet":
         """Import a wallet from previously exported wallet data or a mnemonic seed phrase.
 
         Args:
             data (Union[WalletData, MnemonicSeedPhrase]): Either:
                 - WalletData: The wallet data to import, containing wallet_id and seed
                 - MnemonicSeedPhrase: A valid BIP-39 mnemonic phrase object for importing external wallets
+            network_id (str): The network ID of the wallet. Defaults to "base-sepolia".
 
         Returns:
             Wallet: The imported wallet.
@@ -291,7 +292,7 @@ def import_wallet(cls, data: WalletData | MnemonicSeedPhrase) -> "Wallet":
             seed = seed_bytes.hex()
 
             # Create wallet using the provided seed
-            wallet = cls.create_with_seed(seed=seed)
+            wallet = cls.create_with_seed(seed=seed, network_id=network_id)
             wallet._set_addresses()
             return wallet
 
diff --git a/tests/test_wallet.py b/tests/test_wallet.py
index 342edd5..333a197 100644
--- a/tests/test_wallet.py
+++ b/tests/test_wallet.py
@@ -742,6 +742,68 @@ def test_wallet_import_from_mnemonic_seed_phrase(
     assert wallet.default_address._model.public_key == public_key
 
 
+@patch("cdp.Cdp.use_server_signer", False)
+@patch("cdp.Cdp.api_clients")
+@patch("cdp.wallet.Account")
+def test_wallet_import_from_mnemonic_seed_phrase_specified_network_id(
+    mock_account,
+    mock_api_clients,
+    wallet_factory,
+    address_model_factory,
+):
+    """Test importing a wallet from a mnemonic seed phrase with network ID specified."""
+    # Valid 24-word mnemonic and expected address
+    valid_mnemonic = "crouch cereal notice one canyon kiss tape employ ghost column vanish despair eight razor laptop keen rally gaze riot regret assault jacket risk curve"
+    expected_address = "0x43A0477E658C6e05136e81C576CF02daCEa067bB"
+    public_key = "0x037e6cbdd1d949f60f41d5db7ffa9b3ddce0b77eab35ef7affd3f64cbfd9e33a91"
+
+    # Create mock address model
+    mock_address = address_model_factory(
+        address_id=expected_address,
+        public_key=public_key,
+        wallet_id="new-wallet-id",
+        network_id="base-mainnet",
+        index=0,
+    )
+
+    # Create mock wallet model with the address model
+    mock_wallet = wallet_factory(
+        id="new-wallet-id", network_id="base-mainnet", default_address=mock_address
+    )._model
+
+    # Add debug assertions
+    assert mock_wallet.default_address is not None
+    assert mock_wallet.default_address.address_id == expected_address
+
+    # Mock Account.from_key to return an account with our expected address
+    mock_account_instance = Mock(spec=Account)
+    mock_account_instance.address = expected_address
+    mock_account.from_key = Mock(return_value=mock_account_instance)
+
+    # Mock both API calls to return the same wallet model
+    mock_api_clients.wallets.create_wallet = Mock(return_value=mock_wallet)
+    mock_api_clients.wallets.get_wallet = Mock(return_value=mock_wallet)
+    mock_api_clients.addresses.create_address = Mock(return_value=mock_address)
+
+    # Mock list_addresses call
+    mock_address_list = Mock()
+    mock_address_list.data = [mock_address]
+    mock_api_clients.addresses.list_addresses = Mock(return_value=mock_address_list)
+
+    # Import wallet using mnemonic
+    from cdp.mnemonic_seed_phrase import MnemonicSeedPhrase
+
+    wallet = Wallet.import_wallet(MnemonicSeedPhrase(valid_mnemonic), network_id="base-mainnet")
+
+    # Verify the wallet was created successfully
+    assert isinstance(wallet, Wallet)
+
+    # Verify the default address matches expected address
+    assert wallet.default_address is not None
+    assert wallet.default_address.address_id == expected_address
+    assert wallet.default_address._model.public_key == public_key
+
+
 def test_wallet_import_from_mnemonic_empty_phrase():
     """Test importing a wallet with an empty mnemonic phrase."""
     from cdp.mnemonic_seed_phrase import MnemonicSeedPhrase