Add test for nsfw.go

This commit is contained in:
Theresa Gresch 2020-01-20 17:43:43 +01:00
parent abeb2d16f8
commit 2c3d462390

View file

@ -13,7 +13,7 @@ var modelPath, _ = filepath.Abs("../../assets/resources/nsfw")
var detector = New(modelPath)
func TestNSFW(t *testing.T) {
func TestIsSafe(t *testing.T) {
detect := func(filename string) Labels {
result, err := detector.File(filename)
@ -100,3 +100,14 @@ func TestNSFW(t *testing.T) {
t.Log(err.Error())
}
}
func TestNSFW(t *testing.T) {
porn := Labels{0, 0, 0.11, 0.88, 0}
sexy := Labels{0, 0, 0.2, 0.59, 0.98}
hentai := Labels{0, 0.98, 0.2, 0, 0}
assert.Equal(t, true, porn.NSFW())
assert.Equal(t, true, sexy.NSFW())
assert.Equal(t, true, hentai.NSFW())
}