@@ -237,6 +237,7 @@ type testContext struct {
237237 k8sClient kubernetes.Interface
238238 substrateClient versioned.Interface
239239 persistence * ateredis.Persistence
240+ workerCache * workercache.Cache
240241 fakeAtelet * FakeAteletServer
241242 cleanup func ()
242243 actorTemplateLister listersv1alpha1.ActorTemplateLister
@@ -369,6 +370,7 @@ func setupTest(t *testing.T, ns string) *testContext {
369370 k8sClient : k8sClient ,
370371 substrateClient : substrateClient ,
371372 persistence : persistence ,
373+ workerCache : wc ,
372374 fakeAtelet : fakeAtelet ,
373375 cleanup : cleanup ,
374376 actorTemplateLister : actorTemplateLister ,
@@ -583,13 +585,13 @@ func createWorkerPod(t *testing.T, tc *testContext, ns string, name string, node
583585 t .Fatalf ("failed to update worker pod status: %v" , err )
584586 }
585587
586- // Wait for worker to be registered via API
588+ // Wait for worker to be visible to the scheduler.
587589 err = wait .PollUntilContextTimeout (context .Background (), 100 * time .Millisecond , 5 * time .Second , true , func (ctx context.Context ) (bool , error ) {
588- resp , err := tc .client . ListWorkers ( ctx , & ateapipb. ListWorkersRequest {} )
590+ workers , err := tc .workerCache . Workers ( )
589591 if err != nil {
590- return false , nil // Retry on API error
592+ return false , nil
591593 }
592- for _ , w := range resp . GetWorkers () {
594+ for _ , w := range workers {
593595 if w .GetWorkerNamespace () == ns && w .GetWorkerPod () == name {
594596 return true , nil
595597 }
@@ -608,13 +610,13 @@ func deleteWorkerPod(t *testing.T, tc *testContext, ns string, name string) {
608610 t .Fatalf ("failed to delete worker pod %s: %v" , name , err )
609611 }
610612
611- // Wait for worker to be removed from API
613+ // Wait for worker to be removed from the scheduler.
612614 err = wait .PollUntilContextTimeout (context .Background (), 100 * time .Millisecond , 5 * time .Second , true , func (ctx context.Context ) (bool , error ) {
613- resp , err := tc .client . ListWorkers ( ctx , & ateapipb. ListWorkersRequest {} )
615+ workers , err := tc .workerCache . Workers ( )
614616 if err != nil {
615- return false , nil // Retry on API error
617+ return false , nil
616618 }
617- for _ , w := range resp . GetWorkers () {
619+ for _ , w := range workers {
618620 if w .GetWorkerNamespace () == ns && w .GetWorkerPod () == name {
619621 return false , nil // Still there
620622 }
0 commit comments