From 180fd8deb29d6c1c55e7b25bad07ec1fec76f9a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Granh=C3=A3o?= <daniel.granhao@gmail.com>
Date: Thu, 8 Feb 2024 09:23:43 +0000
Subject: [PATCH] Get wallet pubkey id if not available

---
 honey-badger/src/lib.rs                 | 14 ++++++--
 honey-badger/tests/integration_tests.rs | 47 ++++++++++++++-----------
 2 files changed, 38 insertions(+), 23 deletions(-)

diff --git a/honey-badger/src/lib.rs b/honey-badger/src/lib.rs
index 745eb48..6500128 100644
--- a/honey-badger/src/lib.rs
+++ b/honey-badger/src/lib.rs
@@ -71,8 +71,18 @@ impl Auth {
             .ok_or_permanent_failure("Newly refreshed token is not valid long enough")
     }
 
-    pub fn get_wallet_pubkey_id(&self) -> Option<String> {
-        self.provider.lock().unwrap().get_wallet_pubkey_id()
+    pub fn get_wallet_pubkey_id(&self) -> Result<String> {
+        match self.provider.lock().unwrap().get_wallet_pubkey_id() {
+            Some(id) => Ok(id),
+            None => {
+                self.query_token()?;
+                self.provider
+                    .lock()
+                    .unwrap()
+                    .get_wallet_pubkey_id()
+                    .ok_or_permanent_failure("Failed to get pubkey id for an authenticated wallet")
+            }
+        }
     }
 
     // Not exposed in UDL, used in tests.
diff --git a/honey-badger/tests/integration_tests.rs b/honey-badger/tests/integration_tests.rs
index b978c01..22f2fae 100644
--- a/honey-badger/tests/integration_tests.rs
+++ b/honey-badger/tests/integration_tests.rs
@@ -30,8 +30,14 @@ fn test_invalid_url() {
     )
     .unwrap();
 
-    let id = auth.get_wallet_pubkey_id();
-    assert!(id.is_none());
+    let result = auth.get_wallet_pubkey_id();
+    assert!(matches!(
+        result,
+        Err(Error::RuntimeError {
+            code: GraphQlRuntimeErrorCode::NetworkError,
+            ..
+        })
+    ));
 
     let result = auth.query_token();
     assert!(matches!(
@@ -55,8 +61,14 @@ fn test_502_return() {
     )
     .unwrap();
 
-    let id = auth.get_wallet_pubkey_id();
-    assert!(id.is_none());
+    let result = auth.get_wallet_pubkey_id();
+    assert!(matches!(
+        result,
+        Err(Error::RuntimeError {
+            code: GraphQlRuntimeErrorCode::RemoteServiceUnavailable,
+            ..
+        })
+    ));
 
     let result = auth.query_token();
     assert!(matches!(
@@ -80,21 +92,19 @@ fn test_basic_auth() {
     )
     .unwrap();
 
-    let id = auth.get_wallet_pubkey_id();
-    assert!(id.is_none());
+    let _id = auth.get_wallet_pubkey_id().unwrap();
 
     let token = auth.query_token().unwrap();
     let next_token = auth.query_token().unwrap();
     assert_eq!(token, next_token);
 
-    let id = auth.get_wallet_pubkey_id();
-    assert!(id.is_some());
+    let id = auth.get_wallet_pubkey_id().unwrap();
 
     sleep(Duration::from_secs(1));
     let refreshed_token = auth.refresh_token().unwrap();
     assert_ne!(token, refreshed_token);
 
-    assert_eq!(auth.get_wallet_pubkey_id(), id);
+    assert_eq!(auth.get_wallet_pubkey_id().unwrap(), id);
 }
 
 #[test]
@@ -109,21 +119,19 @@ fn test_owner_auth() {
     )
     .unwrap();
 
-    let id = auth.get_wallet_pubkey_id();
-    assert!(id.is_none());
+    let _id = auth.get_wallet_pubkey_id().unwrap();
 
     let token = auth.query_token().unwrap();
     let next_token = auth.query_token().unwrap();
     assert_eq!(token, next_token);
 
-    let id = auth.get_wallet_pubkey_id();
-    assert!(id.is_some());
+    let id = auth.get_wallet_pubkey_id().unwrap();
 
     sleep(Duration::from_secs(1));
     let refreshed_token = auth.refresh_token().unwrap();
     assert_ne!(token, refreshed_token);
 
-    assert_eq!(auth.get_wallet_pubkey_id(), id);
+    assert_eq!(auth.get_wallet_pubkey_id().unwrap(), id);
 }
 
 #[test]
@@ -138,8 +146,7 @@ fn test_employee_with_no_owner_auth() {
     )
     .unwrap();
 
-    let id = auth.get_wallet_pubkey_id();
-    assert!(id.is_none());
+    let _id = auth.get_wallet_pubkey_id().unwrap();
 
     let result = auth.query_token();
     assert!(matches!(result, Err(Error::InvalidInput { .. })));
@@ -159,21 +166,19 @@ fn test_employee_auth() {
     )
     .unwrap();
 
-    let id = auth.get_wallet_pubkey_id();
-    assert!(id.is_none());
+    let _id = auth.get_wallet_pubkey_id().unwrap();
 
     let token = auth.query_token().unwrap();
     let next_token = auth.query_token().unwrap();
     assert_eq!(token, next_token);
 
-    let id = auth.get_wallet_pubkey_id();
-    assert!(id.is_some());
+    let id = auth.get_wallet_pubkey_id().unwrap();
 
     sleep(Duration::from_secs(1));
     let refreshed_token = auth.refresh_token().unwrap();
     assert_ne!(token, refreshed_token);
 
-    assert_eq!(auth.get_wallet_pubkey_id(), id);
+    assert_eq!(auth.get_wallet_pubkey_id().unwrap(), id);
 }
 
 #[test]