Fix integration tests

GoogleCloudPlatform · Mar 12, 2024 · 1039c27 · 1039c27
1 parent 1188738
commit 1039c27
Showing 1 changed file with 7 additions and 2 deletions.
diff --git 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 {