Update test

cmd/updater.go

@@ -3,16 +3,12 @@ package cmd
 import (
 	"context"
 	"fmt"
-	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/sidecar"
-	"log"
-	"net/http"
-
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/internal/logger"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/internal/metrics"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/chainClient"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/config"
-	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/proofDataFetcher/httpProofDataFetcher"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/services"
+	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/sidecar"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/tracer"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/updater"
 	gethcommon "github.com/ethereum/go-ethereum/common"
@@ -22,6 +18,7 @@ import (
 	"github.com/spf13/viper"
 	"go.uber.org/zap"
 	ddTracer "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"log"
 )
 
 func runUpdater(ctx context.Context, cfg *config.UpdaterConfig, logger *zap.Logger) error {
@@ -34,7 +31,7 @@ func runUpdater(ctx context.Context, cfg *config.UpdaterConfig, logger *zap.Logg
 		return err
 	}
 
-	chainClient, err := chainClient.NewChainClient(ctx, ethClient, cfg.PrivateKey)
+	cc, err := chainClient.NewChainClient(ctx, ethClient, cfg.PrivateKey)
 	if err != nil {
 		logger.Sugar().Errorf("Failed to create new chain client with private key", zap.Error(err))
 		return err
@@ -46,22 +43,19 @@ func runUpdater(ctx context.Context, cfg *config.UpdaterConfig, logger *zap.Logg
 		return err
 	}
 
-	e, _ := config.StringEnvironmentFromEnum(cfg.Environment)
-	dataFetcher := httpProofDataFetcher.NewHttpProofDataFetcher(cfg.ProofStoreBaseUrl, e, cfg.Network, http.DefaultClient, logger)
-
-	transactor, err := services.NewTransactor(chainClient, gethcommon.HexToAddress(cfg.RewardsCoordinatorAddress))
+	transactor, err := services.NewTransactor(cc, gethcommon.HexToAddress(cfg.RewardsCoordinatorAddress))
 	if err != nil {
 		logger.Sugar().Errorf("Failed to initialize transactor", zap.Error(err))
 		return err
 	}
 
-	u, err := updater.NewUpdater(transactor, dataFetcher, sidecarClient, logger)
+	u, err := updater.NewUpdater(transactor, sidecarClient, logger)
 	if err != nil {
 		logger.Sugar().Errorf("Failed to create updater", zap.Error(err))
 		return err
 	}
 
-	err = u.Update(ctx)
+	_, err = u.Update(ctx)
 	if err != nil {
 		logger.Sugar().Infow("Failed to update", zap.Error(err))
 		return nil

pkg/sidecar/client.go

@@ -1,14 +1,20 @@
 package sidecar
 
 import (
+	"context"
 	"crypto/tls"
 	"github.com/Layr-Labs/protocol-apis/gen/protos/eigenlayer/sidecar/v1"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 )
 
+type IRewardsClient interface {
+	GenerateRewards(ctx context.Context, req *v1.GenerateRewardsRequest, opts ...grpc.CallOption) (*v1.GenerateRewardsResponse, error)
+	GenerateRewardsRoot(ctx context.Context, req *v1.GenerateRewardsRootRequest, opts ...grpc.CallOption) (*v1.GenerateRewardsRootResponse, error)
+}
+
 type SidecarClient struct {
-	Rewards v1.RewardsClient
+	Rewards IRewardsClient
 }
 
 func NewSidecarClient(url string) (*SidecarClient, error) {

pkg/updater/updater.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/internal/metrics"
-	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/proofDataFetcher"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/services"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/sidecar"
 	v1 "github.com/Layr-Labs/protocol-apis/gen/protos/eigenlayer/sidecar/v1"
@@ -14,28 +13,30 @@ import (
 )
 
 type Updater struct {
-	transactor       services.Transactor
-	logger           *zap.Logger
-	proofDataFetcher proofDataFetcher.ProofDataFetcher
-	sidecarClient    *sidecar.SidecarClient
+	transactor    services.Transactor
+	logger        *zap.Logger
+	sidecarClient *sidecar.SidecarClient
 }
 
 func NewUpdater(
 	transactor services.Transactor,
-	fetcher proofDataFetcher.ProofDataFetcher,
 	sc *sidecar.SidecarClient,
 	logger *zap.Logger,
 ) (*Updater, error) {
 	return &Updater{
-		transactor:       transactor,
-		logger:           logger,
-		sidecarClient:    sc,
-		proofDataFetcher: fetcher,
+		transactor:    transactor,
+		logger:        logger,
+		sidecarClient: sc,
 	}, nil
 }
 
+type UpdatedRoot struct {
+	SnapshotDate string
+	Root         string
+}
+
 // Update fetches the most recent snapshot and the most recent submitted timestamp from the chain.
-func (u *Updater) Update(ctx context.Context) error {
+func (u *Updater) Update(ctx context.Context) (*UpdatedRoot, error) {
 	span, ctx := ddTracer.StartSpanFromContext(ctx, "updater::Update")
 	defer span.Finish()
 
@@ -44,37 +45,41 @@ func (u *Updater) Update(ctx context.Context) error {
 		RespondWithRewardsData: false,
 	})
 	if err != nil {
-		return fmt.Errorf("failed to generate rewards: %w", err)
+		return nil, fmt.Errorf("failed to generate rewards: %w", err)
 	}
 
 	snapshotDateTime, err := time.Parse(time.DateOnly, res.Snapshot)
 	if err != nil {
-		return fmt.Errorf("failed to parse snapshot date: %w", err)
+		return nil, fmt.Errorf("failed to parse snapshot date: %w", err)
 	}
 
 	rootRes, err := u.sidecarClient.Rewards.GenerateRewardsRoot(ctx, &v1.GenerateRewardsRootRequest{
 		Snapshot: res.Snapshot,
 	})
 
 	if err != nil {
-		return fmt.Errorf("failed to generate rewards root: %w", err)
+		return nil, fmt.Errorf("failed to generate rewards root: %w", err)
 	}
 
-	rootBytes := []byte(rootRes.RewardsRoot)
+	rootBytes := [32]byte{}
+	copy(rootBytes[:], []byte(rootRes.RewardsRoot))
 
 	// send the merkle root to the smart contract
 	u.logger.Sugar().Infow("updating rewards", zap.String("new_root", rootRes.RewardsRoot))
 
 	u.logger.Sugar().Infow("Calculated timestamp", zap.Int64("calculated_until_timestamp", snapshotDateTime.Unix()))
-	if err := u.transactor.SubmitRoot(ctx, [32]byte(rootBytes), uint32(snapshotDateTime.Unix())); err != nil {
+	if err := u.transactor.SubmitRoot(ctx, rootBytes, uint32(snapshotDateTime.Unix())); err != nil {
 		metrics.GetStatsdClient().Incr(metrics.Counter_UpdateFails, nil, 1)
 		metrics.IncCounterUpdateRun(metrics.CounterUpdateRunsFailed)
 		u.logger.Sugar().Errorw("Failed to submit root", zap.Error(err))
-		return err
+		return nil, err
 	} else {
 		metrics.GetStatsdClient().Incr(metrics.Counter_UpdateSuccess, nil, 1)
 		metrics.IncCounterUpdateRun(metrics.CounterUpdateRunsSuccess)
 	}
 
-	return nil
+	return &UpdatedRoot{
+		SnapshotDate: res.Snapshot,
+		Root:         rootRes.RewardsRoot,
+	}, nil
 }

pkg/updater/updater_test.go

@@ -5,16 +5,14 @@ import (
 	"fmt"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/internal/logger"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/internal/metrics"
-	"github.com/Layr-Labs/eigenlayer-rewards-updater/internal/testData"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/mocks"
-	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/proofDataFetcher/httpProofDataFetcher"
+	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/sidecar"
 	"github.com/Layr-Labs/eigenlayer-rewards-updater/pkg/updater"
+	v1 "github.com/Layr-Labs/protocol-apis/gen/protos/eigenlayer/sidecar/v1"
+	"google.golang.org/grpc"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
 	ddTracer "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
-	"io"
 	"net/http"
-	"regexp"
-	"strings"
 	"testing"
 	"time"
 
@@ -30,38 +28,19 @@ func (m *mockHttpClient) Do(req *http.Request) (*http.Response, error) {
 	return m.mockDo(req), nil
 }
 
-func TestUpdaterUpdate(t *testing.T) {
-	env := "preprod"
-	network := "holesky"
-	baseUrl := "https://eigenpayments-dev.s3.us-east-2.amazonaws.com"
-
-	mockClient := &mockHttpClient{
-		mockDo: func(r *http.Request) *http.Response {
-			recentSnapshotsRegex := regexp.MustCompile(`\/recent-snapshots\.json$`)
-			claimAmountsRegex := regexp.MustCompile(`\/(\d{4}-\d{2}-\d{2})\/claim-amounts\.json$`)
-
-			fmt.Printf("request url: %s\n", r.URL.String())
-			if recentSnapshotsRegex.Match([]byte(r.URL.String())) {
-				fmt.Printf("Matched recent snapshots: %s\n", r.URL.String())
-				return &http.Response{
-					StatusCode: 200,
-					Body:       io.NopCloser(strings.NewReader(testData.GetFullSnapshotDatesList())),
-				}
-			} else if claimAmountsRegex.Match([]byte(r.URL.String())) {
-				fmt.Printf("Matched claim amounts: %s\n", r.URL.String())
-				return &http.Response{
-					StatusCode: 200,
-					Body:       io.NopCloser(strings.NewReader(testData.GetFullTestEarnerLines())),
-				}
-			}
-
-			return &http.Response{StatusCode: 400}
-		},
-	}
-	// 2024-05-06
-	currentRewardCalcEndTimestamp := uint32(1714953600)
-	expectedRewardTimestamp := time.Unix(int64(1715040000), 0).UTC()
+type mockRewardsClient struct {
+	mock.Mock
+}
+
+func (m *mockRewardsClient) GenerateRewards(ctx context.Context, req *v1.GenerateRewardsRequest, opts ...grpc.CallOption) (*v1.GenerateRewardsResponse, error) {
+	return &v1.GenerateRewardsResponse{Snapshot: "2024-10-31"}, nil
+}
 
+func (m *mockRewardsClient) GenerateRewardsRoot(ctx context.Context, req *v1.GenerateRewardsRootRequest, opts ...grpc.CallOption) (*v1.GenerateRewardsRootResponse, error) {
+	return &v1.GenerateRewardsRootResponse{RewardsRoot: "0xb4a614cc0bf38dff74822a0744aab5b8897a6868c3b612980436be219a25be21"}, nil
+}
+
+func TestUpdaterUpdate(t *testing.T) {
 	_, err := metrics.InitStatsdClient("", false)
 	fmt.Printf("err: %v\n", err)
 
@@ -71,27 +50,29 @@ func TestUpdaterUpdate(t *testing.T) {
 	span, ctx := ddTracer.StartSpanFromContext(context.Background(), "TestUpdaterUpdate")
 	defer span.Finish()
 
-	logger, _ := logger.NewLogger(&logger.LoggerConfig{Debug: true})
+	l, _ := logger.NewLogger(&logger.LoggerConfig{Debug: true})
 
 	mockTransactor := &mocks.Transactor{}
 
-	fetcher := httpProofDataFetcher.NewHttpProofDataFetcher(baseUrl, env, network, mockClient, logger)
+	mockSidecarClient := &sidecar.SidecarClient{
+		Rewards: &mockRewardsClient{},
+	}
 
-	updater, err := updater.NewUpdater(mockTransactor, fetcher, logger)
+	updater, err := updater.NewUpdater(mockTransactor, mockSidecarClient, l)
 	assert.Nil(t, err)
 
-	// setup data
-	processedData, _ := fetcher.ProcessClaimAmountsFromRawBody(ctx, []byte(testData.GetFullTestEarnerLines()))
+	expectedRoot := "0xb4a614cc0bf38dff74822a0744aab5b8897a6868c3b612980436be219a25be21"
+	expectedSnapshotDate := "2024-10-31"
 
-	rootBytes := processedData.AccountTree.Root()
+	expectedSnapshotDateTime, _ := time.Parse(time.DateOnly, expectedSnapshotDate)
 
-	var root [32]byte
-	copy(root[:], rootBytes)
+	expectedRootBytes := [32]byte{}
+	copy(expectedRootBytes[:], []byte(expectedRoot))
 
-	mockTransactor.On("CurrRewardsCalculationEndTimestamp").Return(currentRewardCalcEndTimestamp, nil)
-	mockTransactor.On("SubmitRoot", mock.Anything, root, uint32(expectedRewardTimestamp.Unix())).Return(nil)
+	mockTransactor.On("SubmitRoot", mock.Anything, expectedRootBytes, uint32(expectedSnapshotDateTime.Unix())).Return(nil)
 
-	accountTree, err := updater.Update(ctx)
+	updatedRoot, err := updater.Update(ctx)
 	assert.Nil(t, err)
-	assert.Equal(t, rootBytes, accountTree.Root())
+	assert.Equal(t, expectedRoot, updatedRoot.Root)
+	assert.Equal(t, "2024-10-31", updatedRoot.SnapshotDate)
 }