Feat: enhance the integration test (#166)

* feat: create a extend private key use private key

* Fix: the fail test

* feat: add WALLET_ADDRESS to .env-example for configuration clarity

* fix: increase sleep duration in transaction mining test to ensure reliability

Author: bonnie57

.env-example

@@ -1,3 +1,4 @@
 WALLET_PRIVATE_KEY=YOUR-PRIVATE-WALLET-KEY
+WALLET_ADDRESS=YOUR-WALLET-ADDRESS
 RPC_PROVIDER_URL=YOUR_RPC_URL
 ETHERSCAN_API_KEY=YOUR-ETHERSCAN-API-KEY-HERE

tests/integration/config/test_config.py

@@ -3,13 +3,16 @@
 from dotenv import load_dotenv
 from web3 import Web3
 
+from tests.integration.config.utils import (
+    create_xprv_from_private_key,
+    get_private_key_from_xprv,
+)
+
 # Load environment variables
 load_dotenv(override=True)
 private_key = os.getenv("WALLET_PRIVATE_KEY")
-private_key_2 = os.getenv("WALLET_PRIVATE_KEY_2")
 rpc_url = os.getenv("RPC_PROVIDER_URL")
 wallet_address = os.getenv("WALLET_ADDRESS")
-wallet_address_2 = os.getenv("WALLET_ADDRESS_2")
 
 if not private_key:
     raise ValueError("WALLET_PRIVATE_KEY environment variable is not set")
@@ -23,8 +26,11 @@
 
 # Set up the account with the private key
 account = web3.eth.account.from_key(private_key)
+# Create the secondary account via extended private key
+xprv = create_xprv_from_private_key(private_key)
+private_key_2 = get_private_key_from_xprv(xprv)
 account_2 = web3.eth.account.from_key(private_key_2)
-
+wallet_address_2 = account_2.address
 # Export all configuration
 __all__ = [
     "web3",

tests/integration/config/utils.py

@@ -1,4 +1,5 @@
 import hashlib
+import hmac
 import os
 
 import base58
@@ -26,6 +27,44 @@
 CORE_METADATA_MODULE = "0x6E81a25C99C6e8430aeC7353325EB138aFE5DC16"
 
 
+def create_xprv_from_private_key(private_key: str) -> str:
+    """
+    Create an extended private key (xprv) using a private key as seed.
+
+    Args:
+        private_key: The private key (hex string with or without 0x prefix)
+
+    Returns:
+        xprv string in base58check format
+    """
+    seed = bytes.fromhex(private_key.removeprefix("0x"))
+    hmac_result = hmac.new(b"Bitcoin seed", seed, hashlib.sha512).digest()
+
+    xprv_bytes = (
+        bytes.fromhex("0488ADE4")
+        + bytes(9)
+        + hmac_result[32:]
+        + bytes([0])
+        + hmac_result[:32]
+    )
+    checksum = hashlib.sha256(hashlib.sha256(xprv_bytes).digest()).digest()[:4]
+    return base58.b58encode(xprv_bytes + checksum).decode()
+
+
+def get_private_key_from_xprv(xprv: str) -> str:
+    """
+    Extract the private key from an xprv.
+
+    Args:
+        xprv: Extended private key in base58check format
+
+    Returns:
+        Private key as hex string with 0x prefix
+    """
+    decoded = base58.b58decode(xprv)
+    return "0x" + decoded[46:78].hex()
+
+
 def get_story_client(web3: Web3, account) -> StoryClient:
     chain_id = 1315  # aeneid chain ID
     return StoryClient(web3, account, chain_id)

tests/integration/conftest.py

@@ -11,7 +11,7 @@
     ROYALTY_POLICY_LAP_ADDRESS,
     ZERO_ADDRESS,
 )
-from tests.integration.config.test_config import account, account_2, web3
+from tests.integration.config.test_config import account, account_2, private_key, web3
 from tests.integration.config.utils import MockERC20, approve, get_story_client
 from tests.integration.setup_for_integration import PIL_LICENSE_TEMPLATE
 
@@ -25,22 +25,25 @@ def story_client() -> StoryClient:
 @pytest.fixture(scope="session")
 def story_client_2() -> StoryClient:
     """Fixture to provide the secondary story client"""
-    return get_story_client(web3, account_2)
+    if not private_key:
+        raise ValueError("Private key is not set")
 
-
-@pytest.fixture(scope="module")
-def nft_collection(story_client: StoryClient):
-    """Fixture to provide the SPG NFT collection"""
-    tx_data = story_client.NFTClient.create_nft_collection(
-        name="test-collection",
-        symbol="TEST",
-        max_supply=100,
-        is_public_minting=True,
-        mint_open=True,
-        contract_uri="test-uri",
-        mint_fee_recipient=account.address,
-    )
-    return tx_data["nft_contract"]
+    story_client_2 = get_story_client(web3, account_2)
+    balance = story_client_2.get_wallet_balance()
+    if balance < web3.to_wei(5, "ether"):
+        tx = {
+            "from": account.address,
+            "to": account_2.address,
+            "value": web3.to_wei(5, "ether"),
+            "nonce": web3.eth.get_transaction_count(account.address),
+            "gas": 21000,
+            "gasPrice": web3.eth.gas_price,
+            "chainId": web3.eth.chain_id,
+        }
+        signed_tx = account.sign_transaction(tx)
+        tx_hash = web3.eth.send_raw_transaction(signed_tx.raw_transaction)
+        web3.eth.wait_for_transaction_receipt(tx_hash)
+    return story_client_2
 
 
 @pytest.fixture(scope="module")
@@ -124,3 +127,17 @@ def mint_and_approve_license_token(
         )
 
     return license_token_ids
+
+
+@pytest.fixture(scope="module")
+def nft_collection(story_client: StoryClient):
+    tx_data = story_client.NFTClient.create_nft_collection(
+        name="test-collection",
+        symbol="TEST",
+        max_supply=100,
+        is_public_minting=True,
+        mint_open=True,
+        contract_uri="test-uri",
+        mint_fee_recipient=account.address,
+    )
+    return tx_data["nft_contract"]

tests/integration/test_integration_transaction_utils.py

@@ -301,3 +301,5 @@ def build_tx(tx_options):
                 build_tx,
                 tx_options={"wait_for_receipt": True, "timeout": 0.001},
             )
+
+        time.sleep(5)  # Wait for the transaction to be mined

tests/unit/resources/test_ip_asset.py

@@ -2061,7 +2061,6 @@ def test_batch_mint_transaction_success(
                 ) as mock_build_multicall_transaction:
                     result = ip_asset.batch_mint_and_register_ip(requests=requests)
 
-                print(mock_build_multicall_transaction.call_args[0][0])
                 assert mock_build_multicall_transaction.call_args[0][0] == [
                     b"encoded_data",
                     b"encoded_data",