From bd5897e74d1a4c42b951c77c5d0cd52cee068f23 Mon Sep 17 00:00:00 2001
From: Ian Shim <100327837+ian-shim@users.noreply.github.com>
Date: Tue, 16 Jan 2024 15:56:18 -0800
Subject: [PATCH] Fix disperser client (#186)

---
 clients/disperser_client.go | 11 +++++------
 tools/traffic/config.go     |  9 ++++++---
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/clients/disperser_client.go b/clients/disperser_client.go
index 44394b5aee..19b5152e38 100644
--- a/clients/disperser_client.go
+++ b/clients/disperser_client.go
@@ -9,8 +9,6 @@ import (
 	disperser_rpc "github.com/Layr-Labs/eigenda/api/grpc/disperser"
 	"github.com/Layr-Labs/eigenda/core"
 	"github.com/Layr-Labs/eigenda/disperser"
-	"github.com/Layr-Labs/eigenda/retriever/flags"
-	"github.com/urfave/cli"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
@@ -23,11 +21,12 @@ type Config struct {
 	UseSecureGrpcFlag bool
 }
 
-func NewConfig(ctx *cli.Context) *Config {
+func NewConfig(hostname, port string, timeout time.Duration, useSecureGrpcFlag bool) *Config {
 	return &Config{
-		Hostname: ctx.GlobalString(flags.HostnameFlag.Name),
-		Port:     ctx.GlobalString(flags.GrpcPortFlag.Name),
-		Timeout:  ctx.Duration(flags.TimeoutFlag.Name),
+		Hostname:          hostname,
+		Port:              port,
+		Timeout:           timeout,
+		UseSecureGrpcFlag: useSecureGrpcFlag,
 	}
 }
 
diff --git a/tools/traffic/config.go b/tools/traffic/config.go
index 3476bf3a1f..a960a05a74 100644
--- a/tools/traffic/config.go
+++ b/tools/traffic/config.go
@@ -20,12 +20,16 @@ type Config struct {
 	LoggingConfig          logging.Config
 	RandomizeBlobs         bool
 	InstanceLaunchInterval time.Duration
-	UseSecureGrpcFlag      bool
 }
 
 func NewConfig(ctx *cli.Context) *Config {
 	return &Config{
-		Config:                 *clients.NewConfig(ctx),
+		Config: *clients.NewConfig(
+			ctx.GlobalString(flags.HostnameFlag.Name),
+			ctx.GlobalString(flags.GrpcPortFlag.Name),
+			ctx.Duration(flags.TimeoutFlag.Name),
+			ctx.GlobalBool(flags.UseSecureGrpcFlag.Name),
+		),
 		NumInstances:           ctx.GlobalUint(flags.NumInstancesFlag.Name),
 		RequestInterval:        ctx.Duration(flags.RequestIntervalFlag.Name),
 		DataSize:               ctx.GlobalUint64(flags.DataSizeFlag.Name),
@@ -34,6 +38,5 @@ func NewConfig(ctx *cli.Context) *Config {
 		LoggingConfig:          logging.ReadCLIConfig(ctx, flags.FlagPrefix),
 		RandomizeBlobs:         ctx.GlobalBool(flags.RandomizeBlobsFlag.Name),
 		InstanceLaunchInterval: ctx.Duration(flags.InstanceLaunchIntervalFlag.Name),
-		UseSecureGrpcFlag:      ctx.GlobalBool(flags.UseSecureGrpcFlag.Name),
 	}
 }