From 1039c27f17ac9ea95adeb6968b895afc234aeadd Mon Sep 17 00:00:00 2001 From: Nick Stogner Date: Fri, 8 Mar 2024 12:06:58 -0500 Subject: [PATCH] Fix integration tests --- tpu-provisioner/internal/controller/provider_test.go | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tpu-provisioner/internal/controller/provider_test.go b/tpu-provisioner/internal/controller/provider_test.go index 63fd2f66d..5e7e25de7 100644 --- a/tpu-provisioner/internal/controller/provider_test.go +++ b/tpu-provisioner/internal/controller/provider_test.go @@ -4,19 +4,24 @@ import ( "sync" "time" + "github.com/GoogleCloudPlatform/ai-on-gke/tpu-provisioner/internal/cloud" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/types" ) +var _ cloud.Provider = &testProvider{} + type testProvider struct { sync.Mutex created map[types.NamespacedName]bool deleted map[string]time.Time + + cloud.Provider } func (p *testProvider) NodePoolLabelKey() string { return "cloud.test.com/test-nodepool" } -func (p *testProvider) EnsureNodePoolForPod(pod *corev1.Pod) error { +func (p *testProvider) EnsureNodePoolForPod(pod *corev1.Pod, _ string) error { p.Lock() defer p.Unlock() p.created[types.NamespacedName{Namespace: pod.Namespace, Name: pod.Name}] = true @@ -29,7 +34,7 @@ func (p *testProvider) getCreated(nn types.NamespacedName) bool { return p.created[nn] } -func (p *testProvider) DeleteNodePoolForNode(node *corev1.Node) error { +func (p *testProvider) DeleteNodePoolForNode(node *corev1.Node, _ string) error { p.Lock() defer p.Unlock() if _, exists := p.deleted[node.Name]; !exists {