feat: Add for_network() method to configure a custom network (#2182)

Signed-off-by: Manish Dait <daitmanish88@gmail.com>

Author: manishdait

src/hiero_sdk_python/client/client.py

@@ -12,13 +12,13 @@
 from dotenv import load_dotenv
 
 from hiero_sdk_python.account.account_id import AccountId
-from hiero_sdk_python.channels import _UserAgentInterceptor
 from hiero_sdk_python.crypto.private_key import PrivateKey
 from hiero_sdk_python.hapi.mirror import (
     consensus_service_pb2_grpc as mirror_consensus_grpc,
 )
 from hiero_sdk_python.hbar import Hbar
 from hiero_sdk_python.logger.logger import Logger, LogLevel
+from hiero_sdk_python.node import _Node
 from hiero_sdk_python.transaction.transaction_id import TransactionId
 
 from .network import Network
@@ -57,9 +57,6 @@ def __init__(self, network: Network = None) -> None:
 
         self.network: Network = network
 
-        self.mirror_channel: grpc.Channel = None
-        self.mirror_stub: mirror_consensus_grpc.ConsensusServiceStub = None
-
         self.max_attempts: int = 10
         self.default_max_query_payment: Hbar = DEFAULT_MAX_QUERY_PAYMENT
 
@@ -69,10 +66,17 @@ def __init__(self, network: Network = None) -> None:
         self._grpc_deadline: float = DEFAULT_GRPC_DEADLINE
         self._request_timeout: float = DEFAULT_REQUEST_TIMEOUT
 
-        self._init_mirror_stub()
-
         self.logger: Logger = Logger(LogLevel.from_env(), "hiero_sdk_python")
 
+    @property
+    def mirror_stub(self) -> mirror_consensus_grpc.ConsensusServiceStub:
+        return self.network.get_mirror_stub()
+
+    @property
+    def mirror_channel(self) -> grpc.Channel:
+        self.network.get_mirror_stub()
+        return self.network._mirror_channel
+
     @classmethod
     def from_env(cls, network: NetworkName | None = None) -> Client:
         """
@@ -153,20 +157,32 @@ def for_previewnet(cls) -> Client:
         """
         return cls(Network("previewnet"))
 
-    def _init_mirror_stub(self) -> None:
+    @classmethod
+    def for_network(cls, network_map: dict[str, AccountId], network_name: str | None = "localhost") -> Client:
         """
-        Connect to a mirror node for topic message subscriptions.
-        Mirror nodes always use TLS (mandatory). We use self.network.get_mirror_address()
-        for a configurable mirror address, which should use port 443 for HTTPS connections.
+        Create a Client with a custom set of nodes and an optional network label.
+
+        Args:
+            network_map (dict[str, AccountId]): A map where keys are "host:port" strings
+                and values are the node's AccountId.
+            network_name (str): A label for the network. Defaults to "localhost".
+
+        Returns:
+            Client: A Client instance configured with the custom network.
         """
-        mirror_address = self.network.get_mirror_address()
-        if mirror_address.endswith(":50212") or mirror_address.endswith(":443"):
-            self.mirror_channel = grpc.secure_channel(mirror_address, grpc.ssl_channel_credentials())
-        else:
-            self.mirror_channel = grpc.insecure_channel(mirror_address)
+        if not network_map:
+            raise ValueError("network_map cannot be empty")
+
+        first_node = next(iter(network_map.values()))
+        shard = first_node.shard
+        realm = first_node.realm
 
-        self.mirror_channel = grpc.intercept_channel(self.mirror_channel, _UserAgentInterceptor())
-        self.mirror_stub = mirror_consensus_grpc.ConsensusServiceStub(self.mirror_channel)
+        for account_id in network_map.values():
+            if shard != account_id.shard or realm != account_id.realm:
+                raise ValueError("network is not valid, all nodes must be in the same shard and realm")
+
+        nodes = [_Node(account_id, address, None) for address, account_id in network_map.items()]
+        return cls(Network(network=network_name, nodes=nodes))
 
     def set_operator(self, account_id: AccountId, private_key: PrivateKey) -> None:
         """Sets the operator credentials (account ID and private key)."""
@@ -203,17 +219,7 @@ def close(self) -> None:
         Closes any open gRPC channels and frees resources.
         Call this when you are done using the Client to ensure a clean shutdown.
         """
-        # Close mirror channel
-        if self.mirror_channel is not None:
-            self.mirror_channel.close()
-            self.mirror_channel = None
-
-        self.mirror_stub = None
-
-        # Fix: Close all consensus node channels
-        if self.network and self.network.nodes:
-            for node in self.network.nodes:
-                node._close()
+        self.network._close()
 
     def set_transport_security(self, enabled: bool) -> Client:
         """

src/hiero_sdk_python/client/network.py

@@ -6,10 +6,12 @@
 import time
 from typing import Any
 
+import grpc
 import requests
 
 from hiero_sdk_python.account.account_id import AccountId
 from hiero_sdk_python.address_book.node_address import NodeAddress
+from hiero_sdk_python.hapi.mirror import consensus_service_pb2_grpc as mirror_consensus_grpc
 from hiero_sdk_python.node import _Node
 
 
@@ -73,7 +75,7 @@ class Network:
 
     def __init__(
         self,
-        network: str = "testnet",
+        network: str | None = None,
         nodes: list[_Node] | None = None,
         mirror_address: str | None = None,
         ledger_id: bytes | None = None,
@@ -97,9 +99,11 @@ def __init__(
             Use Client.set_transport_security() and Client.set_verify_certificates() to customize.
         """
         self.network: str = network or "testnet"
-        self.mirror_address: str = mirror_address or self.MIRROR_ADDRESS_DEFAULT.get(network, "localhost:5600")
+        self._mirror_address: str = mirror_address or self.MIRROR_ADDRESS_DEFAULT.get(self.network, "localhost:5600")
+        self._mirror_channel: grpc.Channel | None = None
+        self._mirror_stub: mirror_consensus_grpc.ConsensusServiceStub | None = None
 
-        self.ledger_id = ledger_id or self.LEDGER_ID.get(network, bytes.fromhex("03"))
+        self.ledger_id = ledger_id or self.LEDGER_ID.get(self.network, bytes.fromhex("03"))
 
         # Default TLS configuration: enabled for hosted networks, disabled for local/custom
         hosted_networks = ("mainnet", "testnet", "previewnet")
@@ -122,13 +126,32 @@ def __init__(
         self._node_index: int = secrets.randbelow(len(self._healthy_nodes))
         self.current_node: _Node = self._healthy_nodes[self._node_index]
 
+    @property
+    def mirror_address(self) -> str:
+        return self._mirror_address
+
+    @mirror_address.setter
+    def mirror_address(self, value: str):
+        """Reset the connection when the address changes."""
+        if not isinstance(value, str):
+            raise TypeError(f"mirror_address must be a string, not {type(value).__name__}")
+
+        value = value.strip()
+        if not value:
+            raise ValueError("mirror_address cannot be empty or just whitespace")
+
+        if self._mirror_address != value:
+            self._mirror_address = value
+            self._close_mirror_node()
+
     def _set_network_nodes(self, nodes: list[_Node] | None = None):
         """Configure the consensus nodes used by this network."""
         final_nodes = self._resolve_nodes(nodes)
 
         # Apply TLS configuration to all nodes
         for node in final_nodes:
-            node._apply_transport_security(self._transport_security)  # pylint: disable=protected-access
+            if self._transport_security:
+                node._apply_transport_security(self._transport_security)  # pylint: disable=protected-access
             node._set_verify_certificates(self._verify_certificates)  # pylint: disable=protected-access
             node._set_root_certificates(self._root_certificates)  # pylint: disable=protected-access
 
@@ -405,3 +428,33 @@ def _mark_node_healthy(self, node: _Node) -> None:
 
         if node not in self._healthy_nodes:
             self._healthy_nodes.append(node)
+
+    def _close_mirror_node(self):
+        """Safely closes the mirror gRPC channel."""
+        if self._mirror_channel is not None:
+            self._mirror_channel.close()
+
+        self._mirror_channel = None
+        self._mirror_stub = None
+
+    def _close(self):
+        """Safely closes the mirror gRPC channel and consensus node."""
+        self._close_mirror_node()
+
+        if self.nodes:
+            for node in self.nodes:
+                node._close()
+
+    def get_mirror_stub(self) -> mirror_consensus_grpc.ConsensusServiceStub:
+        """Returns the mirror stub."""
+        if self._mirror_stub is None:
+            addr = self._mirror_address
+
+            if addr.endswith(":50212") or addr.endswith(":443"):
+                self._mirror_channel = grpc.secure_channel(addr, grpc.ssl_channel_credentials())
+            else:
+                self._mirror_channel = grpc.insecure_channel(addr)
+
+            self._mirror_stub = mirror_consensus_grpc.ConsensusServiceStub(self._mirror_channel)
+
+        return self._mirror_stub

tck/handlers/sdk.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-from hiero_sdk_python import AccountId, Client, Network, PrivateKey
-from hiero_sdk_python.node import _Node
+from hiero_sdk_python import AccountId, Client, PrivateKey
 from tck.handlers.registry import rpc_method
 from tck.param.base import BaseParams
 from tck.param.sdk import SetupParams
@@ -15,11 +14,11 @@ def setup_handler(params: SetupParams) -> SetupResponse:
     operator_private_key = PrivateKey.from_string(params.operatorPrivateKey)
 
     if params.nodeIp and params.nodeAccountId and params.mirrorNetworkIp:
-        client = Client()
-        client.network = Network(
-            nodes=[_Node(AccountId.from_string(params.nodeAccountId), params.nodeIp, None)],
-            mirror_address=params.mirrorNetworkIp,
-        )
+        nodes = {params.nodeIp: AccountId.from_string(params.nodeAccountId)}
+
+        client = Client.for_network(network_map=nodes)
+        client.network.mirror_address = params.mirrorNetworkIp
+
         client.set_operator(operator_account_id, operator_private_key)
 
         client_type = "custom"

tests/unit/client_test.py

@@ -10,8 +10,10 @@
 
 import pytest
 
-from hiero_sdk_python import AccountId, Client, PrivateKey
+from hiero_sdk_python import AccountId, Client
 from hiero_sdk_python.client import client as client_module
+from hiero_sdk_python.client.network import Network
+from hiero_sdk_python.crypto.private_key import PrivateKey
 from hiero_sdk_python.hbar import Hbar
 from hiero_sdk_python.node import _Node
 from hiero_sdk_python.transaction.transaction_id import TransactionId
@@ -527,3 +529,78 @@ def test_get_node_account_ids_raises_when_no_nodes():
         client.get_node_account_ids()
 
     client.close()
+
+
+def test_for_network_initializes_with_custom_map():
+    """Test for_network correctly maps strings to AccountIds and Nodes."""
+    network_map = {"127.0.0.1:50211": AccountId(0, 0, 3), "127.0.0.1:50212": AccountId(0, 0, 4)}
+
+    client = Client.for_network(network_map)
+
+    assert isinstance(client.network, Network)
+    assert len(client.network.nodes) == 2
+
+    node_accounts = [node._account_id.__str__() for node in client.network.nodes]
+    assert "0.0.3" in node_accounts
+    assert "0.0.4" in node_accounts
+
+
+@pytest.mark.parametrize("network", ["mainnet", "testnet", "previewnet"])
+def test_for_network_with_hosted_network_forces_tls(network):
+    """Test that if hosted-net the port 50211 is upgraded to 50212."""
+    network_map = {"34.94.106.61:50211": AccountId(0, 0, 3)}
+    client = Client.for_network(network_map, network_name=network)
+
+    node = client.network.nodes[0]
+    assert str(node._address) == "34.94.106.61:50212"
+    assert node._address._is_transport_security() is True
+    assert client.network.is_transport_security() is True
+
+
+@pytest.mark.parametrize("network", ["local", "localhost", "solo", "custom"])
+def test_for_network_with_non_hosted_network_not_forces_tls(network):
+    """Test that if non hosted-net the port 50211 does not change."""
+    network_map = {"127.0.0.1:50211": AccountId(0, 0, 3)}
+    client = Client.for_network(network_map, network_name=network)
+
+    node = client.network.nodes[0]
+    assert str(node._address) == "127.0.0.1:50211"
+    assert node._address._is_transport_security() is False
+    assert client.network.is_transport_security() is False
+
+
+@pytest.mark.parametrize("network", ["local", "localhost", "solo", "custom"])
+def test_for_network_with_non_hosted_network_not_downgrade_tls(network):
+    """Test that if non hosted-net the port 50212 does not change."""
+    network_map = {"127.0.0.1:50212": AccountId(0, 0, 3)}
+    client = Client.for_network(network_map, network_name=network)
+
+    node = client.network.nodes[0]
+    assert str(node._address) == "127.0.0.1:50212"
+    assert node._address._is_transport_security() is True
+    assert client.network.is_transport_security() is False  # Non hosted network.
+
+
+def test_for_network_with_empty_map_raises_error():
+    """Test that for_network raises ValueError when map is empty."""
+    with pytest.raises(ValueError, match="network_map cannot be empty"):
+        Client.for_network({})
+
+
+@pytest.mark.parametrize(
+    "invalid_map, error_msg",
+    [
+        (
+            {"127.0.0.1:50211": AccountId(0, 0, 3), "127.0.0.1:50212": AccountId(1, 0, 4)},
+            "network is not valid, all nodes must be in the same shard and realm",
+        ),
+        (
+            {"127.0.0.1:50211": AccountId(0, 0, 3), "127.0.0.1:50212": AccountId(0, 1, 4)},
+            "network is not valid, all nodes must be in the same shard and realm",
+        ),
+    ],
+)
+def test_for_network_invalid_shard_realm_raises_error(invalid_map, error_msg):
+    """Test that for_network catches mismatched shards or realms."""
+    with pytest.raises(ValueError, match=error_msg):
+        Client.for_network(invalid_map)