Replace hard wallet types with traits

swap/src/bitcoin.rs

@@ -231,11 +231,26 @@ pub trait GetRawTransaction {
     async fn get_raw_transaction(&self, txid: Txid) -> Result<Transaction>;
 }
 
+#[async_trait]
+pub trait NewAddress {
+    async fn new_address(&self) -> Result<Address>;
+}
+
 #[async_trait]
 pub trait GetNetwork {
     fn get_network(&self) -> Network;
 }
 
+#[async_trait]
+pub trait Balance {
+    async fn balance(&self) -> Result<Amount>;
+}
+
+#[async_trait]
+pub trait FetchTransactionFee {
+    async fn transaction_fee(&self, txid: Txid) -> Result<Amount>;
+}
+
 pub fn recover(S: PublicKey, sig: Signature, encsig: EncryptedSignature) -> Result<SecretKey> {
     let adaptor = Adaptor::<Sha256, Deterministic<Sha256>>::default();
 

swap/src/bitcoin/wallet.rs

@@ -1,8 +1,9 @@
 use crate::{
     bitcoin::{
-        timelocks::BlockHeight, Address, Amount, BroadcastSignedTransaction, BuildTxLockPsbt,
+        timelocks::BlockHeight, Address, Amount, Balance, BroadcastSignedTransaction,
-        GetBlockHeight, GetNetwork, GetRawTransaction, SignTxLock, Transaction,
+        BuildTxLockPsbt, FetchTransactionFee, GetBlockHeight, GetNetwork, GetRawTransaction,
-        TransactionBlockHeight, TxLock, WaitForTransactionFinality, WatchForRawTransaction,
+        NewAddress, SignTxLock, Transaction, TransactionBlockHeight, TxLock,
+        WaitForTransactionFinality, WatchForRawTransaction,
     },
     config::Config,
 };
@@ -30,33 +31,6 @@ impl Wallet {
             network,
         })
     }
-
-    pub async fn balance(&self) -> Result<Amount> {
-        let balance = self.inner.balance().await?;
-        Ok(balance)
-    }
-
-    pub async fn new_address(&self) -> Result<Address> {
-        self.inner.new_address().await.map_err(Into::into)
-    }
-
-    pub async fn transaction_fee(&self, txid: Txid) -> Result<Amount> {
-        let fee = self
-            .inner
-            .get_wallet_transaction(txid)
-            .await
-            .map(|res| {
-                res.fee.map(|signed_amount| {
-                    signed_amount
-                        .abs()
-                        .to_unsigned()
-                        .expect("Absolute value is always positive")
-                })
-            })?
-            .context("Rpc response did not contain a fee")?;
-
-        Ok(fee)
-    }
 }
 
 #[async_trait]
@@ -191,8 +165,44 @@ impl WaitForTransactionFinality for Wallet {
     }
 }
 
+#[async_trait]
+impl NewAddress for Wallet {
+    async fn new_address(&self) -> Result<Address> {
+        self.inner.new_address().await.map_err(Into::into)
+    }
+}
+
 impl GetNetwork for Wallet {
     fn get_network(&self) -> bitcoin::Network {
         self.network
     }
 }
+
+#[async_trait]
+impl Balance for Wallet {
+    async fn balance(&self) -> Result<Amount> {
+        let balance = self.inner.balance().await?;
+        Ok(balance)
+    }
+}
+
+#[async_trait]
+impl FetchTransactionFee for Wallet {
+    async fn transaction_fee(&self, txid: Txid) -> Result<Amount> {
+        let fee = self
+            .inner
+            .get_wallet_transaction(txid)
+            .await
+            .map(|res| {
+                res.fee.map(|signed_amount| {
+                    signed_amount
+                        .abs()
+                        .to_unsigned()
+                        .expect("Absolute value is always positive")
+                })
+            })?
+            .context("Rpc response did not contain a fee")?;
+
+        Ok(fee)
+    }
+}

swap/src/main.rs

@@ -12,7 +12,11 @@
 #![forbid(unsafe_code)]
 #![allow(non_snake_case)]
 
-use crate::cli::{Command, Options, Resume};
+use crate::{
+    bitcoin::Balance as _,
+    cli::{Command, Options, Resume},
+    monero::Balance as _,
+};
 use anyhow::{Context, Result};
 use config::Config;
 use database::Database;
@@ -242,7 +246,7 @@ async fn setup_wallets(
     );
 
     let monero_wallet = monero::Wallet::new(monero_wallet_rpc_url, config.monero_network);
-    let monero_balance = monero_wallet.get_balance().await?;
+    let monero_balance = monero_wallet.balance().await?;
     info!(
         "Connection to Monero wallet succeeded, balance: {}",
         monero_balance

swap/src/monero.rs

@@ -8,6 +8,7 @@ use crate::bitcoin;
 use ::bitcoin::hashes::core::fmt::Formatter;
 use anyhow::Result;
 use async_trait::async_trait;
+use monero_harness::rpc::wallet::BlockHeight;
 use rand::{CryptoRng, RngCore};
 use rust_decimal::{
     prelude::{FromPrimitive, ToPrimitive},
@@ -207,6 +208,22 @@ pub trait CreateWalletForOutput {
     ) -> anyhow::Result<()>;
 }
 
+#[async_trait]
+pub trait Balance {
+    /// Get the balance of the primary account.
+    async fn balance(&self) -> Result<Amount>;
+}
+
+#[async_trait]
+pub trait Refresh {
+    async fn refresh(&self) -> Result<()>;
+}
+
+#[async_trait]
+pub trait FetchBlockHeight {
+    async fn block_height(&self) -> Result<BlockHeight>;
+}
+
 #[derive(thiserror::Error, Debug, Clone, PartialEq)]
 #[error("Overflow, cannot convert {0} to u64")]
 pub struct OverflowError(pub String);

swap/src/monero/wallet.rs

@@ -1,13 +1,13 @@
 use crate::monero::{
-    Amount, CreateWalletForOutput, InsufficientFunds, PrivateViewKey, PublicViewKey, Transfer,
+    Amount, Balance, CreateWalletForOutput, FetchBlockHeight, InsufficientFunds, PrivateViewKey,
-    TransferProof, TxHash, WatchForTransfer,
+    PublicViewKey, Refresh, Transfer, TransferProof, TxHash, WatchForTransfer,
 };
 use ::monero::{Address, Network, PrivateKey, PublicKey};
 use anyhow::Result;
 use async_trait::async_trait;
 use backoff::{backoff::Constant as ConstantBackoff, future::FutureOperation as _};
 use bitcoin::hashes::core::sync::atomic::AtomicU32;
-use monero_harness::rpc::wallet;
+use monero_harness::rpc::{wallet, wallet::BlockHeight};
 use std::{
     str::FromStr,
     sync::{atomic::Ordering, Arc},
@@ -29,13 +29,6 @@ impl Wallet {
             network,
         }
     }
-
-    /// Get the balance of the primary account.
-    pub async fn get_balance(&self) -> Result<Amount> {
-        let amount = self.inner.get_balance(0).await?;
-
-        Ok(Amount::from_piconero(amount))
-    }
 }
 
 #[async_trait]
@@ -166,3 +159,27 @@ impl WatchForTransfer for Wallet {
         Ok(())
     }
 }
+
+#[async_trait]
+impl Balance for Wallet {
+    async fn balance(&self) -> Result<Amount> {
+        let amount = self.inner.get_balance(0).await?;
+
+        Ok(Amount::from_piconero(amount))
+    }
+}
+
+#[async_trait]
+impl Refresh for Wallet {
+    async fn refresh(&self) -> Result<()> {
+        self.inner.refresh().await?;
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl FetchBlockHeight for Wallet {
+    async fn block_height(&self) -> Result<BlockHeight> {
+        self.inner.block_height().await
+    }
+}

swap/src/protocol/alice.rs

@@ -44,17 +44,27 @@ pub mod swap;
 mod swap_response;
 mod transfer_proof;
 
-pub struct Swap {
+pub trait BitcoinWallet: swap::BitcoinWallet + bitcoin::NewAddress {}
+impl BitcoinWallet for bitcoin::Wallet {}
+
+pub trait MoneroWallet: monero::Transfer + monero::CreateWalletForOutput + Send + Sync {}
+impl MoneroWallet for monero::Wallet {}
+
+pub struct Swap<B, M> {
     pub state: AliceState,
     pub event_loop_handle: EventLoopHandle,
-    pub bitcoin_wallet: Arc<bitcoin::Wallet>,
+    pub bitcoin_wallet: Arc<B>,
-    pub monero_wallet: Arc<monero::Wallet>,
+    pub monero_wallet: Arc<M>,
     pub config: Config,
     pub swap_id: Uuid,
     pub db: Database,
 }
 
-pub struct Builder {
+pub struct Builder<B, M>
+where
+    B: BitcoinWallet + bitcoin::NewAddress,
+    M: MoneroWallet,
+{
     swap_id: Uuid,
     identity: Keypair,
     peer_id: PeerId,
@@ -63,8 +73,8 @@ pub struct Builder {
 
     listen_address: Multiaddr,
 
-    bitcoin_wallet: Arc<bitcoin::Wallet>,
+    bitcoin_wallet: Arc<B>,
-    monero_wallet: Arc<monero::Wallet>,
+    monero_wallet: Arc<M>,
 
     init_params: InitParams,
 }
@@ -74,13 +84,17 @@ enum InitParams {
     New { swap_amounts: SwapAmounts },
 }
 
-impl Builder {
+impl<B, M> Builder<B, M>
+where
+    B: BitcoinWallet + bitcoin::NewAddress,
+    M: MoneroWallet,
+{
     pub async fn new(
         seed: Seed,
         config: Config,
         swap_id: Uuid,
-        bitcoin_wallet: Arc<bitcoin::Wallet>,
+        bitcoin_wallet: Arc<B>,
-        monero_wallet: Arc<monero::Wallet>,
+        monero_wallet: Arc<M>,
         db_path: PathBuf,
         listen_address: Multiaddr,
     ) -> Self {
@@ -108,7 +122,7 @@ impl Builder {
         }
     }
 
-    pub async fn build(self) -> Result<(Swap, EventLoop)> {
+    pub async fn build(self) -> Result<(Swap<B, M>, EventLoop)> {
         match self.init_params {
             InitParams::New { swap_amounts } => {
                 let initial_state = self

swap/src/protocol/alice/swap.rs

@@ -2,15 +2,10 @@
 //! Alice holds XMR and wishes receive BTC.
 use crate::{
     bitcoin,
-    bitcoin::{
+    bitcoin::timelocks::ExpiredTimelocks,
-        timelocks::ExpiredTimelocks, TransactionBlockHeight, WaitForTransactionFinality,
-        WatchForRawTransaction,
-    },
     config::Config,
     database,
     database::Database,
-    monero,
-    monero::CreateWalletForOutput,
     protocol::{
         alice,
         alice::{
@@ -22,7 +17,7 @@ use crate::{
                 publish_cancel_transaction, wait_for_bitcoin_encrypted_signature,
                 wait_for_bitcoin_refund, wait_for_locked_bitcoin,
             },
-            AliceState,
+            AliceState, MoneroWallet,
         },
     },
 };
@@ -41,6 +36,20 @@ trait Rng: RngCore + CryptoRng + Send {}
 
 impl<T> Rng for T where T: RngCore + CryptoRng + Send {}
 
+pub trait BitcoinWallet:
+    bitcoin::WatchForRawTransaction
+    + bitcoin::WaitForTransactionFinality
+    + bitcoin::TransactionBlockHeight
+    + bitcoin::GetBlockHeight
+    + bitcoin::BroadcastSignedTransaction
+    + bitcoin::GetRawTransaction
+    + Send
+    + Sync
+{
+}
+
+impl BitcoinWallet for bitcoin::Wallet {}
+
 pub fn is_complete(state: &AliceState) -> bool {
     matches!(
         state,
@@ -51,14 +60,22 @@ pub fn is_complete(state: &AliceState) -> bool {
     )
 }
 
-pub async fn run(swap: alice::Swap) -> Result<AliceState> {
+pub async fn run<B, M>(swap: alice::Swap<B, M>) -> Result<AliceState>
+where
+    B: BitcoinWallet,
+    M: MoneroWallet,
+{
     run_until(swap, is_complete).await
 }
 
-pub async fn run_until(
+pub async fn run_until<B, M>(
-    swap: alice::Swap,
+    swap: alice::Swap<B, M>,
     is_target_state: fn(&AliceState) -> bool,
-) -> Result<AliceState> {
+) -> Result<AliceState>
+where
+    B: BitcoinWallet,
+    M: MoneroWallet,
+{
     run_until_internal(
         swap.state,
         is_target_state,
@@ -75,16 +92,20 @@ pub async fn run_until(
 // State machine driver for swap execution
 #[async_recursion]
 #[allow(clippy::too_many_arguments)]
-async fn run_until_internal(
+async fn run_until_internal<B, M>(
     state: AliceState,
     is_target_state: fn(&AliceState) -> bool,
     mut event_loop_handle: EventLoopHandle,
-    bitcoin_wallet: Arc<bitcoin::Wallet>,
+    bitcoin_wallet: Arc<B>,
-    monero_wallet: Arc<monero::Wallet>,
+    monero_wallet: Arc<M>,
     config: Config,
     swap_id: Uuid,
     db: Database,
-) -> Result<AliceState> {
+) -> Result<AliceState>
+where
+    B: BitcoinWallet,
+    M: MoneroWallet,
+{
     info!("Current state:{}", state);
     if is_target_state(&state) {
         Ok(state)

swap/src/protocol/bob.rs

@@ -42,17 +42,20 @@ pub mod swap;
 mod swap_request;
 mod transfer_proof;
 
-pub struct Swap {
+pub struct Swap<B, M> {
     pub state: BobState,
     pub event_loop_handle: bob::EventLoopHandle,
     pub db: Database,
-    pub bitcoin_wallet: Arc<bitcoin::Wallet>,
+    pub bitcoin_wallet: Arc<B>,
-    pub monero_wallet: Arc<monero::Wallet>,
+    pub monero_wallet: Arc<M>,
     pub config: Config,
     pub swap_id: Uuid,
 }
 
-pub struct Builder {
+pub struct Builder<B, M>
+where
+    B: bitcoin::NewAddress,
+{
     swap_id: Uuid,
     identity: Keypair,
     peer_id: PeerId,
@@ -61,8 +64,8 @@ pub struct Builder {
     alice_address: Multiaddr,
     alice_peer_id: PeerId,
 
-    bitcoin_wallet: Arc<bitcoin::Wallet>,
+    bitcoin_wallet: Arc<B>,
-    monero_wallet: Arc<monero::Wallet>,
+    monero_wallet: Arc<M>,
 
     init_params: InitParams,
     config: Config,
@@ -73,14 +76,17 @@ enum InitParams {
     New { swap_amounts: SwapAmounts },
 }
 
-impl Builder {
+impl<B, M> Builder<B, M>
+where
+    B: bitcoin::NewAddress,
+{
     #[allow(clippy::too_many_arguments)]
     pub fn new(
         seed: Seed,
         db_path: PathBuf,
         swap_id: Uuid,
-        bitcoin_wallet: Arc<bitcoin::Wallet>,
+        bitcoin_wallet: Arc<B>,
-        monero_wallet: Arc<monero::Wallet>,
+        monero_wallet: Arc<M>,
         alice_address: Multiaddr,
         alice_peer_id: PeerId,
         config: Config,
@@ -109,7 +115,7 @@ impl Builder {
         }
     }
 
-    pub async fn build(self) -> Result<(bob::Swap, bob::EventLoop)> {
+    pub async fn build(self) -> Result<(bob::Swap<B, M>, bob::EventLoop)> {
         match self.init_params {
             InitParams::New { swap_amounts } => {
                 let initial_state = self
@@ -164,6 +170,7 @@ impl Builder {
             }
         }
     }
+
     fn init_event_loop(
         &self,
     ) -> Result<(bob::event_loop::EventLoop, bob::event_loop::EventLoopHandle)> {

swap/src/protocol/bob/swap.rs

@@ -17,6 +17,28 @@ use tokio::select;
 use tracing::info;
 use uuid::Uuid;
 
+pub trait BitcoinWallet:
+    bitcoin::BuildTxLockPsbt
+    + bitcoin::GetNetwork
+    + bitcoin::SignTxLock
+    + bitcoin::BroadcastSignedTransaction
+    + bitcoin::WatchForRawTransaction
+    + bitcoin::TransactionBlockHeight
+    + bitcoin::GetBlockHeight
+    + bitcoin::GetRawTransaction
+    + bitcoin::WaitForTransactionFinality
+    + Send
+    + Sync
+{
+}
+impl BitcoinWallet for bitcoin::Wallet {}
+
+pub trait MoneroWallet:
+    monero::WatchForTransfer + monero::FetchBlockHeight + monero::CreateWalletForOutput + Send + Sync
+{
+}
+impl MoneroWallet for monero::Wallet {}
+
 pub fn is_complete(state: &BobState) -> bool {
     matches!(
         state,
@@ -28,14 +50,22 @@ pub fn is_complete(state: &BobState) -> bool {
 }
 
 #[allow(clippy::too_many_arguments)]
-pub async fn run(swap: bob::Swap) -> Result<BobState> {
+pub async fn run<B, M>(swap: bob::Swap<B, M>) -> Result<BobState>
+where
+    B: BitcoinWallet,
+    M: MoneroWallet,
+{
     run_until(swap, is_complete).await
 }
 
-pub async fn run_until(
+pub async fn run_until<B, M>(
-    swap: bob::Swap,
+    swap: bob::Swap<B, M>,
     is_target_state: fn(&BobState) -> bool,
-) -> Result<BobState> {
+) -> Result<BobState>
+where
+    B: BitcoinWallet,
+    M: MoneroWallet,
+{
     run_until_internal(
         swap.state,
         is_target_state,
@@ -53,19 +83,21 @@ pub async fn run_until(
 // State machine driver for swap execution
 #[allow(clippy::too_many_arguments)]
 #[async_recursion]
-async fn run_until_internal<R>(
+async fn run_until_internal<R, B, M>(
     state: BobState,
     is_target_state: fn(&BobState) -> bool,
     mut event_loop_handle: EventLoopHandle,
     db: Database,
-    bitcoin_wallet: Arc<bitcoin::Wallet>,
+    bitcoin_wallet: Arc<B>,
-    monero_wallet: Arc<monero::Wallet>,
+    monero_wallet: Arc<M>,
     mut rng: R,
     swap_id: Uuid,
     config: Config,
 ) -> Result<BobState>
 where
     R: RngCore + CryptoRng + Send,
+    B: BitcoinWallet,
+    M: MoneroWallet,
 {
     info!("Current state: {}", state);
     if is_target_state(&state) {
@@ -139,8 +171,7 @@ where
                     // TODO: This can be optimized further by extracting the block height when
                     //  tx-lock was included. However, scanning a few more blocks won't do any harm
                     //  and is simpler.
-                    let monero_wallet_restore_blockheight =
+                    let monero_wallet_restore_blockheight = monero_wallet.block_height().await?;
-                        monero_wallet.inner.block_height().await?;
 
                     select! {
                         transfer_proof = transfer_proof_watcher => {
@@ -387,15 +418,16 @@ where
     }
 }
 
-pub async fn negotiate<R>(
+pub async fn negotiate<R, W>(
     state0: crate::protocol::bob::state::State0,
     amounts: SwapAmounts,
     swarm: &mut EventLoopHandle,
     mut rng: R,
-    bitcoin_wallet: Arc<crate::bitcoin::Wallet>,
+    bitcoin_wallet: Arc<W>,
 ) -> Result<bob::state::State2>
 where
     R: RngCore + CryptoRng + Send,
+    W: bitcoin::BuildTxLockPsbt + bitcoin::GetNetwork,
 {
     tracing::trace!("Starting negotiate");
     swarm

swap/tests/testutils/mod.rs

@@ -24,18 +24,26 @@ pub struct StartingBalances {
     pub btc: bitcoin::Amount,
 }
 
-struct AliceParams {
+struct AliceParams<B, M>
+where
+    B: swap::protocol::alice::BitcoinWallet,
+    M: swap::protocol::alice::MoneroWallet,
+{
     seed: Seed,
     config: Config,
     swap_id: Uuid,
-    bitcoin_wallet: Arc<bitcoin::Wallet>,
+    bitcoin_wallet: Arc<B>,
-    monero_wallet: Arc<monero::Wallet>,
+    monero_wallet: Arc<M>,
     db_path: PathBuf,
     listen_address: Multiaddr,
 }
 
-impl AliceParams {
+impl<B, M> AliceParams<B, M>
-    pub async fn builder(&self) -> alice::Builder {
+where
+    B: swap::protocol::alice::BitcoinWallet + bitcoin::NewAddress,
+    M: swap::protocol::alice::MoneroWallet,
+{
+    pub async fn builder(&self) -> alice::Builder<B, M> {
         alice::Builder::new(
             self.seed,
             self.config,
@@ -53,19 +61,22 @@ impl AliceParams {
     }
 }
 
-struct BobParams {
+struct BobParams<B, M> {
     seed: Seed,
     db_path: PathBuf,
     swap_id: Uuid,
-    bitcoin_wallet: Arc<bitcoin::Wallet>,
+    bitcoin_wallet: Arc<B>,
-    monero_wallet: Arc<monero::Wallet>,
+    monero_wallet: Arc<M>,
     alice_address: Multiaddr,
     alice_peer_id: PeerId,
     config: Config,
 }
 
-impl BobParams {
+impl<B, M> BobParams<B, M>
-    pub fn builder(&self) -> bob::Builder {
+where
+    B: bitcoin::NewAddress,
+{
+    pub fn builder(&self) -> bob::Builder<B, M> {
         bob::Builder::new(
             self.seed,
             self.db_path.clone(),
@@ -79,22 +90,37 @@ impl BobParams {
     }
 }
 
-pub struct TestContext {
+pub struct TestContext<BA, MA, BB, MB>
+where
+    BA: swap::protocol::alice::BitcoinWallet,
+    MA: swap::protocol::alice::MoneroWallet,
+    BB: swap::protocol::bob::swap::BitcoinWallet,
+    MB: swap::protocol::bob::swap::MoneroWallet,
+{
     swap_amounts: SwapAmounts,
 
-    alice_params: AliceParams,
+    alice_params: AliceParams<BA, MA>,
     alice_starting_balances: StartingBalances,
-    alice_bitcoin_wallet: Arc<bitcoin::Wallet>,
+    alice_bitcoin_wallet: Arc<BA>,
-    alice_monero_wallet: Arc<monero::Wallet>,
+    alice_monero_wallet: Arc<MA>,
 
-    bob_params: BobParams,
+    bob_params: BobParams<BB, MB>,
     bob_starting_balances: StartingBalances,
-    bob_bitcoin_wallet: Arc<bitcoin::Wallet>,
+    bob_bitcoin_wallet: Arc<BB>,
-    bob_monero_wallet: Arc<monero::Wallet>,
+    bob_monero_wallet: Arc<MB>,
 }
 
-impl TestContext {
+impl<BA, MA, BB, MB> TestContext<BA, MA, BB, MB>
-    pub async fn new_swap_as_alice(&mut self) -> alice::Swap {
+where
+    BA: swap::protocol::alice::BitcoinWallet + swap::bitcoin::Balance,
+    MA: swap::protocol::alice::MoneroWallet + swap::monero::Balance + swap::monero::Refresh,
+    BB: swap::protocol::bob::swap::BitcoinWallet
+        + swap::bitcoin::NewAddress
+        + swap::bitcoin::FetchTransactionFee
+        + swap::bitcoin::Balance,
+    MB: swap::protocol::bob::swap::MoneroWallet + swap::monero::Balance + swap::monero::Refresh,
+{
+    pub async fn new_swap_as_alice(&mut self) -> alice::Swap<BA, MA> {
         let (swap, mut event_loop) = self
             .alice_params
             .builder()
@@ -109,7 +135,7 @@ impl TestContext {
         swap
     }
 
-    pub async fn new_swap_as_bob(&mut self) -> bob::Swap {
+    pub async fn new_swap_as_bob(&mut self) -> bob::Swap<BB, MB> {
         let (swap, event_loop) = self
             .bob_params
             .builder()
@@ -123,7 +149,7 @@ impl TestContext {
         swap
     }
 
-    pub async fn recover_alice_from_db(&mut self) -> alice::Swap {
+    pub async fn recover_alice_from_db(&mut self) -> alice::Swap<BA, MA> {
         let (swap, mut event_loop) = self.alice_params.builder().await.build().await.unwrap();
 
         tokio::spawn(async move { event_loop.run().await });
@@ -131,7 +157,7 @@ impl TestContext {
         swap
     }
 
-    pub async fn recover_bob_from_db(&mut self) -> bob::Swap {
+    pub async fn recover_bob_from_db(&mut self) -> bob::Swap<BB, MB> {
         let (swap, event_loop) = self.bob_params.builder().build().await.unwrap();
 
         tokio::spawn(async move { event_loop.run().await });
@@ -149,12 +175,7 @@ impl TestContext {
                 - bitcoin::Amount::from_sat(bitcoin::TX_FEE)
         );
 
-        let xmr_balance_after_swap = self
+        let xmr_balance_after_swap = self.alice_monero_wallet.as_ref().balance().await.unwrap();
-            .alice_monero_wallet
-            .as_ref()
-            .get_balance()
-            .await
-            .unwrap();
         assert!(xmr_balance_after_swap <= self.alice_starting_balances.xmr - self.swap_amounts.xmr);
     }
 
@@ -165,18 +186,8 @@ impl TestContext {
         assert_eq!(btc_balance_after_swap, self.alice_starting_balances.btc);
 
         // Ensure that Alice's balance is refreshed as we use a newly created wallet
-        self.alice_monero_wallet
+        self.alice_monero_wallet.as_ref().refresh().await.unwrap();
-            .as_ref()
+        let xmr_balance_after_swap = self.alice_monero_wallet.as_ref().balance().await.unwrap();
-            .inner
-            .refresh()
-            .await
-            .unwrap();
-        let xmr_balance_after_swap = self
-            .alice_monero_wallet
-            .as_ref()
-            .get_balance()
-            .await
-            .unwrap();
         assert_eq!(xmr_balance_after_swap, self.swap_amounts.xmr);
     }
 
@@ -190,12 +201,7 @@ impl TestContext {
                 - bitcoin::Amount::from_sat(2 * bitcoin::TX_FEE)
         );
 
-        let xmr_balance_after_swap = self
+        let xmr_balance_after_swap = self.alice_monero_wallet.as_ref().balance().await.unwrap();
-            .alice_monero_wallet
-            .as_ref()
-            .get_balance()
-            .await
-            .unwrap();
         assert!(xmr_balance_after_swap <= self.alice_starting_balances.xmr - self.swap_amounts.xmr);
     }
 
@@ -219,13 +225,8 @@ impl TestContext {
         );
 
         // Ensure that Bob's balance is refreshed as we use a newly created wallet
-        self.bob_monero_wallet
+        self.bob_monero_wallet.as_ref().refresh().await.unwrap();
-            .as_ref()
+        let xmr_balance_after_swap = self.bob_monero_wallet.as_ref().balance().await.unwrap();
-            .inner
-            .refresh()
-            .await
-            .unwrap();
-        let xmr_balance_after_swap = self.bob_monero_wallet.as_ref().get_balance().await.unwrap();
         assert_eq!(
             xmr_balance_after_swap,
             self.bob_starting_balances.xmr + self.swap_amounts.xmr
@@ -260,7 +261,7 @@ impl TestContext {
         // Since we cannot be sure who submitted it we have to assert accordingly
         assert!(alice_submitted_cancel || bob_submitted_cancel);
 
-        let xmr_balance_after_swap = self.bob_monero_wallet.as_ref().get_balance().await.unwrap();
+        let xmr_balance_after_swap = self.bob_monero_wallet.as_ref().balance().await.unwrap();
         assert_eq!(xmr_balance_after_swap, self.bob_starting_balances.xmr);
     }
 
@@ -283,14 +284,14 @@ impl TestContext {
             self.bob_starting_balances.btc - self.swap_amounts.btc - lock_tx_bitcoin_fee
         );
 
-        let xmr_balance_after_swap = self.bob_monero_wallet.as_ref().get_balance().await.unwrap();
+        let xmr_balance_after_swap = self.bob_monero_wallet.as_ref().balance().await.unwrap();
         assert_eq!(xmr_balance_after_swap, self.bob_starting_balances.xmr);
     }
 }
 
 pub async fn setup_test<T, F>(testfn: T)
 where
-    T: Fn(TestContext) -> F,
+    T: Fn(TestContext<bitcoin::Wallet, monero::Wallet, bitcoin::Wallet, monero::Wallet>) -> F,
     F: Future<Output = ()>,
 {
     let cli = Cli::default();