From 41e0e6762c91b526cdbf44782496d3705fd73c51 Mon Sep 17 00:00:00 2001
From: Lex Berezhny <lex@damoti.com>
Date: Wed, 30 Oct 2019 17:05:23 -0400
Subject: [PATCH] added wallet management tests and address unsubscribe feature
 on wallet server

---
 lbry/lbry/extras/daemon/Daemon.py              |  6 ++++--
 lbry/tests/integration/test_wallet_commands.py | 13 +++++++++++++
 torba/torba/client/baseledger.py               |  4 ++++
 torba/torba/client/basenetwork.py              |  3 +++
 torba/torba/server/session.py                  | 11 +++++++++++
 5 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/lbry/lbry/extras/daemon/Daemon.py b/lbry/lbry/extras/daemon/Daemon.py
index b3a44ca60..093a03f84 100644
--- a/lbry/lbry/extras/daemon/Daemon.py
+++ b/lbry/lbry/extras/daemon/Daemon.py
@@ -1134,7 +1134,7 @@ class Daemon(metaclass=JSONRPCServerType):
         return wallet
 
     @requires("wallet")
-    def jsonrpc_wallet_remove(self, wallet_id):
+    async def jsonrpc_wallet_remove(self, wallet_id):
         """
         Remove an existing wallet.
 
@@ -1146,8 +1146,10 @@ class Daemon(metaclass=JSONRPCServerType):
 
         Returns: {Wallet}
         """
-        wallet = self.wallet_manager.get_wallet_or_default(wallet_id)
+        wallet = self.wallet_manager.get_wallet_or_error(wallet_id)
         self.wallet_manager.wallets.remove(wallet)
+        for account in wallet.accounts:
+            await self.ledger.unsubscribe_account(account)
         return wallet
 
     @requires("wallet")
diff --git a/lbry/tests/integration/test_wallet_commands.py b/lbry/tests/integration/test_wallet_commands.py
index 9a555c0ad..f41fcb595 100644
--- a/lbry/tests/integration/test_wallet_commands.py
+++ b/lbry/tests/integration/test_wallet_commands.py
@@ -4,6 +4,19 @@ from lbry.testcase import CommandTestCase
 from torba.client.wallet import ENCRYPT_ON_DISK
 
 
+class WalletCommands(CommandTestCase):
+
+    async def test_wallet_create_and_add_subscribe(self):
+        session = next(iter(self.conductor.spv_node.server.session_mgr.sessions))
+        self.assertEqual(len(session.hashX_subs), 27)
+        wallet = await self.daemon.jsonrpc_wallet_create('foo', create_account=True, single_key=True)
+        self.assertEqual(len(session.hashX_subs), 28)
+        await self.daemon.jsonrpc_wallet_remove(wallet.id)
+        self.assertEqual(len(session.hashX_subs), 27)
+        await self.daemon.jsonrpc_wallet_add(wallet.id)
+        self.assertEqual(len(session.hashX_subs), 28)
+
+
 class WalletEncryptionAndSynchronization(CommandTestCase):
 
     SEED = (
diff --git a/torba/torba/client/baseledger.py b/torba/torba/client/baseledger.py
index bbe42af35..fbdddd66e 100644
--- a/torba/torba/client/baseledger.py
+++ b/torba/torba/client/baseledger.py
@@ -392,6 +392,10 @@ class BaseLedger(metaclass=LedgerRegistry):
             await self.subscribe_addresses(address_manager, await address_manager.get_addresses())
         await account.ensure_address_gap()
 
+    async def unsubscribe_account(self, account: baseaccount.BaseAccount):
+        for address in await account.get_addresses():
+            await self.network.unsubscribe_address(address)
+
     async def announce_addresses(self, address_manager: baseaccount.AddressManager, addresses: List[str]):
         await self.subscribe_addresses(address_manager, addresses)
         await self._on_address_controller.add(
diff --git a/torba/torba/client/basenetwork.py b/torba/torba/client/basenetwork.py
index 2ed1e0c95..dfed4d800 100644
--- a/torba/torba/client/basenetwork.py
+++ b/torba/torba/client/basenetwork.py
@@ -261,6 +261,9 @@ class BaseNetwork:
                 self.client.abort()
             raise asyncio.CancelledError()
 
+    def unsubscribe_address(self, address):
+        return self.rpc('blockchain.address.unsubscribe', [address], True)
+
     def get_server_features(self):
         return self.rpc('server.features', (), restricted=True)
 
diff --git a/torba/torba/server/session.py b/torba/torba/server/session.py
index 8cf8fcc6e..edee37d6e 100644
--- a/torba/torba/server/session.py
+++ b/torba/torba/server/session.py
@@ -882,6 +882,9 @@ class ElectrumX(SessionBase):
         self.hashX_subs[hashX] = alias
         return await self.address_status(hashX)
 
+    async def hashX_unsubscribe(self, hashX, alias):
+        del self.hashX_subs[hashX]
+
     def address_to_hashX(self, address):
         try:
             return self.coin.address_to_hashX(address)
@@ -916,6 +919,13 @@ class ElectrumX(SessionBase):
         hashX = self.address_to_hashX(address)
         return await self.hashX_subscribe(hashX, address)
 
+    async def address_unsubscribe(self, address):
+        """Unsubscribe an address.
+
+        address: the address to unsubscribe"""
+        hashX = self.address_to_hashX(address)
+        return await self.hashX_unsubscribe(hashX, address)
+
     async def get_balance(self, hashX):
         utxos = await self.db.all_utxos(hashX)
         confirmed = sum(utxo.value for utxo in utxos)
@@ -1284,6 +1294,7 @@ class ElectrumX(SessionBase):
                 'blockchain.address.get_mempool': self.address_get_mempool,
                 'blockchain.address.listunspent': self.address_listunspent,
                 'blockchain.address.subscribe': self.address_subscribe,
+                'blockchain.address.unsubscribe': self.address_unsubscribe,
             })
 
         self.request_handlers = handlers