From 0fbde6ac92652f837a637554727f46b944b02a50 Mon Sep 17 00:00:00 2001
From: Domino Valdano <2644901+reductionista@users.noreply.github.com>
Date: Fri, 31 Jan 2025 16:36:58 -0800
Subject: [PATCH] Fix tests

---
 pkg/reader/ccip_test.go | 104 ++++++++++++++++++++++------------------
 1 file changed, 57 insertions(+), 47 deletions(-)

diff --git a/pkg/reader/ccip_test.go b/pkg/reader/ccip_test.go
index 3692363a6..0d7b3d069 100644
--- a/pkg/reader/ccip_test.go
+++ b/pkg/reader/ccip_test.go
@@ -462,29 +462,23 @@ func TestCCIPChainReader_DiscoverContracts_HappyPath_Round1(t *testing.T) {
 	mockReaders[destChain] = reader_mocks.NewMockExtended(t)
 	addDestinationContractAssertions(mockReaders[destChain], destNonceMgr, destRMNRemote, destFeeQuoter)
 
-	mockReaders[destChain].EXPECT().ExtendedGetLatestValue(
+	mockReaders[destChain].EXPECT().ExtendedBatchGetLatestValuesSingleMethod(
 		mock.Anything,
 		consts.ContractNameOffRamp,
-		consts.MethodNameOffRampGetAllSourceChainConfigs,
-		primitives.Unconfirmed,
-		map[string]any{},
+		consts.MethodNameGetSourceChainConfig,
 		mock.Anything,
-	).Return(nil).Run(withReturnValueOverridden(func(returnVal interface{}) {
-		v := returnVal.(*selectorsAndConfigs)
-		v.Selectors = []uint64{uint64(sourceChain[0]), uint64(sourceChain[1])}
-		v.SourceChainConfigs = []sourceChainConfig{
-			{
-				OnRamp:    onramps[0],
-				Router:    destRouter,
-				IsEnabled: true,
-			},
-			{
-				OnRamp:    onramps[1],
-				Router:    destRouter,
-				IsEnabled: true,
-			},
-		}
-	}))
+		mock.Anything,
+	).RunAndReturn(withBatchReadSourceChainConfigReturnValues(t,
+		[]any{sourceChainConfig{
+			OnRamp:    onramps[0],
+			Router:    destRouter,
+			IsEnabled: true,
+		}, sourceChainConfig{
+			OnRamp:    onramps[1],
+			Router:    destRouter,
+			IsEnabled: true,
+		}},
+	))
 
 	// mock calls to get fee quoter from onramps and source chain config from offramp.
 	for _, selector := range sourceChain {
@@ -605,29 +599,24 @@ func TestCCIPChainReader_DiscoverContracts_HappyPath_Round2(t *testing.T) {
 	mockReaders[destChain] = reader_mocks.NewMockExtended(t)
 	addDestinationContractAssertions(mockReaders[destChain], destNonceMgr, destRMNRemote, destFeeQuoter)
 
-	mockReaders[destChain].EXPECT().ExtendedGetLatestValue(
+	mockReaders[destChain].EXPECT().ExtendedBatchGetLatestValuesSingleMethod(
 		mock.Anything,
 		consts.ContractNameOffRamp,
-		consts.MethodNameOffRampGetAllSourceChainConfigs,
-		primitives.Unconfirmed,
-		map[string]any{},
+		consts.MethodNameGetSourceChainConfig,
 		mock.Anything,
-	).Return(nil).Run(withReturnValueOverridden(func(returnVal interface{}) {
-		v := returnVal.(*selectorsAndConfigs)
-		v.Selectors = []uint64{uint64(sourceChain[0]), uint64(sourceChain[1])}
-		v.SourceChainConfigs = []sourceChainConfig{
-			{
-				OnRamp:    onramps[0],
-				Router:    destRouter[0],
-				IsEnabled: true,
-			},
-			{
+		mock.Anything,
+	).RunAndReturn(withBatchReadSourceChainConfigReturnValues(t,
+		[]any{sourceChainConfig{
+			OnRamp:    onramps[0],
+			Router:    destRouter[0],
+			IsEnabled: true,
+		},
+			sourceChainConfig{
 				OnRamp:    onramps[1],
 				Router:    destRouter[1],
 				IsEnabled: true,
 			},
-		}
-	}))
+		}))
 
 	// mock calls to get fee quoter from onramps and source chain config from offramp.
 	for i, selector := range sourceChain {
@@ -687,14 +676,13 @@ func TestCCIPChainReader_DiscoverContracts_GetAllSourceChainConfig_Errors(t *tes
 
 	// mock the call for sourceChain2 - failure
 	getLatestValueErr := errors.New("some error")
-	destExtended.EXPECT().ExtendedGetLatestValue(
+	destExtended.EXPECT().ExtendedBatchGetLatestValuesSingleMethod(
 		mock.Anything,
 		consts.ContractNameOffRamp,
-		consts.MethodNameOffRampGetAllSourceChainConfigs,
-		primitives.Unconfirmed,
-		map[string]any{},
+		consts.MethodNameGetSourceChainConfig,
 		mock.Anything,
-	).Return(getLatestValueErr)
+		mock.Anything,
+	).Return(nil, getLatestValueErr)
 
 	// get static config call won't occur because the source chain config call failed.
 
@@ -725,15 +713,16 @@ func TestCCIPChainReader_DiscoverContracts_GetOfframpStaticConfig_Errors(t *test
 	destExtended := reader_mocks.NewMockExtended(t)
 
 	// mock the call for source chain configs
-	destExtended.EXPECT().ExtendedGetLatestValue(
+	destExtended.EXPECT().ExtendedBatchGetLatestValuesSingleMethod(
 		mock.Anything,
 		consts.ContractNameOffRamp,
-		consts.MethodNameOffRampGetAllSourceChainConfigs,
-		primitives.Unconfirmed,
-		map[string]any{},
+		consts.MethodNameGetSourceChainConfig,
 		mock.Anything,
-	).Return(nil) // doesn't matter for this test
-	// mock the call to get the nonce manager - failure
+		sourceChainConfig{},
+	).RunAndReturn(withBatchReadSourceChainConfigReturnValues(t,
+		[]any{sourceChainConfig{}, sourceChainConfig{}}))
+
+	// mock the call to get the static config - failure
 	getLatestValueErr := errors.New("some error")
 	destExtended.EXPECT().ExtendedGetLatestValue(
 		mock.Anything,
@@ -781,6 +770,27 @@ func withReturnValueOverridden(mapper func(returnVal interface{})) func(ctx cont
 	}
 }
 
+func withBatchReadSourceChainConfigReturnValues(
+	t testing.TB,
+	retVals []any) func(context.Context, string, string, []any, any) (types.ContractBatchResults, error) {
+	return func(ctx context.Context,
+		contractName string,
+		methodName string,
+		params []any,
+		retValType any) (types.ContractBatchResults, error) {
+		require.GreaterOrEqual(t, len(retVals), 1)
+		_, ok := retVals[0].(sourceChainConfig)
+		require.True(t, ok)
+		results := make(types.ContractBatchResults, 0, len(retVals))
+		for _, retVal := range retVals {
+			res := types.BatchReadResult{ReadName: methodName}
+			res.SetResult(retVal, nil)
+			results = append(results, res)
+		}
+		return results, nil
+	}
+}
+
 func TestCCIPChainReader_getDestFeeQuoterStaticConfig(t *testing.T) {
 	destCR := reader_mocks.NewMockContractReaderFacade(t)
 	destCR.EXPECT().Bind(mock.Anything, mock.Anything).Return(nil)