gitlab.com/gitlab-org/labkit@v1.21.0/correlation/context_test.go (about)

     1  package correlation
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/require"
     8  )
     9  
    10  func TestExtractFromContext(t *testing.T) {
    11  	require := require.New(t)
    12  
    13  	tests := []struct {
    14  		name string
    15  		ctx  context.Context
    16  		want string
    17  	}{
    18  		{"missing", context.Background(), ""},
    19  		{"set", context.WithValue(context.Background(), keyCorrelationID, "CORRELATION_ID"), "CORRELATION_ID"},
    20  	}
    21  	for _, tt := range tests {
    22  		t.Run(tt.name, func(t *testing.T) {
    23  			require.Equal(tt.want, ExtractFromContext(tt.ctx))
    24  		})
    25  	}
    26  }
    27  
    28  func TestExtractFromContextOrGenerate(t *testing.T) {
    29  	require := require.New(t)
    30  
    31  	tests := []struct {
    32  		name string
    33  		ctx  context.Context
    34  		want string
    35  	}{
    36  		{"missing", context.Background(), ""},
    37  		{"set", context.WithValue(context.Background(), keyCorrelationID, "CORRELATION_ID"), "CORRELATION_ID"},
    38  	}
    39  	for _, tt := range tests {
    40  		t.Run(tt.name, func(t *testing.T) {
    41  			if tt.want == "" { // want a random one
    42  				require.NotEmpty(ExtractFromContextOrGenerate(tt.ctx))
    43  			} else {
    44  				require.Equal(tt.want, ExtractFromContextOrGenerate(tt.ctx))
    45  			}
    46  		})
    47  	}
    48  }
    49  
    50  func TestContextWithCorrelation(t *testing.T) {
    51  	require := require.New(t)
    52  
    53  	tests := []struct {
    54  		name          string
    55  		ctx           context.Context
    56  		correlationID string
    57  		wantValue     string
    58  	}{
    59  		{
    60  			name:          "value",
    61  			ctx:           context.Background(),
    62  			correlationID: "CORRELATION_ID",
    63  			wantValue:     "CORRELATION_ID",
    64  		},
    65  		{
    66  			name:          "empty",
    67  			ctx:           context.Background(),
    68  			correlationID: "",
    69  			wantValue:     "",
    70  		},
    71  	}
    72  	for _, tt := range tests {
    73  		t.Run(tt.name, func(t *testing.T) {
    74  			got := ContextWithCorrelation(tt.ctx, tt.correlationID)
    75  			gotValue := got.Value(keyCorrelationID)
    76  			require.Equal(tt.wantValue, gotValue)
    77  		})
    78  	}
    79  }
    80  
    81  func TestExtractClientNameFromContext(t *testing.T) {
    82  	require := require.New(t)
    83  
    84  	tests := []struct {
    85  		name string
    86  		ctx  context.Context
    87  		want string
    88  	}{
    89  		{"missing", context.Background(), ""},
    90  		{"set", context.WithValue(context.Background(), keyClientName, "CLIENT_NAME"), "CLIENT_NAME"},
    91  	}
    92  	for _, tt := range tests {
    93  		t.Run(tt.name, func(t *testing.T) {
    94  			require.Equal(tt.want, ExtractClientNameFromContext(tt.ctx))
    95  		})
    96  	}
    97  }
    98  
    99  func TestContextWithClientName(t *testing.T) {
   100  	require := require.New(t)
   101  
   102  	tests := []struct {
   103  		name       string
   104  		ctx        context.Context
   105  		clientName string
   106  		wantValue  string
   107  	}{
   108  		{
   109  			name:       "value",
   110  			ctx:        context.Background(),
   111  			clientName: "CLIENT_NAME",
   112  			wantValue:  "CLIENT_NAME",
   113  		},
   114  		{
   115  			name:       "empty",
   116  			ctx:        context.Background(),
   117  			clientName: "",
   118  			wantValue:  "",
   119  		},
   120  	}
   121  	for _, tt := range tests {
   122  		t.Run(tt.name, func(t *testing.T) {
   123  			got := ContextWithClientName(tt.ctx, tt.clientName)
   124  			gotValue := got.Value(keyClientName)
   125  			require.Equal(tt.wantValue, gotValue)
   126  		})
   127  	}
   128  }