diff --git a/util/suffix_tree.go b/util/suffix_tree.go index 5d50d21..182e46c 100644 --- a/util/suffix_tree.go +++ b/util/suffix_tree.go @@ -52,7 +52,7 @@ func (n *node) add(secs []string) { case 0: return case 1: - n.node[secs[length-1]] = &node{node: map[string]*node{}} + n.node[secs[length-1]] = &node{node: map[string]*node{"": &node{}}} default: subNode, ok := n.node[secs[length-1]] if !ok { @@ -70,15 +70,11 @@ func (n *Node) Match(item string) bool { func (n *node) matchSecs(secs []string) bool { length := len(secs) if length == 0 { - switch len(n.node) { - case 0: + if _, ok := n.node[""]; ok { return true - case 1: - _, ok := n.node["*"] - return ok - default: - return false } + _, ok := n.node["*"] + return ok } if n, ok := n.node[secs[length-1]]; ok { diff --git a/util/suffix_tree_test.go b/util/suffix_tree_test.go new file mode 100644 index 0000000..faae54d --- /dev/null +++ b/util/suffix_tree_test.go @@ -0,0 +1,56 @@ +package util + +import ( + "testing" +) + +func TestNode_Match(t *testing.T) { + type test struct { + arg string + want bool + } + tests := []struct { + name string + node *Node + tests []test + }{{ + "simple", + NewNodeFromRules(".", "a.wweir.cc", "b.wweir.cc"), + []test{ + {"a.wweir.cc", true}, + {"b.wweir.cc", true}, + }, + }, { + "parent", + NewNodeFromRules(".", "wweir.cc", "a.wweir.cc"), + []test{ + {"wweir.cc", true}, + {"a.wweir.cc", true}, + {"b.wweir.cc", false}, + }, + }, { + "fuzz1", + NewNodeFromRules(".", "wweir.cc", "*.wweir.cc"), + []test{ + {"wweir.cc", true}, + {"a.wweir.cc", true}, + }, + }, { + "fuzz2", + NewNodeFromRules(".", "a.wweir.cc", "*.wweir.cc"), + []test{ + {"wweir.cc", true}, + {"a.wweir.cc", true}, + {"b.wweir.cc", true}, + }, + }} + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for _, test := range tt.tests { + if got := tt.node.Match(test.arg); got != test.want { + t.Errorf("Node.Match(%s) = %v, want %v", test.arg, got, test.want) + } + } + }) + } +}