@@ -82,7 +82,7 @@ func TestPool(t *testing.T) {
8282 fakeClient := fake .NewClientBuilder ().
8383 WithScheme (scheme ).
8484 Build ()
85- pmf := backendmetrics .NewPodMetricsFactory (& backendmetrics.FakePodMetricsClient {}, time .Second )
85+ pmf := backendmetrics .NewPodMetricsFactory (& backendmetrics.FakePodMetricsClient {}, time .Second , time . Second * 2 )
8686 datastore := NewDatastore (context .Background (), pmf )
8787 _ = datastore .PoolSet (context .Background (), fakeClient , tt .inferencePool )
8888 gotPool , gotErr := datastore .PoolGet ()
@@ -214,7 +214,7 @@ func TestModel(t *testing.T) {
214214 }
215215 for _ , test := range tests {
216216 t .Run (test .name , func (t * testing.T ) {
217- pmf := backendmetrics .NewPodMetricsFactory (& backendmetrics.FakePodMetricsClient {}, time .Second )
217+ pmf := backendmetrics .NewPodMetricsFactory (& backendmetrics.FakePodMetricsClient {}, time .Second , time . Second * 2 )
218218 ds := NewDatastore (t .Context (), pmf )
219219 for _ , m := range test .existingModels {
220220 ds .ObjectiveSetIfOlder (m )
@@ -265,6 +265,7 @@ var (
265265 },
266266 WaitingModels : map [string ]int {},
267267 }
268+
268269 pod1NamespacedName = types.NamespacedName {Name : pod1 .Name , Namespace : pod1 .Namespace }
269270 pod2NamespacedName = types.NamespacedName {Name : pod2 .Name , Namespace : pod2 .Namespace }
270271 inferencePool = & v1.InferencePool {
@@ -314,8 +315,7 @@ func TestMetrics(t *testing.T) {
314315 },
315316 },
316317 storePods : []* corev1.Pod {pod1 , pod2 },
317- want : []* backendmetrics.MetricsState {
318- pod1Metrics ,
318+ want : []* backendmetrics.MetricsState {pod1Metrics ,
319319 // Failed to fetch pod2 metrics so it remains the default values.
320320 {
321321 ActiveModels : map [string ]int {},
@@ -338,14 +338,15 @@ func TestMetrics(t *testing.T) {
338338 fakeClient := fake .NewClientBuilder ().
339339 WithScheme (scheme ).
340340 Build ()
341- pmf := backendmetrics .NewPodMetricsFactory (test .pmc , time .Millisecond )
341+ pmf := backendmetrics .NewPodMetricsFactory (test .pmc , time .Millisecond , time . Second * 2 )
342342 ds := NewDatastore (ctx , pmf )
343343 _ = ds .PoolSet (ctx , fakeClient , inferencePool )
344344 for _ , pod := range test .storePods {
345345 ds .PodUpdateOrAddIfNotExist (pod )
346346 }
347+ time .Sleep (1 * time .Second ) // Give some time for the metrics to be fetched.
347348 assert .EventuallyWithT (t , func (t * assert.CollectT ) {
348- got := ds .PodGetAll ( )
349+ got := ds .PodList ( backendmetrics . AllPodPredicate )
349350 metrics := []* backendmetrics.MetricsState {}
350351 for _ , one := range got {
351352 metrics = append (metrics , one .GetMetrics ())
@@ -431,15 +432,15 @@ func TestPods(t *testing.T) {
431432 for _ , test := range tests {
432433 t .Run (test .name , func (t * testing.T ) {
433434 ctx := context .Background ()
434- pmf := backendmetrics .NewPodMetricsFactory (& backendmetrics.FakePodMetricsClient {}, time .Second )
435+ pmf := backendmetrics .NewPodMetricsFactory (& backendmetrics.FakePodMetricsClient {}, time .Second , time . Second * 2 )
435436 ds := NewDatastore (t .Context (), pmf )
436437 for _ , pod := range test .existingPods {
437438 ds .PodUpdateOrAddIfNotExist (pod )
438439 }
439440
440441 test .op (ctx , ds )
441442 var gotPods []* corev1.Pod
442- for _ , pm := range ds .PodGetAll ( ) {
443+ for _ , pm := range ds .PodList ( backendmetrics . AllPodPredicate ) {
443444 pod := & corev1.Pod {ObjectMeta : metav1.ObjectMeta {Name : pm .GetPod ().NamespacedName .Name , Namespace : pm .GetPod ().NamespacedName .Namespace }, Status : corev1.PodStatus {PodIP : pm .GetPod ().Address }}
444445 gotPods = append (gotPods , pod )
445446 }
0 commit comments