lssa/integration_tests/tests/account.rs

#![expect(
    clippy::tests_outside_test_module,
    reason = "We don't care about these in tests"
)]

use anyhow::{Context as _, Result};
use integration_tests::{TestContext, get_account, new_account, private_mention};
use key_protocol::key_management::KeyChain;
use lee::Data;
use lee_core::account::Nonce;
use log::info;
use sequencer_service_rpc::{ClientError, MAX_ACCOUNTS_PER_REQUEST, RpcClient as _};
use tokio::test;
use wallet::{
    account::{AccountIdWithPrivacy, HumanReadableAccount, Label},
    cli::{
        Command, SubcommandReturnValue,
        account::{AccountSubcommand, ImportSubcommand, NewSubcommand},
        execute_subcommand,
    },
};

#[test]
async fn get_existing_account() -> Result<()> {
    let ctx = TestContext::new().await?;

    let account = get_account(&ctx, ctx.existing_public_accounts()[0]).await?;

    assert_eq!(
        account.program_owner,
        programs::authenticated_transfer().id()
    );
    assert_eq!(account.balance, 10000);
    assert!(account.data.is_empty());
    assert_eq!(account.nonce.0, 1);

    info!("Successfully retrieved account with correct details");

    Ok(())
}

#[test]
async fn get_public_accounts_returns_states_in_input_order() -> Result<()> {
    let ctx = TestContext::new().await?;
    let account_ids = ctx.existing_public_accounts();
    let ordered_account_ids = vec![account_ids[1], account_ids[0], account_ids[1]];

    let empty_accounts = ctx.sequencer_client().get_accounts(Vec::new()).await?;
    let single_account = ctx
        .sequencer_client()
        .get_accounts(vec![account_ids[0]])
        .await?;
    let accounts = ctx
        .sequencer_client()
        .get_accounts(ordered_account_ids)
        .await?;
    let unknown_account = ctx
        .sequencer_client()
        .get_accounts(vec![lee::AccountId::new([0; 32])])
        .await?;
    let oversized_error = ctx
        .sequencer_client()
        .get_accounts(vec![account_ids[0]; MAX_ACCOUNTS_PER_REQUEST + 1])
        .await
        .expect_err("oversized account batches must be rejected");

    assert!(empty_accounts.is_empty());
    assert_eq!(single_account, vec![accounts[1].clone()]);
    assert_eq!(accounts.len(), 3);
    assert_eq!(accounts[0].balance, 20_000);
    assert_eq!(accounts[1].balance, 10_000);
    assert_eq!(accounts[2], accounts[0]);
    assert_eq!(unknown_account, vec![lee::Account::default()]);
    let ClientError::Call(oversized_error) = oversized_error else {
        panic!("expected an RPC call error for oversized batch")
    };
    assert_eq!(oversized_error.code(), -32602);
    assert!(
        oversized_error
            .message()
            .contains("Too many accounts requested")
    );

    Ok(())
}

#[test]
async fn new_public_account_with_label() -> Result<()> {
    let mut ctx = TestContext::new().await?;

    let label = Label::new("my-test-public-account");
    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
        cci: None,
        label: Some(label.clone()),
    }));

    let result = execute_subcommand(ctx.wallet_mut(), command).await?;

    // Extract the account_id from the result
    let wallet::cli::SubcommandReturnValue::RegisterAccount { account_id } = result else {
        panic!("Expected RegisterAccount return value")
    };

    // Verify the label was stored
    let resolved = ctx.wallet().storage().resolve_label(&label);

    assert_eq!(resolved, Some(AccountIdWithPrivacy::Public(account_id)));

    info!("Successfully created public account with label");

    Ok(())
}

#[test]
async fn add_label_to_existing_account() -> Result<()> {
    let mut ctx = TestContext::new().await?;

    let account_id = ctx.existing_private_accounts()[0];
    let label = Label::new("my-test-private-account");
    let command = Command::Account(AccountSubcommand::Label {
        account_id: private_mention(account_id),
        label: label.clone(),
    });

    execute_subcommand(ctx.wallet_mut(), command).await?;

    let resolved = ctx.wallet().storage().resolve_label(&label);

    assert_eq!(resolved, Some(AccountIdWithPrivacy::Private(account_id)));

    info!("Successfully set label on existing private account");

    Ok(())
}

#[test]
async fn new_public_account_without_label() -> Result<()> {
    let mut ctx = TestContext::new().await?;

    let account_id = new_account(&mut ctx, false, None).await?;

    // Verify no label was stored for the account id
    assert!(
        ctx.wallet()
            .storage()
            .labels_for_account(AccountIdWithPrivacy::Public(account_id))
            .next()
            .is_none(),
        "No label should be stored when not provided"
    );

    info!("Successfully created public account without label");

    Ok(())
}

#[test]
async fn import_public_account() -> Result<()> {
    let mut ctx = TestContext::new().await?;

    let private_key = lee::PrivateKey::new_os_random();
    let account_id = lee::AccountId::from(&lee::PublicKey::new_from_private_key(&private_key));

    let command = Command::Account(AccountSubcommand::Import(ImportSubcommand::Public {
        private_key,
    }));
    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::Empty = sub_ret else {
        anyhow::bail!("Expected Empty return value");
    };

    let imported_key = ctx
        .wallet()
        .storage()
        .key_chain()
        .pub_account_signing_key(account_id);
    assert!(
        imported_key.is_some(),
        "Imported public account should be present"
    );

    Ok(())
}

#[test]
async fn import_private_account() -> Result<()> {
    let mut ctx = TestContext::new().await?;

    let key_chain = KeyChain::new_os_random();
    let account_id = lee::AccountId::from((
        &key_chain.nullifier_public_key,
        &key_chain.viewing_public_key,
        0,
    ));
    let account = lee::Account {
        program_owner: programs::authenticated_transfer().id(),
        balance: 777,
        data: Data::default(),
        nonce: Nonce::default(),
    };

    let key_chain_json = serde_json::to_string(&key_chain)
        .context("Failed to serialize key chain for private import")?;
    let account_state = HumanReadableAccount::from(account.clone());

    let command = Command::Account(AccountSubcommand::Import(ImportSubcommand::Private {
        key_chain_json,
        account_state,
        chain_index: None,
        identifier: 0,
    }));
    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
    let SubcommandReturnValue::Empty = sub_ret else {
        anyhow::bail!("Expected Empty return value");
    };

    let imported_acc = ctx
        .wallet()
        .storage()
        .key_chain()
        .private_account(account_id)
        .context("Imported private account should be present")?;

    assert_eq!(
        imported_acc.key_chain.secret_spending_key,
        key_chain.secret_spending_key
    );
    assert_eq!(
        imported_acc.key_chain.nullifier_public_key,
        key_chain.nullifier_public_key
    );
    assert_eq!(
        imported_acc.key_chain.viewing_public_key,
        key_chain.viewing_public_key
    );

    assert_eq!(imported_acc.chain_index, None);

    assert_eq!(imported_acc.kind.identifier(), 0);

    assert_eq!(imported_acc.account, &account);

    Ok(())
}

#[test]
async fn import_private_account_second_time_overrides_account_data() -> Result<()> {
    let mut ctx = TestContext::new().await?;

    let key_chain = KeyChain::new_os_random();
    let account_id = lee::AccountId::from((
        &key_chain.nullifier_public_key,
        &key_chain.viewing_public_key,
        0,
    ));
    let key_chain_json =
        serde_json::to_string(&key_chain).context("Failed to serialize key chain")?;

    let initial_account = lee::Account {
        program_owner: programs::authenticated_transfer().id(),
        balance: 100,
        data: Data::default(),
        nonce: Nonce::default(),
    };

    // First import
    wallet::cli::execute_subcommand(
        ctx.wallet_mut(),
        Command::Account(AccountSubcommand::Import(ImportSubcommand::Private {
            key_chain_json: key_chain_json.clone(),
            account_state: HumanReadableAccount::from(initial_account),
            chain_index: None,
            identifier: 0,
        })),
    )
    .await?;

    let updated_account = lee::Account {
        program_owner: programs::authenticated_transfer().id(),
        balance: 999,
        data: Data::default(),
        nonce: Nonce::default(),
    };

    // Second import with different account data (same key chain)
    wallet::cli::execute_subcommand(
        ctx.wallet_mut(),
        Command::Account(AccountSubcommand::Import(ImportSubcommand::Private {
            key_chain_json,
            account_state: HumanReadableAccount::from(updated_account.clone()),
            chain_index: None,
            identifier: 0,
        })),
    )
    .await?;

    let imported = ctx
        .wallet()
        .storage()
        .key_chain()
        .private_account(account_id)
        .context("Imported private account should be present")?;

    assert_eq!(
        imported.account, &updated_account,
        "Second import should override account data"
    );

    Ok(())
}