PKGBUILD/pkg/testutils/dockerutils/client.go

package dockerutils

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"strings"
	"sync"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"
)

type Client struct {
	*client.Client
	watcher *Watcher
	mutex   *sync.Mutex
}

// Close docker connection
func (client *Client) Close() error {
	client.mutex.Lock()
	defer client.mutex.Unlock()
	if client.watcher != nil {
		client.watcher.stop()
	}
	return client.Close()
}

// ContainerListByLabels returns only containers which match by given labels
func (client *Client) ContainerListByLabels(ctx context.Context, all bool, labels map[string]string) ([]types.Container, error) {
	filterArgs := filters.NewArgs()
	for key, value := range labels {
		filterArgs.Add("label", fmt.Sprintf("%v=%v", key, value))
	}
	containers, err := client.ContainerList(ctx, types.ContainerListOptions{
		All:     all,
		Filters: filterArgs,
	})
	if err != nil {
		return nil, err
	}
	if containers == nil {
		return nil, fmt.Errorf("No containers found by given labels")
	}
	return containers, nil
}

// ContainerListByNames returns only containers which match by given labels
func (client *Client) ContainerListByNames(ctx context.Context, all bool, names ...string) ([]types.Container, error) {
	filterArgs := filters.NewArgs()
	for _, name := range names {
		filterArgs.Add("name", name)
	}
	containers, err := client.ContainerList(ctx, types.ContainerListOptions{
		All:     all,
		Filters: filterArgs,
	})
	if err != nil {
		return nil, err
	}
	if containers == nil {
		return nil, fmt.Errorf("No containers found by given names")
	}
	return containers, nil
}

// ContainerRemoveByIDs deletes all containers which match by their container ids
func (client *Client) ContainerRemoveByIDs(ctx context.Context, containerIDs ...string) error {
	for _, containerID := range containerIDs {
		err := client.ContainerRemove(ctx, containerID, types.ContainerRemoveOptions{Force: true})
		if err != nil {
			return err
		}
	}
	return nil
}

// ContainerRemoveByLabels deletes all containers which match by given labels
func (client *Client) ContainerRemoveByLabels(ctx context.Context, labels map[string]string) error {
	containers, err := client.ContainerListByLabels(ctx, true, labels)
	if err != nil {
		return err
	}
	for _, container := range containers {
		err := client.ContainerRemove(ctx, container.ID, types.ContainerRemoveOptions{Force: true})
		if err != nil {
			return err
		}
	}
	return nil
}

// ContainerRemoveByNames deletes all containers which match by their names
func (client *Client) ContainerRemoveByNames(ctx context.Context, names ...string) error {
	containers, err := client.ContainerListByNames(ctx, true, names...)
	if err != nil {
		return err
	}
	for _, container := range containers {
		err := client.ContainerRemove(ctx, container.ID, types.ContainerRemoveOptions{Force: true})
		if err != nil {
			return err
		}
	}
	return nil
}

// ContainerStopByIDs deletes all containers which match by their container ids
func (client *Client) ContainerStopByIDs(ctx context.Context, timeout time.Duration, containerIDs ...string) error {
	for _, containerID := range containerIDs {
		err := client.ContainerStop(ctx, containerID, &timeout)
		if err != nil {
			return err
		}
	}
	return nil
}

// ContainerStopByLabels shutdown containters which match by given labels
func (client *Client) ContainerStopByLabels(ctx context.Context, timeout time.Duration, labels map[string]string) error {
	containers, err := client.ContainerListByLabels(ctx, true, labels)
	if err != nil {
		return err
	}
	for _, container := range containers {
		err := client.ContainerStop(ctx, container.ID, &timeout)
		if err != nil {
			return err
		}
	}
	return nil
}

// GetWatcher returns a watcher for container health states
func (client *Client) GetWatcher() *Watcher {
	if client.watcher != nil {
		return client.watcher
	}

	client.mutex.Lock()
	defer client.mutex.Unlock()

	client.watcher = &Watcher{
		client:        client,
		errorChannels: make(map[string]chan<- error),
		doneChannels:  make(map[string]chan<- struct{}),
		errorMapper:   make(map[string]ErrorMapper),
		mutex:         new(sync.RWMutex),
	}

	client.watcher.start()
	return client.watcher
}

// NewBuilder returns a new builder for containers
func (client *Client) NewBuilder(image string) *Builder {
	return &Builder{
		client: client,
		containerConfig: &container.Config{
			Image: image,
		},
		hostConfig:     new(container.HostConfig),
		networkConfig:  new(network.NetworkingConfig),
		networks:       make(map[string][]string, 0),
		ports:          make([]string, 0),
		pull:           false,
		waitForHealthy: false,
	}
}

// Pull image
func (client *Client) Pull(ctx context.Context, image string, w io.Writer) error {

	parts := strings.Split(image, "/")
	switch len(parts) {
	case 1:
		image = fmt.Sprintf("docker.io/library/%v", parts[0])
	case 2:
		if strings.Compare(parts[0], "library") == 0 ||
			strings.Compare(parts[0], "docker.io") == 0 {
			image = fmt.Sprintf("docker.io/library/%v", parts[1])
		}
	}

	readCloser, err := client.ImagePull(ctx, image, types.ImagePullOptions{})
	if err != nil {
		return err
	}

	_, err = io.Copy(w, readCloser)
	if err != nil {
		return err
	}

	return nil
}

// PullQuiet image
func (client *Client) PullQuiet(ctx context.Context, image string) error {
	return client.Pull(ctx, image, ioutil.Discard)
}

// New returns a new dockerutil client
func New() (*Client, error) {
	dockerClient, err := client.NewEnvClient()
	if err != nil {
		return nil, err
	}
	return &Client{
		dockerClient,
		nil,
		new(sync.Mutex),
	}, nil
}