From f13a341e205ec06b6bf1317128220166e0183c73 Mon Sep 17 00:00:00 2001
From: Kirill Sibirev <l0kix2@gmail.com>
Date: Wed, 17 Apr 2024 09:37:47 +0200
Subject: [PATCH] Tests

---
 test/e2e/ytsaurus_controller_test.go | 88 ++++++++++++++++++++++++++++
 1 file changed, 88 insertions(+)

diff --git a/test/e2e/ytsaurus_controller_test.go b/test/e2e/ytsaurus_controller_test.go
index 01765e34..3a6b5b36 100644
--- a/test/e2e/ytsaurus_controller_test.go
+++ b/test/e2e/ytsaurus_controller_test.go
@@ -261,6 +261,85 @@ var _ = Describe("Basic test for Ytsaurus controller", func() {
 			"Should run and update Ytsaurus to the next major version",
 			getSimpleUpdateScenario("test-major-update", ytv1.CoreImageNextVer),
 		)
+		It(
+			"Should be updated according to UpdateStrategy",
+			func(ctx context.Context) {
+				namespace := "teststrategy"
+
+				By("Creating a Ytsaurus resource")
+				ytsaurus := ytv1.CreateBaseYtsaurusResource(namespace)
+				DeferCleanup(deleteYtsaurus, ytsaurus)
+				name := types.NamespacedName{Name: ytsaurus.GetName(), Namespace: namespace}
+				deployAndCheck(ytsaurus, namespace)
+
+				By("Run cluster update with strategy blocked")
+				Expect(k8sClient.Get(ctx, name, ytsaurus)).Should(Succeed())
+				ytsaurus.Spec.UpdateStrategy = ytv1.UpdateStrategyBlocked
+				// Discovery count triggers change in all yson configs.
+				ytsaurus.Spec.Discovery.InstanceCount += 1
+				updateAndCheck(ytsaurus, namespace)
+				//Expect(k8sClient.Update(ctx, ytsaurus)).Should(Succeed())
+
+				By("Sleep 5 second and check not updated")
+				time.Sleep(5 * time.Second)
+				ytsaurus = &ytv1.Ytsaurus{}
+				Expect(k8sClient.Get(ctx, name, ytsaurus)).Should(Succeed())
+				Expect(ytsaurus.Status.State).Should(Equal(ytv1.ClusterStateRunning))
+				// TODO: check pods
+
+				By("Run cluster update with strategy stateless only")
+				Expect(k8sClient.Get(ctx, name, ytsaurus)).Should(Succeed())
+				ytsaurus.Spec.UpdateStrategy = ytv1.UpdateStrategyStatelessOnly
+				ytsaurus.Spec.Discovery.InstanceCount += 1
+				//Expect(k8sClient.Update(ctx, ytsaurus)).Should(Succeed())
+				updateAndCheck(ytsaurus, namespace)
+				// UPDATINGCOMPONENTS=["Discovery","DataNode","HttpProxy","ExecNode","Scheduler","ControllerAgent"]
+				EventuallyYtsaurus(ctx, name, reactionTimeout).Should(HaveClusterState(ytv1.ClusterStateUpdating))
+
+				By("Wait cluster update with stateless strategy complete")
+				EventuallyYtsaurus(ctx, name, upgradeTimeout).Should(HaveClusterState(ytv1.ClusterStateRunning))
+				// TODO: check discovery pods timestamps
+
+				By("Run cluster update with strategy tablet nodes only")
+				Expect(k8sClient.Get(ctx, name, ytsaurus)).Should(Succeed())
+				ytsaurus.Spec.UpdateStrategy = ytv1.UpdateStrategyTabletNodesOnly
+				ytsaurus.Spec.Discovery.InstanceCount += 1
+				//Expect(k8sClient.Update(ctx, ytsaurus)).Should(Succeed())
+				updateAndCheck(ytsaurus, namespace)
+				EventuallyYtsaurus(ctx, name, reactionTimeout).Should(HaveClusterState(ytv1.ClusterStateUpdating))
+
+				By("Wait cluster update with tablet nodes strategy complete")
+				EventuallyYtsaurus(ctx, name, upgradeTimeout).Should(HaveClusterState(ytv1.ClusterStateRunning))
+				// TODO: check tnds pods timestamps
+				//             UPDATINGCOMPONENTS
+				// teststrategy   ytsaurus.cluster.ytsaurus.tech/test-ytsaurus   Updating       WaitingForTabletCellsRemoved   ["TabletNode"]
+
+				By("Run cluster update with strategy master only")
+				Expect(k8sClient.Get(ctx, name, ytsaurus)).Should(Succeed())
+				ytsaurus.Spec.UpdateStrategy = ytv1.UpdateStrategyMasterOnly
+				ytsaurus.Spec.Discovery.InstanceCount += 1
+				//Expect(k8sClient.Update(ctx, ytsaurus)).Should(Succeed())
+				updateAndCheck(ytsaurus, namespace)
+				EventuallyYtsaurus(ctx, name, reactionTimeout).Should(HaveClusterState(ytv1.ClusterStateUpdating))
+
+				By("Wait cluster update with master only complete")
+				EventuallyYtsaurus(ctx, name, upgradeTimeout).Should(HaveClusterState(ytv1.ClusterStateRunning))
+				// TODO: check master pods timestamps
+
+				// TODO: strategy full update
+				By("Run cluster update with strategy full")
+				Expect(k8sClient.Get(ctx, name, ytsaurus)).Should(Succeed())
+				ytsaurus.Spec.UpdateStrategy = ytv1.UpdateStrategyFull
+				ytsaurus.Spec.Discovery.InstanceCount += 1
+				//Expect(k8sClient.Update(ctx, ytsaurus)).Should(Succeed())
+				updateAndCheck(ytsaurus, namespace)
+				EventuallyYtsaurus(ctx, name, reactionTimeout).Should(HaveClusterState(ytv1.ClusterStateUpdating))
+
+				By("Wait cluster update with full update complete")
+				EventuallyYtsaurus(ctx, name, upgradeTimeout).Should(HaveClusterState(ytv1.ClusterStateRunning))
+				// check all created recently (or after smth)
+			},
+		)
 
 		// This is a test for specific regression bug when master pods are recreated during PossibilityCheck stage.
 		It("Master shouldn't be recreated before WaitingForPodsCreation state if config changes", func(ctx context.Context) {
@@ -630,6 +709,15 @@ func deployAndCheck(ytsaurus *ytv1.Ytsaurus, namespace string) {
 	checkClusterViability(ytClient)
 }
 
+func updateAndCheck(ytsaurus *ytv1.Ytsaurus, namespace string) {
+	g := ytconfig.NewGenerator(ytsaurus, "local")
+	Expect(k8sClient.Update(ctx, ytsaurus)).Should(Succeed())
+
+	By("Creating ytsaurus client")
+	ytClient := getYtClient(g, namespace)
+	checkClusterBaseViability(ytClient)
+}
+
 func getSimpleUpdateScenario(namespace, newImage string) func(ctx context.Context) {
 	return func(ctx context.Context) {
 		By("Creating a Ytsaurus resource")