account_update_transaction_e2e_test.py

"""
Integration tests for the AccountUpdateTransaction class.
"""

from __future__ import annotations

import pytest

from hiero_sdk_python.account.account_create_transaction import AccountCreateTransaction
from hiero_sdk_python.account.account_id import AccountId
from hiero_sdk_python.account.account_update_transaction import AccountUpdateTransaction
from hiero_sdk_python.crypto.key_list import KeyList
from hiero_sdk_python.crypto.private_key import PrivateKey
from hiero_sdk_python.Duration import Duration
from hiero_sdk_python.hbar import Hbar
from hiero_sdk_python.query.account_info_query import AccountInfoQuery
from hiero_sdk_python.response_code import ResponseCode
from hiero_sdk_python.timestamp import Timestamp


@pytest.mark.integration
def test_integration_account_update_transaction_can_execute(env):
    """Test that the AccountUpdateTransaction can be executed successfully."""
    # Create initial account
    initial_memo = "Initial account memo"

    receipt = (
        AccountCreateTransaction()
        .set_key_without_alias(env.operator_key.public_key())
        .set_initial_balance(Hbar(2))
        .set_account_memo(initial_memo)
        .set_receiver_signature_required(False)
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account creation failed with status: {ResponseCode(receipt.status).name}"
    )

    account_id = receipt.account_id
    assert account_id is not None, "Account ID should not be None"

    # Generate new key for update
    new_private_key = PrivateKey.generate_ed25519()
    new_public_key = new_private_key.public_key()
    new_memo = "Updated account memo"
    new_auto_renew_period = Duration(8000000)  # ~93 days

    # Update account with new key, memo, and other fields
    receipt = (
        AccountUpdateTransaction()
        .set_account_id(account_id)
        .set_key(new_public_key)
        .set_account_memo(new_memo)
        .set_receiver_signature_required(True)
        .set_auto_renew_period(new_auto_renew_period)
        .freeze_with(env.client)
        .sign(new_private_key)  # Sign with new key
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account update failed with status: {ResponseCode(receipt.status).name}"
    )

    # Query account info to verify updates
    info = AccountInfoQuery(account_id).execute(env.client)
    assert str(info.account_id) == str(account_id), "Account ID should match"
    assert info.key.to_bytes_raw() == new_public_key.to_bytes_raw(), "Public key should be updated"
    assert info.account_memo == new_memo, "Account memo should be updated"
    assert info.receiver_signature_required is True, "Receiver signature requirement should be updated"
    assert info.auto_renew_period == new_auto_renew_period, "Auto renew period should be updated"


@pytest.mark.integration
def test_integration_account_update_transaction_set_key_with_threshold_keylist(env):
    """Test rotating to a threshold KeyList and authorizing follow-up updates."""
    # Create an account controlled by the operator key (payer signature is automatic).
    receipt = (
        AccountCreateTransaction()
        .set_key_without_alias(env.operator_key.public_key())
        .set_initial_balance(Hbar(2))
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account creation failed with status: {ResponseCode(receipt.status).name}"
    )

    account_id = receipt.account_id
    assert account_id is not None, "Account ID should not be None"

    # Build a 2-of-2 threshold KeyList and rotate the account key to it.
    key_1_private = PrivateKey.generate_ed25519()
    key_2_private = PrivateKey.generate_ed25519()
    threshold_key = KeyList([key_1_private.public_key(), key_2_private.public_key()], threshold=2)

    tx = AccountUpdateTransaction().set_account_id(account_id).set_key(threshold_key)

    tx.transaction_fee = Hbar.from_hbars(5)
    tx.freeze_with(env.client)
    tx.sign(key_1_private)
    tx.sign(key_2_private)

    receipt = tx.execute(env.client)

    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account key rotation to KeyList failed with status: {ResponseCode(receipt.status).name}"
    )

    # Verify a follow-up update can be authorized with both threshold keys.
    new_memo = "Updated using threshold KeyList"
    receipt = (
        AccountUpdateTransaction()
        .set_account_id(account_id)
        .set_account_memo(new_memo)
        .freeze_with(env.client)
        .sign(key_1_private)
        .sign(key_2_private)
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account update signed by threshold keys failed with status: {ResponseCode(receipt.status).name}"
    )


@pytest.mark.integration
def test_integration_account_update_transaction_fails_with_invalid_account_id(env):
    """Test that AccountUpdateTransaction fails with an invalid account ID."""
    # Create an account ID that doesn't exist on the network
    invalid_account_id = AccountId(0, 0, 999999999)

    receipt = AccountUpdateTransaction().set_account_id(invalid_account_id).execute(env.client)
    assert receipt.status == ResponseCode.INVALID_ACCOUNT_ID, (
        f"Account update should have failed with status INVALID_ACCOUNT_ID, but got {ResponseCode(receipt.status).name}"
    )


@pytest.mark.integration
def test_integration_account_update_transaction_fails_with_invalid_signature(env):
    """Test that AccountUpdateTransaction fails when not signed with the correct key."""
    # Create initial account
    initial_private_key = PrivateKey.generate_ed25519()
    initial_public_key = initial_private_key.public_key()

    receipt = (
        AccountCreateTransaction()
        .set_key_without_alias(initial_public_key)
        .set_initial_balance(Hbar(1))
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account creation failed with status: {ResponseCode(receipt.status).name}"
    )

    account_id = receipt.account_id
    assert account_id is not None, "Account ID should not be None"

    base_info = AccountInfoQuery(account_id).execute(env.client)

    # Try to update without signing with the account's key
    receipt = AccountUpdateTransaction().set_account_id(account_id).set_account_memo("New memo").execute(env.client)

    assert receipt.status == ResponseCode.INVALID_SIGNATURE, (
        f"Account update should have failed with status INVALID_SIGNATURE, but got {ResponseCode(receipt.status).name}"
    )

    # Verify nothing changed on-chain
    info_after = AccountInfoQuery(account_id).execute(env.client)
    assert info_after.account_memo == base_info.account_memo
    assert info_after.key.to_bytes_raw() == initial_public_key.to_bytes_raw()


@pytest.mark.integration
def test_integration_account_update_transaction_partial_update(env):
    """Test that AccountUpdateTransaction can update only specific fields."""
    # Create initial account
    receipt = (
        AccountCreateTransaction()
        .set_key_without_alias(env.operator_key.public_key())
        .set_initial_balance(Hbar(1))
        .set_account_memo("Initial memo")
        .set_receiver_signature_required(False)
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account creation failed with status: {ResponseCode(receipt.status).name}"
    )

    account_id = receipt.account_id
    assert account_id is not None, "Account ID should not be None"

    # Update only the memo, leaving other fields unchanged
    new_memo = "Only memo updated"
    receipt = AccountUpdateTransaction().set_account_id(account_id).set_account_memo(new_memo).execute(env.client)
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account update failed with status: {ResponseCode(receipt.status).name}"
    )

    # Query account info to verify only memo was updated
    info = AccountInfoQuery(account_id).execute(env.client)

    assert str(info.account_id) == str(account_id), "Account ID should match"
    assert info.key.to_bytes_raw() == env.operator_key.public_key().to_bytes_raw(), "Public key should remain unchanged"
    assert info.account_memo == new_memo, "Account memo should be updated"
    assert info.receiver_signature_required is False, "Receiver signature requirement should remain unchanged"


@pytest.mark.integration
def test_integration_account_update_transaction_invalid_auto_renew_period(env):
    """Test that AccountUpdateTransaction fails with invalid auto renew period."""
    # Create initial account
    receipt = (
        AccountCreateTransaction()
        .set_key(env.operator_key.public_key())
        .set_initial_balance(Hbar(1))
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account creation failed with status: {ResponseCode(receipt.status).name}"
    )

    account_id = receipt.account_id
    assert account_id is not None, "Account ID should not be None"

    # Capture existing expiration to ensure it remains unchanged
    original_info = AccountInfoQuery(account_id).execute(env.client)
    invalid_period = Duration(777600000)  # 9000 days

    # Try to update with invalid auto renew period
    receipt = (
        AccountUpdateTransaction().set_account_id(account_id).set_auto_renew_period(invalid_period).execute(env.client)
    )

    assert receipt.status == ResponseCode.AUTORENEW_DURATION_NOT_IN_RANGE, (
        f"Account update should have failed with status AUTORENEW_DURATION_NOT_IN_RANGE, "
        f"but got {ResponseCode(receipt.status).name}"
    )

    # Ensure expiration time was not modified
    info_after = AccountInfoQuery(account_id).execute(env.client)
    assert info_after.expiration_time == original_info.expiration_time


def _apply_tiny_max_fee_if_supported(tx, client) -> bool:
    # Try tx-level setters
    for attr in ("set_max_transaction_fee", "set_max_fee", "set_transaction_fee"):
        if hasattr(tx, attr):
            getattr(tx, attr)(Hbar.from_tinybars(1))
            return True
    # Try client-level default
    for attr in (
        "set_default_max_transaction_fee",
        "set_max_transaction_fee",
        "set_default_max_fee",
        "setMaxTransactionFee",
    ):
        if hasattr(client, attr):
            getattr(client, attr)(Hbar.from_tinybars(1))
            return True
    return False


@pytest.mark.integration
def test_account_update_insufficient_fee_with_valid_expiration_bump(env):
    """If we can cap the fee, a small valid expiration bump should fail with INSUFFICIENT_TX_FEE; otherwise skip."""
    # Create account
    receipt = (
        AccountCreateTransaction()
        .set_key(env.operator_key.public_key())
        .set_initial_balance(Hbar(1))
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS
    account_id = receipt.account_id

    # Use the account's *current* expiration and bump it slightly forward
    info = AccountInfoQuery(account_id).execute(env.client)
    base_expiry_secs = int(info.expiration_time.seconds)

    delta_seconds = 60 * 60 * 24  # +1 day; typically valid
    new_expiry = Timestamp(seconds=base_expiry_secs + delta_seconds, nanos=0)

    tx = AccountUpdateTransaction().set_account_id(account_id).set_expiration_time(new_expiry)

    if not _apply_tiny_max_fee_if_supported(tx, env.client):
        pytest.skip("SDK lacks a max-fee API; cannot deterministically trigger INSUFFICIENT_TX_FEE.")

    receipt = tx.execute(env.client)
    assert receipt.status == ResponseCode.INSUFFICIENT_TX_FEE, (
        f"Expected INSUFFICIENT_TX_FEE but got {ResponseCode(receipt.status).name}"
    )

    # Confirm expiration time did not change
    info_after = AccountInfoQuery(account_id).execute(env.client)
    assert int(info_after.expiration_time.seconds) == base_expiry_secs


@pytest.mark.integration
def test_integration_account_update_transaction_with_only_account_id(env):
    """Test that AccountUpdateTransaction can execute with only account ID set."""
    # Create initial account
    receipt = (
        AccountCreateTransaction()
        .set_key(env.operator_key.public_key())
        .set_initial_balance(Hbar(1))
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account creation failed with status: {ResponseCode(receipt.status).name}"
    )

    account_id = receipt.account_id
    assert account_id is not None, "Account ID should not be None"

    receipt = AccountUpdateTransaction().set_account_id(account_id).execute(env.client)

    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account update failed with status: {ResponseCode(receipt.status).name}"
    )

    # Ensure no fields were unintentionally modified
    info = AccountInfoQuery(account_id).execute(env.client)
    assert str(info.account_id) == str(account_id)
    assert info.key.to_bytes_raw() == env.operator_key.public_key().to_bytes_raw()


@pytest.mark.integration
def test_integration_account_update_transaction_with_max_automatic_token_associations(env):
    """Test updating max_automatic_token_associations and verifying it persists."""
    # Create initial account
    receipt = (
        AccountCreateTransaction()
        .set_key(env.operator_key.public_key())
        .set_initial_balance(Hbar(2))
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account creation failed with status: {ResponseCode(receipt.status).name}"
    )

    account_id = receipt.account_id
    assert account_id is not None, "Account ID should not be None"

    # Update max_automatic_token_associations
    new_max_associations = 100
    receipt = (
        AccountUpdateTransaction()
        .set_account_id(account_id)
        .set_max_automatic_token_associations(new_max_associations)
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account update failed with status: {ResponseCode(receipt.status).name}"
    )

    # Query account info to verify the update persisted
    info = AccountInfoQuery(account_id).execute(env.client)
    assert info.max_automatic_token_associations == new_max_associations, (
        "Max automatic token associations should be updated"
    )


@pytest.mark.integration
def test_integration_account_update_transaction_with_staking_fields(env):
    """Test updating staking fields (staked_account_id, decline_staking_reward)."""
    # Create two accounts - one to stake to
    receipt1 = (
        AccountCreateTransaction()
        .set_key_without_alias(env.operator_key.public_key())
        .set_initial_balance(Hbar(2))
        .execute(env.client)
    )
    assert receipt1.status == ResponseCode.SUCCESS
    staked_account_id = receipt1.account_id

    # Create account to update
    receipt2 = (
        AccountCreateTransaction()
        .set_key_without_alias(env.operator_key.public_key())
        .set_initial_balance(Hbar(2))
        .execute(env.client)
    )
    assert receipt2.status == ResponseCode.SUCCESS
    account_id = receipt2.account_id

    # Update with staking fields
    receipt = (
        AccountUpdateTransaction()
        .set_account_id(account_id)
        .set_staked_account_id(staked_account_id)
        .set_decline_staking_reward(True)
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS, (
        f"Account update with staking fields failed with status: {ResponseCode(receipt.status).name}"
    )

    # Verify staking info reflects the updated values
    info = AccountInfoQuery(account_id).execute(env.client)
    assert info.staking_info.staked_account_id == staked_account_id, "Staked account ID should match"
    assert info.staking_info.staked_node_id is None, "Staked node ID should be cleared when staking to an account"
    assert info.staking_info.decline_reward is True, "Decline staking reward should be true"


@pytest.mark.integration
def test_integration_account_update_transaction_with_staked_node_id(env):
    """Test updating with staked_node_id."""
    # Create account to update
    receipt = (
        AccountCreateTransaction()
        .set_key(env.operator_key.public_key())
        .set_initial_balance(Hbar(2))
        .execute(env.client)
    )
    assert receipt.status == ResponseCode.SUCCESS
    account_id = receipt.account_id

    # Update with staked_node_id (using node 0 as a test value)
    # Note: In a real scenario, you'd use a valid node ID
    receipt = AccountUpdateTransaction().set_account_id(account_id).set_staked_node_id(0).execute(env.client)
    # This might succeed or fail depending on network state, but should not crash
    assert receipt.status in [
        ResponseCode.SUCCESS,
        ResponseCode.INVALID_STAKING_ID,
    ], f"Unexpected status: {ResponseCode(receipt.status).name}"

    if receipt.status == ResponseCode.SUCCESS:
        info = AccountInfoQuery(account_id).execute(env.client)
        assert info.staking_info is not None, "Staking info should be set"
        assert info.staking_info.staked_node_id == 0
        assert info.staking_info.staked_account_id is None