[cmd/telemetrygen] fix issue when passing status-code as string to co…

…mmand - strings were not valid JSON and would therefore cause the JSON Unmarshal to return an error - each status code is now captured explicitly in a switch case statement - mixed case is now supported when passing it as a string also - `TestSpanStatuses` test was updated to run as sub-tests
open-telemetry · Sep 18, 2023 · 5c916ec · 5c916ec
1 parent 8eea663
commit 5c916ec
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 35 deletions.
diff --git a/cmd/telemetrygen/internal/traces/traces.go b/cmd/telemetrygen/internal/traces/traces.go
@@ -6,6 +6,7 @@ package traces
 import (
 	"context"
 	"fmt"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -123,13 +124,18 @@ func Run(c *Config, logger *zap.Logger) error {
 	}
 
 	var statusCode codes.Code
-	if c.StatusCode == "" {
+
+	switch strings.ToLower(c.StatusCode) {
+	case "0", "unset", "":
 		statusCode = codes.Unset
-	} else {
-		if err := statusCode.UnmarshalJSON([]byte(c.StatusCode)); err != nil {
-			return fmt.Errorf("expected `status-code` to be one of (Unset, Error, Ok) or (0, 1, 2), got %q instead", c.StatusCode)
-		}
+	case "1", "error":
+		statusCode = codes.Error
+	case "2", "ok":
+		statusCode = codes.Ok
+	default:
+		return fmt.Errorf("expected `status-code` to be one of (Unset, Error, Ok) or (0, 1, 2), got %q instead", c.StatusCode)
 	}
+
 	wg := sync.WaitGroup{}
 
 	running := &atomic.Bool{}

diff --git a/cmd/telemetrygen/internal/traces/worker_test.go b/cmd/telemetrygen/internal/traces/worker_test.go
@@ -131,46 +131,53 @@ func TestSpanStatuses(t *testing.T) {
 		spanStatus  codes.Code
 		validInput  bool
 	}{
-		{inputStatus: `"Unset"`, spanStatus: codes.Unset, validInput: true},
-		{inputStatus: `"Error"`, spanStatus: codes.Error, validInput: true},
-		{inputStatus: `"Ok"`, spanStatus: codes.Ok, validInput: true},
+		{inputStatus: `Unset`, spanStatus: codes.Unset, validInput: true},
+		{inputStatus: `Error`, spanStatus: codes.Error, validInput: true},
+		{inputStatus: `Ok`, spanStatus: codes.Ok, validInput: true},
+		{inputStatus: `unset`, spanStatus: codes.Unset, validInput: true},
+		{inputStatus: `error`, spanStatus: codes.Error, validInput: true},
+		{inputStatus: `ok`, spanStatus: codes.Ok, validInput: true},
+		{inputStatus: `UNSET`, spanStatus: codes.Unset, validInput: true},
+		{inputStatus: `ERROR`, spanStatus: codes.Error, validInput: true},
+		{inputStatus: `OK`, spanStatus: codes.Ok, validInput: true},
 		{inputStatus: `0`, spanStatus: codes.Unset, validInput: true},
 		{inputStatus: `1`, spanStatus: codes.Error, validInput: true},
 		{inputStatus: `2`, spanStatus: codes.Ok, validInput: true},
-		{inputStatus: `"Foo"`, spanStatus: codes.Unset, validInput: false},
-		{inputStatus: `"UNSET"`, spanStatus: codes.Unset, validInput: false},
+		{inputStatus: `Foo`, spanStatus: codes.Unset, validInput: false},
 		{inputStatus: `-1`, spanStatus: codes.Unset, validInput: false},
 		{inputStatus: `3`, spanStatus: codes.Unset, validInput: false},
+		{inputStatus: `Err`, spanStatus: codes.Unset, validInput: false},
 	}
 
 	for _, tt := range tests {
-		syncer := &mockSyncer{}
-
-		tracerProvider := sdktrace.NewTracerProvider()
-		sp := sdktrace.NewSimpleSpanProcessor(syncer)
-		tracerProvider.RegisterSpanProcessor(sp)
-		otel.SetTracerProvider(tracerProvider)
-
-		cfg := &Config{
-			Config: common.Config{
-				WorkerCount: 1,
-			},
-			NumTraces:  1,
-			StatusCode: tt.inputStatus,
-		}
-
-		// test the program given input, including erroneous inputs
-		if tt.validInput {
-			require.NoError(t, Run(cfg, zap.NewNop()))
-			// verify that the default the span status is set as expected
-			for _, span := range syncer.spans {
-				assert.Equal(t, span.Status().Code, tt.spanStatus, fmt.Sprintf("span status: %v and expected status %v", span.Status().Code, tt.spanStatus))
+		t.Run(fmt.Sprintf("inputStatus=%s", tt.inputStatus), func(t *testing.T) {
+			syncer := &mockSyncer{}
+
+			tracerProvider := sdktrace.NewTracerProvider()
+			sp := sdktrace.NewSimpleSpanProcessor(syncer)
+			tracerProvider.RegisterSpanProcessor(sp)
+			otel.SetTracerProvider(tracerProvider)
+
+			cfg := &Config{
+				Config: common.Config{
+					WorkerCount: 1,
+				},
+				NumTraces:  1,
+				StatusCode: tt.inputStatus,
 			}
-		} else {
-			require.Error(t, Run(cfg, zap.NewNop()))
-		}
-	}
 
+			// test the program given input, including erroneous inputs
+			if tt.validInput {
+				require.NoError(t, Run(cfg, zap.NewNop()))
+				// verify that the default the span status is set as expected
+				for _, span := range syncer.spans {
+					assert.Equal(t, span.Status().Code, tt.spanStatus, fmt.Sprintf("span status: %v and expected status %v", span.Status().Code, tt.spanStatus))
+				}
+			} else {
+				require.Error(t, Run(cfg, zap.NewNop()))
+			}
+		})
+	}
 }
 
 var _ sdktrace.SpanExporter = (*mockSyncer)(nil)