osquery/query_knn_test.go at master · defensestation/osquery · 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
package osquery

import (
	"testing"
)

func TestKNNQuery_Map(t *testing.T) {
	tests := []mapTest{
		{
			"knn query without additional options",
			KNN("vector_field", []float64{1.0, 2.0, 3.0}).K(5),
			map[string]interface{}{
				"knn": map[string]interface{}{
					"vector_field": map[string]interface{}{
						"vector": []float64{1.0, 2.0, 3.0},
						"k":      5,
					},
				},
			},
		},
		{
			"knn query with max_distance option",
			KNN("vector_field", []float64{1.0, 2.0, 3.0}).K(5).MaxDistance(5.5),
			map[string]interface{}{
				"knn": map[string]interface{}{
					"vector_field": map[string]interface{}{
						"vector":       []float64{1.0, 2.0, 3.0},
						"k":            5,
						"max_distance": 5.5,
					},
				},
			},
		},
		{
			"knn query with min_score and filter",
			KNN("vector_field", []float64{0.1, 0.2}).K(5).MinScore(0.75).Filter(Term("status", "active").Map()),
			map[string]interface{}{
				"knn": map[string]interface{}{
					"vector_field": map[string]interface{}{
						"vector":    []float64{0.1, 0.2},
						"k":         5,
						"min_score": 0.75,
						"filter": map[string]interface{}{
							"term": map[string]interface{}{
								"status": map[string]interface{}{
									"value": "active",
								},
							},
						},
					},
				},
			},
		},
		{
			"knn query with method_parameters and expand_nested_docs",
			KNN("vector_field", []float64{4.0, 5.0, 6.0}).K(7).MethodParameters(
				map[string]interface{}{
					"nprobes": 3,
				},
			).ExpandNestedDocs(true),
			map[string]interface{}{
				"knn": map[string]interface{}{
					"vector_field": map[string]interface{}{
						"vector":             []float64{4.0, 5.0, 6.0},
						"k":                  7,
						"method_parameters":  map[string]interface{}{"nprobes": 3},
						"expand_nested_docs": true,
					},
				},
			},
		},
		{
			"knn query with rescore",
			KNN("vector_field", []float64{7, 8, 9}).K(3).Rescore(true),
			map[string]interface{}{
				"knn": map[string]interface{}{
					"vector_field": map[string]interface{}{
						"vector":  []float64{7, 8, 9},
						"k":       3,
						"rescore": true,
					},
				},
			},
		},
	}

	runMapTests(t, tests)
}