Refactor wallet removal and add additional test

chrisguida · Jan 16, 2024 · 6b3d264 · 6b3d264
1 parent 40b361e
commit 6b3d264
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 30 deletions.
diff --git a/src/main.rs b/src/main.rs
@@ -435,36 +435,32 @@ async fn delete(
     let db_dir_path = plugin.state().lock().await.db_dir.clone();
     let wallets = &mut plugin.state().lock().await.wallets;
 
-    if wallets.contains_key(&descriptor_name) {
-        let removed_item = wallets.remove(&descriptor_name);
-        let db_file_path = match removed_item {
-            Some(dw) => {
-                match dw.get_db_file_path(db_dir_path) {
-                    Ok(dw) => dw,
-                    Err(e) => return Err(e),
-                }
-            },
-            None => "Removed wallet was None!".to_string(),
-        };
-        fs::remove_file(db_file_path.clone())?;
-        log::debug!("Deleted smaug db file at {}", db_file_path);
-        let rpc_file = plugin.configuration().rpc_file;
-        let p = Path::new(&rpc_file);
-
-        let mut rpc = ClnRpc::new(p).await?;
-        let _ds_response = rpc
-            .call(Request::Datastore(DatastoreRequest {
-                key: vec!["smaug".to_owned()],
-                string: Some(json!(wallets).to_string()),
-                hex: None,
-                mode: Some(DatastoreMode::CREATE_OR_REPLACE),
-                generation: None,
-            }))
-            .await
-            .map_err(|e| anyhow!("Error calling listdatastore: {:?}", e))?;
-    } else {
-        return Err(anyhow!("can't find wallet {}", descriptor_name));
-    }
+    let removed_item = wallets.remove(&descriptor_name);
+    let db_file_path = match removed_item {
+        Some(dw) => {
+            match dw.get_db_file_path(db_dir_path) {
+                Ok(dw) => dw,
+                Err(e) => return Err(e),
+            }
+        },
+        None => return Err(anyhow!("Removal of wallet '{}' failed.", descriptor_name)),
+    };
+    fs::remove_file(db_file_path.clone())?;
+    log::debug!("Deleted smaug db file at {}", db_file_path);
+    let rpc_file = plugin.configuration().rpc_file;
+    let p = Path::new(&rpc_file);
+
+    let mut rpc = ClnRpc::new(p).await?;
+    let _ds_response = rpc
+        .call(Request::Datastore(DatastoreRequest {
+            key: vec!["smaug".to_owned()],
+            string: Some(json!(wallets).to_string()),
+            hex: None,
+            mode: Some(DatastoreMode::CREATE_OR_REPLACE),
+            generation: None,
+        }))
+        .await
+        .map_err(|e| anyhow!("Error calling listdatastore: {:?}", e))?;
 
     Ok(json!(format!("Deleted wallet: {}", descriptor_name)))
 }

diff --git a/tests/test_rpc.py b/tests/test_rpc.py
@@ -3,6 +3,7 @@
 from conftest import SMAUG_PLUGIN
 from fixtures import *
 from pyln.client import Millisatoshi
+from pyln.client.lightning import RpcError
 from pyln.testing.utils import BITCOIND_CONFIG, only_one, wait_for
 from utils import get_bkpr_smaug_balance, get_only_one_descriptor
 
@@ -146,3 +147,35 @@ def test_rpc_remove(bitcoind, ln_node):
     assert len(smaug_wallets) == 0
     assert result == f"Deleted wallet: {wallet_name}"
     assert not os.path.isfile(db_file_path)
+
+
+def test_rpc_remove_failed(bitcoind, ln_node):
+    """
+    Test RPC remove.
+    """
+
+    # Get external/internal only_one descriptors
+    internal_descriptor = get_only_one_descriptor(bitcoind, "wpkh", True)
+    external_descriptor = get_only_one_descriptor(bitcoind, "wpkh", False)
+
+    # Add wallet to smaug
+    wallet = ln_node.rpc.smaug("add", external_descriptor, internal_descriptor)
+    wallet_name = wallet["name"]
+    db_file_path = f"{str(ln_node.lightning_dir)}/regtest/.smaug/{wallet_name}.db"
+
+    smaug_wallets = ln_node.rpc.smaug("ls")
+    assert len(smaug_wallets) == 1
+    assert wallet_name in smaug_wallets
+    assert os.path.isfile(db_file_path)
+
+    # Try removing wallet from smaug
+    with pytest.raises(RpcError) as e:
+        ln_node.rpc.smaug("remove", "NONEXISTENT_WALLET")
+    assert (
+        "RPC call failed: method: smaug, payload: ('remove', 'NONEXISTENT_WALLET'), error: Removal of wallet 'NONEXISTENT_WALLET' failed."
+        in str(e.value)
+    )
+
+    smaug_wallets = ln_node.rpc.smaug("ls")
+    assert len(smaug_wallets) == 1
+    assert os.path.isfile(db_file_path)