gitlab.com/gitlab-org/labkit@v1.21.0/tracing/env_extractor_test.go

package tracing

import (
	"context"
	"os"
	"reflect"
	"testing"

	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/labkit/correlation"
)

func TestExtractFromEnv(t *testing.T) {
	tests := []struct {
		name              string
		ctx               context.Context
		opts              []ExtractFromEnvOption
		additionalEnv     map[string]string
		wantCorrelationID string
	}{
		{
			name: "no_options",
			ctx:  context.Background(),
			opts: []ExtractFromEnvOption{},
		},
		{
			name: "pass_correlation_id",
			ctx:  context.Background(),
			opts: []ExtractFromEnvOption{},
			additionalEnv: map[string]string{
				envCorrelationIDKey: "abc123",
			},
			wantCorrelationID: "abc123",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			resetEnvironment := addAdditionalEnv(tt.additionalEnv)
			defer resetEnvironment()

			ctx, finished := ExtractFromEnv(tt.ctx, tt.opts...)
			require.NotNil(t, ctx, "ctx is nil")
			require.NotNil(t, finished, "finished is nil")
			gotCorrelationID := correlation.ExtractFromContext(ctx)
			require.Equal(t, tt.wantCorrelationID, gotCorrelationID)
			defer finished()
		})
	}
}

func Test_environAsMap(t *testing.T) {
	tests := []struct {
		name string
		env  []string
		want map[string]string
	}{
		{
			name: "trivial",
			env:  nil,
			want: map[string]string{},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := environAsMap(tt.env); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("environAsMap() = %v, want %v", got, tt.want)
			}
		})
	}
}

// addAdditionalEnv will configure additional environment values
// and return a deferrable function to reset the environment to
// it's original state after the test.
func addAdditionalEnv(envMap map[string]string) func() {
	prevValues := map[string]string{}
	unsetValues := []string{}
	for k, v := range envMap {
		value, exists := os.LookupEnv(k)
		if exists {
			prevValues[k] = value
		} else {
			unsetValues = append(unsetValues, k)
		}
		os.Setenv(k, v)
	}

	return func() {
		for k, v := range prevValues {
			os.Setenv(k, v)
		}

		for _, k := range unsetValues {
			os.Unsetenv(k)
		}
	}
}