refactor: reduce redundancy in avsRegistry writer test (#462)

Layr-Labs · Feb 3, 2025 · fae0cf3 · fae0cf3
1 parent 870f687
commit fae0cf3
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 20 deletions.
diff --git a/chainio/clients/avsregistry/writer.go b/chainio/clients/avsregistry/writer.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"crypto/ecdsa"
 	"crypto/rand"
-	"errors"
 	"math/big"
 
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
@@ -23,6 +22,7 @@ import (
 	"github.com/Layr-Labs/eigensdk-go/crypto/bls"
 	"github.com/Layr-Labs/eigensdk-go/logging"
 	"github.com/Layr-Labs/eigensdk-go/types"
+	"github.com/Layr-Labs/eigensdk-go/utils"
 )
 
 type eLReader interface {
@@ -206,7 +206,7 @@ func (w *ChainWriter) RegisterOperatorInQuorumWithAVSRegistryCoordinator(
 	}
 	receipt, err := w.txMgr.Send(ctx, tx, waitForReceipt)
 	if err != nil {
-		return nil, errors.New("failed to send tx with err: " + err.Error())
+		return nil, utils.WrapError("failed to send tx with err", err.Error())
 	}
 	w.logger.Info(
 		"successfully registered operator with AVS registry coordinator",
@@ -328,7 +328,7 @@ func (w *ChainWriter) RegisterOperator(
 	}
 	receipt, err := w.txMgr.Send(ctx, tx, waitForReceipt)
 	if err != nil {
-		return nil, errors.New("failed to send tx with err: " + err.Error())
+		return nil, utils.WrapError("failed to send tx with err", err.Error())
 	}
 	w.logger.Info(
 		"successfully registered operator with AVS registry coordinator",
@@ -370,7 +370,7 @@ func (w *ChainWriter) UpdateStakesOfEntireOperatorSetForQuorums(
 	}
 	receipt, err := w.txMgr.Send(ctx, tx, waitForReceipt)
 	if err != nil {
-		return nil, errors.New("failed to send tx with err: " + err.Error())
+		return nil, utils.WrapError("failed to send tx with err: ", err.Error())
 	}
 	w.logger.Info(
 		"successfully updated stakes for entire operator set",
@@ -401,7 +401,7 @@ func (w *ChainWriter) UpdateStakesOfOperatorSubsetForAllQuorums(
 	}
 	receipt, err := w.txMgr.Send(ctx, tx, waitForReceipt)
 	if err != nil {
-		return nil, errors.New("failed to send tx with err: " + err.Error())
+		return nil, utils.WrapError("failed to send tx with err", err.Error())
 	}
 	w.logger.Info(
 		"successfully updated stakes of operator subset for all quorums",
@@ -432,7 +432,7 @@ func (w *ChainWriter) DeregisterOperator(
 	}
 	receipt, err := w.txMgr.Send(ctx, tx, waitForReceipt)
 	if err != nil {
-		return nil, errors.New("failed to send tx with err: " + err.Error())
+		return nil, utils.WrapError("failed to send tx with err", err.Error())
 	}
 	w.logger.Info(
 		"successfully deregistered operator with the AVS's registry coordinator",
@@ -464,7 +464,7 @@ func (w *ChainWriter) DeregisterOperatorOperatorSets(
 	}
 	receipt, err := w.txMgr.Send(ctx, tx, waitForReceipt)
 	if err != nil {
-		return nil, errors.New("failed to send tx with err: " + err.Error())
+		return nil, utils.WrapError("failed to send tx with err", err.Error())
 	}
 	w.logger.Info(
 		"successfully deregistered operator with the AVS's registry coordinator",
@@ -491,7 +491,7 @@ func (w *ChainWriter) UpdateSocket(
 	}
 	receipt, err := w.txMgr.Send(ctx, tx, waitForReceipt)
 	if err != nil {
-		return nil, errors.New("failed to send UpdateSocket tx with err: " + err.Error())
+		return nil, utils.WrapError("failed to send UpdateSocket tx with err", err.Error())
 	}
 	return receipt, nil
 }
diff --git a/chainio/clients/avsregistry/writer_test.go b/chainio/clients/avsregistry/writer_test.go
@@ -44,6 +44,9 @@ func TestWriterMethods(t *testing.T) {
 
 	quorumNumbers := types.QuorumNums{0}
 
+	subCtx, cancelFn := context.WithCancel(context.Background())
+	cancelFn()
+
 	t.Run("update socket without being registered", func(t *testing.T) {
 		receipt, err := chainWriter.UpdateSocket(
 			context.Background(),
@@ -122,8 +125,6 @@ func TestWriterMethods(t *testing.T) {
 
 	// Error cases
 	t.Run("fail register operator cancelling context", func(t *testing.T) {
-		subCtx, cancelFn := context.WithCancel(context.Background())
-		cancelFn()
 		receipt, err := chainWriter.RegisterOperator(
 			subCtx,
 			ecdsaPrivateKey,
@@ -137,8 +138,6 @@ func TestWriterMethods(t *testing.T) {
 	})
 
 	t.Run("fail update stake of operator subset cancelling context", func(t *testing.T) {
-		subCtx, cancelFn := context.WithCancel(context.Background())
-		cancelFn()
 		receipt, err := chainWriter.UpdateStakesOfOperatorSubsetForAllQuorums(
 			subCtx,
 			[]gethcommon.Address{addr},
@@ -149,8 +148,6 @@ func TestWriterMethods(t *testing.T) {
 	})
 
 	t.Run("fail update stake of entire operator set cancelling context", func(t *testing.T) {
-		subCtx, cancelFn := context.WithCancel(context.Background())
-		cancelFn()
 		receipt, err := chainWriter.UpdateStakesOfEntireOperatorSetForQuorums(
 			subCtx,
 			[][]gethcommon.Address{{addr}},
@@ -174,8 +171,6 @@ func TestWriterMethods(t *testing.T) {
 	})
 
 	t.Run("fail deregister operator cancelling context", func(t *testing.T) {
-		subCtx, cancelFn := context.WithCancel(context.Background())
-		cancelFn()
 		receipt, err := chainWriter.DeregisterOperator(
 			subCtx,
 			quorumNumbers,
@@ -199,9 +194,6 @@ func TestWriterMethods(t *testing.T) {
 	})
 
 	t.Run("fail update socket cancelling context", func(t *testing.T) {
-		subCtx, cancelFn := context.WithCancel(context.Background())
-
-		cancelFn()
 		receipt, err := chainWriter.UpdateSocket(
 			subCtx,
 			types.Socket(""),

diff --git a/testutils/testclients/testclients.go b/testutils/testclients/testclients.go
@@ -195,7 +195,7 @@ func NewTestTxManager(httpEndpoint string, privateKeyHex string) (*txmgr.SimpleT
 	return txManager, nil
 }
 
-// Creates a testing ChainWriter from an httpEndpoint, private key and config.
+// Creates an avsRegistry testing ChainWriter from an httpEndpoint, private key and config.
 // This is needed because the existing testclients.BuildTestClients returns a
 // ChainWriter with a null rewardsCoordinator, which is required for some of the tests.
 func NewTestAvsRegistryWriterFromConfig(