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/api/types/volume"
	"github.com/docker/docker/client"
)

// Client from the docker API with additional functions
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, containerLabels map[string]string) ([]types.Container, error) {
	filterArgs := filters.NewArgs()
	for key, value := range containerLabels {
		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, containerNames ...string) ([]types.Container, error) {
	filterArgs := filters.NewArgs()
	for _, name := range containerNames {
		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, containerLabels map[string]string) error {
	containers, err := client.ContainerListByLabels(ctx, true, containerLabels)
	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, containerNames ...string) error {
	containers, err := client.ContainerListByNames(ctx, true, containerNames...)
	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, containerLabels map[string]string) error {
	containers, err := client.ContainerListByLabels(ctx, true, containerLabels)
	if err != nil {
		return err
	}
	for _, container := range containers {
		err := client.ContainerStop(ctx, container.ID, &timeout)
		if err != nil {
			return err
		}
	}
	return nil
}

// ContainerStopByNames shutdown containters matching by their names
func (client *Client) ContainerStopByNames(ctx context.Context, timeout time.Duration, containerNames ...string) error {
	containers, err := client.ContainerListByNames(ctx, true, containerNames...)
	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
}

// NetworkListByLabels returns networks which match by given labels
func (client *Client) NetworkListByLabels(ctx context.Context, networkLabels map[string]string) ([]types.NetworkResource, error) {
	args := filters.NewArgs()
	for key, value := range networkLabels {
		args.Add("label", fmt.Sprintf("%v=%v", key, value))
	}

	return client.NetworkList(ctx, types.NetworkListOptions{
		Filters: args,
	})
}

// NetworkListByNames returns networks which match by their names. If a
// network can not be found, the function returns an error
func (client *Client) NetworkListByNames(ctx context.Context, networkNames ...string) ([]types.NetworkResource, error) {
	networks, err := client.NetworkList(ctx, types.NetworkListOptions{})
	if err != nil {
		return nil, err
	}

	foundNetwork := make(map[string]bool, 0)
	for _, networkName := range networkNames {
		foundNetwork[networkName] = false
	}

	filteredNetworks := make([]types.NetworkResource, 0)
	for _, networkName := range networkNames {
		for _, network := range networks {
			if network.Name == networkName {
				filteredNetworks = append(filteredNetworks, network)
				foundNetwork[networkName] = true
			}
		}
	}

	for _, networkName := range networkNames {
		if !foundNetwork[networkName] {
			return nil, fmt.Errorf("Network %v not found", networkName)
		}
	}

	return filteredNetworks, nil
}

// NetworkRemoveByLabels remove all networks which match by given labels
func (client *Client) NetworkRemoveByLabels(ctx context.Context, containerLabels map[string]string) error {
	networks, err := client.NetworkListByLabels(ctx, containerLabels)
	if err != nil {
		return err
	}

	for _, network := range networks {
		err := client.NetworkRemove(ctx, network.ID)
		if err != nil {
			return err
		}
	}

	return nil
}

// NetworkRemoveByNames remove all networks match by their names. If a
// network can not be found, the function returns an error
func (client *Client) NetworkRemoveByNames(ctx context.Context, networkNames ...string) error {
	networks, err := client.NetworkListByNames(ctx, networkNames...)
	if err != nil {
		return err
	}

	for _, network := range networks {
		err := client.NetworkRemove(ctx, network.ID)
		if err != nil {
			return err
		}
	}

	return nil
}

// NetworkRemoveByIDs remove all networks match by their id
func (client *Client) NetworkRemoveByIDs(ctx context.Context, containerIDs ...string) error {
	for _, containerID := range containerIDs {
		err := client.NetworkRemove(ctx, containerID)
		if err != nil {
			return err
		}
	}
	return nil
}

// 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: &network.NetworkingConfig{
			EndpointsConfig: make(map[string]*network.EndpointSettings, 0),
		},
		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)
}

// VolumeListByLabels returns volumes which match by given labels
func (client *Client) VolumeListByLabels(ctx context.Context, volumeLabels map[string]string) (volume.VolumesListOKBody, error) {
	args := filters.NewArgs()
	for key, value := range volumeLabels {
		args.Add("label", fmt.Sprintf("%v=%v", key, value))
	}

	return client.VolumeList(ctx, args)
}

// VolumeListByNames returns volumes which match by their names. If a
// volume can not be found, the function returns an error
func (client *Client) VolumeListByNames(ctx context.Context, volumeNames ...string) (volume.VolumesListOKBody, error) {
	args := filters.NewArgs()
	foundVolumes := make(map[string]bool, 0)
	for _, volumeName := range volumeNames {
		foundVolumes[volumeName] = false
		args.Add("name", volumeName)
	}

	volumes, err := client.VolumeList(ctx, args)
	if err != nil {
		return volume.VolumesListOKBody{}, err
	}

	for _, volume := range volumes.Volumes {
		foundVolumes[volume.Name] = true
	}

	for _, volumeName := range volumeNames {
		if foundVolumes[volumeName] != true {
			return volume.VolumesListOKBody{}, fmt.Errorf("Volume %v not found", volumeName)
		}
	}
	return volumes, err
}

// VolumeRemoveByLabels remove all volumes match by their labels
func (client *Client) VolumeRemoveByLabels(ctx context.Context, volumeLabels map[string]string) error {
	volumes, err := client.VolumeListByLabels(ctx, volumeLabels)
	if err != nil {
		return err
	}

	for _, volume := range volumes.Volumes {
		err := client.VolumeRemove(ctx, volume.Name, true)
		if err != nil {
			return err
		}
	}

	return nil
}

// VolumeRemoveByNames remove all volumes match by their names. If a
// volume can not be found, the function returns an error
func (client *Client) VolumeRemoveByNames(ctx context.Context, volumeNames ...string) error {
	volumes, err := client.VolumeListByNames(ctx, volumeNames...)
	if err != nil {
		return err
	}

	for _, volume := range volumes.Volumes {
		err := client.VolumeRemove(ctx, volume.Name, true)
		if err != nil {
			return err
		}
	}

	return nil
}

// 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
}