From 1ff2c34d23aae1cf1895cb1a2e212746f482ba5c Mon Sep 17 00:00:00 2001
From: lowkeynicc <85139158+lowkeynicc@users.noreply.github.com>
Date: Thu, 22 Aug 2024 13:38:12 -0400
Subject: [PATCH] add emulateAccount fn (#1157)

---
 sdk/src/driftClient.ts | 51 +++++++++++++++++++++++++++++++++++++++---
 1 file changed, 48 insertions(+), 3 deletions(-)

diff --git a/sdk/src/driftClient.ts b/sdk/src/driftClient.ts
index 5fd178dca..0a8625f57 100644
--- a/sdk/src/driftClient.ts
+++ b/sdk/src/driftClient.ts
@@ -714,6 +714,51 @@ export class DriftClient {
 		return success;
 	}
 
+	/**
+	 * Update the subscribed accounts to a given authority, while leaving the
+	 * connected wallet intact. This allows a user to emulate another user's
+	 * account on the UI and sign permissionless transactions with their own wallet.
+	 * @param emulateAuthority
+	 */
+	public async emulateAccount(emulateAuthority: PublicKey): Promise<boolean> {
+		this.skipLoadUsers = false;
+		// Update provider for txSender with new wallet details
+		this.authority = emulateAuthority;
+		this.userStatsAccountPublicKey = undefined;
+		this.includeDelegates = true;
+		const walletSupportsVersionedTxns =
+			//@ts-ignore
+			this.wallet.supportedTransactionVersions?.size ?? 0 > 1;
+		this.txVersion = walletSupportsVersionedTxns ? 0 : 'legacy';
+
+		this.authoritySubAccountMap = new Map<string, number[]>();
+
+		/* Reset user stats account */
+		if (this.userStats?.isSubscribed) {
+			await this.userStats.unsubscribe();
+		}
+
+		this.userStats = undefined;
+
+		this.userStats = new UserStats({
+			driftClient: this,
+			userStatsAccountPublicKey: this.getUserStatsAccountPublicKey(),
+			accountSubscription: this.userStatsAccountSubscriptionConfig,
+		});
+
+		await this.userStats.subscribe();
+
+		let success = true;
+
+		if (this.isSubscribed) {
+			await Promise.all(this.unsubscribeUsers());
+			this.users.clear();
+			success = await this.addAndSubscribeToUsers(emulateAuthority);
+		}
+
+		return success;
+	}
+
 	public async switchActiveUser(subAccountId: number, authority?: PublicKey) {
 		const authorityChanged = authority && !this.authority?.equals(authority);
 
@@ -771,7 +816,7 @@ export class DriftClient {
 	/**
 	 * Adds and subscribes to users based on params set by the constructor or by updateWallet.
 	 */
-	public async addAndSubscribeToUsers(): Promise<boolean> {
+	public async addAndSubscribeToUsers(authority?: PublicKey): Promise<boolean> {
 		// save the rpc calls if driftclient is initialized without a real wallet
 		if (this.skipLoadUsers) return true;
 
@@ -799,12 +844,12 @@ export class DriftClient {
 			let delegatedAccounts = [];
 
 			const userAccountsPromise = this.getUserAccountsForAuthority(
-				this.wallet.publicKey
+				authority ?? this.wallet.publicKey
 			);
 
 			if (this.includeDelegates) {
 				const delegatedAccountsPromise = this.getUserAccountsForDelegate(
-					this.wallet.publicKey
+					authority ?? this.wallet.publicKey
 				);
 				[userAccounts, delegatedAccounts] = await Promise.all([
 					userAccountsPromise,