gitlab.com/gitlab-org/labkit@v1.21.0/tracing/env_injector_test.go (about)

     1  package tracing
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  	"reflect"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  	"gitlab.com/gitlab-org/labkit/correlation"
    11  )
    12  
    13  func TestNewEnvInjector(t *testing.T) {
    14  	tests := []struct {
    15  		name          string
    16  		tracingEnv    string
    17  		correlationID string
    18  		env           []string
    19  		wantEnv       []string
    20  	}{
    21  		{
    22  			name: "trivial",
    23  		},
    24  		{
    25  			name:       "gitlab_tracing_config",
    26  			tracingEnv: "opentracing://null",
    27  			wantEnv: []string{
    28  				"GITLAB_TRACING=opentracing://null",
    29  			},
    30  		},
    31  		{
    32  			name:       "merge",
    33  			tracingEnv: "opentracing://null",
    34  			env: []string{
    35  				"A=1",
    36  			},
    37  			wantEnv: []string{
    38  				"A=1",
    39  				"GITLAB_TRACING=opentracing://null",
    40  			},
    41  		},
    42  		{
    43  			name:          "correlation",
    44  			correlationID: "abc123",
    45  			env:           []string{},
    46  			wantEnv: []string{
    47  				"CORRELATION_ID=abc123",
    48  			},
    49  		},
    50  		{
    51  			name:          "correlation_merge",
    52  			tracingEnv:    "opentracing://null",
    53  			correlationID: "abc123",
    54  			env:           []string{"A=1"},
    55  			wantEnv: []string{
    56  				"A=1",
    57  				"CORRELATION_ID=abc123",
    58  				"GITLAB_TRACING=opentracing://null",
    59  			},
    60  		},
    61  	}
    62  
    63  	for _, tt := range tests {
    64  		t.Run(tt.name, func(t *testing.T) {
    65  			if tt.tracingEnv != "" {
    66  				resetEnvironment := modifyEnvironment(tracingEnvKey, tt.tracingEnv)
    67  				defer resetEnvironment()
    68  			} else {
    69  				resetEnvironment := clearEnvironment(tracingEnvKey)
    70  				defer resetEnvironment()
    71  			}
    72  			envInjector := NewEnvInjector()
    73  
    74  			ctx := context.Background()
    75  			if tt.correlationID != "" {
    76  				ctx = correlation.ContextWithCorrelation(ctx, tt.correlationID)
    77  			}
    78  
    79  			got := envInjector(ctx, tt.env)
    80  			require.Equal(t, tt.wantEnv, got)
    81  		})
    82  	}
    83  }
    84  
    85  // modifyEnvironment will change an environment variable and return a func suitable
    86  // for `defer` to change the value back.
    87  func modifyEnvironment(key string, value string) func() {
    88  	oldValue, hasOldValue := os.LookupEnv(key)
    89  	os.Setenv(key, value)
    90  	return func() {
    91  		if hasOldValue {
    92  			os.Setenv(key, oldValue)
    93  		} else {
    94  			os.Unsetenv(key)
    95  		}
    96  	}
    97  }
    98  
    99  // clearEnvironment will clear an environment variable and return a func suitable
   100  // for `defer` to change the value back.
   101  func clearEnvironment(key string) func() {
   102  	oldValue, hasOldValue := os.LookupEnv(key)
   103  
   104  	if !hasOldValue {
   105  		return func() {
   106  			// Nothing to do
   107  		}
   108  	}
   109  
   110  	os.Unsetenv(key)
   111  	return func() {
   112  		if hasOldValue {
   113  			os.Setenv(key, oldValue)
   114  		}
   115  	}
   116  }
   117  
   118  func Test_appendMapToEnv(t *testing.T) {
   119  	tests := []struct {
   120  		name   string
   121  		env    []string
   122  		envMap map[string]string
   123  		want   []string
   124  	}{
   125  		{
   126  			name:   "empty_case",
   127  			env:    nil,
   128  			envMap: nil,
   129  			want:   nil,
   130  		},
   131  		{
   132  			name:   "no_new_values",
   133  			env:    []string{"A=1"},
   134  			envMap: nil,
   135  			want:   []string{"A=1"},
   136  		},
   137  		{
   138  			name:   "no_original_values",
   139  			env:    nil,
   140  			envMap: map[string]string{"A": "1"},
   141  			want:   []string{"A=1"},
   142  		},
   143  		{
   144  			name:   "merge_values",
   145  			env:    []string{"A=1"},
   146  			envMap: map[string]string{"B": "2"},
   147  			want:   []string{"A=1", "B=2"},
   148  		},
   149  		{
   150  			name:   "merge_conflicts",
   151  			env:    []string{"A=1", "C=3"},
   152  			envMap: map[string]string{"B": "2", "C": "4"},
   153  			want:   []string{"A=1", "C=3", "B=2", "C=4"},
   154  		},
   155  	}
   156  	for _, tt := range tests {
   157  		t.Run(tt.name, func(t *testing.T) {
   158  			if got := appendMapToEnv(tt.env, tt.envMap); !reflect.DeepEqual(got, tt.want) {
   159  				t.Errorf("appendMapToEnv() = %v, want %v", got, tt.want)
   160  			}
   161  		})
   162  	}
   163  }