Refactor ytspec test builders

api/v1/test_helpers.go

@@ -47,20 +47,6 @@ func CreateMinimalYtsaurusResource(namespace string) *Ytsaurus {
 					InstanceCount: 1,
 				},
 			},
-			Bootstrap: &BootstrapSpec{
-				TabletCellBundles: &BundlesBootstrapSpec{
-					Sys: &BundleBootstrapSpec{
-						TabletCellCount:        2,
-						ChangelogPrimaryMedium: ptr.String("default"),
-						SnapshotPrimaryMedium:  ptr.String("default"),
-					},
-					Default: &BundleBootstrapSpec{
-						TabletCellCount:        2,
-						ChangelogPrimaryMedium: ptr.String("default"),
-						SnapshotPrimaryMedium:  ptr.String("default"),
-					},
-				},
-			},
 			PrimaryMasters: MastersSpec{
 				MasterConnectionSpec: MasterConnectionSpec{
 					CellTag: 1,
@@ -112,38 +98,49 @@ func CreateMinimalYtsaurusResource(namespace string) *Ytsaurus {
 			HTTPProxies: []HTTPProxiesSpec{
 				createHTTPProxiesSpec(),
 			},
-			Schedulers: &SchedulersSpec{
-				InstanceSpec: InstanceSpec{
-					InstanceCount: 1,
-				},
-			},
-			ControllerAgents: &ControllerAgentsSpec{
-				InstanceSpec: InstanceSpec{
-					InstanceCount: 1,
-				},
-			},
 		},
 	}
 }
 
 func CreateBaseYtsaurusResource(namespace string) *Ytsaurus {
 	ytsaurus := CreateMinimalYtsaurusResource(namespace)
+	ytsaurus = WithBootstrap(ytsaurus)
+	ytsaurus = WithScheduler(ytsaurus)
+	ytsaurus = WithControllerAgents(ytsaurus)
+	ytsaurus = WithDataNodes(ytsaurus)
+	ytsaurus = WithTabletNodes(ytsaurus)
+	ytsaurus = WithExecNodes(ytsaurus)
+	return ytsaurus
+}
 
-	ytsaurus.Spec.Discovery = DiscoverySpec{
-		InstanceSpec: InstanceSpec{
-			InstanceCount: 1,
-		},
-	}
+// TODO (l0kix2): merge with ytconfig build spec helpers.
+func WithDataNodes(ytsaurus *Ytsaurus) *Ytsaurus {
+	return WithDataNodesCount(ytsaurus, 3)
+}
+
+func WithDataNodesCount(ytsaurus *Ytsaurus, count int) *Ytsaurus {
 	ytsaurus.Spec.DataNodes = []DataNodesSpec{
 		{
-			InstanceSpec: CreateDataNodeInstanceSpec(3),
+			InstanceSpec: CreateDataNodeInstanceSpec(count),
 		},
 	}
+	return ytsaurus
+}
+
+func WithTabletNodes(ytsaurus *Ytsaurus) *Ytsaurus {
+	return WithTabletNodesCount(ytsaurus, 3)
+}
+
+func WithTabletNodesCount(ytsaurus *Ytsaurus, count int) *Ytsaurus {
 	ytsaurus.Spec.TabletNodes = []TabletNodesSpec{
 		{
-			InstanceSpec: CreateTabletNodeSpec(3),
+			InstanceSpec: CreateTabletNodeSpec(count),
 		},
 	}
+	return ytsaurus
+}
+
+func WithExecNodes(ytsaurus *Ytsaurus) *Ytsaurus {
 	ytsaurus.Spec.ExecNodes = []ExecNodesSpec{
 		{
 			InstanceSpec: CreateExecNodeInstanceSpec(),
@@ -152,6 +149,51 @@ func CreateBaseYtsaurusResource(namespace string) *Ytsaurus {
 	return ytsaurus
 }
 
+func WithScheduler(ytsaurus *Ytsaurus) *Ytsaurus {
+	ytsaurus.Spec.Schedulers = &SchedulersSpec{
+		InstanceSpec: InstanceSpec{
+			InstanceCount: 1,
+		},
+	}
+	return ytsaurus
+}
+
+func WithControllerAgents(ytsaurus *Ytsaurus) *Ytsaurus {
+	ytsaurus.Spec.ControllerAgents = &ControllerAgentsSpec{
+		InstanceSpec: InstanceSpec{
+			InstanceCount: 1,
+		},
+	}
+	return ytsaurus
+}
+
+func WithBootstrap(ytsaurus *Ytsaurus) *Ytsaurus {
+	ytsaurus.Spec.Bootstrap = &BootstrapSpec{
+		TabletCellBundles: &BundlesBootstrapSpec{
+			Sys: &BundleBootstrapSpec{
+				TabletCellCount:        2,
+				ChangelogPrimaryMedium: ptr.String("default"),
+				SnapshotPrimaryMedium:  ptr.String("default"),
+			},
+			Default: &BundleBootstrapSpec{
+				TabletCellCount:        2,
+				ChangelogPrimaryMedium: ptr.String("default"),
+				SnapshotPrimaryMedium:  ptr.String("default"),
+			},
+		},
+	}
+	return ytsaurus
+}
+
+func WithQueryTracker(ytsaurus *Ytsaurus) *Ytsaurus {
+	ytsaurus.Spec.QueryTrackers = &QueryTrackerSpec{
+		InstanceSpec: InstanceSpec{
+			InstanceCount: 1,
+		},
+	}
+	return ytsaurus
+}
+
 func createHTTPProxiesSpec() HTTPProxiesSpec {
 	return HTTPProxiesSpec{
 		ServiceType: "NodePort",

test/e2e/ytsaurus_controller_test.go

@@ -33,7 +33,7 @@ const (
 	pollInterval     = time.Millisecond * 250
 	reactionTimeout  = time.Second * 150
 	bootstrapTimeout = time.Minute * 3
-	upgradeTimeout   = time.Minute * 5
+	upgradeTimeout   = time.Minute * 7
 )
 
 var getYtClient = getYtHTTPClient
@@ -54,6 +54,11 @@ func getHTTPProxyAddress(g *ytconfig.Generator, namespace string) string {
 	if proxy != "" {
 		return proxy
 	}
+	// This one is used in real code in YtsaurusClient.
+	proxy = os.Getenv("YTOP_PROXY")
+	if proxy != "" {
+		return proxy
+	}
 
 	if os.Getenv("E2E_YT_PROXY") != "" {
 		panic("E2E_YT_PROXY is deprecated, use E2E_YT_HTTP_PROXY")
@@ -306,7 +311,8 @@ var _ = Describe("Basic test for Ytsaurus controller", func() {
 
 			namespace := "test4"
 
-			ytsaurus := ytv1.CreateBaseYtsaurusResource(namespace)
+			ytsaurus := ytv1.CreateMinimalYtsaurusResource(namespace)
+			ytsaurus = ytv1.WithTabletNodes(ytsaurus)
 
 			g := ytconfig.NewGenerator(ytsaurus, "local")
 
@@ -354,7 +360,8 @@ var _ = Describe("Basic test for Ytsaurus controller", func() {
 
 			namespace := "test5"
 
-			ytsaurus := ytv1.CreateBaseYtsaurusResource(namespace)
+			ytsaurus := ytv1.CreateMinimalYtsaurusResource(namespace)
+			ytsaurus = ytv1.WithDataNodes(ytsaurus)
 			ytsaurus.Spec.TabletNodes = make([]ytv1.TabletNodesSpec, 0)
 
 			g := ytconfig.NewGenerator(ytsaurus, "local")
@@ -405,12 +412,7 @@ var _ = Describe("Basic test for Ytsaurus controller", func() {
 			namespace := "querytrackeraco"
 
 			ytsaurus := ytv1.CreateBaseYtsaurusResource(namespace)
-			ytsaurus.Spec.TabletNodes = []ytv1.TabletNodesSpec{
-				{
-					InstanceSpec: ytv1.InstanceSpec{InstanceCount: 1},
-				},
-			}
-			ytsaurus.Spec.QueryTrackers = &ytv1.QueryTrackerSpec{InstanceSpec: ytv1.InstanceSpec{InstanceCount: 1}}
+			ytsaurus = ytv1.WithQueryTracker(ytsaurus)
 
 			g := ytconfig.NewGenerator(ytsaurus, "local")
 
@@ -479,15 +481,9 @@ var _ = Describe("Basic test for Ytsaurus controller", func() {
 
 			namespace := "remoteexec"
 
-			ytsaurus := ytv1.CreateMinimalYtsaurusResource(namespace)
+			ytsaurus := ytv1.CreateBaseYtsaurusResource(namespace)
 			// Ensure that no local exec nodes exist, only remote ones (which will be created later).
 			ytsaurus.Spec.ExecNodes = []ytv1.ExecNodesSpec{}
-			ytsaurus.Spec.DataNodes = []ytv1.DataNodesSpec{{
-				InstanceSpec: ytv1.CreateDataNodeInstanceSpec(1),
-			}}
-			ytsaurus.Spec.TabletNodes = []ytv1.TabletNodesSpec{{
-				InstanceSpec: ytv1.CreateTabletNodeSpec(1),
-			}}
 			g := ytconfig.NewGenerator(ytsaurus, "local")
 
 			remoteYtsaurus := &ytv1.RemoteYtsaurus{
@@ -548,12 +544,13 @@ var _ = Describe("Basic test for Ytsaurus controller", func() {
 				JobState: &yt.JobCompleted,
 			})
 			Expect(err).ShouldNot(HaveOccurred())
-			Expect(statuses).Should(HaveLen(1))
-			status := statuses[0]
-			Expect(status.Address).Should(
-				ContainSubstring("end-"+ytv1.RemoteResourceName),
-				"actual status: %s", status,
-			)
+			Expect(statuses).Should(Not(BeEmpty()))
+			for _, status := range statuses {
+				Expect(status.Address).Should(
+					ContainSubstring("end-"+ytv1.RemoteResourceName),
+					"actual status: %s", status,
+				)
+			}
 		})
 	})