Skip to content
Merged
62 changes: 34 additions & 28 deletions src/hiero_sdk_python/client/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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
Comment thread
coderabbitai[bot] marked this conversation as resolved.

@classmethod
def from_env(cls, network: NetworkName | None = None) -> Client:
"""
Expand Down Expand Up @@ -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))
Comment thread
coderabbitai[bot] marked this conversation as resolved.

def set_operator(self, account_id: AccountId, private_key: PrivateKey) -> None:
"""Sets the operator credentials (account ID and private key)."""
Expand Down Expand Up @@ -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:
"""
Expand Down
61 changes: 57 additions & 4 deletions src/hiero_sdk_python/client/network.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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,
Expand All @@ -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")
Expand All @@ -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):
Comment thread
manishdait marked this conversation as resolved.
"""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

Expand Down Expand Up @@ -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
Comment thread
exploreriii marked this conversation as resolved.
13 changes: 6 additions & 7 deletions tck/handlers/sdk.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Comment thread
exploreriii marked this conversation as resolved.

client.set_operator(operator_account_id, operator_private_key)

client_type = "custom"
Expand Down
79 changes: 78 additions & 1 deletion tests/unit/client_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Comment thread
manishdait marked this conversation as resolved.


@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)
Loading
Loading