Skip to content

Commit 4a23bbf

Browse files
committed
fix: add ut for modelhub.
Signed-off-by: X1aoZEOuO <nizefeng2002@outlook.com>
1 parent cd84215 commit 4a23bbf

1 file changed

Lines changed: 145 additions & 0 deletions

File tree

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
/*
2+
Copyright 2025 The InftyAI Team.
3+
4+
Licensed under the Apache License,
5+
Version 2.0 (the "License");
6+
you may not use this file except in compliance with the License.
7+
You may obtain a copy of the License at
8+
9+
http://www.apache.org/licenses/LICENSE-2.0
10+
11+
Unless required by applicable law or agreed to in writing, software
12+
distributed under the License is distributed on an "AS IS" BASIS,
13+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
See the License for the specific language governing permissions and
15+
limitations under the License.
16+
*/
17+
18+
package modelSource
19+
20+
import (
21+
"testing"
22+
23+
"github.com/stretchr/testify/assert"
24+
corev1 "k8s.io/api/core/v1"
25+
26+
"github.com/inftyai/llmaz/pkg"
27+
)
28+
29+
func Test_ModelHubProvider_InjectModelLoader(t *testing.T) {
30+
fileName := "weights.gguf"
31+
revision := "v1.2"
32+
allowPatterns := []string{"*.gguf", "*.json"}
33+
ignorePatterns := []string{"*.tmp"}
34+
35+
testCases := []struct {
36+
name string
37+
provider *ModelHubProvider
38+
index int
39+
expectMainModel bool
40+
expectEnvContains []string
41+
}{
42+
{
43+
name: "inject full modelhub with fileName, revision, allow/ignore",
44+
provider: &ModelHubProvider{
45+
modelName: "llama3",
46+
modelID: "meta/llama-3",
47+
modelHub: "Huggingface",
48+
fileName: &fileName,
49+
modelRevision: &revision,
50+
modelAllowPatterns: allowPatterns,
51+
modelIgnorePatterns: ignorePatterns,
52+
},
53+
index: 0,
54+
expectMainModel: true,
55+
expectEnvContains: []string{
56+
"MODEL_SOURCE_TYPE", "MODEL_ID", "MODEL_HUB_NAME", "MODEL_FILENAME",
57+
"REVISION", "MODEL_ALLOW_PATTERNS", "MODEL_IGNORE_PATTERNS",
58+
"HUGGING_FACE_HUB_TOKEN", "HF_TOKEN",
59+
},
60+
},
61+
{
62+
name: "inject with index > 0 skips volume/container mount",
63+
provider: &ModelHubProvider{
64+
modelName: "sub-model",
65+
modelID: "some/model",
66+
modelHub: "Huggingface",
67+
},
68+
index: 1,
69+
expectMainModel: false,
70+
expectEnvContains: []string{
71+
"MODEL_SOURCE_TYPE", "MODEL_ID", "MODEL_HUB_NAME",
72+
"HUGGING_FACE_HUB_TOKEN", "HF_TOKEN",
73+
},
74+
},
75+
}
76+
77+
for _, tc := range testCases {
78+
t.Run(tc.name, func(t *testing.T) {
79+
template := &corev1.PodTemplateSpec{
80+
Spec: corev1.PodSpec{
81+
Containers: []corev1.Container{
82+
{
83+
Name: MODEL_RUNNER_CONTAINER_NAME,
84+
Env: []corev1.EnvVar{
85+
{Name: "HTTP_PROXY", Value: "http://1.1.1.1"},
86+
},
87+
},
88+
},
89+
},
90+
}
91+
92+
tc.provider.InjectModelLoader(template, tc.index)
93+
94+
assert.Len(t, template.Spec.InitContainers, 1)
95+
initContainer := template.Spec.InitContainers[0]
96+
expectedName := MODEL_LOADER_CONTAINER_NAME
97+
if tc.index != 0 {
98+
expectedName += "-" + string(rune('0'+tc.index))
99+
}
100+
assert.Equal(t, expectedName, initContainer.Name)
101+
assert.Equal(t, pkg.LOADER_IMAGE, initContainer.Image)
102+
103+
// Check env vars exist
104+
for _, key := range tc.expectEnvContains {
105+
found := false
106+
for _, env := range initContainer.Env {
107+
if env.Name == key {
108+
found = true
109+
break
110+
}
111+
}
112+
assert.True(t, found, "expected env %s not found", key)
113+
}
114+
115+
// Main model should inject volume & container mount
116+
if tc.expectMainModel {
117+
// Volume should be present
118+
foundVol := false
119+
for _, v := range template.Spec.Volumes {
120+
if v.Name == MODEL_VOLUME_NAME {
121+
foundVol = true
122+
break
123+
}
124+
}
125+
assert.True(t, foundVol, "volume not injected")
126+
127+
// Runner container mount should exist
128+
foundMount := false
129+
for _, m := range template.Spec.Containers[0].VolumeMounts {
130+
if m.Name == MODEL_VOLUME_NAME && m.ReadOnly && m.MountPath == CONTAINER_MODEL_PATH {
131+
foundMount = true
132+
}
133+
}
134+
assert.True(t, foundMount, "volume mount not injected to runner")
135+
} else {
136+
// No volumes or mounts should be injected
137+
assert.Empty(t, template.Spec.Volumes)
138+
assert.Empty(t, template.Spec.Containers[0].VolumeMounts)
139+
}
140+
141+
// Should always carry over container env
142+
assert.Contains(t, initContainer.Env, corev1.EnvVar{Name: "HTTP_PROXY", Value: "http://1.1.1.1"})
143+
})
144+
}
145+
}

0 commit comments

Comments
 (0)