misc: 🍱 add a mock consensus test exercising penumbra_wallet::plan::sweep()

Cargo.toml

@@ -193,6 +193,7 @@ penumbra-test-subscriber         = { path = "crates/test/tracing-subscriber" }
 penumbra-transaction             = { default-features = false, path = "crates/core/transaction" }
 penumbra-txhash                  = { default-features = false, path = "crates/core/txhash" }
 penumbra-view                    = { path = "crates/view" }
+penumbra-wallet                  = { path = "crates/wallet" }
 penumbra-extension               = { path = "crates/penumbra-extension", default-features = false }
 pin-project                      = { version = "1.0.12" }
 pin-project-lite                 = { version = "0.2.9" }

crates/bin/pcli/Cargo.toml

@@ -81,7 +81,7 @@ penumbra-stake = {workspace = true, default-features = false}
 penumbra-tct = {workspace = true, default-features = true}
 penumbra-transaction = {workspace = true, default-features = true}
 penumbra-view = {workspace = true}
-penumbra-wallet = { path = "../../wallet" }
+penumbra-wallet = {workspace = true}
 pin-project = {workspace = true}
 rand = {workspace = true}
 rand_chacha = {workspace = true}

crates/core/app/Cargo.toml

@@ -93,6 +93,7 @@ penumbra-proto                   = { workspace = true, features = ["box-grpc"] }
 penumbra-test-subscriber         = { workspace = true }
 penumbra-mock-tendermint-proxy   = { workspace = true }
 penumbra-view                    = { workspace = true }
+penumbra-wallet                  = { workspace = true }
 rand                             = { workspace = true }
 rand_chacha                      = { workspace = true }
 rand_core                        = { workspace = true }

crates/core/app/tests/app_can_sweep_a_collection_of_small_notes.rs

@@ -0,0 +1,192 @@
+use {
+    anyhow::Context,
+    cnidarium::TempStorage,
+    penumbra_app::{
+        genesis::{AppState, Content},
+        server::consensus::Consensus,
+    },
+    penumbra_asset::{STAKING_TOKEN_ASSET_ID, STAKING_TOKEN_DENOM},
+    penumbra_keys::{keys::AddressIndex, test_keys},
+    penumbra_mock_client::MockClient,
+    penumbra_mock_consensus::TestNode,
+    penumbra_proto::{
+        view::v1::{
+            view_service_client::ViewServiceClient, view_service_server::ViewServiceServer,
+            StatusRequest, StatusResponse,
+        },
+        DomainType,
+    },
+    penumbra_shielded_pool::genesis::Allocation,
+    penumbra_view::ViewClient,
+    penumbra_wallet::plan::SWEEP_COUNT,
+    rand_core::OsRng,
+    std::ops::Deref,
+    tap::{Tap, TapFallible},
+};
+
+mod common;
+
+/// The number of notes placed in the test wallet at genesis.
+//  note: when debugging, it can help to set this to a lower value.
+const COUNT: usize = SWEEP_COUNT + 1;
+
+/// Exercises that the app can process a "sweep", consolidating small notes.
+//  NB: a multi-thread runtime is needed to run both the view server and its client.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn app_can_sweep_a_collection_of_small_notes() -> anyhow::Result<()> {
+    // Install a test logger, and acquire some temporary storage.
+    let guard = common::set_tracing_subscriber_with_env_filter("info".into());
+    let storage = TempStorage::new().await?;
+
+    // Instantiate a mock tendermint proxy, which we will connect to the test node.
+    let proxy = penumbra_mock_tendermint_proxy::TestNodeProxy::new::<Consensus>();
+
+    // Define allocations to the test address, as many small notes.
+    let allocations = {
+        let dust = Allocation {
+            raw_amount: 1_u128.into(),
+            raw_denom: STAKING_TOKEN_DENOM.deref().base_denom().denom,
+            address: test_keys::ADDRESS_0.to_owned(),
+        };
+        std::iter::repeat(dust).take(COUNT).collect()
+    };
+
+    // Define our application state, and start the test node.
+    let mut test_node = {
+        let content = Content {
+            chain_id: TestNode::<()>::CHAIN_ID.to_string(),
+            shielded_pool_content: penumbra_shielded_pool::genesis::Content {
+                allocations,
+                ..Default::default()
+            },
+            ..Default::default()
+        };
+        let app_state = AppState::Content(content);
+        let app_state = serde_json::to_vec(&app_state).unwrap();
+        let consensus = Consensus::new(storage.as_ref().clone());
+        TestNode::builder()
+            .single_validator()
+            .app_state(app_state)
+            .on_block(proxy.on_block_callback())
+            .init_chain(consensus)
+            .await
+            .tap_ok(|e| tracing::info!(hash = %e.last_app_hash_hex(), "finished init chain"))?
+    };
+
+    // Sync the mock client, using the test wallet's spend key, to the latest snapshot.
+    let mut client = MockClient::new(test_keys::SPEND_KEY.clone())
+        .with_sync_to_storage(&storage)
+        .await?
+        .tap(
+            |c| tracing::info!(client.notes = %c.notes.len(), "mock client synced to test storage"),
+        );
+
+    // Jump ahead a few blocks.
+    test_node
+        .fast_forward(10)
+        .tap(|_| tracing::debug!("fast forwarding past genesis"))
+        .await?;
+
+    let grpc_url = "http://127.0.0.1:8081" // see #4517
+        .parse::<url::Url>()?
+        .tap(|url| tracing::debug!(%url, "parsed grpc url"));
+
+    // Spawn the server-side view server.
+    {
+        let make_svc = penumbra_app::rpc::router(
+            storage.as_ref(),
+            proxy,
+            false, /*enable_expensive_rpc*/
+        )?
+        .into_router()
+        .layer(tower_http::cors::CorsLayer::permissive())
+        .into_make_service()
+        .tap(|_| tracing::debug!("initialized rpc service"));
+        let [addr] = grpc_url
+            .socket_addrs(|| None)?
+            .try_into()
+            .expect("grpc url can be turned into a socket address");
+        let server = axum_server::bind(addr).serve(make_svc);
+        tokio::spawn(async { server.await.expect("grpc server returned an error") })
+            .tap(|_| tracing::debug!("grpc server is running"))
+    };
+
+    // Spawn the client-side view server...
+    let view_server = {
+        penumbra_view::ViewServer::load_or_initialize(
+            None::<&camino::Utf8Path>,
+            &*test_keys::FULL_VIEWING_KEY,
+            grpc_url,
+        )
+        .await
+        // TODO(kate): the goal is to communicate with the `ViewServiceServer`.
+        .map(ViewServiceServer::new)
+        .context("initializing view server")?
+    };
+
+    // Create a view client, and get the test wallet's notes.
+    let mut view_client = ViewServiceClient::new(view_server);
+
+    // Sync the view client to the chain.
+    {
+        use futures::StreamExt;
+        let mut status_stream = ViewClient::status_stream(&mut view_client).await?;
+        while let Some(status) = status_stream.next().await.transpose()? {
+            tracing::info!(?status, "view client received status stream response");
+        }
+        // Confirm that the status is as expected: synced up to the 11th block.
+        let status = view_client.status(StatusRequest {}).await?.into_inner();
+        debug_assert_eq!(
+            status,
+            StatusResponse {
+                full_sync_height: 10,
+                partial_sync_height: 10,
+                catching_up: false,
+            }
+        );
+    }
+
+    client.sync_to_latest(storage.latest_snapshot()).await?;
+    debug_assert_eq!(
+        client
+            .notes_by_asset(STAKING_TOKEN_ASSET_ID.deref().to_owned())
+            .count(),
+        COUNT,
+        "client wallet should have {COUNT} notes before sweeping"
+    );
+
+    loop {
+        let plans = penumbra_wallet::plan::sweep(&mut view_client, OsRng)
+            .await
+            .context("constructing sweep plans")?;
+        if plans.is_empty() {
+            break;
+        }
+        for plan in plans {
+            let tx = client.witness_auth_build(&plan).await?;
+            test_node
+                .block()
+                .with_data(vec![tx.encode_to_vec()])
+                .execute()
+                .await?;
+        }
+    }
+
+    let post_sweep_notes = view_client.unspent_notes_by_address_and_asset().await?;
+
+    client.sync_to_latest(storage.latest_snapshot()).await?;
+    assert_eq!(
+        post_sweep_notes
+            .get(&AddressIndex::from(0))
+            .expect("test wallet could not find any notes")
+            .get(&*STAKING_TOKEN_ASSET_ID)
+            .map(Vec::len),
+        Some(2),
+        "destination address should have collected {SWEEP_COUNT} notes into one note"
+    );
+
+    Ok(())
+        .tap(|_| drop(test_node))
+        .tap(|_| drop(storage))
+        .tap(|_| drop(guard))
+}

crates/core/app/tests/view_server_can_be_served_on_localhost.rs

@@ -24,6 +24,7 @@ use {
 
 mod common;
 
+// NB: a multi-thread runtime is needed to run both the view server and its client.
 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
 async fn view_server_can_be_served_on_localhost() -> anyhow::Result<()> {
     // Install a test logger, acquire some temporary storage, and start the test node.
@@ -62,7 +63,7 @@ async fn view_server_can_be_served_on_localhost() -> anyhow::Result<()> {
         .tap(|_| tracing::debug!("fast forwarding past genesis"))
         .await?;
 
-    let grpc_url = "http://127.0.0.1:8080"
+    let grpc_url = "http://127.0.0.1:8080" // see #4517
         .parse::<url::Url>()?
         .tap(|url| tracing::debug!(%url, "parsed grpc url"));
 

crates/test/mock-client/src/lib.rs

@@ -1,10 +1,10 @@
 use anyhow::Error;
 use cnidarium::StateRead;
 use penumbra_compact_block::{component::StateReadExt as _, CompactBlock, StatePayload};
-use penumbra_dex::swap::SwapPlaintext;
+use penumbra_dex::{swap::SwapPlaintext, swap_claim::SwapClaimPlan};
 use penumbra_keys::{keys::SpendKey, FullViewingKey};
 use penumbra_sct::component::{clock::EpochRead, tree::SctRead};
-use penumbra_shielded_pool::{note, Note};
+use penumbra_shielded_pool::{note, Note, SpendPlan};
 use penumbra_tct as tct;
 use penumbra_transaction::{AuthorizationData, Transaction, TransactionPlan, WitnessData};
 use rand_core::OsRng;
@@ -171,20 +171,25 @@ impl MockClient {
     }
 
     pub fn witness_plan(&self, plan: &TransactionPlan) -> Result<WitnessData, Error> {
+        let spend_commitment = |spend: &SpendPlan| spend.note.commit();
+        let spends = plan.spend_plans().map(spend_commitment);
+
+        let swap_claim_commitment = |swap: &SwapClaimPlan| swap.swap_plaintext.swap_commitment();
+        let swap_claims = plan.swap_claim_plans().map(swap_claim_commitment);
+
+        let witness = |commitment| {
+            self.sct
+                .witness(commitment)
+                .ok_or_else(|| anyhow::anyhow!("note commitment {commitment:?} unknown to client"))
+                .map(|proof| (commitment, proof))
+        };
+
         Ok(WitnessData {
             anchor: self.sct.root(),
-            // TODO: this will only witness spends, not other proofs like swaps
-            state_commitment_proofs: plan
-                .spend_plans()
-                .map(|spend| {
-                    let nc = spend.note.commit();
-                    Ok((
-                        nc,
-                        self.sct.witness(nc).ok_or_else(|| {
-                            anyhow::anyhow!("note commitment {:?} unknown to client", nc)
-                        })?,
-                    ))
-                })
+            // TODO: this will only witness spends and swap claims, but not other proofs
+            state_commitment_proofs: spends
+                .chain(swap_claims)
+                .map(witness)
                 .collect::<Result<_, Error>>()?,
         })
     }

crates/wallet/src/plan.rs

@@ -13,6 +13,8 @@ use penumbra_transaction::{TransactionParameters, TransactionPlan};
 pub use penumbra_view::Planner;
 use penumbra_view::{SpendableNoteRecord, ViewClient};
 
+pub const SWEEP_COUNT: usize = 8;
+
 #[instrument(skip(view, rng))]
 pub async fn sweep<V, R>(view: &mut V, mut rng: R) -> anyhow::Result<Vec<TransactionPlan>>
 where
@@ -32,7 +34,7 @@ where
 }
 
 #[instrument(skip(view, rng))]
-pub async fn claim_unclaimed_swaps<V, R>(
+async fn claim_unclaimed_swaps<V, R>(
     view: &mut V,
     mut rng: R,
 ) -> anyhow::Result<Vec<TransactionPlan>>
@@ -84,13 +86,11 @@ where
 }
 
 #[instrument(skip(view, rng))]
-pub async fn sweep_notes<V, R>(view: &mut V, mut rng: R) -> anyhow::Result<Vec<TransactionPlan>>
+async fn sweep_notes<V, R>(view: &mut V, mut rng: R) -> anyhow::Result<Vec<TransactionPlan>>
 where
     V: ViewClient,
     R: RngCore + CryptoRng,
 {
-    const SWEEP_COUNT: usize = 8;
-
     let gas_prices = view.gas_prices().await?;
 
     let all_notes = view