Skip to content

Commit

Permalink
rusk-wallet: Aclimate to new phoenix_balance method
Browse files Browse the repository at this point in the history
  • Loading branch information
Daksh14 committed Sep 5, 2024
1 parent c933c41 commit 0f383bd
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 5 deletions.
14 changes: 10 additions & 4 deletions rusk-wallet/src/wallet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -356,15 +356,21 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
}

let index = addr.index()?;
let notes: Vec<Note> = state
let notes: Vec<NoteLeaf> = state
.fetch_notes(addr.pk())?
.into_iter()
.map(|data| data.note)
.map(|data| NoteLeaf {
note: data.note,
block_height: data.height,
})
.collect();

let seed = self.store.get_seed();

Ok(phoenix_balance(&derive_phoenix_vk(seed, index), notes))
Ok(phoenix_balance(
&derive_phoenix_vk(seed, index),
notes.into_iter(),
))
}

/// Creates a new public address.
Expand Down Expand Up @@ -679,7 +685,7 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
let tx = state.prove_and_propagate(withdraw);

sender_sk.zeroize();
stake_sk.zerorize();
stake_sk.zeroize();

tx
}
Expand Down
18 changes: 17 additions & 1 deletion wallet-core/tests/notes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ fn test_map_owned() {
PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng));

let value = 42;
let note_leaves = vec![
let note_leaves: Vec<Note> = vec![
gen_note(&mut rng, true, &owner_1_pks[0], value), // owner 1
gen_note(&mut rng, true, &owner_1_pks[1], value), // owner 1
gen_note(&mut rng, true, &owner_2_pks[0], value), // owner 2
Expand All @@ -82,6 +82,14 @@ fn test_map_owned() {
gen_note(&mut rng, true, &owner_3_pk, value), // owner 3
];

let note_leaves: Vec<NoteLeaf> = note_leaves
.into_iter()
.map(|note| NoteLeaf {
note,
block_height: 0,
})
.collect();

// notes with idx 0, 1 and 4 are owned by owner_1
let notes_by_1 = map_owned(&owner_1_sks, &note_leaves);
assert_eq!(notes_by_1.len(), 3);
Expand Down Expand Up @@ -143,6 +151,14 @@ fn test_balance() {
spendable: 34,
};

let notes: Vec<NoteLeaf> = notes
.into_iter()
.map(|note| NoteLeaf {
note,
block_height: 0,
})
.collect();

assert_eq!(
phoenix_balance(&(&owner_sk).into(), notes.iter()),
expected_balance
Expand Down

0 comments on commit 0f383bd

Please sign in to comment.