From 47837654ceffbb69a89bc5d203652d6176759e0f Mon Sep 17 00:00:00 2001
From: Tommy Volk <tvolk131@gmail.com>
Date: Sun, 29 Sep 2024 15:00:13 -0500
Subject: [PATCH] chore: cleanup

---
 src/fedimint.rs              | 16 +++++++++++-----
 src/routes/bitcoin_wallet.rs | 29 +++++++++++++++++------------
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/src/fedimint.rs b/src/fedimint.rs
index 51eb486..e5998cc 100644
--- a/src/fedimint.rs
+++ b/src/fedimint.rs
@@ -150,7 +150,10 @@ impl Wallet {
         Ok(())
     }
 
-    pub async fn join_federation(&self, invite_code: InviteCode) -> anyhow::Result<()> {
+    pub async fn join_federation(
+        &self,
+        invite_code: InviteCode,
+    ) -> anyhow::Result<BTreeMap<FederationId, FederationView>> {
         // Note: We're intentionally locking the clients mutex earlier than
         // necessary so that the lock is held while we're accessing the data directory.
         let mut clients = self.clients.lock().await;
@@ -163,7 +166,7 @@ impl Wallet {
 
         // Short-circuit if we're already connected to this federation.
         if federation_data_dir.is_dir() {
-            return Ok(());
+            return Ok(Self::get_current_state(clients).await);
         }
 
         let db: Database = RocksDb::open(federation_data_dir)?.into();
@@ -172,14 +175,17 @@ impl Wallet {
 
         clients.insert(federation_id, client);
 
-        Ok(())
+        Ok(Self::get_current_state(clients).await)
     }
 
     // TODO: Call `ClientModule::leave()` for every module.
     // https://docs.rs/fedimint-client/0.4.2/fedimint_client/module/trait.ClientModule.html#method.leave
     // Currently it isn't implemented for the `LightningClientModule`, so for now we're just checking
     // that the client has a zero balance.
-    pub async fn leave_federation(&self, federation_id: FederationId) -> anyhow::Result<()> {
+    pub async fn leave_federation(
+        &self,
+        federation_id: FederationId,
+    ) -> anyhow::Result<BTreeMap<FederationId, FederationView>> {
         // Note: We're intentionally locking the clients mutex earlier than
         // necessary so that the lock is held while we're accessing the data directory.
         let mut clients = self.clients.lock().await;
@@ -206,7 +212,7 @@ impl Wallet {
             }
         }
 
-        Ok(())
+        Ok(Self::get_current_state(clients).await)
     }
 
     async fn get_current_state(
diff --git a/src/routes/bitcoin_wallet.rs b/src/routes/bitcoin_wallet.rs
index 6e91639..3ee2282 100644
--- a/src/routes/bitcoin_wallet.rs
+++ b/src/routes/bitcoin_wallet.rs
@@ -43,7 +43,7 @@ pub enum Message {
     ConnectedToFederation,
 
     LeaveFederation(FederationId),
-    LeftFederation(FederationId),
+    LeftFederation((FederationId, BTreeMap<FederationId, FederationView>)),
     FailedToLeaveFederation((FederationId, Arc<anyhow::Error>)),
 
     Send(send::Message),
@@ -164,11 +164,12 @@ impl Page {
                 let wallet = self.connected_state.wallet.clone();
 
                 Task::future(async move {
-                    wallet.join_federation(invite_code).await.unwrap();
-                    app::Message::Routes(super::Message::BitcoinWalletPage(
-                        Message::ConnectedToFederation,
-                    ))
+                    let federation_views = wallet.join_federation(invite_code).await.unwrap();
+                    app::Message::UpdateFederationViews(federation_views)
                 })
+                .chain(Task::done(app::Message::Routes(
+                    super::Message::BitcoinWalletPage(Message::ConnectedToFederation),
+                )))
             }
             Message::ConnectedToFederation => {
                 // TODO: Do something here, or remove `ConnectedToFederation` message variant.
@@ -180,21 +181,25 @@ impl Page {
 
                 Task::future(async move {
                     match wallet.leave_federation(federation_id).await {
-                        Ok(()) => app::Message::Routes(super::Message::BitcoinWalletPage(
-                            Message::LeftFederation(federation_id),
-                        )),
+                        Ok(federation_views) => {
+                            app::Message::Routes(super::Message::BitcoinWalletPage(
+                                Message::LeftFederation((federation_id, federation_views)),
+                            ))
+                        }
                         Err(err) => app::Message::Routes(super::Message::BitcoinWalletPage(
                             Message::FailedToLeaveFederation((federation_id, Arc::from(err))),
                         )),
                     }
                 })
             }
-            Message::LeftFederation(_federation_id) => {
+            Message::LeftFederation((_federation_id, federation_views)) => {
                 // TODO: Display toast message.
 
-                Task::done(app::Message::Routes(super::Message::Navigate(
-                    RouteName::BitcoinWallet(SubrouteName::List),
-                )))
+                Task::done(app::Message::UpdateFederationViews(federation_views)).chain(Task::done(
+                    app::Message::Routes(super::Message::Navigate(RouteName::BitcoinWallet(
+                        SubrouteName::List,
+                    ))),
+                ))
             }
             Message::FailedToLeaveFederation((_federation_id, _err)) => {
                 // TODO: Implement this (including displaying a toast message).