feat: improvements from review

Signed-off-by: Smuu <[email protected]>
celestiaorg · Jun 5, 2024 · fa4c769 · fa4c769
1 parent da47cd9
commit fa4c769
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 49 deletions.
diff --git a/e2e/tshark/tshark_test.go b/e2e/tshark/tshark_test.go
@@ -20,25 +20,20 @@ func TestTshark(t *testing.T) {
 	defer cancel()
 
 	knuu, err := knuu.New(ctx)
-	if err != nil {
-		t.Fatalf("Error creating knuu: %v", err)
-	}
+	require.NoError(t, err, "Error creating knuu")
 
 	scope := knuu.Scope()
 	t.Log(scope)
 
 	instance, err := knuu.NewInstance("alpine")
-	if err != nil {
-		t.Fatalf("Error creating instance '%v':", err)
-	}
+	require.NoError(t, err, "Error creating instance")
+
 	err = instance.SetImage(ctx, "docker.io/alpine:latest")
-	if err != nil {
-		t.Fatalf("Error setting image: %v", err)
-	}
+	require.NoError(t, err, "Error setting image")
+
 	err = instance.SetCommand("sleep", "infinity")
-	if err != nil {
-		t.Fatalf("Error setting command: %v", err)
-	}
+	require.NoError(t, err, "Error setting command")
+
 	err = instance.EnableTsharkCollector(
 		"10Gi", // Example volume size
 		os.Getenv("S3_ACCESS_KEY"),
@@ -47,13 +42,10 @@ func TestTshark(t *testing.T) {
 		os.Getenv("S3_BUCKET_NAME"),
 		"tshark/"+scope,
 	)
-	if err != nil {
-		t.Fatalf("Error enabling tshark collector: %v", err)
-	}
+	require.NoError(t, err, "Error enabling tshark collector")
+
 	err = instance.Commit()
-	if err != nil {
-		t.Fatalf("Error committing instance: %v", err)
-	}
+	require.NoError(t, err, "Error committing instance")
 
 	t.Cleanup(func() {
 		require.NoError(t, instance.Destroy(ctx))
@@ -62,17 +54,13 @@ func TestTshark(t *testing.T) {
 	// Test logic
 
 	err = instance.Start(ctx)
-	if err != nil {
-		t.Fatalf("Error starting instance: %v", err)
-	}
+	require.NoError(t, err, "Error starting instance")
+
 	err = instance.WaitInstanceIsRunning(ctx)
-	if err != nil {
-		t.Fatalf("Error waiting for instance to be running: %v", err)
-	}
+	require.NoError(t, err, "Error waiting for instance to be running")
+
 	wget, err := instance.ExecuteCommand(ctx, "echo", "Hello World!")
-	if err != nil {
-		t.Fatalf("Error executing command '%v':", err)
-	}
+	require.NoError(t, err, "Error executing command")
 
 	// wait for 2 minutes to upload network traces to s3
 	time.Sleep(2 * time.Minute)

diff --git a/pkg/instance/tshark.go b/pkg/instance/tshark.go
@@ -4,45 +4,50 @@ import (
 	"context"
 )
 
+const (
+	tsharkCollectorName       = "tshark-collector"
+	tsharkCollectorImage      = "ghcr.io/celestiaorg/tshark-s3:pr-5"
+	tsharkCollectorCPU        = "100m"
+	tsharkCollectorMemory     = "250Mi"
+	tsharkCollectorVolumePath = "/tshark"
+	netAdminCapability        = "NET_ADMIN"
+)
+
 func (i *Instance) createTsharkCollectorInstance(ctx context.Context) (*Instance, error) {
-	tsharkCollector, err := New("tshark-collector", i.SystemDependencies)
+	tsharkCollector, err := New(tsharkCollectorName, i.SystemDependencies)
 	if err != nil {
 		return nil, err
 	}
-	if tsharkCollector.SetImage(ctx, "ghcr.io/celestiaorg/tshark-s3:pr-5") != nil {
+	if tsharkCollector.SetImage(ctx, tsharkCollectorImage) != nil {
 		return nil, err
 	}
 	if tsharkCollector.Commit() != nil {
 		return nil, err
 	}
-	if tsharkCollector.SetCPU("100m") != nil {
-		return nil, err
-	}
-	if tsharkCollector.SetMemory("250Mi", "250Mi") != nil {
+	if tsharkCollector.SetCPU(tsharkCollectorCPU) != nil {
 		return nil, err
 	}
-	if tsharkCollector.AddVolume("/tshark", i.tsharkCollectorConfig.volumeSize) != nil {
+	if tsharkCollector.SetMemory(tsharkCollectorMemory, tsharkCollectorMemory) != nil {
 		return nil, err
 	}
-	if tsharkCollector.SetEnvironmentVariable("STORAGE_ACCESS_KEY_ID", i.tsharkCollectorConfig.s3AccessKey) != nil {
+	if tsharkCollector.AddVolume(tsharkCollectorVolumePath, i.tsharkCollectorConfig.volumeSize) != nil {
 		return nil, err
 	}
-	if tsharkCollector.SetEnvironmentVariable("STORAGE_SECRET_ACCESS_KEY", i.tsharkCollectorConfig.s3SecretKey) != nil {
-		return nil, err
-	}
-	if tsharkCollector.SetEnvironmentVariable("STORAGE_REGION", i.tsharkCollectorConfig.s3Region) != nil {
-		return nil, err
+	envVars := map[string]string{
+		"STORAGE_ACCESS_KEY_ID":     i.tsharkCollectorConfig.s3AccessKey,
+		"STORAGE_SECRET_ACCESS_KEY": i.tsharkCollectorConfig.s3SecretKey,
+		"STORAGE_REGION":            i.tsharkCollectorConfig.s3Region,
+		"STORAGE_BUCKET_NAME":       i.tsharkCollectorConfig.s3Bucket,
+		"STORAGE_KEY_PREFIX":        i.tsharkCollectorConfig.s3KeyPrefix,
+		"CAPTURE_FILE_NAME":         i.k8sName + ".pcapng",
 	}
-	if tsharkCollector.SetEnvironmentVariable("STORAGE_BUCKET_NAME", i.tsharkCollectorConfig.s3Bucket) != nil {
-		return nil, err
-	}
-	if tsharkCollector.SetEnvironmentVariable("STORAGE_KEY_PREFIX", i.tsharkCollectorConfig.s3KeyPrefix) != nil {
-		return nil, err
-	}
-	if tsharkCollector.SetEnvironmentVariable("CAPTURE_FILE_NAME", i.k8sName+".pcapng") != nil {
-		return nil, err
+
+	for key, value := range envVars {
+		if tsharkCollector.SetEnvironmentVariable(key, value) != nil {
+			return nil, err
+		}
 	}
-	if tsharkCollector.AddCapability("NET_ADMIN") != nil {
+	if tsharkCollector.AddCapability(netAdminCapability) != nil {
 		return nil, err
 	}
 	return tsharkCollector, nil