diff options
Diffstat (limited to 'modules')
| -rw-r--r-- | modules/ratelimit/ratelimit_test.go | 8 | ||||
| -rw-r--r-- | modules/starturl/starturl_test.go | 4 |
2 files changed, 12 insertions, 0 deletions
diff --git a/modules/ratelimit/ratelimit_test.go b/modules/ratelimit/ratelimit_test.go index 23cc8c8..9aa8acb 100644 --- a/modules/ratelimit/ratelimit_test.go +++ b/modules/ratelimit/ratelimit_test.go @@ -6,6 +6,7 @@ package ratelimit_test import ( "net/http" + "sync" "testing" "time" @@ -19,6 +20,7 @@ import ( func TestRatelimit(t *testing.T) { var times []time.Time + var mu sync.Mutex mods := []flyscrape.Module{ &starturl.Module{URL: "http://www.example.com"}, @@ -28,7 +30,9 @@ func TestRatelimit(t *testing.T) { return flyscrape.MockTransport(200, `<a href="foo">foo</a>`) }, ReceiveResponseFn: func(r *flyscrape.Response) { + mu.Lock() times = append(times, time.Now()) + mu.Unlock() }, }, &ratelimit.Module{ @@ -53,6 +57,7 @@ func TestRatelimit(t *testing.T) { func TestRatelimitConcurrency(t *testing.T) { var times []time.Time + var mu sync.Mutex mods := []flyscrape.Module{ &starturl.Module{URL: "http://www.example.com"}, @@ -60,7 +65,10 @@ func TestRatelimitConcurrency(t *testing.T) { hook.Module{ AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper { return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) { + mu.Lock() times = append(times, time.Now()) + mu.Unlock() + time.Sleep(10 * time.Millisecond) return flyscrape.MockResponse(200, ` <a href="foo"></a> diff --git a/modules/starturl/starturl_test.go b/modules/starturl/starturl_test.go index 54f899a..4b6a554 100644 --- a/modules/starturl/starturl_test.go +++ b/modules/starturl/starturl_test.go @@ -6,6 +6,7 @@ package starturl_test import ( "net/http" + "sync" "testing" "github.com/philippta/flyscrape" @@ -100,6 +101,7 @@ func TestStartURL_MultipleStartingURLs(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { urls := []string{} + mu := sync.Mutex{} mods := []flyscrape.Module{ tc.startURLModFn(), @@ -108,7 +110,9 @@ func TestStartURL_MultipleStartingURLs(t *testing.T) { return flyscrape.MockTransport(http.StatusOK, "") }, BuildRequestFn: func(r *flyscrape.Request) { + mu.Lock() urls = append(urls, r.URL) + mu.Unlock() }, }, } |