From e463f98a777e1287275c22ed18a3f660c8033dbc Mon Sep 17 00:00:00 2001
From: fmhoeger <fmhoeger@mirsem.org>
Date: Sat, 6 Jan 2024 15:31:15 -0600
Subject: [PATCH] Refactor Smaug wallet db file deletion; add additional test

---
 src/main.rs       | 52 +++++++++++++++++++++++++----------------------
 tests/test_rpc.py | 33 ++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 24 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 0bd2164..5e398eb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -432,31 +432,35 @@ async fn delete(
     plugin: Plugin<State>,
     descriptor_name: String,
 ) -> Result<serde_json::Value, Error> {
-    let db_dir = plugin.state().lock().await.db_dir.clone();
+    let db_dir_path = plugin.state().lock().await.db_dir.clone();
     let wallets = &mut plugin.state().lock().await.wallets;
-    let _removed_item: Option<DescriptorWallet>;
-    if wallets.contains_key(&descriptor_name) {
-        let removed_item = wallets.remove(&descriptor_name);
-        let db_path = removed_item.unwrap().get_db_file_path(db_dir).unwrap();
-        fs::remove_file(db_path.clone())?;
-        log::debug!("Deleted smaug db file at {}", db_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!("Can't find wallet '{}'.", 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
index 1a73310..e38ac5b 100644
--- 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 failing 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 nonexistent 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: Can't find wallet 'NONEXISTENT_WALLET'."
+        in str(e.value)
+    )
+
+    smaug_wallets = ln_node.rpc.smaug("ls")
+    assert len(smaug_wallets) == 1
+    assert os.path.isfile(db_file_path)