added optional paramter blockNumber

RobinNagpal · Oct 21, 2024 · 119ebdb · 119ebdb
1 parent 8ee8fde
commit 119ebdb
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 10 deletions.
diff --git a/plugins/scenario/utils/TokenSourcer.ts b/plugins/scenario/utils/TokenSourcer.ts
@@ -12,6 +12,7 @@ interface SourceTokenParameters {
   asset: string;
   address: string;
   blacklist: string[];
+  blockNumber?: number;
 }
 
 export async function fetchQuery(
@@ -55,14 +56,20 @@ export async function sourceTokens({
   asset,
   address,
   blacklist,
+  blockNumber,
 }: SourceTokenParameters) {
   let amount = BigNumber.from(amount_);
   if (amount.isZero()) {
     return;
   } else if (amount.isNegative()) {
     await removeTokens(dm, amount.abs(), asset, address);
   } else {
-    await addTokens(dm, amount, asset, address, [address].concat(blacklist));
+    if(blockNumber) {
+      await addTokens(dm, amount, asset, address, [address].concat(blacklist), blockNumber);
+    } else {
+      await addTokens(dm, amount, asset, address, [address].concat(blacklist));
+    }
+
   }
 }
 

diff --git a/scenario/RewardsScenario.ts b/scenario/RewardsScenario.ts
@@ -260,7 +260,12 @@ async function testScalingReward(properties: CometProperties, context: CometCont
     [albert.address]
   );
   await newRewards.connect(albert.signer).setRewardConfigWithMultiplier(comet.address, rewardTokenAddress, multiplier);
-  await context.sourceTokens(100000, rewardTokenAddress, newRewards.address);
+  await context.sourceTokens(
+    100000,
+    rewardTokenAddress,
+    newRewards.address,
+    2751700
+  );
 
   await baseAsset.approve(albert, comet.address);
   await albert.safeSupplyAsset({ asset: baseAssetAddress, amount: 100n * baseScale });

diff --git a/scenario/context/CometContext.ts b/scenario/context/CometContext.ts
@@ -254,7 +254,7 @@ export class CometContext {
     throw new Error(`Unable to find asset by address ${address}`);
   }
 
-  async sourceTokens(amount: number | bigint, asset: CometAsset | string, recipient: AddressLike) {
+  async sourceTokens(amount: number | bigint, asset: CometAsset | string, recipient: AddressLike, blockNumber?: number) {
     const { world } = this;
     const recipientAddress = resolveAddress(recipient);
     const cometAsset = typeof asset === 'string' ? this.getAssetByAddress(asset) : asset;
@@ -281,13 +281,24 @@ export class CometContext {
     if (amountRemaining != 0n) {
       // Source from logs (expensive, in terms of node API limits)
       debug('Source Tokens: sourcing from logs...', amountRemaining, cometAsset.address);
-      await sourceTokens({
-        dm: this.world.deploymentManager,
-        amount: amountRemaining,
-        asset: cometAsset.address,
-        address: recipientAddress,
-        blacklist: [comet.address],
-      });
+      if (blockNumber) {
+        await sourceTokens({
+          dm: this.world.deploymentManager,
+          amount: amountRemaining,
+          asset: cometAsset.address,
+          address: recipientAddress,
+          blacklist: [comet.address],
+          blockNumber,
+        });
+      } else {
+        await sourceTokens({
+          dm: this.world.deploymentManager,
+          amount: amountRemaining,
+          asset: cometAsset.address,
+          address: recipientAddress,
+          blacklist: [comet.address],
+        });
+      }
     }
   }