Refactoring

smartcontractkit · Nov 19, 2024 · 99c3452 · 99c3452
1 parent 1619dff
commit 99c3452
Showing 1 changed file with 134 additions and 80 deletions.
diff --git a/integration-tests/smoke/ccip_test.go b/integration-tests/smoke/ccip_test.go
@@ -111,63 +111,18 @@ func TestInitialDeployOnLocal(t *testing.T) {
 	// TODO: Apply the proposal.
 }
 
-func setupEnvironment(t *testing.T) (ccdeploy.DeployedEnv, ccdeploy.CCIPOnChainState) {
-	lggr := logger.TestLogger(t)
-	tenv, _, _ := testsetups.NewLocalDevEnvironmentWithDefaultPrice(t, lggr)
-	state, err := ccdeploy.LoadOnchainState(tenv.Env)
-	require.NoError(t, err)
-	return tenv, state
-}
-
-func deployAndApproveTokens(t *testing.T, e ccdeploy.DeployedEnv, state ccdeploy.CCIPOnChainState, tokenName string) (*burn_mint_erc677.BurnMintERC677, *burn_mint_erc677.BurnMintERC677) {
-	srcToken, _, dstToken, _, err := ccdeploy.DeployTransferableToken(
-		logger.TestLogger(t),
-		e.Env.Chains,
-		e.HomeChainSel,
-		e.FeedChainSel,
-		state,
-		e.Env.ExistingAddresses,
-		tokenName,
-	)
-	require.NoError(t, err)
-
-	tenTokens := new(big.Int).Mul(big.NewInt(1e18), big.NewInt(10))
-	mintAndApprove(t, e, state, srcToken, e.HomeChainSel, tenTokens)
-	mintAndApprove(t, e, state, dstToken, e.FeedChainSel, tenTokens)
-
-	return srcToken, dstToken
-}
-
-func mintAndApprove(t *testing.T, e ccdeploy.DeployedEnv, state ccdeploy.CCIPOnChainState, token *burn_mint_erc677.BurnMintERC677, chainSel uint64, amount *big.Int) {
-	tx, err := token.Mint(
-		e.Env.Chains[chainSel].DeployerKey,
-		e.Env.Chains[chainSel].DeployerKey.From,
-		amount,
-	)
-	require.NoError(t, err)
-	_, err = e.Env.Chains[chainSel].Confirm(tx)
-	require.NoError(t, err)
-
-	tx, err = token.Approve(
-		e.Env.Chains[chainSel].DeployerKey,
-		state.Chains[chainSel].Router.Address(),
-		amount,
-	)
-	require.NoError(t, err)
-	_, err = e.Env.Chains[chainSel].Confirm(tx)
-	require.NoError(t, err)
-}
-
 func TestTokenTransfer(t *testing.T) {
 	t.Parallel()
 	tenv, state := setupEnvironment(t)
 
 	output, err := changeset.DeployPrerequisites(tenv.Env, changeset.DeployPrerequisiteConfig{
 		ChainSelectors: tenv.Env.AllChainSelectors(),
 	})
+
 	require.NoError(t, err)
 	require.NoError(t, tenv.Env.ExistingAddresses.Merge(output.AddressBook))
 
+	// Deploy CCIP contracts
 	output, err = changeset.InitialDeploy(tenv.Env, ccdeploy.DeployCCIPContractConfig{
 		HomeChainSel:   tenv.HomeChainSel,
 		FeedChainSel:   tenv.FeedChainSel,
@@ -181,11 +136,10 @@ func TestTokenTransfer(t *testing.T) {
 	state, err = ccdeploy.LoadOnchainState(tenv.Env)
 	require.NoError(t, err)
 
-	srcToken1, dstToken1 := deployAndApproveTokens(t, tenv, state, "TOKEN1")
-	srcToken2, dstToken2 := deployAndApproveTokens(t, tenv, state, "TOKEN2")
-
+	// Replay logs
 	ccdeploy.ReplayLogs(t, tenv.Env.Offchain, tenv.ReplayBlocks)
 
+	// Apply the jobs
 	for nodeID, jobs := range output.JobSpecs {
 		for _, job := range jobs {
 			_, err := tenv.Env.Offchain.ProposeJob(ccdeploy.Context(t),
@@ -197,8 +151,14 @@ func TestTokenTransfer(t *testing.T) {
 		}
 	}
 
+	// Add all lanes
 	require.NoError(t, ccdeploy.AddLanesForAll(tenv.Env, state))
 
+	// Deploy and approve tokens
+	srcToken1, dstToken1 := deployAndApproveTokens(t, tenv, state, "Token1")
+	srcToken2, dstToken2 := deployAndApproveTokens(t, tenv, state, "Token2")
+
+	// Define your scenarios
 	scenarios := []struct {
 		name         string
 		srcChain     uint64
@@ -222,15 +182,15 @@ func TestTokenTransfer(t *testing.T) {
 		},
 		{
 			name:     "Send token to contract",
-			srcChain: tenv.FeedChainSel,
-			dstChain: tenv.HomeChainSel,
+			srcChain: tenv.HomeChainSel, // Will be set in the test
+			dstChain: tenv.FeedChainSel, // Will be set in the test
 			tokenAmounts: []router.ClientEVMTokenAmount{
 				{
-					Token:  dstToken1.Address(),
+					Token:  srcToken1.Address(), // Will be set in the test
 					Amount: big.NewInt(1e18),
 				},
 			},
-			receiver: state.Chains[tenv.HomeChainSel].Receiver.Address(),
+			receiver: state.Chains[tenv.FeedChainSel].Receiver.Address(), // Will be set in the test
 			data:     []byte(""),
 		},
 		{
@@ -252,8 +212,8 @@ func TestTokenTransfer(t *testing.T) {
 		},
 		{
 			name:     "Send N tokens to contract",
-			srcChain: tenv.FeedChainSel,
-			dstChain: tenv.HomeChainSel,
+			srcChain: tenv.HomeChainSel,
+			dstChain: tenv.FeedChainSel,
 			tokenAmounts: []router.ClientEVMTokenAmount{
 				{
 					Token:  srcToken1.Address(),
@@ -268,14 +228,16 @@ func TestTokenTransfer(t *testing.T) {
 					Amount: big.NewInt(3e18),
 				},
 			},
-			receiver: state.Chains[tenv.HomeChainSel].Receiver.Address(),
+			receiver: state.Chains[tenv.FeedChainSel].Receiver.Address(),
 			data:     []byte(""),
 		},
 	}
 
 	for _, scenario := range scenarios {
-		scenario := scenario
+		scenario := scenario // Capture range variable
 		t.Run(scenario.name, func(t *testing.T) {
+
+			// Prepare for message sending
 			startBlocks := make(map[uint64]*uint64)
 			expectedSeqNum := make(map[uint64]uint64)
 
@@ -285,6 +247,29 @@ func TestTokenTransfer(t *testing.T) {
 			block := latesthdr.Number.Uint64()
 			startBlocks[scenario.dstChain] = &block
 
+			// Fetch initial balances and aggregate total amounts
+			initialBalances := make(map[common.Address]*big.Int)
+			totalAmountsTransferred := make(map[common.Address]*big.Int)
+
+			for _, tokenAmount := range scenario.tokenAmounts {
+				dstToken := getDestinationToken(scenario, tenv, srcToken1, dstToken1, srcToken2, dstToken2, tokenAmount.Token)
+				require.NotNil(t, dstToken, "Destination token not found")
+
+				// Fetch initial balance
+				if _, exists := initialBalances[dstToken.Address()]; !exists {
+					balance, err := dstToken.BalanceOf(nil, scenario.receiver)
+					require.NoError(t, err)
+					initialBalances[dstToken.Address()] = balance
+				}
+
+				// Initialize and sum up total amounts
+				if _, exists := totalAmountsTransferred[dstToken.Address()]; !exists {
+					totalAmountsTransferred[dstToken.Address()] = big.NewInt(0)
+				}
+				totalAmountsTransferred[dstToken.Address()].Add(totalAmountsTransferred[dstToken.Address()], tokenAmount.Amount)
+			}
+
+			// Prepare message
 			msg := router.ClientEVM2AnyMessage{
 				Receiver:     common.LeftPadBytes(scenario.receiver.Bytes(), 32),
 				Data:         scenario.data,
@@ -293,39 +278,108 @@ func TestTokenTransfer(t *testing.T) {
 				ExtraArgs:    nil,
 			}
 
+			// Send the message
 			msgSentEvent := ccdeploy.TestSendRequest(t, tenv.Env, state, scenario.srcChain, scenario.dstChain, false, msg)
 			expectedSeqNum[scenario.dstChain] = msgSentEvent.SequenceNumber
 
+			// Wait for commit and execution
 			ccdeploy.ConfirmCommitForAllWithExpectedSeqNums(t, tenv.Env, state, expectedSeqNum, startBlocks)
 			ccdeploy.ConfirmExecWithSeqNrForAll(t, tenv.Env, state, expectedSeqNum, startBlocks)
 
-			for _, tokenAmount := range scenario.tokenAmounts {
-				var dstToken *burn_mint_erc677.BurnMintERC677
-				if scenario.srcChain == tenv.HomeChainSel && scenario.dstChain == tenv.FeedChainSel {
-					if tokenAmount.Token == srcToken1.Address() {
-						dstToken = dstToken1
-					} else if tokenAmount.Token == srcToken2.Address() {
-						dstToken = dstToken2
-					}
-				} else if scenario.srcChain == tenv.FeedChainSel && scenario.dstChain == tenv.HomeChainSel {
-					if tokenAmount.Token == srcToken1.Address() {
-						dstToken = dstToken1
-					} else if tokenAmount.Token == srcToken2.Address() {
-						dstToken = dstToken2
-					}
-				}
-				require.NotNil(t, dstToken, "Destination token not found")
+			// Fetch final balances and assert
+			for tokenAddress, totalAmount := range totalAmountsTransferred {
+				dstToken := getTokenByAddress(dstToken1, dstToken2, tokenAddress)
+				require.NotNil(t, dstToken, "Destination token not found for address %s", tokenAddress.Hex())
 
-				balance, err := dstToken.BalanceOf(nil, scenario.receiver)
+				finalBalance, err := dstToken.BalanceOf(nil, scenario.receiver)
 				require.NoError(t, err)
 
-				expectedBalance := tokenAmount.Amount
-				if scenario.name == "Send N tokens to contract" && tokenAmount.Token == srcToken1.Address() {
-					expectedBalance = new(big.Int).Add(big.NewInt(1e18), big.NewInt(3e18))
-				}
+				initialBalance := initialBalances[dstToken.Address()]
+				expectedBalance := new(big.Int).Add(initialBalance, totalAmount)
 
-				require.Equal(t, expectedBalance, balance, "Incorrect balance for token %s", dstToken.Address().Hex())
+				require.Equal(t, expectedBalance, finalBalance, "Incorrect balance for token %s", dstToken.Address().Hex())
 			}
 		})
 	}
 }
+
+func getTokenByAddress(dstToken1, dstToken2 *burn_mint_erc677.BurnMintERC677, tokenAddress common.Address) *burn_mint_erc677.BurnMintERC677 {
+	if dstToken1.Address() == tokenAddress {
+		return dstToken1
+	} else if dstToken2.Address() == tokenAddress {
+		return dstToken2
+	}
+	return nil
+}
+
+// Helper function to determine the destination token
+func getDestinationToken(scenario struct {
+	name         string
+	srcChain     uint64
+	dstChain     uint64
+	tokenAmounts []router.ClientEVMTokenAmount
+	receiver     common.Address
+	data         []byte
+}, tenv ccdeploy.DeployedEnv, srcToken1, dstToken1, srcToken2, dstToken2 *burn_mint_erc677.BurnMintERC677, tokenAddress common.Address) *burn_mint_erc677.BurnMintERC677 {
+	if scenario.srcChain == tenv.HomeChainSel && scenario.dstChain == tenv.FeedChainSel {
+		if tokenAddress == srcToken1.Address() {
+			return dstToken1
+		} else if tokenAddress == srcToken2.Address() {
+			return dstToken2
+		}
+	} else if scenario.srcChain == tenv.FeedChainSel && scenario.dstChain == tenv.HomeChainSel {
+		if tokenAddress == dstToken1.Address() {
+			return srcToken1
+		} else if tokenAddress == dstToken2.Address() {
+			return srcToken2
+		}
+	}
+	return nil
+}
+
+func setupEnvironment(t *testing.T) (ccdeploy.DeployedEnv, ccdeploy.CCIPOnChainState) {
+	lggr := logger.TestLogger(t)
+	tenv := ccdeploy.NewMemoryEnvironment(t, lggr, 2, 4, ccdeploy.MockLinkPrice, ccdeploy.MockWethPrice)
+	state, err := ccdeploy.LoadOnchainState(tenv.Env)
+	require.NoError(t, err)
+	return tenv, state
+}
+
+func deployAndApproveTokens(t *testing.T, e ccdeploy.DeployedEnv, state ccdeploy.CCIPOnChainState, tokenName string) (*burn_mint_erc677.BurnMintERC677, *burn_mint_erc677.BurnMintERC677) {
+	srcToken, _, dstToken, _, err := ccdeploy.DeployTransferableToken(
+		logger.TestLogger(t),
+		e.Env.Chains,
+		e.HomeChainSel,
+		e.FeedChainSel,
+		state,
+		e.Env.ExistingAddresses,
+		tokenName,
+	)
+	require.NoError(t, err)
+
+	tenTokens := new(big.Int).Mul(big.NewInt(1e18), big.NewInt(10))
+	mintAndApprove(t, e, state, srcToken, e.HomeChainSel, tenTokens)
+	mintAndApprove(t, e, state, dstToken, e.FeedChainSel, tenTokens)
+
+	return srcToken, dstToken
+}
+
+func mintAndApprove(t *testing.T, e ccdeploy.DeployedEnv, state ccdeploy.CCIPOnChainState, token *burn_mint_erc677.BurnMintERC677, chainSel uint64, amount *big.Int) {
+	tx, err := token.Mint(
+		e.Env.Chains[chainSel].DeployerKey,
+		e.Env.Chains[chainSel].DeployerKey.From,
+		amount,
+	)
+	require.NoError(t, err)
+	_, err = e.Env.Chains[chainSel].Confirm(tx)
+	require.NoError(t, err)
+
+	tx, err = token.Approve(
+		e.Env.Chains[chainSel].DeployerKey,
+		state.Chains[chainSel].Router.Address(),
+		amount,
+	)
+	require.NoError(t, err)
+	_, err = e.Env.Chains[chainSel].Confirm(tx)
+	require.NoError(t, err)
+}