1 files changed, 26 insertions, 23 deletions
diff --git a/modules/depth/depth_test.go b/modules/depth/depth_test.go
index 10b67e9..a596eb4 100644
--- a/modules/depth/depth_test.go
+++ b/modules/depth/depth_test.go
@@ -21,31 +21,34 @@ func TestDepth(t *testing.T) {
 	var urls []string
 	var mu sync.Mutex
 
-	scraper := flyscrape.NewScraper()
-	scraper.LoadModule(&starturl.Module{URL: "http://www.example.com"})
-	scraper.LoadModule(&followlinks.Module{})
-	scraper.LoadModule(&depth.Module{Depth: 2})
-	scraper.LoadModule(hook.Module{
-		AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
-			return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) {
-				switch r.URL.String() {
-				case "http://www.example.com":
-					return flyscrape.MockResponse(200, `<a href="http://www.google.com">Google</a>`)
-				case "http://www.google.com":
-					return flyscrape.MockResponse(200, `<a href="http://www.duckduckgo.com">DuckDuckGo</a>`)
-				case "http://www.duckduckgo.com":
-					return flyscrape.MockResponse(200, `<a href="http://www.example.com">Example</a>`)
-				}
-				return flyscrape.MockResponse(200, "")
-			})
-		},
-		ReceiveResponseFn: func(r *flyscrape.Response) {
-			mu.Lock()
-			urls = append(urls, r.Request.URL)
-			mu.Unlock()
+	mods := []flyscrape.Module{
+		&starturl.Module{URL: "http://www.example.com"},
+		&followlinks.Module{},
+		&depth.Module{Depth: 2},
+		hook.Module{
+			AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
+				return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) {
+					switch r.URL.String() {
+					case "http://www.example.com":
+						return flyscrape.MockResponse(200, `<a href="http://www.google.com">Google</a>`)
+					case "http://www.google.com":
+						return flyscrape.MockResponse(200, `<a href="http://www.duckduckgo.com">DuckDuckGo</a>`)
+					case "http://www.duckduckgo.com":
+						return flyscrape.MockResponse(200, `<a href="http://www.example.com">Example</a>`)
+					}
+					return flyscrape.MockResponse(200, "")
+				})
+			},
+			ReceiveResponseFn: func(r *flyscrape.Response) {
+				mu.Lock()
+				urls = append(urls, r.Request.URL)
+				mu.Unlock()
+			},
 		},
-	})
+	}
 
+	scraper := flyscrape.NewScraper()
+	scraper.Modules = mods
 	scraper.Run()
 
 	require.Len(t, urls, 3)