Add logs and test

Layr-Labs · Jun 7, 2024 · fa38314 · fa38314
1 parent 2635c32
commit fa38314
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 2 deletions.
diff --git a/tools/traffic/cmd/main.go b/tools/traffic/cmd/main.go
@@ -39,6 +39,7 @@ func trafficGeneratorMain(ctx *cli.Context) error {
 
 	var signer core.BlobRequestSigner
 	if config.SignerPrivateKey != "" {
+		log.Println("Using signer private key")
 		signer = auth.NewSigner(config.SignerPrivateKey)
 	}
 

diff --git a/tools/traffic/generator.go b/tools/traffic/generator.go
@@ -105,15 +105,15 @@ func (g *TrafficGenerator) sendRequest(ctx context.Context, data []byte) error {
 			return err
 		}
 
-		g.Logger.Info("successfully dispersed new blob,", "key", hex.EncodeToString(key), "status", blobStatus.String())
+		g.Logger.Info("successfully dispersed new blob", "authenticated", true, "key", hex.EncodeToString(key), "status", blobStatus.String())
 		return nil
 	} else {
 		blobStatus, key, err := g.DisperserClient.DisperseBlob(ctxTimeout, data, g.Config.CustomQuorums)
 		if err != nil {
 			return err
 		}
 
-		g.Logger.Info("successfully dispersed new blob,", "key", hex.EncodeToString(key), "status", blobStatus.String())
+		g.Logger.Info("successfully dispersed new blob", "authenticated", false, "key", hex.EncodeToString(key), "status", blobStatus.String())
 		return nil
 	}
 

diff --git a/tools/traffic/generator_test.go b/tools/traffic/generator_test.go
@@ -40,3 +40,32 @@ func TestTrafficGenerator(t *testing.T) {
 	cancel()
 	disperserClient.AssertNumberOfCalls(t, "DisperseBlob", 2)
 }
+
+func TestTrafficGeneratorAuthenticated(t *testing.T) {
+	disperserClient := clientsmock.NewMockDisperserClient()
+	logger := logging.NewNoopLogger()
+
+	trafficGenerator := &traffic.TrafficGenerator{
+		Logger: logger,
+		Config: &traffic.Config{
+			Config: clients.Config{
+				Timeout: 1 * time.Second,
+			},
+			DataSize:         1000_000,
+			RequestInterval:  2 * time.Second,
+			SignerPrivateKey: "Hi",
+		},
+		DisperserClient: disperserClient,
+	}
+
+	processing := disperser.Processing
+	disperserClient.On("DisperseBlobAuthenticated", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+		Return(&processing, []byte{1}, nil)
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		_ = trafficGenerator.StartTraffic(ctx)
+	}()
+	time.Sleep(5 * time.Second)
+	cancel()
+	disperserClient.AssertNumberOfCalls(t, "DisperseBlobAuthenticated", 2)
+}