github.com/lingyao2333/mo-zero@v1.4.1/core/stores/redis/hook_test.go (about)

     1  package redis
     2  
     3  import (
     4  	"context"
     5  	"log"
     6  	"strings"
     7  	"testing"
     8  	"time"
     9  
    10  	red "github.com/go-redis/redis/v8"
    11  	"github.com/lingyao2333/mo-zero/core/logx"
    12  	ztrace "github.com/lingyao2333/mo-zero/core/trace"
    13  	"github.com/stretchr/testify/assert"
    14  	tracesdk "go.opentelemetry.io/otel/trace"
    15  )
    16  
    17  func TestHookProcessCase1(t *testing.T) {
    18  	ztrace.StartAgent(ztrace.Config{
    19  		Name:     "go-zero-test",
    20  		Endpoint: "http://localhost:14268/api/traces",
    21  		Batcher:  "jaeger",
    22  		Sampler:  1.0,
    23  	})
    24  	defer ztrace.StopAgent()
    25  
    26  	writer := log.Writer()
    27  	var buf strings.Builder
    28  	log.SetOutput(&buf)
    29  	defer log.SetOutput(writer)
    30  
    31  	ctx, err := durationHook.BeforeProcess(context.Background(), red.NewCmd(context.Background()))
    32  	if err != nil {
    33  		t.Fatal(err)
    34  	}
    35  
    36  	assert.Nil(t, durationHook.AfterProcess(ctx, red.NewCmd(context.Background())))
    37  	assert.False(t, strings.Contains(buf.String(), "slow"))
    38  	assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
    39  }
    40  
    41  func TestHookProcessCase2(t *testing.T) {
    42  	ztrace.StartAgent(ztrace.Config{
    43  		Name:     "go-zero-test",
    44  		Endpoint: "http://localhost:14268/api/traces",
    45  		Batcher:  "jaeger",
    46  		Sampler:  1.0,
    47  	})
    48  	defer ztrace.StopAgent()
    49  
    50  	w, restore := injectLog()
    51  	defer restore()
    52  
    53  	ctx, err := durationHook.BeforeProcess(context.Background(), red.NewCmd(context.Background()))
    54  	if err != nil {
    55  		t.Fatal(err)
    56  	}
    57  	assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
    58  
    59  	time.Sleep(slowThreshold.Load() + time.Millisecond)
    60  
    61  	assert.Nil(t, durationHook.AfterProcess(ctx, red.NewCmd(context.Background(), "foo", "bar")))
    62  	assert.True(t, strings.Contains(w.String(), "slow"))
    63  	assert.True(t, strings.Contains(w.String(), "trace"))
    64  	assert.True(t, strings.Contains(w.String(), "span"))
    65  }
    66  
    67  func TestHookProcessCase3(t *testing.T) {
    68  	writer := log.Writer()
    69  	var buf strings.Builder
    70  	log.SetOutput(&buf)
    71  	defer log.SetOutput(writer)
    72  
    73  	assert.Nil(t, durationHook.AfterProcess(context.Background(), red.NewCmd(context.Background())))
    74  	assert.True(t, buf.Len() == 0)
    75  }
    76  
    77  func TestHookProcessCase4(t *testing.T) {
    78  	writer := log.Writer()
    79  	var buf strings.Builder
    80  	log.SetOutput(&buf)
    81  	defer log.SetOutput(writer)
    82  
    83  	ctx := context.WithValue(context.Background(), startTimeKey, "foo")
    84  	assert.Nil(t, durationHook.AfterProcess(ctx, red.NewCmd(context.Background())))
    85  	assert.True(t, buf.Len() == 0)
    86  }
    87  
    88  func TestHookProcessPipelineCase1(t *testing.T) {
    89  	writer := log.Writer()
    90  	var buf strings.Builder
    91  	log.SetOutput(&buf)
    92  	defer log.SetOutput(writer)
    93  
    94  	ctx, err := durationHook.BeforeProcessPipeline(context.Background(), []red.Cmder{
    95  		red.NewCmd(context.Background()),
    96  	})
    97  	assert.NoError(t, err)
    98  	assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
    99  
   100  	assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
   101  		red.NewCmd(context.Background()),
   102  	}))
   103  	assert.False(t, strings.Contains(buf.String(), "slow"))
   104  }
   105  
   106  func TestHookProcessPipelineCase2(t *testing.T) {
   107  	ztrace.StartAgent(ztrace.Config{
   108  		Name:     "go-zero-test",
   109  		Endpoint: "http://localhost:14268/api/traces",
   110  		Batcher:  "jaeger",
   111  		Sampler:  1.0,
   112  	})
   113  	defer ztrace.StopAgent()
   114  
   115  	w, restore := injectLog()
   116  	defer restore()
   117  
   118  	ctx, err := durationHook.BeforeProcessPipeline(context.Background(), []red.Cmder{
   119  		red.NewCmd(context.Background()),
   120  	})
   121  	assert.NoError(t, err)
   122  	assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
   123  
   124  	time.Sleep(slowThreshold.Load() + time.Millisecond)
   125  
   126  	assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
   127  		red.NewCmd(context.Background(), "foo", "bar"),
   128  	}))
   129  	assert.True(t, strings.Contains(w.String(), "slow"))
   130  	assert.True(t, strings.Contains(w.String(), "trace"))
   131  	assert.True(t, strings.Contains(w.String(), "span"))
   132  }
   133  
   134  func TestHookProcessPipelineCase3(t *testing.T) {
   135  	w, restore := injectLog()
   136  	defer restore()
   137  
   138  	assert.Nil(t, durationHook.AfterProcessPipeline(context.Background(), []red.Cmder{
   139  		red.NewCmd(context.Background()),
   140  	}))
   141  	assert.True(t, len(w.String()) == 0)
   142  }
   143  
   144  func TestHookProcessPipelineCase4(t *testing.T) {
   145  	w, restore := injectLog()
   146  	defer restore()
   147  
   148  	ctx := context.WithValue(context.Background(), startTimeKey, "foo")
   149  	assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
   150  		red.NewCmd(context.Background()),
   151  	}))
   152  	assert.True(t, len(w.String()) == 0)
   153  }
   154  
   155  func TestHookProcessPipelineCase5(t *testing.T) {
   156  	writer := log.Writer()
   157  	var buf strings.Builder
   158  	log.SetOutput(&buf)
   159  	defer log.SetOutput(writer)
   160  
   161  	ctx := context.WithValue(context.Background(), startTimeKey, "foo")
   162  	assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
   163  		red.NewCmd(context.Background()),
   164  	}))
   165  	assert.True(t, buf.Len() == 0)
   166  }
   167  
   168  func TestLogDuration(t *testing.T) {
   169  	w, restore := injectLog()
   170  	defer restore()
   171  
   172  	logDuration(context.Background(), []red.Cmder{
   173  		red.NewCmd(context.Background(), "get", "foo"),
   174  	}, 1*time.Second)
   175  	assert.True(t, strings.Contains(w.String(), "get foo"))
   176  
   177  	logDuration(context.Background(), []red.Cmder{
   178  		red.NewCmd(context.Background(), "get", "foo"),
   179  		red.NewCmd(context.Background(), "set", "bar", 0),
   180  	}, 1*time.Second)
   181  	assert.True(t, strings.Contains(w.String(), `get foo\nset bar 0`))
   182  }
   183  
   184  func injectLog() (r *strings.Builder, restore func()) {
   185  	var buf strings.Builder
   186  	w := logx.NewWriter(&buf)
   187  	o := logx.Reset()
   188  	logx.SetWriter(w)
   189  
   190  	return &buf, func() {
   191  		logx.Reset()
   192  		logx.SetWriter(o)
   193  	}
   194  }