From c04d3978da7882776d1abda477cc303204da7b9d Mon Sep 17 00:00:00 2001
From: Tianjie Wei <weitianjie@gmail.com>
Date: Wed, 13 Apr 2022 21:01:59 -0700
Subject: [PATCH] Allow users to chain staking calls after claimNOTE

---
 contracts/mainnet/connectors/notional/main.sol | 8 ++++++--
 test/mainnet/notional/notional.helpers.ts      | 4 ++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/contracts/mainnet/connectors/notional/main.sol b/contracts/mainnet/connectors/notional/main.sol
index 575cc68f..a7b85c76 100644
--- a/contracts/mainnet/connectors/notional/main.sol
+++ b/contracts/mainnet/connectors/notional/main.sol
@@ -681,12 +681,14 @@ abstract contract NotionalResolver is Events, Helpers {
 	function mintSNoteFromETH(
 		uint256 noteAmount,
 		uint256 ethAmount,
-		uint256 minBPT
+		uint256 minBPT,
+		uint256 getId
 	)
 		external
 		payable
 		returns (string memory _eventName, bytes memory _eventParam)
 	{
+		noteAmount = getUint(getId, noteAmount);
 		if (noteAmount == type(uint256).max)
 			noteAmount = note.balanceOf(address(this));
 
@@ -708,12 +710,14 @@ abstract contract NotionalResolver is Events, Helpers {
 	function mintSNoteFromWETH(
 		uint256 noteAmount,
 		uint256 wethAmount,
-		uint256 minBPT
+		uint256 minBPT,
+		uint256 getId
 	)
 		external
 		payable
 		returns (string memory _eventName, bytes memory _eventParam)
 	{
+		noteAmount = getUint(getId, noteAmount);
 		if (noteAmount == type(uint256).max)
 			noteAmount = note.balanceOf(address(this));
 
diff --git a/test/mainnet/notional/notional.helpers.ts b/test/mainnet/notional/notional.helpers.ts
index ad57c8c1..d5d04bfa 100644
--- a/test/mainnet/notional/notional.helpers.ts
+++ b/test/mainnet/notional/notional.helpers.ts
@@ -212,7 +212,7 @@ const mintSNoteFromETH = async (
         {
             connector: "NOTIONAL-TEST-A",
             method: "mintSNoteFromETH",
-            args: [noteAmount, ethAmount, minBPT]
+            args: [noteAmount, ethAmount, minBPT, 0]
         }        
     ]
 
@@ -232,7 +232,7 @@ const mintSNoteFromWETH = async (
         {
             connector: "NOTIONAL-TEST-A",
             method: "mintSNoteFromWETH",
-            args: [noteAmount, wethAmount, minBPT]
+            args: [noteAmount, wethAmount, minBPT, 0]
         }        
     ]