fix: TokenBalance and TokenRelationship query response mapping

Signed-off-by: Ricky Saechao <ricky@launchbadge.com>

Author: RickyLB

Sources/Hedera/Account/AccountBalance.swift

@@ -21,7 +21,7 @@
 import Foundation
 import HederaProtobufs
 
-private struct TokenBalance {
+internal struct TokenBalance {
     fileprivate let id: TokenId
     fileprivate let balance: UInt64
     fileprivate let decimals: UInt32
@@ -59,6 +59,16 @@ public struct AccountBalance: Sendable {
 
     private let tokensInner: [TokenBalance]
 
+    internal init(
+        accountId: AccountId,
+        hbars: Hbar,
+        tokensInner: [TokenBalance]
+    ) {
+        self.accountId = accountId
+        self.hbars = hbars
+        self.tokensInner = tokensInner
+    }
+
     // hack to work around deprecated warning
     private var tokenBalancesInner: [TokenId: UInt64] {
         Dictionary(uniqueKeysWithValues: tokensInner.map { ($0.id, $0.balance) })

Sources/Hedera/Account/AccountBalanceQuery.swift

@@ -88,11 +88,25 @@ public final class AccountBalanceQuery: Query<AccountBalance> {
     internal override func makeQueryResponse(_ context: Context, _ response: Proto_Response.OneOf_Response) async throws
         -> Response
     {
+        let mirrorNodeGateway = try MirrorNodeGateway.forNetwork(context.mirrorNetworkNodes, context.ledgerId)
+        let mirrorNodeService = MirrorNodeService.init(mirrorNodeGateway)
+
         guard case .cryptogetAccountBalance(let proto) = response else {
             throw HError.fromProtobuf("unexpected \(response) received, expected `cryptogetAccountBalance`")
         }
 
-        return try .fromProtobuf(proto)
+        let accountId = try AccountId.fromProtobuf(proto.accountID)
+        let tokenBalanceProto = try await mirrorNodeService.getTokenBalancesForAccount(String(accountId.num))
+
+        // var tokenBalances: [TokenId: UInt64] = [:]
+
+        // for balance in tokenBalanceProto {
+        //     tokenBalances[.fromProtobuf(balance.tokenID)] = balance.balance
+        // }
+
+        return AccountBalance(
+            accountId: accountId, hbars: .fromTinybars(Int64(proto.balance)),
+            tokensInner: .fromProtobuf(tokenBalanceProto))
     }
 
     public override func validateChecksums(on ledgerId: LedgerId) throws {

Sources/Hedera/Account/AccountId.swift

@@ -135,6 +135,15 @@ public struct AccountId: Sendable, EntityId, ValidateChecksums {
     public func toBytes() -> Data {
         toProtobufBytes()
     }
+
+    public func populateAccountIdNum(_ client: Client) async throws -> Self {
+        let mirrorNodeGateway = try MirrorNodeGateway.forClient(client)
+        let mirrorNodeService = MirrorNodeService(mirrorNodeGateway)
+
+        let accountNumFromMirror = try await mirrorNodeService.getAccountNum(self.evmAddress!.toString())
+
+        return Self(shard: shard, realm: realm, num: accountNumFromMirror)
+    }
 }
 
 extension AccountId: TryProtobufCodable {

Sources/Hedera/Account/AccountInfoQuery.swift

@@ -69,10 +69,9 @@ public final class AccountInfoQuery: Query<AccountInfo> {
             throw HError.fromProtobuf("unexpected \(response) received, expected `cryptoGetInfo`")
         }
 
-        let accountInfoProto = proto.accountInfo
-        let accountId = try AccountId.fromProtobuf(accountInfoProto.accountID)
+        let accountInfo = try AccountInfo.fromProtobuf(proto.accountInfo)
         let tokenRelationshipsProto = try await mirrorNodeService.getTokenRelationshipsForAccount(
-            String(describing: accountId.num))
+            String(describing: accountInfo.accountId.num))
 
         var tokenRelationships: [TokenId: TokenRelationship] = [:]
 
@@ -81,26 +80,26 @@ public final class AccountInfoQuery: Query<AccountInfo> {
         }
 
         return AccountInfo(
-            accountId: try AccountId.fromProtobuf(accountInfoProto.accountID),
-            contractAccountId: accountInfoProto.contractAccountID,
-            isDeleted: accountInfoProto.deleted,
-            proxyAccountId: try .fromProtobuf(accountInfoProto.proxyAccountID),
-            proxyReceived: Hbar.fromTinybars(accountInfoProto.proxyReceived),
-            key: try .fromProtobuf(accountInfoProto.key),
-            balance: .fromTinybars(Int64(accountInfoProto.balance)),
-            sendRecordThreshold: Hbar.fromTinybars(Int64(accountInfoProto.generateSendRecordThreshold)),
-            receiveRecordThreshold: Hbar.fromTinybars(Int64(accountInfoProto.generateReceiveRecordThreshold)),
-            isReceiverSignatureRequired: accountInfoProto.receiverSigRequired,
-            expirationTime: .fromProtobuf(accountInfoProto.expirationTime),
-            autoRenewPeriod: .fromProtobuf(accountInfoProto.autoRenewPeriod),
-            accountMemo: accountInfoProto.memo,
-            ownedNfts: UInt64(accountInfoProto.ownedNfts),
-            maxAutomaticTokenAssociations: UInt32(accountInfoProto.maxAutomaticTokenAssociations),
-            aliasKey: try .fromAliasBytes(accountInfoProto.alias),
-            ethereumNonce: UInt64(accountInfoProto.ethereumNonce),
+            accountId: accountInfo.accountId,
+            contractAccountId: accountInfo.contractAccountId,
+            isDeleted: accountInfo.isDeleted,
+            proxyAccountId: accountInfo.proxyAccountId,
+            proxyReceived: accountInfo.proxyReceived,
+            key: accountInfo.key,
+            balance: accountInfo.balance,
+            sendRecordThreshold: accountInfo.sendRecordThreshold,
+            receiveRecordThreshold: accountInfo.receiveRecordThreshold,
+            isReceiverSignatureRequired: accountInfo.isReceiverSignatureRequired,
+            expirationTime: accountInfo.expirationTime,
+            autoRenewPeriod: accountInfo.autoRenewPeriod,
+            accountMemo: accountInfo.accountMemo,
+            ownedNfts: accountInfo.ownedNfts,
+            maxAutomaticTokenAssociations: accountInfo.maxAutomaticTokenAssociations,
+            aliasKey: accountInfo.aliasKey,
+            ethereumNonce: accountInfo.ethereumNonce,
             tokenRelationships: tokenRelationships,
-            ledgerId: .fromBytes(accountInfoProto.ledgerID),
-            staking: try .fromProtobuf(accountInfoProto.stakingInfo)
+            ledgerId: accountInfo.ledgerId,
+            staking: accountInfo.staking
         )
     }
 

Sources/Hedera/Contract/ContractInfo.swift

@@ -104,7 +104,7 @@ extension ContractInfo: TryProtobufCodable {
 
         self.init(
             contractId: try .fromProtobuf(proto.contractID),
-            accountId: try .fromProtobuf(proto.accountID),
+            accountId: try AccountId.fromProtobuf(proto.accountID),
             contractAccountId: proto.contractAccountID,
             adminKey: try .fromProtobuf(adminKey),
             expirationTime: .fromProtobuf(expirationTime),

Sources/Hedera/Contract/ContractInfoQuery.swift

@@ -65,10 +65,10 @@ public final class ContractInfoQuery: Query<ContractInfo> {
         guard case .contractGetInfo(let proto) = response else {
             throw HError.fromProtobuf("unexpected \(response) received, expected `contractGetInfo`")
         }
-        let contractInfoProto = proto.contractInfo
-        let contractId = try ContractId.fromProtobuf(contractInfoProto.contractID)
+
+        let contractInfo = try ContractInfo.fromProtobuf(proto.contractInfo)
         let tokenRelationshipsProto = try await mirrorNodeService.getTokenRelationshipsForAccount(
-            String(describing: contractId.num))
+            String(describing: contractInfo.contractId.num))
 
         var tokenRelationships: [TokenId: TokenRelationship] = [:]
 
@@ -77,21 +77,22 @@ public final class ContractInfoQuery: Query<ContractInfo> {
         }
 
         return ContractInfo(
-            contractId: try ContractId.fromProtobuf(contractInfoProto.contractID),
-            accountId: try AccountId.fromProtobuf(contractInfoProto.accountID),
-            contractAccountId: contractInfoProto.contractAccountID,
-            adminKey: try .fromProtobuf(contractInfoProto.adminKey),
-            expirationTime: .fromProtobuf(contractInfoProto.expirationTime),
-            autoRenewPeriod: .fromProtobuf(contractInfoProto.autoRenewPeriod),
-            storage: UInt64(contractInfoProto.storage),
-            contractMemo: contractInfoProto.memo,
-            balance: .fromTinybars(Int64(contractInfoProto.balance)),
-            isDeleted: contractInfoProto.deleted,
-            autoRenewAccountId: try .fromProtobuf(contractInfoProto.autoRenewAccountID),
-            maxAutomaticTokenAssociations: UInt32(contractInfoProto.maxAutomaticTokenAssociations),
+            contractId: contractInfo.contractId,
+            accountId: contractInfo.accountId,
+            contractAccountId: contractInfo.contractAccountId,
+            adminKey: contractInfo.adminKey,
+            expirationTime: contractInfo.expirationTime,
+            autoRenewPeriod: contractInfo.autoRenewPeriod,
+            storage: contractInfo.storage,
+            contractMemo: contractInfo.contractMemo,
+            balance: contractInfo.balance,
+            isDeleted: contractInfo.isDeleted,
+            autoRenewAccountId: contractInfo.autoRenewAccountId,
+            maxAutomaticTokenAssociations: contractInfo.maxAutomaticTokenAssociations,
             tokenRelationships: tokenRelationships,
-            ledgerId: .fromBytes(contractInfoProto.ledgerID),
-            stakingInfo: try .fromProtobuf(contractInfoProto.stakingInfo))
+            ledgerId: contractInfo.ledgerId,
+            stakingInfo: contractInfo.stakingInfo
+        )
     }
 
     internal override func validateChecksums(on ledgerId: LedgerId) throws {

Tests/HederaE2ETests/Contract/ContractInfo.swift

@@ -22,7 +22,7 @@ import Hedera
 import XCTest
 
 internal final class ContractInfo: XCTestCase {
-    internal func testQueryBoo() async throws {
+    internal func testQuery() async throws {
         let testEnv = try TestEnvironment.nonFree
 
         let contractId = try await ContractHelpers.makeContract(testEnv, operatorAdminKey: true)