Skip to content

Commit da80048

Browse files
dido18mirkoCrobu
dido18
and
mirkoCrobu
authored
test(internal/orchestrator/modelsindex): add unit test for the model list index (#22)
* feat(models): add model_labels field and update test data for models * refactor(tests): update model index tests for clarity and accuracy * fix(tests): correct model retrieval order in TestGenerateModelsIndexFromFile * Update internal/orchestrator/modelsindex/modelsindex_test.go Co-authored-by: mirkoCrobu <m.crobu@ext.arduino.cc> * refactor(tests): streamline model index test structure for clarity * refactor(tests): simplify model retrieval tests for clarity and consistency * refactor(tests): rename TestGenerateModelsIndexFromFile to TestModelsIndex for consistency --------- Co-authored-by: mirkoCrobu <m.crobu@ext.arduino.cc>
1 parent df2d206 commit da80048

File tree

3 files changed

+184
-0
lines changed

3 files changed

+184
-0
lines changed

‎internal/orchestrator/modelsindex/models_index.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ type AIModel struct {
4848
ModuleDescription string `yaml:"description"`
4949
Runner string `yaml:"runner"`
5050
Bricks []string `yaml:"bricks,omitempty"`
51+
ModelLabels []string `yaml:"model_labels,omitempty"`
5152
Metadata map[string]string `yaml:"metadata,omitempty"`
5253
ModelConfiguration map[string]string `yaml:"model_configuration,omitempty"`
5354
}

‎internal/orchestrator/modelsindex/modelsindex_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package modelsindex
2+
3+
import (
4+
"testing"
5+
6+
"github.com/arduino/go-paths-helper"
7+
"github.com/stretchr/testify/assert"
8+
"github.com/stretchr/testify/require"
9+
)
10+
11+
func TestModelsIndex(t *testing.T) {
12+
modelsIndex, err := GenerateModelsIndexFromFile(paths.New("testdata"))
13+
require.NoError(t, err)
14+
require.NotNil(t, modelsIndex)
15+
16+
t.Run("it parses a valid model-list.yaml", func(t *testing.T) {
17+
models := modelsIndex.GetModels()
18+
assert.Len(t, models, 2, "Expected 2 models to be parsed")
19+
})
20+
21+
t.Run("it gets a model by ID", func(t *testing.T) {
22+
model, found := modelsIndex.GetModelByID("not-existing-model")
23+
assert.False(t, found)
24+
assert.Nil(t, model)
25+
26+
model, found = modelsIndex.GetModelByID("face-detection")
27+
assert.Equal(t, "brick", model.Runner)
28+
require.True(t, found, "face-detection should be found")
29+
assert.Equal(t, "face-detection", model.ID)
30+
assert.Equal(t, "Lightweight-Face-Detection", model.Name)
31+
assert.Equal(t, "Face bounding box detection. This model is trained on the WIDER FACE dataset and can detect faces in images.", model.ModuleDescription)
32+
assert.Equal(t, []string{"face"}, model.ModelLabels)
33+
assert.Equal(t, "/models/ootb/ei/lw-face-det.eim", model.ModelConfiguration["EI_OBJ_DETECTION_MODEL"])
34+
assert.Equal(t, []string{"arduino:object_detection", "arduino:video_object_detection"}, model.Bricks)
35+
assert.Equal(t, "qualcomm-ai-hub", model.Metadata["source"])
36+
assert.Equal(t, "false", model.Metadata["ei-gpu-mode"])
37+
assert.Equal(t, "face-det-lite", model.Metadata["source-model-id"])
38+
assert.Equal(t, "https://aihub.qualcomm.com/models/face_det_lite", model.Metadata["source-model-url"])
39+
})
40+
41+
t.Run("it fails if model-list.yaml does not exist", func(t *testing.T) {
42+
nonExistentPath := paths.New("nonexistentdir")
43+
modelsIndex, err := GenerateModelsIndexFromFile(nonExistentPath)
44+
assert.Error(t, err)
45+
assert.Nil(t, modelsIndex)
46+
})
47+
48+
t.Run("it gets models by a brick", func(t *testing.T) {
49+
model := modelsIndex.GetModelsByBrick("not-existing-brick")
50+
assert.Nil(t, model)
51+
52+
model = modelsIndex.GetModelsByBrick("arduino:object_detection")
53+
assert.Len(t, model, 1)
54+
assert.Equal(t, "face-detection", model[0].ID)
55+
})
56+
57+
t.Run("it gets models by bricks", func(t *testing.T) {
58+
models := modelsIndex.GetModelsByBricks([]string{"arduino:non_existing"})
59+
assert.Len(t, models, 0)
60+
assert.Nil(t, models)
61+
62+
models = modelsIndex.GetModelsByBricks([]string{"arduino:video_object_detection"})
63+
assert.Len(t, models, 2)
64+
assert.Equal(t, "face-detection", models[0].ID)
65+
assert.Equal(t, "yolox-object-detection", models[1].ID)
66+
67+
models = modelsIndex.GetModelsByBricks([]string{"arduino:object_detection", "arduino:video_object_detection"})
68+
assert.Len(t, models, 2)
69+
assert.Equal(t, "face-detection", models[0].ID)
70+
assert.Equal(t, "yolox-object-detection", models[1].ID)
71+
})
72+
}

‎internal/orchestrator/modelsindex/testdata/models-list.yaml‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
models:
2+
- face-detection:
3+
runner: brick
4+
name : "Lightweight-Face-Detection"
5+
description: "Face bounding box detection. This model is trained on the WIDER FACE dataset and can detect faces in images."
6+
model_configuration:
7+
"EI_OBJ_DETECTION_MODEL": "/models/ootb/ei/lw-face-det.eim"
8+
model_labels:
9+
- face
10+
bricks:
11+
- arduino:object_detection
12+
- arduino:video_object_detection
13+
metadata:
14+
source: "qualcomm-ai-hub"
15+
ei-gpu-mode: false
16+
source-model-id: "face-det-lite"
17+
source-model-url: "https://aihub.qualcomm.com/models/face_det_lite"
18+
- yolox-object-detection:
19+
runner: brick
20+
name : "General purpose object detection - YoloX"
21+
description: "General purpose object detection model based on YoloX Nano. This model is trained on the COCO dataset and can detect 80 different object classes."
22+
model_configuration:
23+
"EI_OBJ_DETECTION_MODEL": "/models/ootb/ei/yolo-x-nano.eim"
24+
model_labels:
25+
- airplane
26+
- apple
27+
- backpack
28+
- banana
29+
- baseball bat
30+
- baseball glove
31+
- bear
32+
- bed
33+
- bench
34+
- bicycle
35+
- bird
36+
- boat
37+
- book
38+
- bottle
39+
- bowl
40+
- broccoli
41+
- bus
42+
- cake
43+
- car
44+
- carrot
45+
- cat
46+
- cell phone
47+
- chair
48+
- clock
49+
- couch
50+
- cow
51+
- cup
52+
- dining table
53+
- dog
54+
- donut
55+
- elephant
56+
- fire hydrant
57+
- fork
58+
- frisbee
59+
- giraffe
60+
- hair drier
61+
- handbag
62+
- hot dog
63+
- horse
64+
- keyboard
65+
- kite
66+
- knife
67+
- laptop
68+
- microwave
69+
- motorcycle
70+
- mouse
71+
- orange
72+
- oven
73+
- parking meter
74+
- person
75+
- pizza
76+
- potted plant
77+
- refrigerator
78+
- remote
79+
- sandwich
80+
- scissors
81+
- sheep
82+
- sink
83+
- skateboard
84+
- skis
85+
- snowboard
86+
- spoon
87+
- sports ball
88+
- stop sign
89+
- suitcase
90+
- surfboard
91+
- teddy bear
92+
- tennis racket
93+
- tie
94+
- toaster
95+
- toilet
96+
- toothbrush
97+
- traffic light
98+
- train
99+
- truck
100+
- tv
101+
- umbrella
102+
- vase
103+
- wine glass
104+
- zebra
105+
metadata:
106+
source: "edgeimpulse"
107+
ei-project-id: 717280
108+
source-model-id: "YOLOX-Nano"
109+
source-model-url: "https://github.com/Megvii-BaseDetection/YOLOX"
110+
bricks:
111+
- arduino:video_object_detection

0 commit comments

Comments
 (0)