fix flaky test

Author: linkdata

dialer_test.go

@@ -1,6 +1,7 @@
 package bwlimit
 
 import (
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -21,7 +22,9 @@ func TestDialer_Dial(t *testing.T) {
 		ContextDialer: l2.Wrap(d1),
 	}
 
-	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write([]byte("Hello world!"))
+	}))
 	defer srv.Close()
 
 	client1 := &http.Client{
@@ -34,23 +37,33 @@ func TestDialer_Dial(t *testing.T) {
 			Dial: d2.Dial,
 		},
 	}
-	<-l1.WaitCh()
 	resp1, err1 := client1.Get(srv.URL)
 	resp2, err2 := client2.Get(srv.URL)
-	<-l1.WaitCh()
 
 	if err1 != nil {
 		t.Fatal(err1)
 	}
 	t.Log(resp1.Status)
+	if _, err := io.ReadAll(resp1.Body); err != nil {
+		t.Fatal(err)
+	}
+	if err := resp1.Body.Close(); err != nil {
+		t.Fatal(err)
+	}
 
 	if err2 != nil {
 		t.Fatal(err2)
 	}
 	t.Log(resp2.Status)
+	if _, err := io.ReadAll(resp2.Body); err != nil {
+		t.Fatal(err)
+	}
+	if err := resp2.Body.Close(); err != nil {
+		t.Fatal(err)
+	}
 
-	r1 := l1.Reads.Count.Load()
-	r2 := l2.Reads.Count.Load()
+	r1 := l1.Reads.Count.Load() + l1.Reads.count.Load()
+	r2 := l2.Reads.Count.Load() + l2.Reads.count.Load()
 	if r1 != r2*2 {
 		t.Error(r1, r2)
 	}

ticker.go

@@ -11,7 +11,13 @@ type Ticker struct {
 	ch chan struct{}
 }
 
-var DefaultTicker *Ticker
+var DefaultTicker *Ticker = NewTicker()
+
+func NewTicker() (ot *Ticker) {
+	ot = &Ticker{ch: make(chan struct{})}
+	go ot.run()
+	return
+}
 
 // NewLimiter returns a new Limiter using this Ticker.
 // If you provide limits, the first will set
@@ -49,10 +55,3 @@ func (ot *Ticker) run() {
 		close(oldCh)
 	}
 }
-
-func init() {
-	DefaultTicker = &Ticker{
-		ch: make(chan struct{}),
-	}
-	go DefaultTicker.run()
-}

ticker_test.go

@@ -0,0 +1,25 @@
+package bwlimit
+
+import (
+	"bytes"
+	"testing"
+	"time"
+)
+
+func TestTicker_zeroValue_NewLimiter(t *testing.T) {
+	l := (&Ticker{}).NewLimiter(1000)
+	defer l.Stop()
+
+	done := make(chan struct{})
+	go func() {
+		_, _ = l.Reads.io(bytes.NewReader(make([]byte, 1)).Read, make([]byte, 1))
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("read stalled with zero-value ticker")
+	}
+}
+