Allow configuration of links to follow

1 files changed, 96 insertions, 0 deletions

diff --git a/modules/followlinks/followlinks_test.go b/modules/followlinks/followlinks_test.go
index 0a628c3..f3eb4fe 100644
--- a/modules/followlinks/followlinks_test.go
+++ b/modules/followlinks/followlinks_test.go
@@ -47,3 +47,99 @@ func TestFollowLinks(t *testing.T) {
 	require.Contains(t, urls, "http://www.google.com")
 	require.Contains(t, urls, "http://www.google.com/baz")
 }
+
+func TestFollowSelector(t *testing.T) {
+	var urls []string
+	var mu sync.Mutex
+
+	scraper := flyscrape.NewScraper()
+	scraper.LoadModule(&starturl.Module{URL: "http://www.example.com/foo/bar"})
+	scraper.LoadModule(&followlinks.Module{
+		Follow: []string{".next a[href]"},
+	})
+
+	scraper.LoadModule(hook.Module{
+		AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
+			return flyscrape.MockTransport(200, `
+				<a href="/baz">Baz</a>
+				<a href="baz">Baz</a>
+                <div class="next">
+				    <a href="http://www.google.com">Google</a>
+                </div>`)
+		},
+		ReceiveResponseFn: func(r *flyscrape.Response) {
+			mu.Lock()
+			urls = append(urls, r.Request.URL)
+			mu.Unlock()
+		},
+	})
+
+	scraper.Run()
+
+	require.Len(t, urls, 2)
+	require.Contains(t, urls, "http://www.example.com/foo/bar")
+	require.Contains(t, urls, "http://www.google.com")
+}
+
+func TestFollowDataAttr(t *testing.T) {
+	var urls []string
+	var mu sync.Mutex
+
+	scraper := flyscrape.NewScraper()
+	scraper.LoadModule(&starturl.Module{URL: "http://www.example.com/foo/bar"})
+	scraper.LoadModule(&followlinks.Module{
+		Follow: []string{"[data-url]"},
+	})
+
+	scraper.LoadModule(hook.Module{
+		AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
+			return flyscrape.MockTransport(200, `
+				<a href="/baz">Baz</a>
+				<a href="baz">Baz</a>
+				<div data-url="http://www.google.com">Google</div>`)
+		},
+		ReceiveResponseFn: func(r *flyscrape.Response) {
+			mu.Lock()
+			urls = append(urls, r.Request.URL)
+			mu.Unlock()
+		},
+	})
+
+	scraper.Run()
+
+	require.Len(t, urls, 2)
+	require.Contains(t, urls, "http://www.example.com/foo/bar")
+	require.Contains(t, urls, "http://www.google.com")
+}
+
+func TestFollowMultiple(t *testing.T) {
+	var urls []string
+	var mu sync.Mutex
+
+	scraper := flyscrape.NewScraper()
+	scraper.LoadModule(&starturl.Module{URL: "http://www.example.com/foo/bar"})
+	scraper.LoadModule(&followlinks.Module{
+		Follow: []string{"a.prev", "a.next"},
+	})
+
+	scraper.LoadModule(hook.Module{
+		AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
+			return flyscrape.MockTransport(200, `
+				<a href="/baz">Baz</a>
+				<a class="prev" href="a">a</a>
+				<a class="next" href="b">b</a>`)
+		},
+		ReceiveResponseFn: func(r *flyscrape.Response) {
+			mu.Lock()
+			urls = append(urls, r.Request.URL)
+			mu.Unlock()
+		},
+	})
+
+	scraper.Run()
+
+	require.Len(t, urls, 3)
+	require.Contains(t, urls, "http://www.example.com/foo/bar")
+	require.Contains(t, urls, "http://www.example.com/foo/a")
+	require.Contains(t, urls, "http://www.example.com/foo/b")
+}