From 3832432930b406f9e8f4663e75cce8a2e7aa5cd6 Mon Sep 17 00:00:00 2001
From: refactor-ring <refactor_ring0@proton.me>
Date: Fri, 30 Jun 2023 16:32:25 +0200
Subject: [PATCH] Added support for the "get_attribute" and "set_attribute"
 methods

---
 src/lib.rs                            | 29 +++++++++++++++++++++++++++
 tests/clients_tests/basic_wallet.rs   | 16 +++++++++++++++
 tests/clients_tests/helpers/wallet.rs | 14 +++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index a6272cd..9d46bbc 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1327,6 +1327,35 @@ impl WalletClient {
 
         Ok((u16::try_from(major)?, u16::try_from(minor)?))
     }
+
+    /// Returns an attribute as a string or an error when there is no attribute for the given key
+    pub async fn get_attribute(&self, key: String) -> anyhow::Result<String> {
+        let params = empty().chain(once(("key", key.into())));
+
+        #[derive(Deserialize)]
+        struct Rsp {
+            value: String,
+        }
+
+        Ok(self
+            .inner
+            .request::<Rsp>("get_attribute", RpcParams::map(params))
+            .await?
+            .value)
+    }
+
+    /// Set an arbitrary attribute which is saved in the wallet
+    pub async fn set_attribute(&self, key: String, value: String) -> anyhow::Result<()> {
+        let params = empty()
+            .chain(once(("key", key.into())))
+            .chain(once(("value", value.into())));
+
+        self.inner
+            .request::<IgnoredAny>("set_attribute", RpcParams::map(params))
+            .await?;
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
diff --git a/tests/clients_tests/basic_wallet.rs b/tests/clients_tests/basic_wallet.rs
index 520d41d..50db0f4 100644
--- a/tests/clients_tests/basic_wallet.rs
+++ b/tests/clients_tests/basic_wallet.rs
@@ -297,5 +297,21 @@ pub async fn run() {
 
     helpers::wallet::sign_and_verify_assert_ok(&wallet, "test message").await;
 
+    helpers::wallet::get_attribute_error(&wallet, "nonexistingattribute".to_string()).await;
+
+    helpers::wallet::set_attribute_assert_ok(
+        &wallet,
+        "attribute_key".to_string(),
+        "attribute_value".to_string(),
+    )
+    .await;
+
+    helpers::wallet::get_attribute_assert_ok(
+        &wallet,
+        "attribute_key".to_string(),
+        "attribute_value".to_string(),
+    )
+    .await;
+
     helpers::wallet::close_wallet_assert_ok(&wallet).await;
 }
diff --git a/tests/clients_tests/helpers/wallet.rs b/tests/clients_tests/helpers/wallet.rs
index bb5edc8..f15ac7b 100644
--- a/tests/clients_tests/helpers/wallet.rs
+++ b/tests/clients_tests/helpers/wallet.rs
@@ -741,3 +741,17 @@ pub async fn sign_and_verify_assert_ok(wallet: &WalletClient, message: &str) {
         .unwrap();
     assert!(signature_ok);
 }
+
+pub async fn get_attribute_error(wallet: &WalletClient, key: String) {
+    let err = wallet.get_attribute(key).await.unwrap_err();
+    assert_eq!(err.to_string(), "Server error: Attribute not found.");
+}
+
+pub async fn get_attribute_assert_ok(wallet: &WalletClient, key: String, expected_value: String) {
+    let value_from_wallet = wallet.get_attribute(key).await.unwrap();
+    assert_eq!(value_from_wallet, expected_value);
+}
+
+pub async fn set_attribute_assert_ok(wallet: &WalletClient, key: String, value: String) {
+    wallet.set_attribute(key, value).await.unwrap()
+}