chore: added wait_for_mirro_node util to e2e test

Signed-off-by: Manish Dait <[email protected]>

Author: manishdait

tests/integration/account_id_population_e2e_test.py

@@ -12,7 +12,7 @@
     TransactionGetReceiptQuery,
 )
 from hiero_sdk_python.transaction.transfer_transaction import TransferTransaction
-from tests.integration.utils import env
+from tests.integration.utils import env, wait_for_mirror_node
 
 
 @pytest.fixture
@@ -45,24 +45,26 @@ def test_populate_account_id_num(env, evm_address):
     )
 
     assert transfer_receipt is not None
-    assert len(transfer_receipt.children) > 0
+    assert (
+        len(transfer_receipt.children) > 0
+    ), "Expected child transaction for auto-account creation"
 
-    account_id = transfer_receipt.children[0].account_id
-
-    assert account_id is not None
-
-    # Wait for mirrornode to update
-    time.sleep(5)
+    created_account_id = transfer_receipt.children[0].account_id
+    assert created_account_id is not None
 
     mirror_account_id = AccountId.from_evm_address(evm_address, 0, 0)
-    new_account_id = mirror_account_id.populate_account_num(env.client)
-
     assert mirror_account_id.num == 0
 
-    assert new_account_id.evm_address == mirror_account_id.evm_address
-    assert new_account_id.shard == account_id.shard
-    assert new_account_id.realm == account_id.realm
-    assert new_account_id.num == account_id.num
+    # Wait for mirrornode to update
+    resolved_account_id = wait_for_mirror_node(
+        fn=lambda: mirror_account_id.populate_account_num(env.client),
+        predicate=lambda acc: acc.num != 0,
+    )
+
+    assert resolved_account_id.evm_address == mirror_account_id.evm_address
+    assert resolved_account_id.shard == created_account_id.shard
+    assert resolved_account_id.realm == created_account_id.realm
+    assert resolved_account_id.num == created_account_id.num
 
 
 @pytest.mark.integration
@@ -86,17 +88,21 @@ def test_populate_account_id_evm_address(env, evm_address):
     )
 
     assert transfer_receipt is not None
-    assert len(transfer_receipt.children) > 0
+    assert (
+        len(transfer_receipt.children) > 0
+    ), "Expected child transaction for auto-account creation"
 
-    account_id = transfer_receipt.children[0].account_id
+    created_account_id = transfer_receipt.children[0].account_id
+    assert created_account_id is not None
+    assert created_account_id.evm_address is None
 
     # Wait for mirror_node to update
-    time.sleep(5)
-
-    new_account_id = account_id.populate_evm_address(env.client)
-    assert account_id.evm_address is None
+    resolved_account_id = wait_for_mirror_node(
+        fn=lambda: created_account_id.populate_evm_address(env.client),
+        predicate=lambda acc: acc.evm_address is not None,
+    )
 
-    assert new_account_id.evm_address == evm_address
-    assert new_account_id.shard == account_id.shard
-    assert new_account_id.realm == account_id.realm
-    assert new_account_id.num == account_id.num
+    assert resolved_account_id.shard == created_account_id.shard
+    assert resolved_account_id.realm == created_account_id.realm
+    assert resolved_account_id.num == created_account_id.num
+    assert resolved_account_id.evm_address == evm_address

tests/integration/contract_id_population_e2e_test.py

@@ -19,7 +19,7 @@
 from hiero_sdk_python.contract.contract_info_query import ContractInfoQuery
 from hiero_sdk_python.file.file_create_transaction import FileCreateTransaction
 from hiero_sdk_python.response_code import ResponseCode
-from tests.integration.utils import env
+from tests.integration.utils import env, wait_for_mirror_node
 
 
 @pytest.mark.integration
@@ -49,25 +49,22 @@ def test_populate_contract_id_num(env):
         .execute(env.client)
     )
     assert contract_receipt.status == ResponseCode.SUCCESS
-    contract_id_with_num = contract_receipt.contract_id
 
-    assert contract_id_with_num is not None
+    created_contract_id = contract_receipt.contract_id
+    assert created_contract_id is not None
 
     # Query contract info to get evm_address
-    info = ContractInfoQuery().set_contract_id(contract_id_with_num).execute(env.client)
-    contract_id_with_evm_addr = ContractId.from_evm_address(
-        0, 0, info.contract_account_id
-    )
+    info = ContractInfoQuery().set_contract_id(created_contract_id).execute(env.client)
+    contract_with_evm = ContractId.from_evm_address(0, 0, info.contract_account_id)
+    assert contract_with_evm.contract == 0
 
     # Wait for mirror_node to update
-    time.sleep(5)
-
-    final_contract_id = contract_id_with_evm_addr.populate_contract_num(env.client)
-
-    # should not update
-    assert contract_id_with_evm_addr.contract == 0
+    resolved_contract_id = wait_for_mirror_node(
+        fn=lambda: contract_with_evm.populate_contract_num(env.client),
+        predicate=lambda contract: contract.contract != 0,
+    )
 
-    assert final_contract_id.shard == contract_id_with_num.shard
-    assert final_contract_id.realm == contract_id_with_num.realm
-    assert final_contract_id.contract == contract_id_with_num.contract
-    assert final_contract_id.evm_address == contract_id_with_evm_addr.evm_address
+    assert resolved_contract_id.shard == created_contract_id.shard
+    assert resolved_contract_id.realm == created_contract_id.realm
+    assert resolved_contract_id.contract == created_contract_id.contract
+    assert resolved_contract_id.evm_address == contract_with_evm.evm_address

tests/integration/utils.py

@@ -1,8 +1,9 @@
 import os
+import time
 from pytest import fixture
 from dotenv import load_dotenv
 from dataclasses import dataclass
-from typing import Optional
+from typing import Callable, Optional, TypeVar
 from hiero_sdk_python.account.account_id import AccountId
 from hiero_sdk_python.client.client import Client
 from hiero_sdk_python.client.network import Network
@@ -11,43 +12,54 @@
 from hiero_sdk_python.logger.log_level import LogLevel
 from hiero_sdk_python.response_code import ResponseCode
 from hiero_sdk_python.tokens.supply_type import SupplyType
-from hiero_sdk_python.tokens.token_create_transaction import TokenCreateTransaction, TokenKeys, TokenParams
-from hiero_sdk_python.tokens.token_associate_transaction import TokenAssociateTransaction
+from hiero_sdk_python.tokens.token_create_transaction import (
+    TokenCreateTransaction,
+    TokenKeys,
+    TokenParams,
+)
+from hiero_sdk_python.tokens.token_associate_transaction import (
+    TokenAssociateTransaction,
+)
 from hiero_sdk_python.account.account_create_transaction import AccountCreateTransaction
 from hiero_sdk_python.transaction.transfer_transaction import TransferTransaction
-from hiero_sdk_python.hbar                    import Hbar
+from hiero_sdk_python.hbar import Hbar
+
+T = TypeVar("T")
 
 load_dotenv(override=True)
 
+
 @fixture
 def env():
     """Integration test environment with client/operator set up."""
     e = IntegrationTestEnv()
     yield e
     e.close()
 
+
 @dataclass
 class Account:
-    id:    AccountId
-    key:   PrivateKey
+    id: AccountId
+    key: PrivateKey
+
 
 class IntegrationTestEnv:
 
     def __init__(self) -> None:
-        network = Network(os.getenv('NETWORK'))
+        network = Network(os.getenv("NETWORK"))
         self.client = Client(network)
         self.operator_id: Optional[AccountId] = None
         self.operator_key: Optional[PrivateKey] = None
-        operator_id = os.getenv('OPERATOR_ID')
-        operator_key = os.getenv('OPERATOR_KEY')
+        operator_id = os.getenv("OPERATOR_ID")
+        operator_key = os.getenv("OPERATOR_KEY")
         if operator_id and operator_key:
             self.operator_id = AccountId.from_string(operator_id)
             self.operator_key = PrivateKey.from_string(operator_key)
             self.client.set_operator(self.operator_id, self.operator_key)
 
         self.client.logger.set_level(LogLevel.ERROR)
         self.public_operator_key = self.operator_key.public_key()
-        
+
     def close(self):
         self.client.close()
 
@@ -56,8 +68,8 @@ def create_account(self, initial_hbar: float = 1.0) -> Account:
         key = PrivateKey.generate()
         tx = (
             AccountCreateTransaction()
-                .set_key_without_alias(key.public_key())
-                .set_initial_balance(Hbar(initial_hbar))
+            .set_key_without_alias(key.public_key())
+            .set_initial_balance(Hbar(initial_hbar))
         )
         receipt = tx.execute(self.client)
         if receipt.status != ResponseCode.SUCCESS:
@@ -66,18 +78,20 @@ def create_account(self, initial_hbar: float = 1.0) -> Account:
             )
         return Account(id=receipt.account_id, key=key)
 
-    def associate_and_transfer(self, receiver: AccountId, receiver_key: PrivateKey, token_id, amount: int):
+    def associate_and_transfer(
+        self, receiver: AccountId, receiver_key: PrivateKey, token_id, amount: int
+    ):
         """
         Associate the token with `receiver`, then transfer `amount` of the token
         from the operator to that receiver.
         """
         assoc_receipt = (
             TokenAssociateTransaction()
-                .set_account_id(receiver)
-                .add_token_id(token_id)
-                .freeze_with(self.client)
-                .sign(receiver_key)
-                .execute(self.client)
+            .set_account_id(receiver)
+            .add_token_id(token_id)
+            .freeze_with(self.client)
+            .sign(receiver_key)
+            .execute(self.client)
         )
         if assoc_receipt.status != ResponseCode.SUCCESS:
             raise AssertionError(
@@ -86,15 +100,16 @@ def associate_and_transfer(self, receiver: AccountId, receiver_key: PrivateKey,
 
         transfer_receipt = (
             TransferTransaction()
-                .add_token_transfer(token_id, self.operator_id, -amount)
-                .add_token_transfer(token_id, receiver, amount)
-                .execute(self.client) # auto-signs with operator’s key
+            .add_token_transfer(token_id, self.operator_id, -amount)
+            .add_token_transfer(token_id, receiver, amount)
+            .execute(self.client)  # auto-signs with operator’s key
         )
         if transfer_receipt.status != ResponseCode.SUCCESS:
             raise AssertionError(
                 f"Transfer failed: {ResponseCode(transfer_receipt.status).name}"
             )
 
+
 def create_fungible_token(env, opts=[]):
     """
     Create a fungible token with the given options.
@@ -106,36 +121,39 @@ def create_fungible_token(env, opts=[]):
              lambda tx: tx.set_treasury_account_id(custom_treasury_id).freeze_with(client)
     """
     token_params = TokenParams(
-            token_name="PTokenTest34",
-            token_symbol="PTT34",
-            decimals=2,
-            initial_supply=1000,
-            treasury_account_id=env.operator_id,
-            token_type=TokenType.FUNGIBLE_COMMON,
-            supply_type=SupplyType.FINITE,
-            max_supply=10000
-        )
-    
+        token_name="PTokenTest34",
+        token_symbol="PTT34",
+        decimals=2,
+        initial_supply=1000,
+        treasury_account_id=env.operator_id,
+        token_type=TokenType.FUNGIBLE_COMMON,
+        supply_type=SupplyType.FINITE,
+        max_supply=10000,
+    )
+
     token_keys = TokenKeys(
-            admin_key=env.operator_key,
-            supply_key=env.operator_key,
-            freeze_key=env.operator_key,
-            wipe_key=env.operator_key
-            # pause_key=  None  # implicitly “no pause key” use opts to add one
-        )
-        
+        admin_key=env.operator_key,
+        supply_key=env.operator_key,
+        freeze_key=env.operator_key,
+        wipe_key=env.operator_key,
+        # pause_key=  None  # implicitly “no pause key” use opts to add one
+    )
+
     token_transaction = TokenCreateTransaction(token_params, token_keys)
-    
+
     # Apply optional functions to the token creation transaction
     for opt in opts:
         opt(token_transaction)
-    
+
     token_receipt = token_transaction.execute(env.client)
-    
-    assert token_receipt.status == ResponseCode.SUCCESS, f"Token creation failed with status: {ResponseCode(token_receipt.status).name}"
-    
+
+    assert (
+        token_receipt.status == ResponseCode.SUCCESS
+    ), f"Token creation failed with status: {ResponseCode(token_receipt.status).name}"
+
     return token_receipt.token_id
 
+
 def create_nft_token(env, opts=[]):
     """
     Create a non-fungible token (NFT) with the given options.
@@ -154,15 +172,14 @@ def create_nft_token(env, opts=[]):
         treasury_account_id=env.operator_id,
         token_type=TokenType.NON_FUNGIBLE_UNIQUE,
         supply_type=SupplyType.FINITE,
-        max_supply=10000  
+        max_supply=10000,
     )
-    
+
     token_keys = TokenKeys(
         admin_key=env.operator_key,
         supply_key=env.operator_key,
-        freeze_key=env.operator_key
+        freeze_key=env.operator_key,
         # pause_key=  None  # implicitly “no pause key” use opts to add one
-
     )
 
     transaction = TokenCreateTransaction(token_params, token_keys)
@@ -172,7 +189,51 @@ def create_nft_token(env, opts=[]):
         opt(transaction)
 
     token_receipt = transaction.execute(env.client)
-    
-    assert token_receipt.status == ResponseCode.SUCCESS, f"Token creation failed with status: {ResponseCode(token_receipt.status).name}"
-    
-    return token_receipt.token_id
+
+    assert (
+        token_receipt.status == ResponseCode.SUCCESS
+    ), f"Token creation failed with status: {ResponseCode(token_receipt.status).name}"
+
+    return token_receipt.token_id
+
+
+def wait_for_mirror_node(
+    fn: Callable[[], T],
+    predicate: Callable[[T], bool],
+    timeout: float = 5,
+    interval: float = 1,
+) -> T:
+    """
+    Polls fn until predicate(result) returns True or timeout is reached
+
+    Args:
+        fn: Function that fetches data from mirror node.
+        predicate: Condition that determines success.
+        timeout: Max time to wait (seconds).
+        interval: Sleep interval between retries (seconds).
+
+    Returns:
+        T: The successful result.
+    """
+    deadline = time.monotonic() + timeout
+    last_response = None
+    last_exception = None
+
+    while time.monotonic() < deadline:
+        try:
+            last_response = fn()
+            if predicate(last_response):
+                return last_response
+        except Exception as e:
+            last_exception = e
+
+        time.sleep(interval)
+
+    if last_exception is not None:
+        raise TimeoutError(
+            "Timed out waiting for mirror node, Last call raised an exception"
+        ) from last_exception
+
+    raise TimeoutError(
+        f"Timed out waiting for mirror node. Last response: {last_response}"
+    )