diff --git a/src/lib.rs b/src/lib.rs
index 4fc3898..1ae4b74 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,7 +8,8 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use bdk_chain::{
-    local_chain, miniscript, tx_graph, Anchor, ConfirmationBlockTime, DescriptorExt, Merge,
+    local_chain, miniscript, tx_graph, Anchor, ConfirmationBlockTime, DescriptorExt, DescriptorId,
+    Merge,
 };
 use bdk_wallet::bitcoin::{
     self,
@@ -216,30 +217,23 @@ impl Store {
 
         if let Some(ref descriptor) = changeset.descriptor {
             insert_descriptor(&mut tx, wallet_name, descriptor, External).await?;
-            if let Some(last_revealed) = changeset
-                .indexer
-                .last_revealed
-                .get(&descriptor.descriptor_id())
-            {
-                update_last_revealed(&mut tx, wallet_name, *last_revealed, External).await?;
-            }
         }
 
-        if let Some(ref change_descriptor) = changeset.clone().change_descriptor {
+        if let Some(ref change_descriptor) = changeset.change_descriptor {
             insert_descriptor(&mut tx, wallet_name, change_descriptor, Internal).await?;
-            if let Some(last_revealed) = changeset
-                .indexer
-                .last_revealed
-                .get(&change_descriptor.descriptor_id())
-            {
-                update_last_revealed(&mut tx, wallet_name, *last_revealed, Internal).await?;
-            }
         }
 
         if let Some(network) = changeset.network {
             insert_network(&mut tx, wallet_name, network).await?;
         }
 
+        let last_revealed_indices = &changeset.indexer.last_revealed;
+        if !last_revealed_indices.is_empty() {
+            for (desc_id, index) in last_revealed_indices {
+                update_last_revealed(&mut tx, wallet_name, *desc_id, *index).await?;
+            }
+        }
+
         local_chain_changeset_persist_to_postgres(&mut tx, wallet_name, &changeset.local_chain)
             .await?;
         tx_graph_changeset_persist_to_postgres(&mut tx, wallet_name, &changeset.tx_graph).await?;
@@ -302,20 +296,17 @@ async fn insert_network(
 async fn update_last_revealed(
     tx: &mut Transaction<'_, Postgres>,
     wallet_name: &str,
+    descriptor_id: DescriptorId,
     last_revealed: u32,
-    keychain: KeychainKind,
 ) -> Result<(), BdkSqlxError> {
     info!("update last revealed");
-    let keychain = match keychain {
-        External => "External",
-        Internal => "Internal",
-    };
+
     sqlx::query(
-        "UPDATE keychain SET last_revealed = $1 WHERE wallet_name = $2 AND keychainkind = $3",
+        "UPDATE keychain SET last_revealed = $1 WHERE wallet_name = $2 AND descriptor_id = $3",
     )
     .bind(last_revealed as i32)
     .bind(wallet_name)
-    .bind(keychain)
+    .bind(descriptor_id.to_byte_array())
     .execute(&mut **tx)
     .await?;
 
diff --git a/src/test.rs b/src/test.rs
index 356ae28..594e7f1 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -1,8 +1,11 @@
 use crate::{drop_all, Store};
 use bdk_chain::bitcoin::secp256k1::Secp256k1;
 use bdk_chain::bitcoin::Network;
-use bdk_chain::bitcoin::Network::Signet;
-use bdk_wallet::{wallet_name_from_descriptor, KeychainKind, Wallet};
+use bdk_wallet::{
+    wallet_name_from_descriptor,
+    KeychainKind::{self, *},
+    Wallet,
+};
 use better_panic::Settings;
 use sqlx::PgPool;
 use std::env;
@@ -15,7 +18,7 @@ pub fn get_test_tr_single_sig_xprv_with_change_desc() -> (&'static str, &'static
         "tr(tprv8ZgxMBicQKsPdDArR4xSAECuVxeX1jwwSXR4ApKbkYgZiziDc4LdBy2WvJeGDfUSE4UT4hHhbgEwbdq8ajjUHiKDegkwrNU6V55CxcxonVN/1/*)")
 }
 
-const NETWORK: Network = Signet;
+const NETWORK: Network = Network::Signet;
 
 #[tracing::instrument]
 #[tokio::test]
@@ -34,16 +37,9 @@ async fn wallet_is_persisted() -> anyhow::Result<()> {
 
     // Set up the database URL (you might want to use a test-specific database)
     let url = env::var("DATABASE_TEST_URL").expect("DATABASE_TEST_URL must be set for tests");
-
     let pg = PgPool::connect(&url.clone()).await?;
-    match drop_all(pg).await {
-        Ok(_) => {
-            dbg!("tables dropped")
-        }
-        Err(_) => {
-            dbg!("Error dropping tables")
-        }
-    };
+    drop_all(pg).await?;
+    println!("tables dropped");
 
     // Define descriptors (you may need to adjust these based on your exact requirements)
     let (external_desc, internal_desc) = get_test_tr_single_sig_xprv_with_change_desc();
@@ -56,26 +52,24 @@ async fn wallet_is_persisted() -> anyhow::Result<()> {
     )?;
 
     // Create a new wallet
-    let wallet_spk_index = {
-        let mut store = Store::new_with_url(url.clone(), Some(wallet_name.clone())).await?;
-        let mut wallet = Wallet::create(external_desc, internal_desc)
-            .network(NETWORK)
-            .create_wallet_async(&mut store)
-            .await?;
-
-        let deposit_address = wallet.reveal_next_address(KeychainKind::External);
-        let change_address = wallet.reveal_next_address(KeychainKind::Internal);
-        dbg!(deposit_address.address);
-        dbg!(change_address.address);
+    let mut store = Store::new_with_url(url.clone(), Some(wallet_name.clone())).await?;
+    let mut wallet = Wallet::create(external_desc, internal_desc)
+        .network(NETWORK)
+        .create_wallet_async(&mut store)
+        .await?;
+
+    let external_addr0 = wallet.reveal_next_address(KeychainKind::External);
+    for keychain in [External, Internal] {
+        let _ = wallet.reveal_addresses_to(keychain, 2);
+    }
 
-        assert!(wallet.persist_async(&mut store).await?);
-        wallet.spk_index().clone()
-    };
+    assert!(wallet.persist_async(&mut store).await?);
+    let wallet_spk_index = wallet.spk_index();
 
     {
         // Recover the wallet
-        let mut store = Store::new_with_url(url.clone(), Some(wallet_name.clone())).await?;
-        let mut wallet = Wallet::load()
+        let mut store = Store::new_with_url(url.clone(), Some(wallet_name)).await?;
+        let wallet = Wallet::load()
             .descriptor(KeychainKind::External, Some(external_desc))
             .descriptor(KeychainKind::Internal, Some(internal_desc))
             .load_wallet_async(&mut store)
@@ -92,8 +86,8 @@ async fn wallet_is_persisted() -> anyhow::Result<()> {
             wallet_spk_index.last_revealed_indices()
         );
 
-        let recovered_address = wallet.reveal_next_address(KeychainKind::External);
-        println!("Recovered next address: {}", recovered_address.address);
+        let recovered_addr = wallet.peek_address(KeychainKind::External, 0);
+        assert_eq!(recovered_addr, external_addr0, "failed to recover address");
 
         assert_eq!(
             wallet.public_descriptor(KeychainKind::External).to_string(),
@@ -104,7 +98,7 @@ async fn wallet_is_persisted() -> anyhow::Result<()> {
     // Clean up (optional, depending on your test database strategy)
     // You might want to delete the test wallet from the database here
     let db = PgPool::connect(&url).await?;
-    drop_all(db).await.expect("hope its not mainet");
+    drop_all(db).await.expect("hope its not mainnet");
 
     Ok(())
 }