ctxreader_test.go 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. package ctxreader
  2. import (
  3. "context"
  4. "crypto/rand"
  5. "testing"
  6. "time"
  7. "github.com/stretchr/testify/require"
  8. "github.com/stretchr/testify/suite"
  9. )
  10. type testReader struct {
  11. closed bool
  12. }
  13. func (r *testReader) Read(p []byte) (int, error) {
  14. return rand.Reader.Read(p)
  15. }
  16. func (r *testReader) Close() error {
  17. r.closed = true
  18. return nil
  19. }
  20. type CtxReaderTestSuite struct {
  21. suite.Suite
  22. }
  23. func (s *CtxReaderTestSuite) TestReadUntilCanceled() {
  24. ctx, cancel := context.WithCancel(context.Background())
  25. r := New(ctx, &testReader{}, false)
  26. p := make([]byte, 1024)
  27. _, err := r.Read(p)
  28. require.Nil(s.T(), err)
  29. cancel()
  30. time.Sleep(time.Second)
  31. _, err = r.Read(p)
  32. require.Equal(s.T(), err, context.Canceled)
  33. }
  34. func (s *CtxReaderTestSuite) TestReturnOriginalOnBackgroundContext() {
  35. rr := &testReader{}
  36. r := New(context.Background(), rr, false)
  37. require.Equal(s.T(), rr, r)
  38. }
  39. func (s *CtxReaderTestSuite) TestClose() {
  40. ctx, cancel := context.WithCancel(context.Background())
  41. defer cancel()
  42. rr := &testReader{}
  43. New(ctx, rr, true).Close()
  44. require.True(s.T(), rr.closed)
  45. }
  46. func (s *CtxReaderTestSuite) TestCloseOnCancel() {
  47. ctx, cancel := context.WithCancel(context.Background())
  48. rr := &testReader{}
  49. New(ctx, rr, true)
  50. cancel()
  51. time.Sleep(time.Second)
  52. require.True(s.T(), rr.closed)
  53. }
  54. func (s *CtxReaderTestSuite) TestDontCloseOnCancel() {
  55. ctx, cancel := context.WithCancel(context.Background())
  56. rr := &testReader{}
  57. New(ctx, rr, false)
  58. cancel()
  59. time.Sleep(time.Second)
  60. require.False(s.T(), rr.closed)
  61. }
  62. func TestCtxReader(t *testing.T) {
  63. suite.Run(t, new(CtxReaderTestSuite))
  64. }
  65. func BenchmarkRawReader(b *testing.B) {
  66. r := testReader{}
  67. b.ResetTimer()
  68. p := make([]byte, 1024)
  69. for i := 0; i < b.N; i++ {
  70. r.Read(p)
  71. }
  72. }
  73. func BenchmarkCtxReader(b *testing.B) {
  74. ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
  75. defer cancel()
  76. r := New(ctx, &testReader{}, true)
  77. b.ResetTimer()
  78. p := make([]byte, 1024)
  79. for i := 0; i < b.N; i++ {
  80. r.Read(p)
  81. }
  82. }