From fb84ca746e92e371161f1e1de3b01a048a9ae979 Mon Sep 17 00:00:00 2001
From: Philipp Tanlak <philipp.tanlak@gmail.com>
Date: Wed, 11 Oct 2023 20:20:30 +0200
Subject: Implement file based caching

---
 modules/ratelimit/ratelimit.go      | 16 ++++++++++------
 modules/ratelimit/ratelimit_test.go |  8 ++++----
 2 files changed, 14 insertions(+), 10 deletions(-)

(limited to 'modules/ratelimit')

diff --git a/modules/ratelimit/ratelimit.go b/modules/ratelimit/ratelimit.go
index 9588db3..b23cd7a 100644
--- a/modules/ratelimit/ratelimit.go
+++ b/modules/ratelimit/ratelimit.go
@@ -5,6 +5,7 @@
 package ratelimit
 
 import (
+	"net/http"
 	"time"
 
 	"github.com/philippta/flyscrape"
@@ -45,11 +46,14 @@ func (m *Module) Provision(v flyscrape.Context) {
 	}()
 }
 
-func (m *Module) BuildRequest(_ *flyscrape.Request) {
+func (m *Module) AdaptTransport(t http.RoundTripper) http.RoundTripper {
 	if m.disabled() {
-		return
+		return t
 	}
-	<-m.semaphore
+	return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) {
+		<-m.semaphore
+		return t.RoundTrip(r)
+	})
 }
 
 func (m *Module) Finalize() {
@@ -64,7 +68,7 @@ func (m *Module) disabled() bool {
 }
 
 var (
-	_ flyscrape.RequestBuilder = (*Module)(nil)
-	_ flyscrape.Provisioner    = (*Module)(nil)
-	_ flyscrape.Finalizer      = (*Module)(nil)
+	_ flyscrape.TransportAdapter = (*Module)(nil)
+	_ flyscrape.Provisioner      = (*Module)(nil)
+	_ flyscrape.Finalizer        = (*Module)(nil)
 )
diff --git a/modules/ratelimit/ratelimit_test.go b/modules/ratelimit/ratelimit_test.go
index ffd061c..1fe22b1 100644
--- a/modules/ratelimit/ratelimit_test.go
+++ b/modules/ratelimit/ratelimit_test.go
@@ -23,17 +23,17 @@ func TestRatelimit(t *testing.T) {
 	scraper := flyscrape.NewScraper()
 	scraper.LoadModule(&starturl.Module{URL: "http://www.example.com"})
 	scraper.LoadModule(&followlinks.Module{})
-	scraper.LoadModule(&ratelimit.Module{
-		Rate: 100,
-	})
 	scraper.LoadModule(hook.Module{
 		AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
 			return flyscrape.MockTransport(200, `<a href="foo">foo</a>`)
 		},
-		BuildRequestFn: func(r *flyscrape.Request) {
+		ReceiveResponseFn: func(r *flyscrape.Response) {
 			times = append(times, time.Now())
 		},
 	})
+	scraper.LoadModule(&ratelimit.Module{
+		Rate: 100,
+	})
 
 	start := time.Now()
 	scraper.Run()
-- 
cgit v1.2.3