refactor: use StakingInfo wrapper in ContractInfo and remove duplicate account_info changes

Signed-off-by: Mounil <mounilkankhara@gmail.com>

Author: Mounil2005

CHANGELOG.md

@@ -21,6 +21,7 @@ This changelog is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.
 - Fixed duplication in GitHub bot next issue recommendations by parsing actual issue descriptions instead of blind truncation (#1658)
 
 ### Src
+- Add `staking_info` field to `ContractInfo` class to expose staking metadata using the `StakingInfo` wrapper. (#1365)
 - Fix `TopicInfo.__str__()` to format `expiration_time` in UTC so unit tests pass in non-UTC environments. (#1800)
 - Resolve CodeQL `reflected-XSS` warning in TCK JSON-RPC endpoint
 - Improve `keccak256` docstring formatting for better readability and consistency (#1624)
@@ -31,7 +32,6 @@ This changelog is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.
 - Refactored contract_delete_transaction example to use Client.from_env. (#1823)
 
 ### Added
-- Add staking_info fields (`staked_account_id`, `staked_node_id`, `decline_staking_reward`) to ContractInfo class to expose staking metadata from protobuf. (#1365)
 
 ### Docs
 

src/hiero_sdk_python/account/account_info.py

@@ -9,9 +9,9 @@
 from hiero_sdk_python.account.account_id import AccountId
 from hiero_sdk_python.crypto.public_key import PublicKey
 from hiero_sdk_python.Duration import Duration
+from hiero_sdk_python.hapi.services.basic_types_pb2 import StakingInfo
 from hiero_sdk_python.hapi.services.crypto_get_info_pb2 import CryptoGetInfoResponse
 from hiero_sdk_python.hbar import Hbar
-from hiero_sdk_python.staking_info import StakingInfo
 from hiero_sdk_python.timestamp import Timestamp
 from hiero_sdk_python.tokens.token_relationship import TokenRelationship
 
@@ -38,7 +38,9 @@ class AccountInfo:
             associated with this account.
         account_memo (Optional[str]): The memo associated with this account.
         owned_nfts (Optional[int]): The number of NFTs owned by this account.
-        staking_info (Optional[StakingInfo]): The staking information for this account.
+        staked_account_id (Optional[AccountId]): The account to which this account is staked.
+        staked_node_id (Optional[int]): The node to which this account is staked.
+        decline_staking_reward (bool): Whether this account declines receiving staking rewards. 
     """
 
     account_id: Optional[AccountId] = None
@@ -54,7 +56,9 @@ class AccountInfo:
     account_memo: Optional[str] = None
     owned_nfts: Optional[int] = None
     max_automatic_token_associations: Optional[int] = None
-    staking_info: Optional[StakingInfo] = None
+    staked_account_id: Optional[AccountId] = None
+    staked_node_id: Optional[int] = None
+    decline_staking_reward: Optional[bool] = None
 
     @classmethod
     def _from_proto(cls, proto: CryptoGetInfoResponse.AccountInfo) -> "AccountInfo":
@@ -96,13 +100,21 @@ def _from_proto(cls, proto: CryptoGetInfoResponse.AccountInfo) -> "AccountInfo":
             account_memo=proto.memo,
             owned_nfts=proto.ownedNfts,
             max_automatic_token_associations=proto.max_automatic_token_associations,
-            staking_info=(
-                StakingInfo._from_proto(proto.staking_info)
-                if proto.HasField('staking_info')
-                else None
-            ),
         )
 
+        staking_info = proto.staking_info if proto.HasField('staking_info') else None
+
+        if staking_info:
+            account_info.staked_account_id = (
+                AccountId._from_proto(staking_info.staked_account_id) 
+                if staking_info.HasField('staked_account_id') else None
+            )
+            account_info.staked_node_id = (
+                staking_info.staked_node_id 
+                if staking_info.HasField('staked_node_id') else None
+            )
+            account_info.decline_staking_reward = staking_info.decline_reward
+
         return account_info
 
     def _to_proto(self) -> CryptoGetInfoResponse.AccountInfo:
@@ -135,7 +147,11 @@ def _to_proto(self) -> CryptoGetInfoResponse.AccountInfo:
             memo=self.account_memo,
             ownedNfts=self.owned_nfts,
             max_automatic_token_associations=self.max_automatic_token_associations,
-            staking_info=self.staking_info._to_proto() if self.staking_info else None,
+            staking_info=StakingInfo(
+                staked_account_id=self.staked_account_id._to_proto() if self.staked_account_id else None,
+                staked_node_id=self.staked_node_id if self.staked_node_id else None,
+                decline_reward=self.decline_staking_reward
+            ),
         )
 
     def __str__(self) -> str:
@@ -150,10 +166,11 @@ def __str__(self) -> str:
             (self.account_memo, "Memo"),
             (self.owned_nfts, "Owned NFTs"),
             (self.max_automatic_token_associations, "Max Automatic Token Associations"),
+            (self.staked_account_id, "Staked Account ID"),
+            (self.staked_node_id, "Staked Node ID"),
             (self.proxy_received, "Proxy Received"),
             (self.expiration_time, "Expiration Time"),
             (self.auto_renew_period, "Auto Renew Period"),
-            (self.staking_info, "Staking Info"),
         ]
 
         # Use a list comprehension to process simple fields (reduces complexity score)
@@ -165,6 +182,9 @@ def __str__(self) -> str:
 
         if self.receiver_signature_required is not None:
             lines.append(f"Receiver Signature Required: {self.receiver_signature_required}")
+            
+        if self.decline_staking_reward is not None:
+            lines.append(f"Decline Staking Reward: {self.decline_staking_reward}")
 
         if self.token_relationships:
             lines.append(f"Token Relationships: {len(self.token_relationships)}")
@@ -182,6 +202,7 @@ def __repr__(self) -> str:
             f"receiver_signature_required={self.receiver_signature_required!r}, "
             f"owned_nfts={self.owned_nfts!r}, "
             f"account_memo={self.account_memo!r}, "
-            f"staking_info={self.staking_info!r}"
+            f"staked_node_id={self.staked_node_id!r}, "
+            f"staked_account_id={self.staked_account_id!r}"
             f")"
         )

tests/integration/contract_info_query_e2e_test.py

@@ -82,6 +82,17 @@ def test_integration_contract_info_query_can_execute(env):
     assert info.auto_renew_account_id == env.operator_id, "Auto renew account ID mismatch"
     assert info.auto_renew_period == auto_renew_period, "Auto renew period mismatch"
 
+    # Verify staking_info is populated (contracts default to no staking)
+    staking = info.staking_info
+    if staking is None:
+        raise AssertionError("staking_info should not be None")
+    if staking.staked_account_id is not None:
+        raise AssertionError("staked_account_id should be None by default")
+    if staking.staked_node_id is not None:
+        raise AssertionError("staked_node_id should be None by default")
+    if staking.decline_reward is not False:
+        raise AssertionError("decline_reward should be False by default")
+
 
 @pytest.mark.integration
 def test_integration_contract_info_query_get_cost(env):

tests/unit/account_info_test.py

@@ -9,8 +9,6 @@
 from hiero_sdk_python.tokens.token_relationship import TokenRelationship
 from hiero_sdk_python.tokens.token_id import TokenId
 from hiero_sdk_python.hapi.services.crypto_get_info_pb2 import CryptoGetInfoResponse
-from hiero_sdk_python.hapi.services.basic_types_pb2 import StakingInfo as StakingInfoProto
-from hiero_sdk_python.staking_info import StakingInfo
 
 pytestmark = pytest.mark.unit
 
@@ -213,154 +211,3 @@ def test_str_and_repr(account_info):
     assert "contract_account_id='0.0.100'" in info_repr
     assert "account_memo='Test account memo'" in info_repr
 
-
-def test_from_proto_with_staking_info():
-    """Test from_proto with staking account info"""
-    public_key = PrivateKey.generate_ed25519().public_key()
-    proto = CryptoGetInfoResponse.AccountInfo(
-        accountID=AccountId(0, 0, 100)._to_proto(),
-        key=public_key._to_proto(),
-        balance=5000000,
-        staking_info=StakingInfoProto(
-            staked_account_id=AccountId(0, 0, 500)._to_proto(),
-            decline_reward=True,
-        ),
-    )
-    
-    account_info = AccountInfo._from_proto(proto)
-    
-    assert account_info.staking_info is not None
-    assert account_info.staking_info.staked_account_id == AccountId(0, 0, 500)
-    assert account_info.staking_info.decline_reward is True
-
-
-def test_from_proto_with_staked_node_id():
-    """Test from_proto with staked_node_id"""
-    public_key = PrivateKey.generate_ed25519().public_key()
-    proto = CryptoGetInfoResponse.AccountInfo(
-        accountID=AccountId(0, 0, 100)._to_proto(),
-        key=public_key._to_proto(),
-        balance=5000000,
-        staking_info=StakingInfoProto(
-            staked_node_id=3,
-            decline_reward=False,
-        ),
-    )
-    
-    account_info = AccountInfo._from_proto(proto)
-    
-    assert account_info.staking_info is not None
-    assert account_info.staking_info.staked_node_id == 3
-    assert account_info.staking_info.decline_reward is False
-
-
-def test_from_proto_with_no_staking_info():
-    """Test from_proto without staking info"""
-    public_key = PrivateKey.generate_ed25519().public_key()
-    proto = CryptoGetInfoResponse.AccountInfo(
-        accountID=AccountId(0, 0, 100)._to_proto(),
-        key=public_key._to_proto(),
-        balance=5000000,
-    )
-    
-    account_info = AccountInfo._from_proto(proto)
-    
-    assert account_info.staking_info is None
-
-
-def test_to_proto_with_staking_info():
-    """Test to_proto with staking info"""
-    account_info = AccountInfo(
-        account_id=AccountId(0, 0, 100),
-        balance=Hbar.from_tinybars(5000000),
-        staking_info=StakingInfo(
-            staked_account_id=AccountId(0, 0, 500),
-            decline_reward=True,
-        ),
-    )
-    
-    proto = account_info._to_proto()
-    
-    assert proto.HasField('staking_info')
-    assert proto.staking_info.HasField('staked_account_id')
-    assert proto.staking_info.staked_account_id == AccountId(0, 0, 500)._to_proto()
-    assert proto.staking_info.decline_reward is True
-
-
-def test_to_proto_with_staked_node_id():
-    """Test to_proto with staked_node_id"""
-    account_info = AccountInfo(
-        account_id=AccountId(0, 0, 100),
-        balance=Hbar.from_tinybars(5000000),
-        staking_info=StakingInfo(
-            staked_node_id=5,
-            decline_reward=False,
-        ),
-    )
-    
-    proto = account_info._to_proto()
-    
-    assert proto.HasField('staking_info')
-    assert proto.staking_info.staked_node_id == 5
-    assert proto.staking_info.decline_reward is False
-
-
-def test_proto_conversion_staking_node_round_trip():
-    """Test proto conversion round trip with staked_node_id"""
-    account_info = AccountInfo(
-        account_id=AccountId(0, 0, 100),
-        key=PrivateKey.generate_ed25519().public_key(),
-        balance=Hbar.from_tinybars(5000000),
-        staking_info=StakingInfo(
-            staked_node_id=7,
-            decline_reward=False,
-        ),
-    )
-    
-    converted = AccountInfo._from_proto(account_info._to_proto())
-    
-    assert converted.account_id == account_info.account_id
-    assert converted.balance.to_tinybars() == account_info.balance.to_tinybars()
-    assert converted.staking_info.staked_account_id is None
-    assert converted.staking_info.staked_node_id == 7
-    assert converted.staking_info.decline_reward is False
-
-
-def test_proto_conversion_staking_account_round_trip():
-    """Test proto conversion round trip with staked_account_id"""
-    account_info = AccountInfo(
-        account_id=AccountId(0, 0, 100),
-        key=PrivateKey.generate_ed25519().public_key(),
-        balance=Hbar.from_tinybars(5000000),
-        staking_info=StakingInfo(
-            staked_account_id=AccountId(0, 0, 600),
-            decline_reward=True,
-        ),
-    )
-    
-    converted = AccountInfo._from_proto(account_info._to_proto())
-    
-    assert converted.account_id == account_info.account_id
-    assert converted.balance.to_tinybars() == account_info.balance.to_tinybars()
-    assert converted.staking_info.staked_account_id == AccountId(0, 0, 600)
-    assert converted.staking_info.staked_node_id is None
-    assert converted.staking_info.decline_reward is True
-
-
-def test_proto_conversion_with_staked_node_zero():
-    """Test proto conversion with staked_node_id set to 0"""
-    account_info = AccountInfo(
-        account_id=AccountId(0, 0, 100),
-        key=PrivateKey.generate_ed25519().public_key(),
-        balance=Hbar.from_tinybars(5000000),
-        staking_info=StakingInfo(
-            staked_node_id=0,
-            decline_reward=True,
-        ),
-    )
-    
-    converted = AccountInfo._from_proto(account_info._to_proto())
-    
-    assert converted.staking_info is not None
-    assert converted.staking_info.staked_node_id == 0
-    assert converted.staking_info.decline_reward is True

tests/unit/contract_info_test.py

@@ -350,7 +350,6 @@ def test_from_proto_with_no_staking_info():
 
 def test_from_proto_with_staked_node_id():
     """Test from_proto with staked_node_id (staked to node)"""
-    public_key = PrivateKey.generate_ed25519().public_key()
     proto = ContractGetInfoResponse.ContractInfo(
         contractID=ContractId(0, 0, 200)._to_proto(),
         accountID=AccountId(0, 0, 300)._to_proto(),
@@ -370,7 +369,7 @@ def test_from_proto_with_staked_node_id():
     assert contract_info.staking_info.decline_reward is True
 
 
-def test_to_proto_with_staked_account_id(token_relationship):
+def test_to_proto_with_staked_account_id():
     """Test to_proto with staked_account_id"""
     contract_info = ContractInfo(
         contract_id=ContractId(0, 0, 200),
@@ -428,6 +427,7 @@ def test_proto_conversion_staking_node_round_trip():
     assert converted.balance == contract_info.balance
     assert converted.staking_info.staked_account_id is None
     assert converted.staking_info.staked_node_id == 7
+    assert isinstance(converted.staking_info.staked_node_id, int)
     assert converted.staking_info.decline_reward is False
 
 
@@ -449,5 +449,43 @@ def test_proto_conversion_staking_account_round_trip():
     assert converted.account_id == contract_info.account_id
     assert converted.balance == contract_info.balance
     assert converted.staking_info.staked_account_id == AccountId(0, 0, 600)
+    assert isinstance(converted.staking_info.staked_account_id, AccountId)
     assert converted.staking_info.staked_node_id is None
     assert converted.staking_info.decline_reward is True
+
+
+def test_proto_conversion_with_staked_node_zero():
+    """Test proto conversion with staked_node_id set to 0 (valid edge case)"""
+    contract_info = ContractInfo(
+        contract_id=ContractId(0, 0, 200),
+        account_id=AccountId(0, 0, 300),
+        balance=5000000,
+        staking_info=StakingInfo(
+            staked_node_id=0,
+            decline_reward=True,
+        ),
+    )
+    
+    converted = ContractInfo._from_proto(contract_info._to_proto())
+    
+    assert converted.staking_info is not None
+    assert converted.staking_info.staked_node_id == 0
+    assert isinstance(converted.staking_info.staked_node_id, int)
+    assert converted.staking_info.decline_reward is True
+
+
+def test_proto_conversion_no_staking_info_round_trip():
+    """Test proto conversion round trip with no staking info (None should stay None)"""
+    contract_info = ContractInfo(
+        contract_id=ContractId(0, 0, 200),
+        account_id=AccountId(0, 0, 300),
+        balance=5000000,
+        staking_info=None,
+    )
+    
+    converted = ContractInfo._from_proto(contract_info._to_proto())
+    
+    assert converted.contract_id == contract_info.contract_id
+    assert converted.account_id == contract_info.account_id
+    assert converted.balance == contract_info.balance
+    assert converted.staking_info is None