server_timing_redis_test.go 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. package repository
  2. import (
  3. "context"
  4. "errors"
  5. "strings"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
  9. "github.com/redis/go-redis/v9"
  10. )
  11. func TestServerTimingRedisHookRecordsCommands(t *testing.T) {
  12. collector := servertiming.New(time.Now())
  13. ctx := servertiming.WithCollector(context.Background(), collector)
  14. hook := serverTimingRedisHook{}
  15. process := hook.ProcessHook(func(context.Context, redis.Cmder) error {
  16. time.Sleep(time.Millisecond)
  17. return errors.New("redis failure")
  18. })
  19. if err := process(ctx, redis.NewStringCmd(ctx, "get", "sensitive-key")); err == nil {
  20. t.Fatal("ProcessHook did not return the underlying error")
  21. }
  22. pipeline := hook.ProcessPipelineHook(func(context.Context, []redis.Cmder) error {
  23. time.Sleep(time.Millisecond)
  24. return nil
  25. })
  26. commands := []redis.Cmder{
  27. redis.NewStringCmd(ctx, "get", "first-secret"),
  28. redis.NewStringCmd(ctx, "get", "second-secret"),
  29. redis.NewStatusCmd(ctx, "set", "third-secret", "value"),
  30. }
  31. if err := pipeline(ctx, commands); err != nil {
  32. t.Fatal(err)
  33. }
  34. header := collector.HeaderValue(time.Now(), "bypass")
  35. if !strings.Contains(header, `commands=4`) {
  36. t.Fatalf("header %q does not report one command and a three-command pipeline", header)
  37. }
  38. if strings.Contains(header, "secret") || strings.Contains(header, "get") {
  39. t.Fatalf("Redis command details leaked into header: %q", header)
  40. }
  41. }
  42. func TestServerTimingRedisHookSkipsInactiveContext(t *testing.T) {
  43. called := false
  44. hook := serverTimingRedisHook{}
  45. process := hook.ProcessHook(func(context.Context, redis.Cmder) error {
  46. called = true
  47. return nil
  48. })
  49. ctx := context.Background()
  50. if err := process(ctx, redis.NewStringCmd(ctx, "ping")); err != nil {
  51. t.Fatal(err)
  52. }
  53. if !called {
  54. t.Fatal("inactive Redis command did not reach the next hook")
  55. }
  56. }