hook_test.go 5.8 KB

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