flashtext/flashtext_test.go at master · linkerlin/flashtext · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
package flashtext

import (
	"encoding/json"
	"io/ioutil"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/assert"
)

var (
	extractCases = make([]*ExtractorTestCase, 0, 100)
	removeCases  = make([]*RemoveTestCase, 0, 100)
)

func TestExtract(t *testing.T) {
	Init(t)
	//case insensitive
	for _, c := range extractCases {
		p := NewKeywordProcessor()
		p.SetCaseSenstive(false)
		for cleanName, keywords := range c.KeywordDict {
			for _, keyword := range keywords {
				p.AddKeywordAndName(keyword, cleanName)
			}
		}
		res := p.ExtractKeywords(c.Sentence, &Option{Longest: true})
		resultArray := []string{}
		for _, result := range res {
			resultArray = append(resultArray, result.Keyword)
		}
		assert.EqualValues(t, c.Keywords, resultArray, "insensitive keywords should match at sentence:"+c.Sentence)
	}

	//case sensitive
	for _, c := range extractCases {
		p := NewKeywordProcessor()
		p.SetCaseSenstive(true)
		for cleanName, keywords := range c.KeywordDict {
			for _, keyword := range keywords {
				p.AddKeywordAndName(keyword, cleanName)
			}
		}
		res := p.ExtractKeywords(c.Sentence, &Option{Longest: true})
		resultArray := []string{}
		for _, result := range res {
			resultArray = append(resultArray, result.Keyword)
		}
		assert.EqualValues(t, c.KeywordsCaseSensitive, resultArray, "sensitive keywords should match at sentence:"+c.Sentence)
	}
}

func TestRemove(t *testing.T) {
	Init(t)
	//case insensitive
	for _, c := range removeCases {
		p := NewKeywordProcessor()
		p.SetCaseSenstive(false)
		for cleanName, keywords := range c.KeywordDict {
			for _, keyword := range keywords {
				p.AddKeywordAndName(keyword, cleanName)
			}
		}
		for _, keywords := range c.RemoveKeywordDict {
			p.RemoveKeywords(keywords...)
		}
		res := p.ExtractKeywords(c.Sentence, &Option{Longest: true})
		resultArray := []string{}
		for _, result := range res {
			resultArray = append(resultArray, result.Keyword)
		}
		assert.EqualValues(t, c.Keywords, resultArray, "insensitive keywords should match at sentence:"+c.Sentence)
	}

	//case sensitive
	for _, c := range removeCases {
		p := NewKeywordProcessor()
		p.SetCaseSenstive(true)
		for cleanName, keywords := range c.KeywordDict {
			for _, keyword := range keywords {
				p.AddKeywordAndName(keyword, cleanName)
			}
		}
		for _, keywords := range c.RemoveKeywordDict {
			p.RemoveKeywords(keywords...)
		}
		res := p.ExtractKeywords(c.Sentence, &Option{Longest: true})
		resultArray := []string{}
		for _, result := range res {
			resultArray = append(resultArray, result.Keyword)
		}
		assert.EqualValues(t, c.KeywordsCaseSensitive, resultArray, "sensitive keywords should match at sentence:"+c.Sentence)
	}
}

// read test_cases json files
func Init(t *testing.T) {
	err := json.Unmarshal(helperLoadBytes(t, "extracts.json"), &extractCases)
	if err != nil {
		t.Fatal(err)
	}
}

func helperLoadBytes(t *testing.T, name string) []byte {
	path := filepath.Join("test_cases", name) // relative path
	bytes, err := ioutil.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	return bytes
}

type ExtractorTestCase struct {
	Sentence              string
	KeywordDict           map[string][]string `json:"keyword_dict"`
	Explaination          string
	Keywords              []string
	KeywordsCaseSensitive []string `json:"keywords_case_sensitive"`
}

type RemoveTestCase struct {
	Sentence              string              `json:"sentence"`
	KeywordDict           map[string][]string `json:"remove_keyword_dict"`
	RemoveKeywordDict     map[string][]string `json:"remove_keyword_dict"`
	Keywords              []string            `json:"keywords"`
	KeywordsCaseSensitive []string            `json:"keywords_case_sensitive"`
}