| 106 | } |
| 107 | |
| 108 | func TestBurst(t *testing.T) { |
| 109 | sampler := &BurstSampler{Burst: 1, Period: time.Second} |
| 110 | |
| 111 | t0 := time.Now() |
| 112 | now := t0 |
| 113 | mockedTime := func() time.Time { |
| 114 | return now |
| 115 | } |
| 116 | |
| 117 | TimestampFunc = mockedTime |
| 118 | defer func() { TimestampFunc = time.Now }() |
| 119 | |
| 120 | scenario := []struct { |
| 121 | tm time.Time |
| 122 | want bool |
| 123 | }{ |
| 124 | {t0, true}, |
| 125 | {t0.Add(time.Second - time.Nanosecond), false}, |
| 126 | {t0.Add(time.Second), true}, |
| 127 | {t0.Add(time.Second + time.Nanosecond), false}, |
| 128 | } |
| 129 | |
| 130 | for i, step := range scenario { |
| 131 | now = step.tm |
| 132 | got := sampler.Sample(NoLevel) |
| 133 | if got != step.want { |
| 134 | t.Errorf("step %d (t=%s): expect %t got %t", i, step.tm, step.want, got) |
| 135 | } |
| 136 | } |
| 137 | } |
| 138 | |
| 139 | func TestLevelSampler(t *testing.T) { |
| 140 | // Create mock samplers that return true for specific levels |