code grouming, missing test, reduce code smells (#30)

s0rg · Oct 3, 2022 · b47e9c5 · b47e9c5
1 parent 5de5e8d
commit b47e9c5
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 16 deletions.
diff --git a/pkg/crawler/crawler.go b/pkg/crawler/crawler.go
@@ -281,7 +281,12 @@ func (c *Crawler) fetch(
 
 	switch {
 	case isHTML(hdrs.Get(contentType)):
-		links.ExtractHTML(base, body, c.cfg.Brute, c.filter, c.linkHandler)
+		links.ExtractHTML(body, links.ExtractArgs{
+			Base:    base,
+			Brute:   c.cfg.Brute,
+			Filter:  c.filter,
+			Handler: c.linkHandler,
+		})
 	case isSitemap(uri):
 		links.ExtractSitemap(base, body, c.sitemapHandler)
 	}

diff --git a/pkg/crawler/util_test.go b/pkg/crawler/util_test.go
@@ -0,0 +1,31 @@
+package crawler
+
+import (
+	"testing"
+
+	"golang.org/x/net/html"
+	"golang.org/x/net/html/atom"
+)
+
+func TestPrepareFilter(t *testing.T) {
+	t.Parallel()
+
+	f := prepareFilter([]string{"a", "form", "foo", "div"})
+	d := []html.Token{
+		{DataAtom: atom.A},
+		{DataAtom: atom.Form},
+		{DataAtom: atom.Div},
+	}
+
+	for _, k := range d {
+		if !f(k) {
+			t.Fatalf("not allowed: %v", k)
+		}
+	}
+
+	v := html.Token{DataAtom: atom.Video}
+
+	if f(v) {
+		t.Fatalf("allowed: %v", v)
+	}
+}
diff --git a/pkg/links/html.go b/pkg/links/html.go
@@ -25,16 +25,17 @@ const (
 type HTMLHandler func(atom.Atom, string)
 type TokenFilter func(html.Token) bool
 
+type ExtractArgs struct {
+	Base    *url.URL
+	Brute   bool
+	Filter  TokenFilter
+	Handler HTMLHandler
+}
+
 func AllowALL(_ html.Token) bool { return true }
 
 // ExtractHTML run `handler` for every link found inside html from `r`, rebasing them to `b` (if need).
-func ExtractHTML(
-	b *url.URL,
-	r io.Reader,
-	brute bool,
-	allowed TokenFilter,
-	handler HTMLHandler,
-) {
+func ExtractHTML(r io.Reader, a ExtractArgs) {
 	var (
 		tkns = html.NewTokenizer(r)
 		key  = keySRC
@@ -46,12 +47,12 @@ func ExtractHTML(
 		case html.ErrorToken:
 			return
 		case html.StartTagToken, html.SelfClosingTagToken:
-			if t = tkns.Token(); allowed(t) {
-				extractToken(b, t, &key, handler)
+			if t = tkns.Token(); a.Filter(t) {
+				extractToken(a.Base, t, &key, a.Handler)
 			}
 		case html.CommentToken:
-			if brute {
-				extractComment(tkns.Token().Data, handler)
+			if a.Brute {
+				extractComment(tkns.Token().Data, a.Handler)
 			}
 		}
 	}

diff --git a/pkg/links/html_test.go b/pkg/links/html_test.go
@@ -195,8 +195,13 @@ func TestExtractURLS(t *testing.T) {
 
 			var res string
 
-			ExtractHTML(testBase, io.NopCloser(buf), true, AllowALL, func(_ atom.Atom, s string) {
-				res = s
+			ExtractHTML(io.NopCloser(buf), ExtractArgs{
+				Base:   testBase,
+				Brute:  true,
+				Filter: AllowALL,
+				Handler: func(_ atom.Atom, s string) {
+					res = s
+				},
 			})
 
 			if tc.hasLink {
@@ -247,8 +252,13 @@ func TestExtractAllowed(t *testing.T) {
 		return tkn.DataAtom == atom.A
 	}
 
-	ExtractHTML(testBase, io.NopCloser(buf), true, filter, func(_ atom.Atom, s string) {
-		res = append(res, s)
+	ExtractHTML(io.NopCloser(buf), ExtractArgs{
+		Base:   testBase,
+		Brute:  true,
+		Filter: filter,
+		Handler: func(_ atom.Atom, s string) {
+			res = append(res, s)
+		},
 	})
 
 	if len(res) != 1 {