Add tests

Cargo.toml

@@ -61,5 +61,8 @@ features = ["ff_derive"]
 [build-dependencies]
 tower-grpc-build = { git = "https://github.com/tower-rs/tower-grpc",  features = ["tower-hyper"]  }
 
+[dev-dependencies]
+rand_core = "0.5.1"
+
 [profile.release]
 debug = true

src/lightwallet/mod.rs

@@ -640,10 +640,10 @@ impl LightWallet {
                                 continue;
                             }
 
+                            // Update the WalletTx 
+                            info!("A sapling output was sent in {}", tx.txid());
                             {
-                                // Update the WalletTx 
                                 // Do it in a short scope because of the write lock.
-                                info!("A sapling output was sent in {}", tx.txid());
                                 
                                 let mut txs = self.txs.write().unwrap();
                                 if txs.get(&tx.txid()).unwrap().outgoing_metadata.iter()
@@ -1078,3 +1078,220 @@ impl LightWallet {
         Some(raw_tx.into_boxed_slice())
     }
 }
+
+
+
+#[cfg(test)]
+pub mod tests {
+    use std::convert::TryInto;
+    use ff::{Field, PrimeField, PrimeFieldRepr};
+    use pairing::bls12_381::Bls12;
+    use rand_core::{RngCore, OsRng};
+    use protobuf::Message;
+    use zcash_client_backend::{encoding::encode_payment_address,
+        proto::compact_formats::{
+            CompactBlock, CompactOutput, CompactSpend, CompactTx,
+        }
+    };
+    use zcash_primitives::{
+        block::BlockHash,
+        jubjub::fs::Fs,
+        note_encryption::{Memo, SaplingNoteEncryption},
+        primitives::{Note, PaymentAddress},
+        transaction::components::Amount,
+        zip32::{ExtendedFullViewingKey, ExtendedSpendingKey},
+        JUBJUB,
+    };
+
+    use crate::lightwallet::LightWallet;
+    use crate::LightClientConfig;
+
+    /// Create a fake CompactBlock at the given height, containing a single output paying
+    /// the given address. Returns the CompactBlock and the nullifier for the new note.
+    pub(crate) fn fake_compact_block(
+        height: i32,
+        prev_hash: BlockHash,
+        extfvk: ExtendedFullViewingKey,
+        value: Amount,
+    ) -> (CompactBlock, Vec<u8>) {
+        let to = extfvk.default_address().unwrap().1;
+
+        // Create a fake Note for the account
+        let mut rng = OsRng;
+        let note = Note {
+            g_d: to.diversifier.g_d::<Bls12>(&JUBJUB).unwrap(),
+            pk_d: to.pk_d.clone(),
+            value: value.into(),
+            r: Fs::random(&mut rng),
+        };
+        let encryptor = SaplingNoteEncryption::new(
+            extfvk.fvk.ovk,
+            note.clone(),
+            to.clone(),
+            Memo::default(),
+            &mut rng,
+        );
+        let mut cmu = vec![];
+        note.cm(&JUBJUB).into_repr().write_le(&mut cmu).unwrap();
+        let mut epk = vec![];
+        encryptor.epk().write(&mut epk).unwrap();
+        let enc_ciphertext = encryptor.encrypt_note_plaintext();
+
+        // Create a fake CompactBlock containing the note
+        let mut cout = CompactOutput::new();
+        cout.set_cmu(cmu);
+        cout.set_epk(epk);
+        cout.set_ciphertext(enc_ciphertext[..52].to_vec());
+        let mut ctx = CompactTx::new();
+        let mut txid = vec![0; 32];
+        rng.fill_bytes(&mut txid);
+        ctx.set_hash(txid);
+        ctx.outputs.push(cout);
+        let mut cb = CompactBlock::new();
+        cb.set_height(height as u64);
+        cb.hash.resize(32, 0);
+        rng.fill_bytes(&mut cb.hash);
+        cb.prevHash.extend_from_slice(&prev_hash.0);
+        cb.vtx.push(ctx);
+        (cb, note.nf(&extfvk.fvk.vk, 0, &JUBJUB))
+    }
+
+    /// Create a fake CompactBlock at the given height, spending a single note from the
+    /// given address.
+    pub(crate) fn fake_compact_block_spending(
+        height: i32,
+        prev_hash: BlockHash,
+        (nf, in_value): (Vec<u8>, Amount),
+        extfvk: ExtendedFullViewingKey,
+        to: PaymentAddress<Bls12>,
+        value: Amount,
+    ) -> CompactBlock {
+        let mut rng = OsRng;
+
+        // Create a fake CompactBlock containing the note
+        let mut cspend = CompactSpend::new();
+        cspend.set_nf(nf);
+        let mut ctx = CompactTx::new();
+        let mut txid = vec![0; 32];
+        rng.fill_bytes(&mut txid);
+        ctx.set_hash(txid);
+        ctx.spends.push(cspend);
+
+        // Create a fake Note for the payment
+        ctx.outputs.push({
+            let note = Note {
+                g_d: to.diversifier.g_d::<Bls12>(&JUBJUB).unwrap(),
+                pk_d: to.pk_d.clone(),
+                value: value.into(),
+                r: Fs::random(&mut rng),
+            };
+            let encryptor = SaplingNoteEncryption::new(
+                extfvk.fvk.ovk,
+                note.clone(),
+                to,
+                Memo::default(),
+                &mut rng,
+            );
+            let mut cmu = vec![];
+            note.cm(&JUBJUB).into_repr().write_le(&mut cmu).unwrap();
+            let mut epk = vec![];
+            encryptor.epk().write(&mut epk).unwrap();
+            let enc_ciphertext = encryptor.encrypt_note_plaintext();
+
+            let mut cout = CompactOutput::new();
+            cout.set_cmu(cmu);
+            cout.set_epk(epk);
+            cout.set_ciphertext(enc_ciphertext[..52].to_vec());
+            cout
+        });
+
+        // Create a fake Note for the change
+        ctx.outputs.push({
+            let change_addr = extfvk.default_address().unwrap().1;
+            let note = Note {
+                g_d: change_addr.diversifier.g_d::<Bls12>(&JUBJUB).unwrap(),
+                pk_d: change_addr.pk_d.clone(),
+                value: (in_value - value).into(),
+                r: Fs::random(&mut rng),
+            };
+            let encryptor = SaplingNoteEncryption::new(
+                extfvk.fvk.ovk,
+                note.clone(),
+                change_addr,
+                Memo::default(),
+                &mut rng,
+            );
+            let mut cmu = vec![];
+            note.cm(&JUBJUB).into_repr().write_le(&mut cmu).unwrap();
+            let mut epk = vec![];
+            encryptor.epk().write(&mut epk).unwrap();
+            let enc_ciphertext = encryptor.encrypt_note_plaintext();
+
+            let mut cout = CompactOutput::new();
+            cout.set_cmu(cmu);
+            cout.set_epk(epk);
+            cout.set_ciphertext(enc_ciphertext[..52].to_vec());
+            cout
+        });
+
+        let mut cb = CompactBlock::new();
+        cb.set_height(height as u64);
+        cb.hash.resize(32, 0);
+        rng.fill_bytes(&mut cb.hash);
+        cb.prevHash.extend_from_slice(&prev_hash.0);
+        cb.vtx.push(ctx);
+        cb
+    }
+
+    #[test]
+    fn z_balances() {
+        let wallet = LightWallet::new(None, &LightClientConfig {
+            server: "0.0.0.0:0".to_string(),
+            chain_name: "test".to_string(),
+            sapling_activation_height: 0
+        }).unwrap();
+
+        const AMOUNT1:u64 = 5;
+        // Address is encoded in bech32
+        let address = Some(encode_payment_address(wallet.config.hrp_sapling_address(), 
+                                            &wallet.extfvks[0].default_address().unwrap().1));
+
+        let (cb1, _) = fake_compact_block(
+            0,
+            BlockHash([0; 32]),
+            wallet.extfvks[0].clone(),
+            Amount::from_u64(AMOUNT1).unwrap(),
+        );
+        
+
+        // Make sure that the intial state is empty
+        assert_eq!(wallet.txs.read().unwrap().len(), 0);
+        assert_eq!(wallet.blocks.read().unwrap().len(), 0);
+        assert_eq!(wallet.zbalance(None), 0);
+        assert_eq!(wallet.zbalance(address.clone()), 0);
+
+        wallet.scan_block(&cb1.write_to_bytes().unwrap()).unwrap();
+        
+        assert_eq!(wallet.txs.read().unwrap().len(), 1);
+        assert_eq!(wallet.blocks.read().unwrap().len(), 1);
+        assert_eq!(wallet.zbalance(None), AMOUNT1);
+        assert_eq!(wallet.zbalance(address.clone()), AMOUNT1);
+
+        const AMOUNT2:u64 = 10;
+
+        // Add a second block
+        let (cb2, _) = fake_compact_block(
+            1,  // Block number 1
+            BlockHash{0: cb1.hash[..].try_into().unwrap()},
+            wallet.extfvks[0].clone(),
+            Amount::from_u64(AMOUNT2).unwrap(),
+        );
+
+        wallet.scan_block(&cb2.write_to_bytes().unwrap()).unwrap();
+        
+        assert_eq!(wallet.txs.read().unwrap().len(), 2);
+        assert_eq!(wallet.blocks.read().unwrap().len(), 2);
+        assert_eq!(wallet.zbalance(None), AMOUNT1 + AMOUNT2);
+        assert_eq!(wallet.zbalance(address.clone()), AMOUNT1 + AMOUNT2);
+    }
+}