From c22199138559c85eb132b9bf04408fb282e6b61b Mon Sep 17 00:00:00 2001
From: lowkeynicc <85139158+lowkeynicc@users.noreply.github.com>
Date: Fri, 20 Oct 2023 16:28:55 -0400
Subject: [PATCH] =?UTF-8?q?sdk:=20make=20it=20possible=20to=20update=20mul?=
 =?UTF-8?q?tiple=20accounts=20margin=20settings=20in=20=E2=80=A6=20(#659)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* sdk: make it possible to update multiple accounts margin settings in single tx

* update tests and changelog
---
 CHANGELOG.md                 |  1 +
 sdk/src/driftClient.ts       | 68 ++++++++++++++++++++++--------------
 tests/spotWithdrawUtil100.ts |  4 ++-
 tests/subaccounts.ts         |  7 ++--
 4 files changed, 48 insertions(+), 32 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f856f5b0e..0660eeb92 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Breaking
 
 - sdk: remove getMakerLimitBids/Asks from DLOB
+- sdk: updateUserMarginEnabled and updateUserCustomMarginRatio now take in an array of params to allow multiple subaccounts to be update in a single tx
 
 ## [2.41.0] - 2023-10-05
 
diff --git a/sdk/src/driftClient.ts b/sdk/src/driftClient.ts
index b86def378..597f2635b 100644
--- a/sdk/src/driftClient.ts
+++ b/sdk/src/driftClient.ts
@@ -875,21 +875,48 @@ export class DriftClient {
 	}
 
 	public async updateUserCustomMarginRatio(
+		updates: { marginRatio: number; subAccountId: number }[]
+	): Promise<TransactionSignature> {
+		const ixs = await Promise.all(
+			updates.map(async ({ marginRatio, subAccountId }) => {
+				const ix = await this.getUpdateUserCustomMarginRatioIx(
+					marginRatio,
+					subAccountId
+				);
+				return ix;
+			})
+		);
+
+		const tx = await this.buildTransaction(ixs, this.txParams);
+
+		const { txSig } = await this.sendTransaction(tx, [], this.opts);
+		return txSig;
+	}
+
+	public async getUpdateUserCustomMarginRatioIx(
 		marginRatio: number,
 		subAccountId = 0
-	): Promise<TransactionSignature> {
-		const tx = await this.program.transaction.updateUserCustomMarginRatio(
+	): Promise<TransactionInstruction> {
+		const userAccountPublicKey = getUserAccountPublicKeySync(
+			this.program.programId,
+			this.wallet.publicKey,
+			subAccountId
+		);
+
+		await this.addUser(subAccountId, this.wallet.publicKey);
+
+		const ix = this.program.instruction.updateUserCustomMarginRatio(
 			subAccountId,
 			marginRatio,
 			{
 				accounts: {
-					user: await this.getUserAccountPublicKey(),
+					user: userAccountPublicKey,
 					authority: this.wallet.publicKey,
 				},
 			}
 		);
-		const { txSig } = await this.sendTransaction(tx, [], this.opts);
-		return txSig;
+
+		return ix;
 	}
 
 	public async getUpdateUserMarginTradingEnabledIx(
@@ -930,31 +957,18 @@ export class DriftClient {
 	}
 
 	public async updateUserMarginTradingEnabled(
-		marginTradingEnabled: boolean,
-		subAccountId = 0
+		updates: { marginTradingEnabled: boolean; subAccountId: number }[]
 	): Promise<TransactionSignature> {
-		const userAccountPublicKey = getUserAccountPublicKeySync(
-			this.program.programId,
-			this.wallet.publicKey,
-			subAccountId
+		const ixs = await Promise.all(
+			updates.map(async ({ marginTradingEnabled, subAccountId }) => {
+				return await this.getUpdateUserMarginTradingEnabledIx(
+					marginTradingEnabled,
+					subAccountId
+				);
+			})
 		);
 
-		await this.addUser(subAccountId, this.wallet.publicKey);
-		const remainingAccounts = this.getRemainingAccounts({
-			userAccounts: [this.getUserAccount(subAccountId)],
-		});
-
-		const tx = await this.program.transaction.updateUserMarginTradingEnabled(
-			subAccountId,
-			marginTradingEnabled,
-			{
-				accounts: {
-					user: userAccountPublicKey,
-					authority: this.wallet.publicKey,
-				},
-				remainingAccounts,
-			}
-		);
+		const tx = await this.buildTransaction(ixs, this.txParams);
 
 		const { txSig } = await this.sendTransaction(tx, [], this.opts);
 		return txSig;
diff --git a/tests/spotWithdrawUtil100.ts b/tests/spotWithdrawUtil100.ts
index f00a8f9a0..cc7afb768 100644
--- a/tests/spotWithdrawUtil100.ts
+++ b/tests/spotWithdrawUtil100.ts
@@ -547,7 +547,9 @@ describe('test function when spot market at >= 100% util', () => {
 
 		const marketIndex = 1;
 
-		await firstUserDriftClient.updateUserMarginTradingEnabled(true, 0);
+		const updates = [{ marginTradingEnabled: true, subAccountId: 0 }];
+
+		await firstUserDriftClient.updateUserMarginTradingEnabled(updates);
 
 		const takerDriftClientUser = new User({
 			driftClient: firstUserDriftClient,
diff --git a/tests/subaccounts.ts b/tests/subaccounts.ts
index ffcf31bc3..808677b31 100644
--- a/tests/subaccounts.ts
+++ b/tests/subaccounts.ts
@@ -174,10 +174,9 @@ describe('subaccounts', () => {
 	it('Update custom margin ratio', async () => {
 		const subAccountId = 0;
 		const customMarginRatio = MARGIN_PRECISION.toNumber() * 2;
-		await driftClient.updateUserCustomMarginRatio(
-			customMarginRatio,
-			subAccountId
-		);
+
+		const updates = [{ marginRatio: customMarginRatio, subAccountId }];
+		await driftClient.updateUserCustomMarginRatio(updates);
 
 		await driftClient.fetchAccounts();
 		assert(driftClient.getUserAccount().maxMarginRatio === customMarginRatio);