From c5d431dc368a07800eac744f64c0dad5c8515579 Mon Sep 17 00:00:00 2001
From: Kirill Sibirev <sibirev@nebius.com>
Date: Tue, 16 Apr 2024 12:46:43 +0200
Subject: [PATCH] Add local update tests (#234)

* Local update test

* Local update test

* Local update test
---
 controllers/ytsaurus_local_test.go | 58 +++++++++++++++++++++++-------
 pkg/testutil/testhelper.go         | 13 +++++++
 2 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/controllers/ytsaurus_local_test.go b/controllers/ytsaurus_local_test.go
index 2b1127f4..dd3ba501 100644
--- a/controllers/ytsaurus_local_test.go
+++ b/controllers/ytsaurus_local_test.go
@@ -1,9 +1,11 @@
 package controllers_test
 
 import (
+	"fmt"
 	"path/filepath"
 	"testing"
 
+	"github.com/stretchr/testify/require"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	ctrl "sigs.k8s.io/controller-runtime"
@@ -20,8 +22,7 @@ const (
 	dndsNameOne       = "dn-1"
 )
 
-func TestYtsaurusFromScratch(t *testing.T) {
-	namespace := "ytsaurus-from-scratch"
+func prepareTest(t *testing.T, namespace string) *testutil.TestHelper {
 	h := testutil.NewTestHelper(t, namespace, filepath.Join("..", "config", "crd", "bases"))
 	reconcilerSetup := func(mgr ctrl.Manager) error {
 		return (&controllers.YtsaurusReconciler{
@@ -31,6 +32,26 @@ func TestYtsaurusFromScratch(t *testing.T) {
 		}).SetupWithManager(mgr)
 	}
 	h.Start(reconcilerSetup)
+	return h
+}
+
+func waitClusterState(h *testutil.TestHelper, expectedState ytv1.ClusterState) {
+	h.Logf("[ Wait for YTsaurus %s state ]", expectedState)
+	testutil.FetchAndCheckEventually(
+		h,
+		ytsaurusName,
+		&ytv1.Ytsaurus{},
+		fmt.Sprintf("cluster state is %s", expectedState),
+		func(obj client.Object) bool {
+			state := obj.(*ytv1.Ytsaurus).Status.State
+			return state == expectedState
+		},
+	)
+}
+
+func TestYtsaurusFromScratch(t *testing.T) {
+	namespace := "ytsaurus-from-scratch"
+	h := prepareTest(t, namespace)
 	defer h.Stop()
 
 	ytsaurusResource := testutil.BuildMinimalYtsaurus(namespace, ytsaurusName)
@@ -78,15 +99,28 @@ func TestYtsaurusFromScratch(t *testing.T) {
 			return len(secret.Data["YT_TOKEN"]) != 0
 		},
 	)
+	waitClusterState(h, ytv1.ClusterStateRunning)
+}
 
-	testutil.FetchAndCheckEventually(
-		h,
-		ytsaurusName,
-		&ytv1.Ytsaurus{},
-		"cluster state is running",
-		func(obj client.Object) bool {
-			state := obj.(*ytv1.Ytsaurus).Status.State
-			return state == ytv1.ClusterStateRunning
-		},
-	)
+func TestYtsaurusUpdateStatelessComponent(t *testing.T) {
+	namespace := "upd-discovery"
+	h := prepareTest(t, namespace)
+	defer h.Stop()
+
+	ytsaurusResource := testutil.BuildMinimalYtsaurus(namespace, ytsaurusName)
+	testutil.DeployObject(h, &ytsaurusResource)
+
+	waitClusterState(h, ytv1.ClusterStateRunning)
+
+	imageUpdated := testYtsaurusImage + "-updated"
+	ytsaurusResource.Spec.Discovery.Image = &imageUpdated
+	t.Log("[ Updating discovery with disabled full update ]")
+	ytsaurusResource.Spec.EnableFullUpdate = false
+	testutil.UpdateObject(h, &ytv1.Ytsaurus{}, &ytsaurusResource)
+
+	waitClusterState(h, ytv1.ClusterStateRunning)
+
+	sts := appsv1.StatefulSet{}
+	testutil.GetObject(h, "ds", &sts)
+	require.Equal(t, imageUpdated, sts.Spec.Template.Spec.Containers[0].Image)
 }
diff --git a/pkg/testutil/testhelper.go b/pkg/testutil/testhelper.go
index cc2930ca..df1ba00a 100644
--- a/pkg/testutil/testhelper.go
+++ b/pkg/testutil/testhelper.go
@@ -138,12 +138,25 @@ func (h *TestHelper) createNamespace() {
 	require.NoError(h.t, err)
 }
 
+func (h *TestHelper) Log(args ...any) {
+	h.t.Log(args...)
+}
+
+func (h *TestHelper) Logf(format string, args ...any) {
+	h.t.Logf(format, args...)
+}
+
 // helpers
 func GetObject(h *TestHelper, key string, emptyObject client.Object) {
 	k8sCli := h.GetK8sClient()
 	err := k8sCli.Get(context.Background(), h.GetObjectKey(key), emptyObject)
 	require.NoError(h.t, err)
 }
+func ListObjects(h *TestHelper, emptyList client.ObjectList) {
+	k8sCli := h.GetK8sClient()
+	err := k8sCli.List(context.Background(), emptyList)
+	require.NoError(h.t, err)
+}
 func DeployObject(h *TestHelper, object client.Object) {
 	k8sCli := h.GetK8sClient()
 	err := k8sCli.Create(context.Background(), object)