From 2cba89787741428e7cbeea73332074d8f515b1e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82?= <pawel@uber.com>
Date: Wed, 13 Jul 2022 15:03:20 -0700
Subject: [PATCH] Add a test verifying initial startup sequence (#97)

* Add a test verifying initial startup sequence

See https://github.com/uber-go/ratelimit/pull/95#discussion_r915251700

From that discussion I wasn't sure whether the proposed the initial
startup sequence of the limiter - i.e. whether at startup we always
block, or always allow.

Since we didn't seem to have that codified (perhaps apart from the
`example_test.go`) this PR adds a test to verify this.

This is still slightly (2/1000) flaky, but I think that's good enough
to add this in - should be valuable anyway.

* channels are great
---
 ratelimit_test.go | 67 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/ratelimit_test.go b/ratelimit_test.go
index 4c06e20..7268f87 100644
--- a/ratelimit_test.go
+++ b/ratelimit_test.go
@@ -22,6 +22,8 @@ type testRunner interface {
 	// afterFunc executes a func at a given time.
 	// not using clock.AfterFunc because andres-erbsen/clock misses a nap there.
 	afterFunc(d time.Duration, fn func())
+	// some tests want raw access to the clock.
+	getClock() *clock.Mock
 }
 
 type runnerImpl struct {
@@ -89,6 +91,10 @@ func (r *runnerImpl) createLimiter(rate int, opts ...Option) Limiter {
 	return r.constructor(rate, opts...)
 }
 
+func (r *runnerImpl) getClock() *clock.Mock {
+	return r.clock
+}
+
 // startTaking tries to Take() on passed in limiters in a loop/goroutine.
 func (r *runnerImpl) startTaking(rls ...Limiter) {
 	r.goWait(func() {
@@ -202,6 +208,67 @@ func TestPer(t *testing.T) {
 	})
 }
 
+// TestInitial verifies that the initial sequence is scheduled as expected.
+func TestInitial(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		msg  string
+		opts []Option
+	}{
+		{
+			msg: "With Slack",
+		},
+		{
+			msg:  "Without Slack",
+			opts: []Option{WithoutSlack},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.msg, func(t *testing.T) {
+			runTest(t, func(r testRunner) {
+				rl := r.createLimiter(10, tt.opts...)
+
+				var (
+					clk  = r.getClock()
+					prev = clk.Now()
+
+					results = make(chan time.Time)
+					have    []time.Duration
+					startWg sync.WaitGroup
+				)
+				startWg.Add(3)
+
+				for i := 0; i < 3; i++ {
+					go func() {
+						startWg.Done()
+						results <- rl.Take()
+					}()
+				}
+
+				startWg.Wait()
+				clk.Add(time.Second)
+
+				for i := 0; i < 3; i++ {
+					ts := <-results
+					have = append(have, ts.Sub(prev))
+					prev = ts
+				}
+
+				assert.Equal(t,
+					[]time.Duration{
+						0,
+						time.Millisecond * 100,
+						time.Millisecond * 100,
+					},
+					have,
+					"bad timestamps for inital takes",
+				)
+			})
+		})
+	}
+}
+
 func TestSlack(t *testing.T) {
 	t.Parallel()
 	// To simulate slack, we combine two limiters.