chore(PSDK-640): add network_id to WalletData

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add `network_id` to `WalletData` so that it is saved with the seed data and surfaced via the export function
+
 ### [0.12.1] - 2024-12-10
 
 ### Added

README.md

@@ -189,7 +189,7 @@ list(address.trades())
 The SDK creates wallets with [Developer-Managed (1-1)](https://docs.cdp.coinbase.com/mpc-wallet/docs/wallets#developer-managed-wallets) keys by default, which means you are responsible for securely storing the keys required to re-instantiate wallets. The below code walks you through how to export a wallet and store it in a secure location.
 
 ```python
-# Export the data required to re-instantiate the wallet. The data contains the seed and the ID of the wallet.
+# Export the data required to re-instantiate the wallet. The data contains the seed, the ID of the wallet, and the network ID.
 data = wallet.export_data()
 ```
 

cdp/wallet.py

@@ -492,7 +492,7 @@ def export_data(self) -> WalletData:
         if self._master is None or self._seed is None:
             raise ValueError("Wallet does not have seed loaded")
 
-        return WalletData(self.id, self._seed)
+        return WalletData(self.id, self._seed, self.network_id)
 
     def save_seed(self, file_path: str, encrypt: bool | None = False) -> None:
         """Save the wallet seed to a file.
@@ -530,6 +530,7 @@ def save_seed(self, file_path: str, encrypt: bool | None = False) -> None:
             "encrypted": encrypt,
             "auth_tag": auth_tag,
             "iv": iv,
+            "network_id": self.network_id,
         }
 
         with open(file_path, "w") as f:

cdp/wallet_data.py

@@ -1,16 +1,18 @@
 class WalletData:
     """A class representing wallet data required to recreate a wallet."""
 
-    def __init__(self, wallet_id: str, seed: str) -> None:
+    def __init__(self, wallet_id: str, seed: str, network_id: str | None = None) -> None:
         """Initialize the WalletData class.
 
         Args:
             wallet_id (str): The ID of the wallet.
             seed (str): The seed of the wallet.
+            network_id (str | None): The network ID of the wallet. Defaults to None.
 
         """
         self._wallet_id = wallet_id
         self._seed = seed
+        self._network_id = network_id
 
     @property
     def wallet_id(self) -> str:
@@ -32,14 +34,27 @@ def seed(self) -> str:
         """
         return self._seed
 
+    @property
+    def network_id(self) -> str | None:
+        """Get the network ID of the wallet.
+
+        Returns:
+            str: The network ID of the wallet.
+
+        """
+        return self._network_id
+
     def to_dict(self) -> dict[str, str]:
         """Convert the wallet data to a dictionary.
 
         Returns:
             dict[str, str]: The dictionary representation of the wallet data.
 
         """
-        return {"wallet_id": self.wallet_id, "seed": self.seed}
+        result = {"wallet_id": self.wallet_id, "seed": self.seed}
+        if self._network_id is not None:
+            result["network_id"] = self.network_id
+        return result
 
     @classmethod
     def from_dict(cls, data: dict[str, str]) -> "WalletData":
@@ -52,7 +67,11 @@ def from_dict(cls, data: dict[str, str]) -> "WalletData":
             WalletData: The wallet data.
 
         """
-        return cls(data["wallet_id"], data["seed"])
+        return cls(
+            data["wallet_id"],
+            data["seed"],
+            data.get("network_id")
+        )
 
     def __str__(self) -> str:
         """Return a string representation of the WalletData object.
@@ -61,7 +80,7 @@ def __str__(self) -> str:
             str: A string representation of the wallet data.
 
         """
-        return f"WalletData: (wallet_id: {self.wallet_id}, seed: {self.seed})"
+        return f"WalletData: (wallet_id: {self.wallet_id}, seed: {self.seed}, network_id: {self.network_id})"
 
     def __repr__(self) -> str:
         """Return a string representation of the WalletData object.

tests/test_wallet.py

@@ -658,3 +658,22 @@ def test_wallet_quote_fund_no_default_address(wallet_factory):
 
         with pytest.raises(ValueError, match="Default address does not exist"):
             wallet.quote_fund(amount="1.0", asset_id="eth")
+
+@patch("cdp.Cdp.use_server_signer", False)
+@patch("cdp.wallet.os")
+@patch("cdp.wallet.Bip32Slip10Secp256k1")
+def test_wallet_export_data(mock_bip32, mock_os, wallet_factory, master_key_factory):
+    """Test Wallet export_data method."""
+    seed = b"\x00" * 64
+    mock_urandom = Mock(return_value=seed)
+    mock_os.urandom = mock_urandom
+    mock_from_seed = Mock(return_value=master_key_factory(seed))
+    mock_bip32.FromSeed = mock_from_seed
+
+    wallet = wallet_factory()
+
+    exported = wallet.export_data()
+
+    assert exported.wallet_id == wallet.id
+    assert exported.seed == seed.hex()
+    assert exported.network_id == wallet.network_id