github.com/jxskiss/gopkg/v2@v2.14.9-0.20240514120614-899f3e7952b4/zlog/context_test.go (about)

     1  package zlog
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"sync"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  	"go.uber.org/zap"
    14  	"go.uber.org/zap/zapcore"
    15  )
    16  
    17  func TestAddFields(t *testing.T) {
    18  	ctx := context.Background()
    19  	assert.Equal(t, ctx, AddFields(ctx))
    20  
    21  	ctx = AddFields(ctx,
    22  		zap.String("a", "aValue"),
    23  		zap.Int("b", 1234),
    24  	)
    25  	fs := GetFields(ctx)
    26  	require.Len(t, fs, 2)
    27  	assert.Equal(t, fs[0].Key, "a")
    28  	assert.Equal(t, fs[1].Key, "b")
    29  
    30  	ctx = AddFields(ctx,
    31  		zap.Duration("b", time.Minute), // override "b":"1234"
    32  	)
    33  	fs = GetFields(ctx)
    34  	require.Len(t, fs, 2)
    35  	assert.Equal(t, fs[1].Key, "b")
    36  	assert.Equal(t, fs[1].Integer, int64(time.Minute))
    37  }
    38  
    39  func TestAddFieldsConcurrently(t *testing.T) {
    40  	ctx := context.Background()
    41  	ctx = AddFields(ctx,
    42  		zap.String("a", "aValue"),
    43  		zap.String("b", "bValue"),
    44  	)
    45  	var wg sync.WaitGroup
    46  	for i := 0; i < 10; i++ {
    47  		value := fmt.Sprintf("value %d", i)
    48  		wg.Add(1)
    49  		go func() {
    50  			defer wg.Done()
    51  			ctx1 := AddFields(ctx,
    52  				zap.String("a", value),
    53  				zap.String("b", value),
    54  				zap.String("c", value),
    55  				zap.String("d", value),
    56  			)
    57  			fields := GetFields(ctx1)
    58  			assert.Len(t, fields, 4)
    59  			assert.Equal(t, value, fields[0].String)
    60  			assert.Equal(t, value, fields[1].String)
    61  			assert.Equal(t, value, fields[2].String)
    62  			assert.Equal(t, value, fields[3].String)
    63  		}()
    64  	}
    65  	wg.Wait()
    66  
    67  	fields := GetFields(ctx)
    68  	assert.Len(t, fields, 2)
    69  	assert.Equal(t, "aValue", fields[0].String)
    70  	assert.Equal(t, "bValue", fields[1].String)
    71  }
    72  
    73  func TestWithCtx(t *testing.T) {
    74  	ctx := context.Background()
    75  	helperReplace := func() (*bytes.Buffer, func()) {
    76  		var buf = &bytes.Buffer{}
    77  		l, p, err := NewWithOutput(&Config{Development: false, Level: "info"}, zapcore.AddSync(buf))
    78  		if err != nil {
    79  			panic(err)
    80  		}
    81  		resetFunc := ReplaceGlobals(l, p)
    82  		return buf, resetFunc
    83  	}
    84  
    85  	t.Run("with logger", func(t *testing.T) {
    86  		buf, rf := helperReplace()
    87  		defer rf()
    88  
    89  		ctx2 := WithLogger(ctx, L().With(zap.String("a", "aValue")))
    90  		WithCtx(ctx2).Info("test message")
    91  		got := buf.String()
    92  		assert.Contains(t, got, "test message")
    93  		assert.Contains(t, got, `"a":"aValue"`)
    94  	})
    95  
    96  	t.Run("with sugared logger", func(t *testing.T) {
    97  		buf, rf := helperReplace()
    98  		defer rf()
    99  
   100  		ctx2 := WithLogger(ctx, S().With(zap.Int64("a", 12345), "b", "bValue"))
   101  		WithCtx(ctx2).Info("test message")
   102  		got := buf.String()
   103  		assert.Contains(t, got, "test message")
   104  		assert.Contains(t, got, `"a":12345`)
   105  		assert.Contains(t, got, `"b":"bValue"`)
   106  	})
   107  
   108  	t.Run("no logger / with fields", func(t *testing.T) {
   109  		buf, rf := helperReplace()
   110  		defer rf()
   111  
   112  		ctx2 := AddFields(ctx, zap.String("a", "aValue"))
   113  		WithCtx(ctx2).Info("test message")
   114  		got := buf.String()
   115  		assert.Contains(t, got, "test message")
   116  		assert.Contains(t, got, `"a":"aValue"`)
   117  	})
   118  
   119  	t.Run("no logger / with fields / with extra", func(t *testing.T) {
   120  		buf, rf := helperReplace()
   121  		defer rf()
   122  
   123  		ctx2 := AddFields(ctx, zap.String("a", "aValue"))
   124  		WithCtx(ctx2, zap.Int("b", 12345)).
   125  			Info("test message", zap.Int32("c", 12345))
   126  		got := buf.String()
   127  		assert.Contains(t, got, "test message")
   128  		assert.Contains(t, got, `"a":"aValue"`)
   129  		assert.Contains(t, got, `"b":12345`)
   130  		assert.Contains(t, got, `"c":12345`)
   131  	})
   132  
   133  	t.Run("no logger / no fields / no extra", func(t *testing.T) {
   134  		buf, rf := helperReplace()
   135  		defer rf()
   136  
   137  		WithCtx(ctx).Info("test message")
   138  		got := buf.String()
   139  		assert.Contains(t, got, "test message")
   140  		assert.NotContains(t, got, `"a":"aValue"`)
   141  		assert.NotContains(t, got, `"b":12345`)
   142  	})
   143  }