From 5bef4ea63bc88cbb9b0b4ffc487235dbc583432d Mon Sep 17 00:00:00 2001
From: Sargun Dhillon <sargun@sargun.me>
Date: Mon, 16 Nov 2020 03:50:55 -0800
Subject: [PATCH] Add some more envtests

This tests if the lease is created
---
 internal/test/vk_envtest/env_test.go | 68 +++++++++++++++++++++++-----
 1 file changed, 57 insertions(+), 11 deletions(-)

diff --git a/internal/test/vk_envtest/env_test.go b/internal/test/vk_envtest/env_test.go
index baa3720ac..193b4bc30 100644
--- a/internal/test/vk_envtest/env_test.go
+++ b/internal/test/vk_envtest/env_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"flag"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -13,6 +14,7 @@ import (
 	logruslogger "github.com/virtual-kubelet/virtual-kubelet/log/logrus"
 	"github.com/virtual-kubelet/virtual-kubelet/node"
 	"gotest.tools/assert"
+	is "gotest.tools/assert/cmp"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -46,12 +48,23 @@ func TestEnvtest(t *testing.T) {
 	}()
 
 	t.Log("Env test environment ready")
-	t.Run("E2ERun", func(t *testing.T) {
-		testNodeE2ERun(t, env)
+	t.Run("E2ERunWithoutLeases", func(t *testing.T) {
+		testNodeE2ERun(t, env, false)
+	})
+	t.Run("E2ERunWithLeases", func(t *testing.T) {
+		testNodeE2ERun(t, env, true)
 	})
 }
 
-func testNodeE2ERun(t *testing.T, env *envtest.Environment) {
+func nodeNameForTest(t *testing.T) string {
+	name := t.Name()
+	name = strings.ToLower(name)
+	name = strings.ReplaceAll(name, "/", "-")
+	name = strings.ReplaceAll(name, "_", "-")
+	return name
+}
+
+func testNodeE2ERun(t *testing.T, env *envtest.Environment, withLeases bool) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -71,13 +84,18 @@ func testNodeE2ERun(t *testing.T, env *envtest.Environment) {
 
 	testNode := &corev1.Node{
 		ObjectMeta: metav1.ObjectMeta{
-			Name: "e2e-envtest",
+			Name: nodeNameForTest(t),
 		},
 	}
 
 	testNodeCopy := testNode.DeepCopy()
 
-	node, err := node.NewNodeController(testProvider, testNode, nodes)
+	opts := []node.NodeControllerOpt{}
+	leasesClient := clientset.CoordinationV1beta1().Leases(corev1.NamespaceNodeLease)
+	if withLeases {
+		opts = append(opts, node.WithNodeEnableLeaseV1Beta1(leasesClient, nil))
+	}
+	node, err := node.NewNodeController(testProvider, testNode, nodes, opts...)
 	assert.NilError(t, err)
 
 	chErr := make(chan error, 1)
@@ -85,22 +103,50 @@ func testNodeE2ERun(t *testing.T, env *envtest.Environment) {
 		chErr <- node.Run(ctx)
 	}()
 
-	<-node.Ready()
+	log.G(ctx).Debug("Waiting for node ready")
+	select {
+	case <-node.Ready():
+	case err = <-chErr:
+		t.Fatal(err)
+	case <-ctx.Done():
+		t.Fatal(ctx.Err())
+	}
 
 	now := time.Now()
+	var n *corev1.Node
 	for time.Since(now) < time.Minute*5 {
-		n, err := nodes.Get(ctx, testNodeCopy.Name, metav1.GetOptions{})
+		n, err = nodes.Get(ctx, testNodeCopy.Name, metav1.GetOptions{})
 		if errors.IsNotFound(err) {
 			continue
 		}
 		if err == nil {
 			t.Log(n)
-			cancel()
-			err := <-chErr
-			assert.NilError(t, err)
-			return
+			goto node_found
 		}
 		time.Sleep(100 * time.Millisecond)
 	}
 	t.Fatal("Node never found")
+
+node_found:
+	if withLeases {
+		for time.Since(now) < time.Minute*5 {
+			l, err := leasesClient.Get(ctx, testNodeCopy.Name, metav1.GetOptions{})
+			if errors.IsNotFound(err) {
+				continue
+			}
+			if err == nil {
+				t.Log(l)
+				assert.Assert(t, is.Len(l.OwnerReferences, 1))
+				assert.Assert(t, is.Equal(l.OwnerReferences[0].Name, n.Name))
+				assert.Assert(t, is.Equal(l.OwnerReferences[0].UID, n.UID))
+				goto lease_found
+			}
+			time.Sleep(100 * time.Millisecond)
+		}
+	}
+
+lease_found:
+	cancel()
+	err = <-chErr
+	assert.NilError(t, err)
 }