Adding Restart support

cmd/main.go

@@ -275,6 +275,7 @@ func initializeCLICommands() []cli.Command {
 	return []cli.Command{
 		*cmd.NewKillCLICommand(topContext),
 		*cmd.NewExecCLICommand(topContext),
+		*cmd.NewRestartCLICommand(topContext),
 		*cmd.NewStopCLICommand(topContext),
 		*cmd.NewPauseCLICommand(topContext),
 		*cmd.NewRemoveCLICommand(topContext),

pkg/chaos/docker/cmd/restart.go

@@ -0,0 +1,66 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/urfave/cli"
+
+	"github.com/alexei-led/pumba/pkg/chaos"
+	"github.com/alexei-led/pumba/pkg/chaos/docker"
+)
+
+type restartContext struct {
+	context context.Context
+}
+
+// NewRestartCLICommand initialize CLI restart command and bind it to the restartContext
+func NewRestartCLICommand(ctx context.Context) *cli.Command {
+	cmdContext := &restartContext{context: ctx}
+	return &cli.Command{
+		Name: "restart",
+		Flags: []cli.Flag{
+			cli.StringFlag{
+				Name:  "command, s",
+				Usage: "shell command, that will be sent by Pumba to the target container(s)",
+				Value: "kill 1",
+			},
+			cli.IntFlag{
+				Name:  "limit, l",
+				Usage: "limit number of container to restart (0: restart all matching)",
+				Value: 0,
+			},
+		},
+		Usage:       "restart specified containers",
+		ArgsUsage:   fmt.Sprintf("containers (name, list of names, or RE2 regex if prefixed with %q)", chaos.Re2Prefix),
+		Description: "send command to target container(s)",
+		Action:      cmdContext.restart,
+	}
+}
+
+// RESTART Command
+func (cmd *restartContext) restart(c *cli.Context) error {
+	// get random
+	random := c.GlobalBool("random")
+	// get labels
+	labels := c.GlobalStringSlice("label")
+	// get dry-run mode
+	dryRun := c.GlobalBool("dry-run")
+	// get skip error flag
+	skipError := c.GlobalBool("skip-error")
+	// get interval
+	interval := c.GlobalString("interval")
+	// get names or pattern
+	names, pattern := chaos.GetNamesOrPattern(c)
+	// get command
+	command := c.String("command")
+	// get limit for number of containers to restart
+	limit := c.Int("limit")
+	// init restart command
+	restartCommand, err := docker.NewRestartCommand(chaos.DockerClient, names, pattern, labels, command, limit, dryRun)
+	if err != nil {
+		return err
+	}
+	// run restart command
+	return chaos.RunChaosCommand(cmd.context, restartCommand, interval, random, skipError)
+}

pkg/chaos/docker/restart.go

@@ -0,0 +1,70 @@
+package docker
+
+import (
+	"context"
+
+	"github.com/alexei-led/pumba/pkg/chaos"
+	"github.com/alexei-led/pumba/pkg/container"
+	"github.com/pkg/errors"
+	log "github.com/sirupsen/logrus"
+)
+
+// RestartCommand `docker restart` command
+type RestartCommand struct {
+	client  container.Client
+	names   []string
+	pattern string
+	labels  []string
+	command string
+	limit   int
+	dryRun  bool
+}
+
+// NewRestartCommand create new Restart Command instance
+func NewRestartCommand(client container.Client, names []string, pattern string, labels []string, command string, limit int, dryRun bool) (chaos.Command, error) {
+	restart := &RestartCommand{client, names, pattern, labels, command, limit, dryRun}
+	if restart.command == "" {
+		restart.command = "kill 1"
+	}
+	return restart, nil
+}
+
+// Run restart command
+func (k *RestartCommand) Run(ctx context.Context, random bool) error {
+	log.Debug("restarting all matching containers")
+	log.WithFields(log.Fields{
+		"names":   k.names,
+		"pattern": k.pattern,
+		"labels":  k.labels,
+		"limit":   k.limit,
+		"random":  random,
+	}).Debug("listing matching containers")
+	containers, err := container.ListNContainers(ctx, k.client, k.names, k.pattern, k.labels, k.limit)
+	if err != nil {
+		return err
+	}
+	if len(containers) == 0 {
+		log.Warning("no containers to restart")
+		return nil
+	}
+
+	// select single random container from matching container and replace list with selected item
+	if random {
+		if c := container.RandomContainer(containers); c != nil {
+			containers = []*container.Container{c}
+		}
+	}
+
+	for _, container := range containers {
+		log.WithFields(log.Fields{
+			"container": container,
+			"command":   k.command,
+		}).Debug("restarting container")
+		c := container
+		err = k.client.RestartContainer(ctx, c, k.command, k.dryRun)
+		if err != nil {
+			return errors.Wrap(err, "failed to restart container")
+		}
+	}
+	return nil
+}

pkg/chaos/docker/restart_test.go

@@ -0,0 +1,223 @@
+package docker
+
+import (
+	"context"
+	"errors"
+	"reflect"
+	"testing"
+
+	"github.com/alexei-led/pumba/pkg/chaos"
+	"github.com/alexei-led/pumba/pkg/container"
+	"github.com/stretchr/testify/mock"
+)
+
+//nolint:funlen
+func TestRestartCommand_Run(t *testing.T) {
+	type wantErrors struct {
+		listError bool
+		restartError bool
+	}
+	type fields struct {
+		names   []string
+		pattern string
+		labels  []string
+		command string
+		limit   int
+		dryRun  bool
+	}
+	type args struct {
+		ctx    context.Context
+		random bool
+	}
+	tests := []struct {
+		name     string
+		fields   fields
+		args     args
+		expected []*container.Container
+		wantErr  bool
+		errs     wantErrors
+	}{
+		{
+			name: "restart matching containers by names",
+			fields: fields{
+				names:   []string{"c1", "c2", "c3"},
+				command: "kill 1",
+			},
+			args: args{
+				ctx: context.TODO(),
+			},
+			expected: container.CreateTestContainers(3),
+		},
+		{
+			name: "restart matching labeled containers by names",
+			fields: fields{
+				names:   []string{"c1", "c2", "c3"},
+				labels:  []string{"key=value"},
+				command: "kill 1",
+			},
+			args: args{
+				ctx: context.TODO(),
+			},
+			expected: container.CreateLabeledTestContainers(3, map[string]string{"key": "value"}),
+		},
+		{
+			name: "restart matching containers by filter with limit",
+			fields: fields{
+				pattern: "^c?",
+				command: "kill -STOP 1",
+				limit:   2,
+			},
+			args: args{
+				ctx: context.TODO(),
+			},
+			expected: container.CreateTestContainers(3),
+		},
+		{
+			name: "restart random matching container by names",
+			fields: fields{
+				names:   []string{"c1", "c2", "c3"},
+				command: "kill 1",
+			},
+			args: args{
+				ctx:    context.TODO(),
+				random: true,
+			},
+			expected: container.CreateTestContainers(3),
+		},
+		{
+			name: "no matching containers by names",
+			fields: fields{
+				names:   []string{"c1", "c2", "c3"},
+				command: "kill 1",
+			},
+			args: args{
+				ctx: context.TODO(),
+			},
+		},
+		{
+			name: "error listing containers",
+			fields: fields{
+				names:   []string{"c1", "c2", "c3"},
+				command: "kill 1",
+			},
+			args: args{
+				ctx: context.TODO(),
+			},
+			wantErr: true,
+			errs:    wantErrors{listError: true},
+		},
+		{
+			name: "error restarting container",
+			fields: fields{
+				names:   []string{"c1", "c2", "c3"},
+				command: "kill 1",
+			},
+			args: args{
+				ctx: context.TODO(),
+			},
+			expected: container.CreateTestContainers(3),
+			wantErr:  true,
+			errs:     wantErrors{restartError: true},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockClient := new(container.MockClient)
+			k := &RestartCommand{
+				client:  mockClient,
+				names:   tt.fields.names,
+				pattern: tt.fields.pattern,
+				labels:  tt.fields.labels,
+				command: tt.fields.command,
+				limit:   tt.fields.limit,
+				dryRun:  tt.fields.dryRun,
+			}
+			opts := container.ListOpts{Labels: tt.fields.labels}
+			call := mockClient.On("ListContainers", tt.args.ctx, mock.AnythingOfType("container.FilterFunc"), opts)
+			if tt.errs.listError {
+				call.Return(tt.expected, errors.New("ERROR"))
+				goto Invoke
+			} else {
+				call.Return(tt.expected, nil)
+				if tt.expected == nil {
+					goto Invoke
+				}
+			}
+			if tt.args.random {
+				mockClient.On("RestartContainer", tt.args.ctx, mock.AnythingOfType("*container.Container"), tt.fields.command, tt.fields.dryRun).Return(nil)
+			} else {
+				for i := range tt.expected {
+					if tt.fields.limit == 0 || i < tt.fields.limit {
+						call = mockClient.On("RestartContainer", tt.args.ctx, mock.AnythingOfType("*container.Container"), tt.fields.command, tt.fields.dryRun)
+						if tt.errs.restartError {
+							call.Return(errors.New("ERROR"))
+							goto Invoke
+						} else {
+							call.Return(nil)
+						}
+					}
+				}
+			}
+		Invoke:
+			if err := k.Run(tt.args.ctx, tt.args.random); (err != nil) != tt.wantErr {
+				t.Errorf("RestartCommand.Run() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			mockClient.AssertExpectations(t)
+		})
+	}
+}
+
+func TestNewRestartCommand(t *testing.T) {
+	type args struct {
+		client  container.Client
+		names   []string
+		pattern string
+		labels  []string
+		command string
+		limit   int
+		dryRun  bool
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    chaos.Command
+		wantErr bool
+	}{
+		{
+			name: "create new restart command",
+			args: args{
+				names:   []string{"c1", "c2"},
+				command: "kill -TERM 1",
+				limit:   10,
+			},
+			want: &RestartCommand{
+				names:   []string{"c1", "c2"},
+				command: "kill -TERM 1",
+				limit:   10,
+			},
+		},
+		{
+			name: "empty command",
+			args: args{
+				names:   []string{"c1", "c2"},
+				command: "",
+			},
+			want: &RestartCommand{
+				names:   []string{"c1", "c2"},
+				command: "kill 1",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := NewRestartCommand(tt.args.client, tt.args.names, tt.args.pattern, tt.args.labels, tt.args.command, tt.args.limit, tt.args.dryRun)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("NewRestartCommand() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("NewRestartCommand() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}