ce7cb702218693d08ea6d83464d2457c9250027c

services/short/tunnel/lease_test.go

@@ -0,0 +1,135 @@
+package tunnel
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	coordination "k8s.io/api/coordination/v1"
+	core "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/client-go/kubernetes/fake"
+	ktest "k8s.io/client-go/testing"
+)
+
+func testKube() *fake.Clientset {
+	kube := fake.NewClientset(&coordination.Lease{ObjectMeta: metav1.ObjectMeta{Name: "tunnel", Namespace: "short", ResourceVersion: "1"}})
+	for _, name := range []string{"tunnel-0", "tunnel-1"} {
+		kube.Tracker().Add(&core.Pod{ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: "short", Labels: map[string]string{"app.kubernetes.io/name": "tunnel", "app.kubernetes.io/instance": "short"}}})
+	}
+	return kube
+}
+
+func TestClaim(t *testing.T) {
+	kube := testKube()
+	ctx := context.Background()
+	if err := label(ctx, kube, "short", "tunnel-0", "active"); err != nil {
+		t.Fatal(err)
+	}
+	if err := claim(ctx, kube, "short", "tunnel-1"); err != nil {
+		t.Fatal(err)
+	}
+	pods, _ := kube.CoreV1().Pods("short").List(ctx, metav1.ListOptions{})
+	for _, pod := range pods.Items {
+		if (pod.Labels["role"] == "active") != (pod.Name == "tunnel-1") {
+			t.Fatalf("wrong backend: %s %v", pod.Name, pod.Labels)
+		}
+	}
+	if err := label(ctx, kube, "short", "tunnel-0", "active"); err != nil {
+		t.Fatal(err)
+	}
+	if err := label(ctx, kube, "short", "tunnel-1", nil); err != nil {
+		t.Fatal(err)
+	}
+	kube.PrependReactor("patch", "pods", func(a ktest.Action) (bool, runtime.Object, error) {
+		if a.(ktest.PatchAction).GetName() == "tunnel-0" {
+			return true, nil, fmt.Errorf("API unavailable")
+		}
+		return false, nil, nil
+	})
+	if err := claim(ctx, kube, "short", "tunnel-1"); err == nil {
+		t.Fatal("claimed without removing old backend")
+	}
+	pod, _ := kube.CoreV1().Pods("short").Get(ctx, "tunnel-1", metav1.GetOptions{})
+	if pod.Labels["role"] == "active" {
+		t.Fatal("two active backends after API failure")
+	}
+}
+
+func TestLeaseHandoff(t *testing.T) {
+	kube := testKube()
+	var api sync.Mutex
+	version := 1
+	kube.PrependReactor("update", "leases", func(a ktest.Action) (bool, runtime.Object, error) {
+		api.Lock()
+		defer api.Unlock()
+		u := a.(ktest.UpdateAction).GetObject().(*coordination.Lease).DeepCopy()
+		gvr := schema.GroupVersionResource{Group: "coordination.k8s.io", Version: "v1", Resource: "leases"}
+		old, err := kube.Tracker().Get(gvr, "short", "tunnel")
+		if err != nil {
+			return true, nil, err
+		}
+		if u.ResourceVersion != old.(*coordination.Lease).ResourceVersion {
+			return true, nil, apierrors.NewConflict(gvr.GroupResource(), "tunnel", fmt.Errorf("version changed"))
+		}
+		version++
+		u.ResourceVersion = fmt.Sprint(version)
+		return true, u, kube.Tracker().Update(gvr, u, "short")
+	})
+	ctx, cancel := context.WithTimeout(context.Background(), 75*time.Second)
+	defer cancel()
+	first, stopFirst := context.WithCancel(ctx)
+	defer stopFirst()
+	var count atomic.Int32
+	started := make(chan string, 2)
+	done := make(chan error, 2)
+	work := func(name string) func(context.Context) error {
+		return func(ctx context.Context) error {
+			if count.Add(1) != 1 {
+				t.Error("overlapping tunnel owners")
+			}
+			defer count.Add(-1)
+			if err := claim(ctx, kube, "short", name); err != nil {
+				return err
+			}
+			started <- name
+			<-ctx.Done()
+			return nil
+		}
+	}
+	go func() { done <- ownership(first, kube, "short", "tunnel-0", work("tunnel-0")) }()
+	select {
+	case name := <-started:
+		if name != "tunnel-0" {
+			t.Fatal(name)
+		}
+	case <-ctx.Done():
+		t.Fatal("first owner did not start")
+	}
+	go func() { done <- ownership(ctx, kube, "short", "tunnel-1", work("tunnel-1")) }()
+	stopFirst()
+	if err := <-done; err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case name := <-started:
+		if name != "tunnel-1" {
+			t.Fatal(name)
+		}
+	case <-ctx.Done():
+		t.Fatal("standby did not acquire expired lease")
+	}
+	cancel()
+	if err := <-done; err != nil {
+		t.Fatal(err)
+	}
+	if count.Load() != 0 {
+		t.Fatal("owner still running")
+	}
+}