Implement manual following (#82)

2 files changed, 49 insertions, 7 deletions

diff --git a/modules/followlinks/followlinks.go b/modules/followlinks/followlinks.go
index c1448be..3ce2797 100644
--- a/modules/followlinks/followlinks.go
+++ b/modules/followlinks/followlinks.go
@@ -18,7 +18,7 @@ func init() {
 }
 
 type Module struct {
-	Follow []string `json:"follow"`
+	Follow *[]string `json:"follow"`
 }
 
 func (Module) ModuleInfo() flyscrape.ModuleInfo {
@@ -29,18 +29,26 @@ func (Module) ModuleInfo() flyscrape.ModuleInfo {
 }
 
 func (m *Module) Provision(ctx flyscrape.Context) {
-	if len(m.Follow) == 0 {
-		m.Follow = []string{"a[href]"}
+	if m.Follow == nil {
+		m.Follow = &[]string{"a[href]"}
 	}
 }
 
 func (m *Module) ReceiveResponse(resp *flyscrape.Response) {
+	if m.Follow == nil {
+		return
+	}
+
 	for _, link := range m.parseLinks(string(resp.Body), resp.Request.URL) {
 		resp.Visit(link)
 	}
 }
 
 func (m *Module) parseLinks(html string, origin string) []string {
+	if m.Follow == nil {
+		return nil
+	}
+
 	var links []string
 	doc, err := goquery.NewDocumentFromReader(strings.NewReader(html))
 	if err != nil {
@@ -54,7 +62,7 @@ func (m *Module) parseLinks(html string, origin string) []string {
 
 	uniqueLinks := make(map[string]bool)
 
-	for _, selector := range m.Follow {
+	for _, selector := range *m.Follow {
 		attr := parseSelectorAttr(selector)
 		doc.Find(selector).Each(func(i int, s *goquery.Selection) {
 			link, _ := s.Attr(attr)
diff --git a/modules/followlinks/followlinks_test.go b/modules/followlinks/followlinks_test.go
index af186f9..b09b0cd 100644
--- a/modules/followlinks/followlinks_test.go
+++ b/modules/followlinks/followlinks_test.go
@@ -57,7 +57,7 @@ func TestFollowSelector(t *testing.T) {
 	mods := []flyscrape.Module{
 		&starturl.Module{URL: "http://www.example.com/foo/bar"},
 		&followlinks.Module{
-			Follow: []string{".next a[href]"},
+			Follow: &[]string{".next a[href]"},
 		},
 		hook.Module{
 			AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
@@ -92,7 +92,7 @@ func TestFollowDataAttr(t *testing.T) {
 	mods := []flyscrape.Module{
 		&starturl.Module{URL: "http://www.example.com/foo/bar"},
 		&followlinks.Module{
-			Follow: []string{"[data-url]"},
+			Follow: &[]string{"[data-url]"},
 		},
 		hook.Module{
 			AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
@@ -125,7 +125,7 @@ func TestFollowMultiple(t *testing.T) {
 	mods := []flyscrape.Module{
 		&starturl.Module{URL: "http://www.example.com/foo/bar"},
 		&followlinks.Module{
-			Follow: []string{"a.prev", "a.next"},
+			Follow: &[]string{"a.prev", "a.next"},
 		},
 		hook.Module{
 			AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
@@ -151,3 +151,37 @@ func TestFollowMultiple(t *testing.T) {
 	require.Contains(t, urls, "http://www.example.com/foo/a")
 	require.Contains(t, urls, "http://www.example.com/foo/b")
 }
+
+func TestFollowNoFollow(t *testing.T) {
+	var urls []string
+	var mu sync.Mutex
+
+	mods := []flyscrape.Module{
+		&starturl.Module{URL: "http://www.example.com/foo/bar"},
+		&followlinks.Module{
+			Follow: &[]string{},
+		},
+		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 := flyscrape.NewScraper()
+	scraper.Modules = mods
+	scraper.Run()
+
+	require.Len(t, urls, 1)
+	require.Contains(t, urls, "http://www.example.com/foo/bar")
+}