From dd84f877f92510feaa4cc9f130b8d3385d7891bf Mon Sep 17 00:00:00 2001
From: BenRey <br@massa.net>
Date: Wed, 27 Nov 2024 10:18:47 +0100
Subject: [PATCH 1/2] Add balancesOf method to MRC20 contract and corresponding
 tests

---
 src/contracts-wrappers/token.ts | 20 +++++++++++++++++++-
 test/integration/MRC20.spec.ts  | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/src/contracts-wrappers/token.ts b/src/contracts-wrappers/token.ts
index db900a19..2c2a67a4 100644
--- a/src/contracts-wrappers/token.ts
+++ b/src/contracts-wrappers/token.ts
@@ -1,4 +1,4 @@
-import { Args, bytesToStr, U256, U8 } from '../basicElements'
+import { Args, bytesToStr, strToBytes, U256, U8 } from '../basicElements'
 import { Operation } from '../operation'
 import { CallSCOptions, ReadSCOptions, SmartContract } from '../smartContracts'
 
@@ -85,6 +85,24 @@ export class MRC20 extends SmartContract {
     return U256.fromBytes(res.value)
   }
 
+  async balancesOf(addresses: string[]): Promise<
+    {
+      address: string
+      balance: bigint
+    }[]
+  > {
+    const res = await this.provider.readStorage(
+      this.address,
+      addresses.map((a) => strToBytes(`BALANCE${a}`)),
+      true
+    )
+
+    return res.map((v, i) => ({
+      address: addresses[i],
+      balance: v.length ? U256.fromBytes(v) : 0n,
+    }))
+  }
+
   async transfer(
     to: string,
     amount: bigint,
diff --git a/test/integration/MRC20.spec.ts b/test/integration/MRC20.spec.ts
index 8ce2e238..ad9fe922 100644
--- a/test/integration/MRC20.spec.ts
+++ b/test/integration/MRC20.spec.ts
@@ -1,3 +1,4 @@
+import { account, Account } from '../../src'
 import { MRC20 } from '../../src/contracts-wrappers'
 import { provider } from './setup'
 
@@ -80,4 +81,35 @@ describe('Generic token wrapper tests', () => {
     )
     expect(newAllowance).toBe(previousAllowance)
   })
+
+  test('balancesOf', async () => {
+    const amounts = [1_000n, 2_000n]
+    const recipientAddresses = [
+      await (await Account.generate()).address.toString(),
+      await (await Account.generate()).address.toString(),
+    ]
+
+    const b = await usdcContract.balancesOf(recipientAddresses)
+    for (const { address, balance } of b) {
+      expect(balance).toBe(0n)
+    }
+
+    const operation = await usdcContract.transfer(
+      recipientAddresses[0],
+      amounts[0]
+    )
+
+    const operation2 = await usdcContract.transfer(
+      recipientAddresses[1],
+      amounts[1]
+    )
+    await operation2.waitFinalExecution()
+
+    const recipientBalance = await usdcContract.balancesOf(recipientAddresses)
+
+    expect(recipientBalance[0].balance).toBe(amounts[0])
+    expect(recipientBalance[0].address).toBe(recipientAddresses[0])
+    expect(recipientBalance[1].balance).toBe(amounts[1])
+    expect(recipientBalance[1].address).toBe(recipientAddresses[1])
+  })
 })

From 61c17219c7d61d871e46e9fcd96cad768aedf53e Mon Sep 17 00:00:00 2001
From: BenRey <br@massa.net>
Date: Wed, 27 Nov 2024 11:09:00 +0100
Subject: [PATCH 2/2] Update balancesOf method in MRC20 contract to accept
 final parameter and adjust related tests

---
 src/contracts-wrappers/token.ts | 11 +++++++----
 test/integration/MRC20.spec.ts  | 14 +++++++-------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/contracts-wrappers/token.ts b/src/contracts-wrappers/token.ts
index 2c2a67a4..636a8a4d 100644
--- a/src/contracts-wrappers/token.ts
+++ b/src/contracts-wrappers/token.ts
@@ -1,4 +1,4 @@
-import { Args, bytesToStr, strToBytes, U256, U8 } from '../basicElements'
+import { Args, bytesToStr, U256, U8 } from '../basicElements'
 import { Operation } from '../operation'
 import { CallSCOptions, ReadSCOptions, SmartContract } from '../smartContracts'
 
@@ -85,7 +85,10 @@ export class MRC20 extends SmartContract {
     return U256.fromBytes(res.value)
   }
 
-  async balancesOf(addresses: string[]): Promise<
+  async balancesOf(
+    addresses: string[],
+    final = true
+  ): Promise<
     {
       address: string
       balance: bigint
@@ -93,8 +96,8 @@ export class MRC20 extends SmartContract {
   > {
     const res = await this.provider.readStorage(
       this.address,
-      addresses.map((a) => strToBytes(`BALANCE${a}`)),
-      true
+      addresses.map((a) => `BALANCE${a}`),
+      final
     )
 
     return res.map((v, i) => ({
diff --git a/test/integration/MRC20.spec.ts b/test/integration/MRC20.spec.ts
index ad9fe922..8c42ca33 100644
--- a/test/integration/MRC20.spec.ts
+++ b/test/integration/MRC20.spec.ts
@@ -1,4 +1,4 @@
-import { account, Account } from '../../src'
+import { Account } from '../../src'
 import { MRC20 } from '../../src/contracts-wrappers'
 import { provider } from './setup'
 
@@ -94,18 +94,18 @@ describe('Generic token wrapper tests', () => {
       expect(balance).toBe(0n)
     }
 
-    const operation = await usdcContract.transfer(
-      recipientAddresses[0],
-      amounts[0]
-    )
+    await usdcContract.transfer(recipientAddresses[0], amounts[0])
 
     const operation2 = await usdcContract.transfer(
       recipientAddresses[1],
       amounts[1]
     )
-    await operation2.waitFinalExecution()
+    await operation2.waitSpeculativeExecution()
 
-    const recipientBalance = await usdcContract.balancesOf(recipientAddresses)
+    const recipientBalance = await usdcContract.balancesOf(
+      recipientAddresses,
+      false
+    )
 
     expect(recipientBalance[0].balance).toBe(amounts[0])
     expect(recipientBalance[0].address).toBe(recipientAddresses[0])