From 0e99a9e8f308a12a03da4817b5c331b9e9da7f87 Mon Sep 17 00:00:00 2001
From: Ivan Vershigora <ivan.vershigora@gmail.com>
Date: Fri, 29 Nov 2024 11:18:15 +0000
Subject: [PATCH] feat: add wallet stop method

---
 package-lock.json     |  4 ++--
 package.json          |  2 +-
 src/wallet/index.ts   | 22 ++++++++++++++++++++++
 tests/boost.test.ts   |  2 +-
 tests/receive.test.ts |  2 +-
 tests/storage.test.ts | 27 +++++++++++++++++++++++----
 6 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 566fbc6..7ca7370 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "beignet",
-  "version": "0.0.48",
+  "version": "0.0.49",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "beignet",
-      "version": "0.0.48",
+      "version": "0.0.49",
       "license": "MIT",
       "dependencies": {
         "@bitcoinerlab/secp256k1": "1.0.5",
diff --git a/package.json b/package.json
index f799248..52b8817 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "beignet",
-  "version": "0.0.48",
+  "version": "0.0.49",
   "description": "A self-custodial, JS Bitcoin wallet management library.",
   "main": "dist/index.js",
   "scripts": {
diff --git a/src/wallet/index.ts b/src/wallet/index.ts
index e39ae8b..6d34644 100644
--- a/src/wallet/index.ts
+++ b/src/wallet/index.ts
@@ -281,6 +281,28 @@ export class Wallet {
 		}
 	}
 
+	/**
+	 * Stops the wallet. Use this method to prepare the wallet to be de
+	 * @returns {Promise<Result<string>>}
+	 */
+	public async stop(): Promise<Result<string>> {
+		try {
+			// if we are refreshing, we need to wait for it to finish
+			if (this.isRefreshing) {
+				await this.refreshWallet();
+			}
+			// disable onMessage callback
+			this.disableMessages = true;
+			// disable saving to storage
+			this._setData = undefined;
+			// disconnect from Electrum
+			await this.electrum.disconnect();
+			return ok('Wallet stopped.');
+		} catch (e) {
+			return err(e);
+		}
+	}
+
 	public async switchNetwork(
 		network: EAvailableNetworks,
 		servers?: TServer | TServer[]
diff --git a/tests/boost.test.ts b/tests/boost.test.ts
index 32e4486..81a32fc 100644
--- a/tests/boost.test.ts
+++ b/tests/boost.test.ts
@@ -10,7 +10,7 @@ import {
 	EProtocol,
 	generateMnemonic,
 	Wallet
-} from '../src';
+} from '../';
 import {
 	bitcoinURL,
 	electrumHost,
diff --git a/tests/receive.test.ts b/tests/receive.test.ts
index bc3d460..6f846a2 100644
--- a/tests/receive.test.ts
+++ b/tests/receive.test.ts
@@ -10,7 +10,7 @@ import {
 	generateMnemonic,
 	sleep,
 	Wallet
-} from '../src';
+} from '../';
 import {
 	bitcoinURL,
 	electrumHost,
diff --git a/tests/storage.test.ts b/tests/storage.test.ts
index 9025167..174dc53 100644
--- a/tests/storage.test.ts
+++ b/tests/storage.test.ts
@@ -2,6 +2,7 @@ import { validateMnemonic } from 'bip39';
 import { expect } from 'chai';
 import net from 'net';
 import tls from 'tls';
+import sinon from 'sinon';
 
 import { Wallet } from '../';
 import { deleteDirectory, getData, servers, setData } from '../example/helpers';
@@ -20,6 +21,10 @@ const testTimeout = 60000;
 let wallet: Wallet;
 const WALLET_NAME = 'storagetestwallet0';
 
+const storage = { getData, setData };
+const storageSpy = sinon.spy(storage, 'setData');
+const messageSpy = sinon.spy();
+
 describe('Storage Test', async function (): Promise<void> {
 	this.timeout(testTimeout);
 
@@ -31,10 +36,8 @@ describe('Storage Test', async function (): Promise<void> {
 			network: EAvailableNetworks.testnet,
 			name: WALLET_NAME,
 			addressType: EAddressType.p2wpkh,
-			storage: {
-				getData,
-				setData
-			},
+			storage,
+			onMessage: messageSpy,
 			electrumOptions: {
 				servers: servers[EAvailableNetworks.testnet],
 				net,
@@ -43,6 +46,7 @@ describe('Storage Test', async function (): Promise<void> {
 		});
 		if (res.isErr()) throw res.error;
 		wallet = res.value;
+		await wallet.refreshWallet({});
 	});
 
 	after(async function () {
@@ -146,6 +150,21 @@ describe('Storage Test', async function (): Promise<void> {
 		);
 	});
 
+	it('Should successfully stop the wallet', async () => {
+		wallet.refreshWallet(); // start wallet refresh in the background
+		const r = await wallet.stop();
+		if (r.isErr()) throw r.error;
+		storageSpy.resetHistory();
+		messageSpy.resetHistory();
+		// try to force the wallet to update it's data
+		wallet.feeEstimates.timestamp = 0;
+		await wallet.updateFeeEstimates(true);
+		await wallet.refreshWallet();
+		// make sure that the wallet did not call setData or onMessage
+		expect(storageSpy.called).to.equal(false);
+		expect(messageSpy.called).to.equal(false);
+	});
+
 	// it('Attempts to create a new wallet using the same name of an existing wallet in storage.', async () => {
 	// 	const createRes = await Wallet.create({
 	// 		mnemonic: generateMnemonic(),