Added tests for the go/trace package (#15052)

Signed-off-by: VaibhavMalik4187 <vaibhavmalik2018@gmail.com>

go/trace/fake_test.go

@@ -0,0 +1,33 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package trace
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNoopTracingServer(t *testing.T) {
+	factoryFunc := tracingBackendFactories["noop"]
+	tracingSvc, closer, err := factoryFunc("value")
+	require.NoError(t, err)
+	require.NoError(t, closer.Close())
+	span, err := tracingSvc.NewFromString("parent", "label")
+	require.NoError(t, err)
+	require.Empty(t, span)
+}

go/trace/logger_test.go

@@ -0,0 +1,72 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package trace
+
+import (
+	"io"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// If captureStdout is false, it will capture the outut of
+// os.Stderr
+func captureOutput(t *testing.T, f func(), captureStdout bool) string {
+	oldVal := os.Stderr
+	if captureStdout {
+		oldVal = os.Stdout
+	}
+	t.Cleanup(func() {
+		// Ensure reset even if deferred function panics
+		if captureStdout {
+			os.Stdout = oldVal
+		} else {
+			os.Stderr = oldVal
+		}
+	})
+
+	r, w, _ := os.Pipe()
+	if captureStdout {
+		os.Stdout = w
+	} else {
+		os.Stderr = w
+	}
+
+	f()
+
+	w.Close()
+	got, _ := io.ReadAll(r)
+
+	return string(got)
+}
+
+func TestLoggerLogAndError(t *testing.T) {
+	logger := traceLogger{}
+
+	// Test Error() output
+	output := captureOutput(t, func() {
+		logger.Error("This is an error message")
+	}, false)
+	assert.Contains(t, output, "This is an error message")
+
+	// Test Log() output
+	output = captureOutput(t, func() {
+		logger.Log("This is an log message")
+	}, false)
+	assert.Contains(t, output, "This is an log message")
+}

go/trace/opentracing_test.go

@@ -17,12 +17,14 @@ limitations under the License.
 package trace
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"testing"
 
 	"github.com/opentracing/opentracing-go"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestExtractMapFromString(t *testing.T) {
@@ -47,3 +49,29 @@ func TestErrorConditions(t *testing.T) {
 	_, err = extractMapFromString("this is not base64") // malformed base64
 	assert.Error(t, err)
 }
+
+func TestNewClientSpan(t *testing.T) {
+	svc := openTracingService{
+		Tracer: &fakeTracer{},
+	}
+	clientSpan := svc.NewClientSpan(nil, "test-svc", "test-label")
+	require.NotEmpty(t, clientSpan)
+
+	clientSpan = svc.New(clientSpan, "client-span")
+	require.NotEmpty(t, clientSpan)
+
+	spanFromCtx, ok := svc.FromContext(context.Background())
+	require.False(t, ok)
+	require.Nil(t, spanFromCtx)
+
+	ctx := svc.NewContext(context.TODO(), clientSpan)
+	require.NotNil(t, ctx)
+	clientSpan.Finish()
+
+	spanFromCtx, ok = svc.FromContext(ctx)
+	require.True(t, ok)
+	require.NotEmpty(t, spanFromCtx)
+
+	ctx = svc.NewContext(context.TODO(), &mockSpan{})
+	require.Nil(t, ctx)
+}

go/trace/plugin_datadog_test.go

@@ -0,0 +1,39 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package trace
+
+import (
+	"testing"
+
+	"github.com/opentracing/opentracing-go"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetOpenTracingTracer(t *testing.T) {
+	tracer := datadogTracer{
+		actual: opentracing.GlobalTracer(),
+	}
+	require.Equal(t, opentracing.GlobalTracer(), tracer.GetOpenTracingTracer())
+}
+
+func TestNewDataDogTracerHostAndPortNotSet(t *testing.T) {
+	tracingSvc, closer, err := newDatadogTracer("svc")
+	expectedErr := "need host and port to datadog agent to use datadog tracing"
+	require.ErrorContains(t, err, expectedErr)
+	require.Nil(t, tracingSvc)
+	require.Nil(t, closer)
+}

go/trace/plugin_jaeger_test.go

@@ -0,0 +1,35 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package trace
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewJaegerTracerFromEnv(t *testing.T) {
+	tracingSvc, closer, err := newJagerTracerFromEnv("noop")
+	require.NoError(t, err)
+	require.NotEmpty(t, tracingSvc)
+	require.NotEmpty(t, closer)
+
+	tracingSvc, closer, err = newJagerTracerFromEnv("")
+	require.ErrorContains(t, err, "no service name provided")
+	require.Empty(t, tracingSvc)
+	require.Empty(t, closer)
+}

go/trace/trace_test.go

@@ -22,7 +22,9 @@ import (
 	"io"
 	"testing"
 
+	"github.com/opentracing/opentracing-go"
 	"github.com/spf13/viper"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 
 	"vitess.io/vitess/go/viperutil/vipertest"
@@ -68,13 +70,80 @@ func TestRegisterService(t *testing.T) {
 	}
 }
 
+func TestNewFromString(t *testing.T) {
+	tests := []struct {
+		parent      string
+		label       string
+		context     context.Context
+		expectedLog string
+		isPresent   bool
+		expectedErr string
+	}{
+		{
+			parent:      "",
+			label:       "empty parent",
+			context:     context.TODO(),
+			expectedLog: "",
+			isPresent:   true,
+			expectedErr: "parent is empty",
+		},
+		{
+			parent:      "parent",
+			label:       "non-empty parent",
+			expectedLog: "[key: sql-statement-type values:non-empty parent]\n",
+			context:     context.Background(),
+			isPresent:   false,
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.label, func(t *testing.T) {
+			span, ctx, err := NewFromString(context.Background(), tt.parent, tt.label)
+			if tt.expectedErr == "" {
+				require.NoError(t, err)
+				require.NotEmpty(t, span)
+				require.Equal(t, tt.context, ctx)
+
+				got := captureOutput(t, func() {
+					AnnotateSQL(span, &fakeStringer{tt.label})
+				}, true)
+
+				require.Equal(t, tt.expectedLog, got)
+			} else {
+				require.ErrorContains(t, err, tt.expectedErr)
+				require.Nil(t, span)
+				require.Nil(t, ctx)
+			}
+
+			copySpan := CopySpan(context.TODO(), tt.context)
+			if tt.isPresent {
+				require.Equal(t, tt.context, copySpan)
+			} else {
+				require.Equal(t, context.TODO(), copySpan)
+			}
+		})
+	}
+}
+
+func TestNilCloser(t *testing.T) {
+	nc := nilCloser{}
+	require.Nil(t, nc.Close())
+}
+
 type fakeTracer struct {
 	name string
 	log  []string
 }
 
+func (f *fakeTracer) GetOpenTracingTracer() opentracing.Tracer {
+	return opentracing.GlobalTracer()
+}
+
 func (f *fakeTracer) NewFromString(parent, label string) (Span, error) {
-	panic("implement me")
+	if parent == "" {
+		return &mockSpan{tracer: f}, fmt.Errorf("parent is empty")
+	}
+	return &mockSpan{tracer: f}, nil
 }
 
 func (f *fakeTracer) New(parent Span, label string) Span {
@@ -84,7 +153,10 @@ func (f *fakeTracer) New(parent Span, label string) Span {
 }
 
 func (f *fakeTracer) FromContext(ctx context.Context) (Span, bool) {
-	return nil, false
+	if ctx == context.Background() {
+		return nil, false
+	}
+	return &mockSpan{}, true
 }
 
 func (f *fakeTracer) NewContext(parent context.Context, span Span) context.Context {
@@ -113,4 +185,13 @@ func (m *mockSpan) Finish() {
 
 func (m *mockSpan) Annotate(key string, value any) {
 	m.tracer.log = append(m.tracer.log, fmt.Sprintf("key: %v values:%v", key, value))
+	fmt.Println(m.tracer.log)
+}
+
+type fakeStringer struct {
+	str string
+}
+
+func (fs *fakeStringer) String() string {
+	return fs.str
 }

go/trace/utils_test.go

@@ -0,0 +1,41 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package trace
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestLogErrorsWhenClosing(t *testing.T) {
+	logFunc := LogErrorsWhenClosing(&fakeCloser{})
+
+	got := captureOutput(t, func() {
+		logFunc()
+	}, false)
+
+	require.Contains(t, string(got), "test error")
+}
+
+type fakeCloser struct {
+}
+
+func (fc *fakeCloser) Close() error {
+	return fmt.Errorf("test error")
+}