From 6dda27b46504fc54dcc20c663c987b33bdd8f287 Mon Sep 17 00:00:00 2001
From: Maximo Palopoli <96491141+maximopalopoli@users.noreply.github.com>
Date: Fri, 7 Feb 2025 18:06:09 -0300
Subject: [PATCH] feat: add set ejection cooldown function (#545)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Tomás Grüner <47506558+MegaRedHand@users.noreply.github.com>
---
 chainio/clients/avsregistry/writer.go      | 24 +++++++++++++++
 chainio/clients/avsregistry/writer_test.go | 34 ++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/chainio/clients/avsregistry/writer.go b/chainio/clients/avsregistry/writer.go
index 079dfc3b..cebffb9b 100644
--- a/chainio/clients/avsregistry/writer.go
+++ b/chainio/clients/avsregistry/writer.go
@@ -598,3 +598,27 @@ func (w *ChainWriter) SetAccountIdentifier(
 	}
 	return receipt, nil
 }
+
+// Sets the ejection cooldown with the value received by parameter. The ejection cooldown is the time an operator has to
+// wait to join any quorum after being rejected. Returns the receipt of the transaction in case of success.
+func (w *ChainWriter) SetEjectionCooldown(
+	ctx context.Context,
+	ejectionCooldown *big.Int,
+	waitForReceipt bool,
+) (*gethtypes.Receipt, error) {
+	w.logger.Info("setting ejection cooldown with value ", ejectionCooldown)
+
+	noSendTxOpts, err := w.txMgr.GetNoSendTxOpts()
+	if err != nil {
+		return nil, err
+	}
+	tx, err := w.registryCoordinator.SetEjectionCooldown(noSendTxOpts, ejectionCooldown)
+	if err != nil {
+		return nil, err
+	}
+	receipt, err := w.txMgr.Send(ctx, tx, waitForReceipt)
+	if err != nil {
+		return nil, utils.WrapError("failed to send SetEjectionCooldown tx with err", err.Error())
+	}
+	return receipt, nil
+}
diff --git a/chainio/clients/avsregistry/writer_test.go b/chainio/clients/avsregistry/writer_test.go
index 13618a3b..3306d69f 100644
--- a/chainio/clients/avsregistry/writer_test.go
+++ b/chainio/clients/avsregistry/writer_test.go
@@ -603,3 +603,37 @@ func TestSetAccountIdentifier(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, newAccountIdentifier.String(), testutils.ANVIL_SECOND_ADDRESS)
 }
+
+func TestSetEjectionCooldown(t *testing.T) {
+	// Test set up
+	clients, anvilHttpEndpoint := testclients.BuildTestClients(t)
+
+	contractAddrs := testutils.GetContractAddressesFromContractRegistry(anvilHttpEndpoint)
+
+	chainWriter := clients.AvsRegistryChainWriter
+
+	ejectionCooldown := big.NewInt(2873)
+
+	ethHttpClient := clients.EthHttpClient
+
+	registryCoordinatorContract, err := regcoord.NewContractRegistryCoordinator(
+		contractAddrs.RegistryCoordinator,
+		ethHttpClient,
+	)
+	require.NoError(t, err)
+
+	// At first, ejectionCooldown is zero
+	cooldown, err := registryCoordinatorContract.EjectionCooldown(&bind.CallOpts{})
+	require.NoError(t, err)
+	assert.Equal(t, cooldown.Int64(), int64(0))
+
+	// Set a new ejectionCooldown
+	receipt, err := chainWriter.SetEjectionCooldown(context.Background(), ejectionCooldown, true)
+	require.NoError(t, err)
+	require.Equal(t, receipt.Status, gethtypes.ReceiptStatusSuccessful)
+
+	// After change, ejectionCooldown is the value set
+	newCooldown, err := registryCoordinatorContract.EjectionCooldown(&bind.CallOpts{})
+	require.NoError(t, err)
+	assert.Equal(t, newCooldown, ejectionCooldown)
+}