From 3243233ebe6580d59c7718c0bdea19636a04c586 Mon Sep 17 00:00:00 2001
From: Chris Duncan <chris@zoso.dev>
Date: Thu, 3 Jul 2025 16:26:14 -0700
Subject: [PATCH] Test account lock and unlock.

---
 src/lib/account.ts               | 33 +++++++++++++++++++++------
 test/test.lock-unlock-wallet.mjs | 38 ++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 7 deletions(-)

diff --git a/src/lib/account.ts b/src/lib/account.ts
index f74c950..f68d6ee 100644
--- a/src/lib/account.ts
+++ b/src/lib/account.ts
@@ -141,12 +141,23 @@ export class Account {
 			throw new Error('Failed to unlock wallet')
 		}
 		try {
+			const data: { id: string, privateKey: string | null } = {
+				id: this.#pub,
+				privateKey: null
+			}
 			if (this.#prv != null) {
-				await Account.#poolSafe.assign({
-					method: 'put',
-					name: this.#prv,
-					password
-				})
+				data.privateKey = this.#prv
+			}
+			console.log(this.#prv)
+			const response = (await Account.#poolSafe.assign({
+				method: 'put',
+				name: this.#pub,
+				password,
+				data
+			}))[0]
+			const success = response.result
+			if (!success) {
+				throw null
 			}
 		} catch (err) {
 			console.error(`Failed to lock account ${this.address}`, err)
@@ -166,11 +177,19 @@ export class Account {
 			throw new Error('Failed to unlock wallet')
 		}
 		try {
-			this.#prv = await Account.#poolSafe.assign({
+			const response = (await Account.#poolSafe.assign({
 				method: 'get',
 				name: this.#pub,
 				password
-			})
+			}))[0]
+			console.log(response)
+			const { id, privateKey } = response.result
+			if (id !== this.#pub) {
+				throw null
+			}
+			if (privateKey != null) {
+				this.#prv = privateKey
+			}
 		} catch (err) {
 			console.error(`Failed to unlock account ${this.address}`, err)
 			return false
diff --git a/test/test.lock-unlock-wallet.mjs b/test/test.lock-unlock-wallet.mjs
index a22c5d9..5062183 100644
--- a/test/test.lock-unlock-wallet.mjs
+++ b/test/test.lock-unlock-wallet.mjs
@@ -51,6 +51,25 @@ await suite('Lock and unlock wallets', async () => {
 		await wallet.destroy()
 	})
 
+	await test('locking and unlocking a Bip44Wallet Account with a password', async () => {
+		const wallet = await Bip44Wallet.fromMnemonic(NANO_TEST_VECTORS.PASSWORD, NANO_TEST_VECTORS.MNEMONIC, NANO_TEST_VECTORS.PASSWORD)
+		await wallet.unlock(NANO_TEST_VECTORS.PASSWORD)
+		const account = await wallet.account()
+		const lockResult = await account.lock(NANO_TEST_VECTORS.PASSWORD)
+
+		assert.equals(lockResult, true)
+		assert.ok('privateKey' in account)
+		assert.nullish(account.privateKey)
+
+		const unlockResult = await account.unlock(NANO_TEST_VECTORS.PASSWORD)
+
+		assert.equals(unlockResult, true)
+		assert.ok('privateKey' in account)
+		assert.equals(account.privateKey, NANO_TEST_VECTORS.PRIVATE_0)
+
+		await wallet.destroy()
+	})
+
 	await test('fail to unlock a Bip44Wallet with different passwords', async () => {
 		const wallet = await Bip44Wallet.fromMnemonic(NANO_TEST_VECTORS.PASSWORD, NANO_TEST_VECTORS.MNEMONIC, NANO_TEST_VECTORS.PASSWORD)
 		await wallet.unlock(NANO_TEST_VECTORS.PASSWORD)
@@ -185,6 +204,25 @@ await suite('Lock and unlock wallets', async () => {
 		await wallet.destroy()
 	})
 
+	await test('locking and unlocking a Blake2bWallet Account with a password', async () => {
+		const wallet = await Blake2bWallet.fromSeed(NANO_TEST_VECTORS.PASSWORD, TREZOR_TEST_VECTORS.ENTROPY_0)
+		await wallet.unlock(NANO_TEST_VECTORS.PASSWORD)
+		const account = await wallet.account()
+		const lockResult = await account.lock(NANO_TEST_VECTORS.PASSWORD)
+
+		assert.equals(lockResult, true)
+		assert.ok('privateKey' in account)
+		assert.nullish(account.privateKey)
+
+		const unlockResult = await account.unlock(NANO_TEST_VECTORS.PASSWORD)
+
+		assert.equals(unlockResult, true)
+		assert.ok('privateKey' in account)
+		assert.equals(account.privateKey, NANO_TEST_VECTORS.PRIVATE_0)
+
+		await wallet.destroy()
+	})
+
 	await test('fail to unlock a Blake2bWallet with different passwords', async () => {
 		const wallet = await Blake2bWallet.fromSeed(NANO_TEST_VECTORS.PASSWORD, TREZOR_TEST_VECTORS.ENTROPY_1)
 
-- 
2.47.3