kubectl-persistent-logger/logs/watcher_test.go

package logs_test

import (
	"bytes"
	"context"
	"errors"
	"io"
	"log"
	"strings"
	"sync"
	"testing"
	"time"

	"git.netflux.io/rob/kubectl-persistent-logger/logs"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/watch"
	dynamicclient "k8s.io/client-go/dynamic/fake"
	testclient "k8s.io/client-go/kubernetes/fake"
	k8stest "k8s.io/client-go/testing"
)

// concurrentBuffer is a buffer which implements io.Writer and fmt.Stringer. It
// is useful in tests to allow the output to be read back from the destination
// buffer.
type concurrentBuffer struct {
	bytes.Buffer
	mu sync.Mutex
}

func (cb *concurrentBuffer) Write(p []byte) (int, error) {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	return cb.Buffer.Write(p)
}

func (cb *concurrentBuffer) String() string {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	return cb.Buffer.String()
}

type mockPodWatcher struct{ err error }

func (m *mockPodWatcher) WatchPods(ctx context.Context) error { return m.err }
func (m *mockPodWatcher) Close()                              {}

func mockPodwatcherFunc(err error) logs.PodWatcherFunc {
	return func(logs.KubernetesClient, string, string, labels.Selector, io.Writer, *log.Logger) logs.PodWatcherInterface {
		return &mockPodWatcher{err: err}
	}
}

func discardLogger() *log.Logger {
	return log.New(io.Discard, "", 0)
}

func buildDeployment(t *testing.T, name string) *unstructured.Unstructured {
	deployment := new(unstructured.Unstructured)
	deployment.SetAPIVersion("v1")
	deployment.SetKind("deployment")
	deployment.SetName("mydeployment")
	deployment.SetNamespace("default")
	deployment.SetUID(types.UID("foo"))
	deployment.SetResourceVersion("1")
	require.NoError(t, unstructured.SetNestedField(deployment.Object, map[string]any{"app": "myapp"}, "spec", "selector", "matchLabels"))
	return deployment
}

func TestWatcherStrictExist(t *testing.T) {
	client := logs.KubernetesClient{Untyped: dynamicclient.NewSimpleDynamicClient(runtime.NewScheme())}

	var buf bytes.Buffer
	params := logs.WatcherParams{Name: "mydeployment", Type: "deployments", Namespace: "default", StrictExist: true}
	watcher := logs.NewWatcher(params, client, mockPodwatcherFunc(nil), &buf, discardLogger())

	err := watcher.Watch(context.Background())
	assert.EqualError(t, err, `deployments.apps "mydeployment" not found`)
}

func TestWatcherPodWatcherError(t *testing.T) {
	deploymentsWatcher := watch.NewFake()

	untypedClient := dynamicclient.NewSimpleDynamicClient(runtime.NewScheme())
	untypedClient.PrependWatchReactor("deployments", k8stest.DefaultWatchReactor(deploymentsWatcher, nil))
	client := logs.KubernetesClient{Untyped: untypedClient}

	var buf bytes.Buffer
	wantErr := errors.New("foo")
	params := logs.WatcherParams{Name: "mydeployment", Type: "deployments", Namespace: "default"}
	watcher := logs.NewWatcher(params, client, mockPodwatcherFunc(wantErr), &buf, discardLogger())

	go func() {
		defer deploymentsWatcher.Stop()

		deployment := buildDeployment(t, "mydeployment")
		deploymentsWatcher.Add(deployment)
	}()

	err := watcher.Watch(context.Background())
	assert.EqualError(t, err, wantErr.Error())
}

func TestWatcherClosedChannel(t *testing.T) {
	deploymentsWatcher := watch.NewFake()

	untypedClient := dynamicclient.NewSimpleDynamicClient(runtime.NewScheme())
	untypedClient.PrependWatchReactor("deployments", k8stest.DefaultWatchReactor(deploymentsWatcher, nil))
	client := logs.KubernetesClient{Untyped: untypedClient}

	var buf bytes.Buffer
	params := logs.WatcherParams{Name: "mydeployment", Type: "deployments", Namespace: "default"}
	watcher := logs.NewWatcher(params, client, nil, &buf, discardLogger())
	// Immediately stop the watcher, which closes the ResultChan.
	// This should be expected to be handled.
	deploymentsWatcher.Stop()

	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*500)
	defer cancel()

	err := watcher.Watch(ctx)
	require.Equal(t, context.DeadlineExceeded, err)
	assert.Equal(t, "", buf.String())
}

func TestWatcherWithPodWatcher(t *testing.T) {
	deploymentsWatcher := watch.NewFake()
	defer deploymentsWatcher.Stop()

	podsWatcher := watch.NewFake()
	defer podsWatcher.Stop()

	typedClient := testclient.NewSimpleClientset()
	typedClient.PrependWatchReactor("pods", k8stest.DefaultWatchReactor(podsWatcher, nil))
	untypedClient := dynamicclient.NewSimpleDynamicClient(runtime.NewScheme())
	untypedClient.PrependWatchReactor("deployments", k8stest.DefaultWatchReactor(deploymentsWatcher, nil))
	client := logs.KubernetesClient{Typed: typedClient, Untyped: untypedClient}

	var cb concurrentBuffer
	params := logs.WatcherParams{Name: "mydeployment", Type: "deployments", Namespace: "default"}
	watcher := logs.NewWatcher(params, client, logs.NewPodWatcher, &cb, discardLogger())

	go func() {
		deployment := buildDeployment(t, "mydeployment")
		deploymentsWatcher.Add(deployment)
		time.Sleep(time.Millisecond * 250)

		pods := []*corev1.Pod{
			{ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default", ResourceVersion: "1"}, Status: corev1.PodStatus{Phase: corev1.PodRunning}},
			{ObjectMeta: metav1.ObjectMeta{Name: "bar", Namespace: "default", ResourceVersion: "1"}, Status: corev1.PodStatus{Phase: corev1.PodRunning}},
			{ObjectMeta: metav1.ObjectMeta{Name: "baz", Namespace: "default", ResourceVersion: "1"}, Status: corev1.PodStatus{Phase: corev1.PodPending}},
		}
		for _, pod := range pods {
			podsWatcher.Add(pod)
			time.Sleep(time.Millisecond * 250)
		}

		watcher.Close()
	}()

	err := watcher.Watch(context.Background())
	require.NoError(t, err)
	lines := bufToLines(&cb)
	require.Len(t, lines, 2)
	assert.ElementsMatch(t, []string{"[foo] fake logs", "[bar] fake logs"}, lines)
}

// bufToLines splits the output buffer removing newline characters.
func bufToLines(buf *concurrentBuffer) []string {
	return strings.Split(strings.TrimSpace(buf.String()), "\n")
}