Merge pull request #1638 from dbachrach/f-edit-set-replicas

Add support for kustomize edit set replicas

kustomize/internal/commands/edit/set/all.go

@@ -30,6 +30,7 @@ func NewCmdSet(fSys filesys.FileSystem, v ifc.Validator) *cobra.Command {
 		newCmdSetNameSuffix(fSys),
 		newCmdSetNamespace(fSys, v),
 		newCmdSetImage(fSys),
+		newCmdSetReplicas(fSys),
 	)
 	return c
 }

kustomize/internal/commands/edit/set/setimage_test.go

@@ -207,7 +207,7 @@ func TestSetImage(t *testing.T) {
 
 			// assert
 			if err != tc.expected.err {
-				t.Errorf("Unexpedted error from set image command. Actual: %v\nExpected: %v", err, tc.expected.err)
+				t.Errorf("Unexpected error from set image command. Actual: %v\nExpected: %v", err, tc.expected.err)
 				t.FailNow()
 			}
 

kustomize/internal/commands/edit/set/setreplicas.go

@@ -0,0 +1,132 @@
+// Copyright 2019 The Kubernetes Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package set
+
+import (
+	"errors"
+	"sort"
+	"strconv"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"sigs.k8s.io/kustomize/kustomize/v3/internal/commands/kustfile"
+	"sigs.k8s.io/kustomize/v3/filesys"
+	"sigs.k8s.io/kustomize/v3/pkg/types"
+)
+
+type setReplicasOptions struct {
+	replicasMap map[string]types.Replica
+}
+
+// errors
+
+var (
+	errReplicasNoArgs      = errors.New("no replicas specified")
+	errReplicasInvalidArgs = errors.New(`invalid format of replica, use the following format: <name>=<count>`)
+)
+
+const replicasSeparator = "="
+
+// newCmdSetReplicas sets the new replica count for a resource in the kustomization.
+func newCmdSetReplicas(fSys filesys.FileSystem) *cobra.Command {
+	var o setReplicasOptions
+
+	cmd := &cobra.Command{
+		Use:   "replicas",
+		Short: `Sets replicas count for resources in the kustomization file`,
+		Example: `
+The command
+  set replicas my-app=3 other-app=1
+will add
+
+replicas:
+- name: my-app
+  count: 3
+- name: other-app
+  count: 1
+
+to the kustomization file if it doesn't exist,
+and overwrite the previous ones if the replicas name exists.
+`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			err := o.Validate(args)
+			if err != nil {
+				return err
+			}
+			return o.RunSetReplicas(fSys)
+		},
+	}
+	return cmd
+}
+
+// Validate validates setImage command.
+func (o *setReplicasOptions) Validate(args []string) error {
+	if len(args) == 0 {
+		return errReplicasNoArgs
+	}
+
+	o.replicasMap = make(map[string]types.Replica)
+
+	for _, arg := range args {
+
+		replica, err := parseReplicasArg(arg)
+		if err != nil {
+			return err
+		}
+		o.replicasMap[replica.Name] = replica
+	}
+	return nil
+}
+
+// RunSetReplicas runs setReplicas command.
+func (o *setReplicasOptions) RunSetReplicas(fSys filesys.FileSystem) error {
+	mf, err := kustfile.NewKustomizationFile(fSys)
+	if err != nil {
+		return err
+	}
+	m, err := mf.Read()
+	if err != nil {
+		return err
+	}
+
+	// append only new replicas from kustomize file
+	for _, rep := range m.Replicas {
+		if _, ok := o.replicasMap[rep.Name]; ok {
+			continue
+		}
+
+		o.replicasMap[rep.Name] = rep
+	}
+
+	var replicas []types.Replica
+	for _, v := range o.replicasMap {
+		replicas = append(replicas, v)
+	}
+
+	sort.Slice(replicas, func(i, j int) bool {
+		return replicas[i].Name < replicas[j].Name
+	})
+
+	m.Replicas = replicas
+	return mf.Write(m)
+}
+
+func parseReplicasArg(arg string) (types.Replica, error) {
+
+	// matches a name and a replica count
+	// <name>=<count>
+	if s := strings.Split(arg, replicasSeparator); len(s) == 2 {
+		count, err := strconv.ParseInt(s[1], 10, 64)
+		if err != nil {
+			return types.Replica{}, errReplicasInvalidArgs
+		}
+
+		return types.Replica{
+			Name:  s[0],
+			Count: count,
+		}, nil
+	}
+
+	return types.Replica{}, errReplicasInvalidArgs
+}

kustomize/internal/commands/edit/set/setreplicas_test.go

@@ -0,0 +1,145 @@
+// Copyright 2019 The Kubernetes Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package set
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+
+	"sigs.k8s.io/kustomize/kustomize/v3/internal/commands/testutils"
+	"sigs.k8s.io/kustomize/v3/filesys"
+)
+
+func TestSetReplicas(t *testing.T) {
+	type given struct {
+		args         []string
+		infileReplicas []string
+	}
+	type expected struct {
+		fileOutput []string
+		err        error
+	}
+	testCases := []struct {
+		description string
+		given       given
+		expected    expected
+	}{
+		{
+			given: given{
+				args: []string{"app=5"},
+			},
+			expected: expected{
+				fileOutput: []string{
+					"replicas:",
+					"- count: 5",
+					"  name: app",
+				}},
+		},
+		{
+			description: "override file",
+			given: given{
+				args: []string{"app=5"},
+				infileReplicas: []string{
+					"replicas:",
+					"- count: 1",
+					"  name: app",
+					"- count: 2",
+					"  name: other-app",
+				},
+			},
+			expected: expected{
+				fileOutput: []string{
+					"replicas:",
+					"- count: 5",
+					"  name: app",
+					"- count: 2",
+					"  name: other-app",
+				}},
+		},
+		{
+			description: "multiple args with multiple overrides",
+			given: given{
+				args: []string{
+					"app=100",
+					"other-app=200",
+				},
+				infileReplicas: []string{
+					"replicas:",
+					"- count: 1",
+					"  name: app",
+					"- count: 2",
+					"  name: other-app",
+				},
+			},
+			expected: expected{
+				fileOutput: []string{
+					"replicas:",
+					"- count: 100",
+					"  name: app",
+					"- count: 200",
+					"  name: other-app",
+				}},
+		},
+		{
+			description: "error: no args",
+			expected: expected{
+				err: errReplicasNoArgs,
+			},
+		},
+		{
+			description: "error: invalid args -- no =",
+			given: given{
+				args: []string{"bad", "args"},
+			},
+			expected: expected{
+				err: errReplicasInvalidArgs,
+			},
+		},
+		{
+			description: "error: invalid args -- non-integer count",
+			given: given{
+				args: []string{"app=bad"},
+			},
+			expected: expected{
+				err: errReplicasInvalidArgs,
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("%s%v", tc.description, tc.given.args), func(t *testing.T) {
+			fSys := filesys.MakeFsInMemory()
+			cmd := newCmdSetReplicas(fSys)
+
+			if len(tc.given.infileReplicas) > 0 {
+				// write file with infileReplicas
+				testutils.WriteTestKustomizationWith(
+					fSys,
+					[]byte(strings.Join(tc.given.infileReplicas, "\n")))
+			} else {
+				testutils.WriteTestKustomization(fSys)
+			}
+
+			// act
+			err := cmd.RunE(cmd, tc.given.args)
+
+			// assert
+			if err != tc.expected.err {
+				t.Errorf("Unexpected error from set replicas command. Actual: %v\nExpected: %v", err, tc.expected.err)
+				t.FailNow()
+			}
+
+			content, err := testutils.ReadTestKustomization(fSys)
+			if err != nil {
+				t.Errorf("unexpected read error: %v", err)
+				t.FailNow()
+			}
+			expectedStr := strings.Join(tc.expected.fileOutput, "\n")
+			if !strings.Contains(string(content), expectedStr) {
+				t.Errorf("unexpected replicas in kustomization file. \nActual:\n%s\nExpected:\n%s", content, expectedStr)
+			}
+		})
+	}
+}