add some tests

Layr-Labs · Sep 5, 2024 · 5d5e854 · 5d5e854
1 parent 925b333
commit 5d5e854
Show file tree

Hide file tree

Showing 6 changed files with 332 additions and 37 deletions.
diff --git a/pkg/operator/allocations/initializedelay.go b/pkg/operator/allocations/initializedelay.go
@@ -50,7 +50,7 @@ func initializeDelayAction(cCtx *cli.Context, p utils.Prompter) error {
 		return eigenSdkUtils.WrapError("failed to create new eth client", err)
 	}
 
-	// Temp to test modify allocations
+	// Temp to test modify Allocations
 	config.avsDirectoryAddress = gethcommon.HexToAddress("0x8BffE5a668DB26bc5Ce8dC9C0096fB634747b62A")
 
 	if config.broadcast {

diff --git a/pkg/operator/allocations/testdata/allocations1.csv b/pkg/operator/allocations/testdata/allocations1.csv
@@ -0,0 +1,5 @@
+avs_address,operator_set_id,strategy_address,bips
+0x2222AAC0C980Cc029624b7ff55B88Bc6F63C538f,1,0x49989b32351Eb9b8ab2d5623cF22E7F7C23e5630,2000
+0x2222AAC0C980Cc029624b7ff55B88Bc6F63C538f,3,0x49989b32351Eb9b8ab2d5623cF22E7F7C23e5630,1000
+0x111116fE4F8C2f83E3eB2318F090557b7CD0BF76,4,0x232326fE4F8C2f83E3eB2318F090557b7CD02222,3000
+0x111116fE4F8C2f83E3eB2318F090557b7CD0BF76,5,0x545456fE4F8C2f83E3eB2318F090557b7CD04567,4000
diff --git a/pkg/operator/allocations/testdata/allocations_duplicate.csv b/pkg/operator/allocations/testdata/allocations_duplicate.csv
@@ -0,0 +1,6 @@
+avs_address,operator_set_id,strategy_address,bips
+0x2222AAC0C980Cc029624b7ff55B88Bc6F63C538f,1,0x49989b32351Eb9b8ab2d5623cF22E7F7C23e5630,2000
+0x2222AAC0C980Cc029624b7ff55B88Bc6F63C538f,3,0x49989b32351Eb9b8ab2d5623cF22E7F7C23e5630,1000
+0x111116fE4F8C2f83E3eB2318F090557b7CD0BF76,4,0x232326fE4F8C2f83E3eB2318F090557b7CD02222,3000
+0x111116fE4F8C2f83E3eB2318F090557b7CD0BF76,5,0x545456fE4F8C2f83E3eB2318F090557b7CD04567,4000
+0x111116fE4F8C2f83E3eB2318F090557b7CD0BF76,5,0x545456fE4F8C2f83E3eB2318F090557b7CD04567,5000
diff --git a/pkg/operator/allocations/types.go b/pkg/operator/allocations/types.go
@@ -11,18 +11,18 @@ import (
 	gethcommon "github.com/ethereum/go-ethereum/common"
 )
 
-type bulkModifyAllocations struct {
-	allocations           []contractIAVSDirectory.IAVSDirectoryMagnitudeAllocation
-	allocatableMagnitudes map[gethcommon.Address]uint64
+type BulkModifyAllocations struct {
+	Allocations           []contractIAVSDirectory.IAVSDirectoryMagnitudeAllocation
+	AllocatableMagnitudes map[gethcommon.Address]uint64
 }
 
-func (b *bulkModifyAllocations) Print() {
-	for _, a := range b.allocations {
+func (b *BulkModifyAllocations) Print() {
+	for _, a := range b.Allocations {
 		fmt.Printf(
 			"Strategy: %s, Expected Total Magnitude: %d, Allocatable Magnitude %d\n",
 			a.Strategy.Hex(),
 			a.ExpectedTotalMagnitude,
-			b.allocatableMagnitudes[a.Strategy],
+			b.AllocatableMagnitudes[a.Strategy],
 		)
 		for i, opSet := range a.OperatorSets {
 			fmt.Printf(
@@ -37,21 +37,20 @@ func (b *bulkModifyAllocations) Print() {
 }
 
 type updateConfig struct {
-	network               string
-	rpcUrl                string
-	environment           string
-	chainID               *big.Int
-	output                string
-	outputType            string
-	broadcast             bool
-	operatorAddress       gethcommon.Address
-	avsAddress            gethcommon.Address
-	strategyAddress       gethcommon.Address
-	operatorSetId         uint32
-	bipsToAllocate        uint64
-	bulkModifyAllocations *bulkModifyAllocations
-	signerConfig          *types.SignerConfig
-	csvFilePath           string
+	network         string
+	rpcUrl          string
+	environment     string
+	chainID         *big.Int
+	output          string
+	outputType      string
+	broadcast       bool
+	operatorAddress gethcommon.Address
+	avsAddress      gethcommon.Address
+	strategyAddress gethcommon.Address
+	operatorSetId   uint32
+	bipsToAllocate  uint64
+	signerConfig    *types.SignerConfig
+	csvFilePath     string
 }
 
 type allocation struct {

diff --git a/pkg/operator/allocations/update.go b/pkg/operator/allocations/update.go
@@ -94,7 +94,7 @@ func updateAllocations(cCtx *cli.Context, p utils.Prompter) error {
 
 	allocationsToUpdate, err := generateAllocationsParams(ctx, elReader, config, logger)
 	if err != nil {
-		return eigenSdkUtils.WrapError("failed to generate allocations params", err)
+		return eigenSdkUtils.WrapError("failed to generate Allocations params", err)
 	}
 
 	if config.broadcast {
@@ -120,7 +120,7 @@ func updateAllocations(cCtx *cli.Context, p utils.Prompter) error {
 		receipt, err := eLWriter.ModifyAllocations(
 			ctx,
 			config.operatorAddress,
-			allocationsToUpdate.allocations,
+			allocationsToUpdate.Allocations,
 			contractIAVSDirectory.ISignatureUtilsSignatureWithSaltAndExpiry{
 				Expiry: big.NewInt(0),
 			},
@@ -143,14 +143,14 @@ func updateAllocations(cCtx *cli.Context, p utils.Prompter) error {
 		// to claim. So we hardcode the gas limit to 150_000 so that we can
 		// create unsigned tx without gas limit estimation from contract bindings
 		if common.IsSmartContractAddress(config.operatorAddress, ethClient) {
-			// Claimer is a smart contract
+			// address is a smart contract
 			noSendTxOpts.GasLimit = 150_000
 		}
 
 		unsignedTx, err := contractBindings.AvsDirectory.ModifyAllocations(
 			noSendTxOpts,
 			config.operatorAddress,
-			allocationsToUpdate.allocations,
+			allocationsToUpdate.Allocations,
 			contractIAVSDirectory.ISignatureUtilsSignatureWithSaltAndExpiry{
 				Expiry: big.NewInt(0),
 			})
@@ -211,7 +211,7 @@ func generateAllocationsParams(
 	elReader elChainReader,
 	config *updateConfig,
 	logger logging.Logger,
-) (*bulkModifyAllocations, error) {
+) (*BulkModifyAllocations, error) {
 	allocations := make([]contractIAVSDirectory.IAVSDirectoryMagnitudeAllocation, 0)
 	var allocatableMagnitudes map[gethcommon.Address]uint64
 
@@ -257,9 +257,9 @@ func generateAllocationsParams(
 		}
 	}
 
-	return &bulkModifyAllocations{
-		allocations:           allocations,
-		allocatableMagnitudes: allocatableMagnitudes,
+	return &BulkModifyAllocations{
+		Allocations:           allocations,
+		AllocatableMagnitudes: allocatableMagnitudes,
 	}, nil
 }
 
@@ -272,19 +272,47 @@ func computeAllocations(
 	if err != nil {
 		return nil, nil, eigenSdkUtils.WrapError("failed to parse allocations csv", err)
 	}
+
+	err = validateDataFromCSV(allocations)
+	if err != nil {
+		return nil, nil, eigenSdkUtils.WrapError("failed to validate data from csv", err)
+	}
+
 	strategies := getUniqueStrategies(allocations)
 	strategyTotalMagnitudes, err := parallelGetMagnitudes(strategies, operatorAddress, elReader)
 	if err != nil {
 		return nil, nil, eigenSdkUtils.WrapError("failed to get total magnitudes", err)
 	}
+
 	allocatableMagnitudePerStrategy, err := parallelGetAllocatableMagnitudes(strategies, operatorAddress, elReader)
 	if err != nil {
 		return nil, nil, eigenSdkUtils.WrapError("failed to get allocatable magnitudes", err)
 	}
+
 	magnitudeAllocations := convertAllocationsToMagnitudeAllocations(allocations, strategyTotalMagnitudes)
 	return magnitudeAllocations, allocatableMagnitudePerStrategy, nil
 }
 
+func validateDataFromCSV(allocations []allocation) error {
+	// check for duplicated (avs_address,operator_set_id,strategy_address)
+	tuples := make(map[string]struct{})
+
+	for _, alloc := range allocations {
+		tuple := fmt.Sprintf("%s_%d_%s", alloc.AvsAddress.Hex(), alloc.OperatorSetId, alloc.StrategyAddress.Hex())
+		if _, exists := tuples[tuple]; exists {
+			return fmt.Errorf(
+				"duplicate combination found: avs_address=%s, operator_set_id=%d, strategy_address=%s",
+				alloc.AvsAddress.Hex(),
+				alloc.OperatorSetId,
+				alloc.StrategyAddress.Hex(),
+			)
+		}
+		tuples[tuple] = struct{}{}
+	}
+
+	return nil
+}
+
 func parallelGetAllocatableMagnitudes(
 	strategies []gethcommon.Address,
 	operatorAddress gethcommon.Address,
@@ -451,7 +479,6 @@ func readAndValidateUpdateFlags(cCtx *cli.Context, logger logging.Logger) (*upda
 		logger.Debugf("Failed to get signer config: %s", err)
 	}
 
-	// TODO(shrimalmadhur): Read bulk modify allocations
 	csvFilePath := cCtx.String(flags.CSVFileFlag.Name)
 	chainId := utils.NetworkNameToChainId(network)