test(1781): assert token balance via mirror node (HIP-367)

sdk/account_balance_query_e2e_test.go

@@ -39,15 +39,20 @@ func TestIntegrationAccountBalanceQueryCanGetTokenBalance(t *testing.T) {
 	tokenID, err := createFungibleToken(&env)
 	require.NoError(t, err)
 
-	balance, err := NewAccountBalanceQuery().
+	operatorID := env.Client.GetOperatorAccountID()
+
+	_, err = NewAccountBalanceQuery().
 		SetNodeAccountIDs(env.NodeAccountIDs).
-		SetAccountID(env.Client.GetOperatorAccountID()).
+		SetAccountID(operatorID).
 		Execute(env.Client)
 	require.NoError(t, err)
 
-	assert.Equal(t, balance, balance)
-	// TODO: assert.Equal(t, uint64(1000000), balance.Tokens.Get(*tokenID))
-	// TODO: assert.Equal(t, uint64(3), balance.TokenDecimals.Get(*tokenID))
+	// HIP-367 dropped token balances from the consensus response; assert via the mirror node.
+	balance, decimals, err := getTokenBalanceFromMirror(env.Client, operatorID, tokenID)
+	require.NoError(t, err)
+	assert.Equal(t, uint64(1_000_000), balance)
+	assert.Equal(t, uint64(18), decimals)
+
 	err = CloseIntegrationTestEnv(env, &tokenID)
 	require.NoError(t, err)
 }

sdk/utilities_for_test.go

@@ -5,7 +5,9 @@ package hiero
 // SPDX-License-Identifier: Apache-2.0
 
 import (
+	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"net/url"
 	"os"
@@ -288,6 +290,75 @@ func createFungibleToken(env *IntegrationTestEnv, opts ...TokenCreateTransaction
 	return *receipt.TokenID, err
 }
 
+// mirrorTokenBalance is one entry from GET /api/v1/accounts/{id}/tokens.
+type mirrorTokenBalance struct {
+	TokenID  string `json:"token_id"`
+	Balance  uint64 `json:"balance"`
+	Decimals uint64 `json:"decimals"`
+}
+
+type mirrorAccountTokensResponse struct {
+	Tokens []mirrorTokenBalance `json:"tokens"`
+}
+
+// getTokenBalanceFromMirror returns accountID's balance and decimals for tokenID from the
+// mirror node (the source of truth since HIP-367), polling with bounded retries to absorb
+// propagation delay and returning as soon as the relationship appears.
+func getTokenBalanceFromMirror(client *Client, accountID AccountID, tokenID TokenID) (uint64, uint64, error) {
+	baseURL, err := client.GetMirrorRestApiBaseUrl()
+	if err != nil {
+		return 0, 0, err
+	}
+	requestURL := fmt.Sprintf("%s/accounts/%s/tokens?token.id=%s", baseURL, accountID.String(), tokenID.String())
+
+	const maxAttempts = 10
+	const retryDelay = 2 * time.Second
+
+	httpClient := &http.Client{Timeout: 30 * time.Second}
+	var lastErr error
+	for attempt := 0; attempt < maxAttempts; attempt++ {
+		if attempt > 0 {
+			time.Sleep(retryDelay)
+		}
+
+		resp, err := httpClient.Get(requestURL) // #nosec
+		if err != nil {
+			lastErr = err
+			continue
+		}
+
+		if resp.StatusCode != http.StatusOK {
+			lastErr = fmt.Errorf("mirror node returned status %d for %s", resp.StatusCode, requestURL)
+			drainAndClose(resp.Body)
+			continue
+		}
+
+		var parsed mirrorAccountTokensResponse
+		decodeErr := json.NewDecoder(resp.Body).Decode(&parsed)
+		drainAndClose(resp.Body)
+		if decodeErr != nil {
+			lastErr = decodeErr
+			continue
+		}
+
+		for _, rel := range parsed.Tokens {
+			if rel.TokenID == tokenID.String() {
+				return rel.Balance, rel.Decimals, nil
+			}
+		}
+		lastErr = fmt.Errorf("token %s not yet reflected on mirror node for account %s", tokenID.String(), accountID.String())
+	}
+
+	return 0, 0, fmt.Errorf("failed to fetch token balance from mirror node after %d attempts: %w", maxAttempts, lastErr)
+}
+
+// drainAndClose exhausts and closes an HTTP response body so the underlying
+// connection can be reused across retries.
+func drainAndClose(body io.ReadCloser) {
+	_, _ = io.Copy(io.Discard, body)
+	_ = body.Close()
+}
+
 type AccountCreateTransactionCustomizer func(transaction *AccountCreateTransaction)
 
 func createAccount(env *IntegrationTestEnv, opts ...AccountCreateTransactionCustomizer) (AccountID, PrivateKey, error) {