From 9f3e3af21ed2b302d99864202fbaf530c9b74489 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ti=E1=BA=BFn=20Nguy=E1=BB=85n=20Kh=E1=BA=AFc?=
 <tien.nguyenkhac@icloud.com>
Date: Wed, 13 Nov 2024 23:54:30 +1300
Subject: [PATCH] refactor: wallet aggregation

---
 .changeset/selfish-keys-know.md               |  5 ++
 .../src/actions/aggregate-wallets.spec.ts     | 66 +++++++++++++++++++
 .../core/src/actions/aggregate-wallets.ts     |  8 ++-
 packages/core/src/wallets/provider.ts         |  9 ---
 4 files changed, 78 insertions(+), 10 deletions(-)
 create mode 100644 .changeset/selfish-keys-know.md
 create mode 100644 packages/core/src/actions/aggregate-wallets.spec.ts

diff --git a/.changeset/selfish-keys-know.md b/.changeset/selfish-keys-know.md
new file mode 100644
index 00000000..a4cb08ec
--- /dev/null
+++ b/.changeset/selfish-keys-know.md
@@ -0,0 +1,5 @@
+---
+"@reactive-dot/core": patch
+---
+
+Refactored wallet aggregation logic.
diff --git a/packages/core/src/actions/aggregate-wallets.spec.ts b/packages/core/src/actions/aggregate-wallets.spec.ts
new file mode 100644
index 00000000..cc1e3124
--- /dev/null
+++ b/packages/core/src/actions/aggregate-wallets.spec.ts
@@ -0,0 +1,66 @@
+import type { PolkadotSignerAccount } from "../wallets/account.js";
+import { WalletProvider } from "../wallets/provider.js";
+import { Wallet } from "../wallets/wallet.js";
+import { aggregateWallets } from "./aggregate-wallets.js";
+import { Observable } from "rxjs";
+import { it, expect } from "vitest";
+
+class MockWallet extends Wallet {
+  id = "mock";
+
+  name = "Mock";
+
+  initialize() {
+    throw new Error("Method not implemented.");
+  }
+
+  connected$ = new Observable<boolean>();
+
+  connect() {
+    throw new Error("Method not implemented.");
+  }
+
+  disconnect() {
+    throw new Error("Method not implemented.");
+  }
+
+  accounts$ = new Observable<PolkadotSignerAccount[]>();
+}
+
+class MockProvider extends WalletProvider {
+  getWallets() {
+    return Array.from({ length: 10 }).map(() => new MockWallet());
+  }
+}
+
+it("aggregates wallets from provider", async () => {
+  const wallets = await aggregateWallets([
+    new MockWallet(),
+    new MockProvider(),
+  ]);
+
+  expect(wallets).toHaveLength(11);
+
+  for (const wallet of wallets) {
+    expect(wallet).toBeInstanceOf(Wallet);
+  }
+});
+
+it("only aggregate once", async () => {
+  const walletsAndProviders = [
+    new MockWallet(),
+    new MockProvider(),
+    new MockWallet(),
+    new MockProvider(),
+  ];
+
+  const firstRunWallets = await aggregateWallets(walletsAndProviders);
+
+  const secondRunWallets = await aggregateWallets(walletsAndProviders);
+
+  expect(firstRunWallets).toHaveLength(secondRunWallets.length);
+
+  for (const [index, wallet] of firstRunWallets.entries()) {
+    expect(wallet).toBe(secondRunWallets.at(index));
+  }
+});
diff --git a/packages/core/src/actions/aggregate-wallets.ts b/packages/core/src/actions/aggregate-wallets.ts
index 46d6a823..bc03773e 100644
--- a/packages/core/src/actions/aggregate-wallets.ts
+++ b/packages/core/src/actions/aggregate-wallets.ts
@@ -1,5 +1,7 @@
 import { Wallet, type WalletProvider } from "../wallets/index.js";
 
+const providerWallets = new WeakMap<WalletProvider, Wallet[]>();
+
 export function aggregateWallets(
   providersOrWallets: ReadonlyArray<Wallet | WalletProvider>,
 ) {
@@ -7,7 +9,11 @@ export function aggregateWallets(
     providersOrWallets.map((walletOrProvider) =>
       walletOrProvider instanceof Wallet
         ? [walletOrProvider]
-        : walletOrProvider.getOrFetchWallets(),
+        : (async () =>
+            providerWallets.get(walletOrProvider) ??
+            providerWallets
+              .set(walletOrProvider, await walletOrProvider.getWallets())
+              .get(walletOrProvider)!)(),
     ),
   ).then((wallets) => wallets.flat());
 }
diff --git a/packages/core/src/wallets/provider.ts b/packages/core/src/wallets/provider.ts
index 5ed3ec4b..d39d7d30 100644
--- a/packages/core/src/wallets/provider.ts
+++ b/packages/core/src/wallets/provider.ts
@@ -1,14 +1,5 @@
 import type { Wallet } from "./wallet.js";
 
 export abstract class WalletProvider {
-  #wallets: Wallet[] | undefined;
-
   abstract getWallets(): Wallet[] | Promise<Wallet[]>;
-
-  /**
-   * @internal
-   */
-  async getOrFetchWallets() {
-    return (this.#wallets ??= await this.getWallets());
-  }
 }