quorate/services/short/tunnel/lease_test.go

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")
	}
}