PKGBUILD/pkg/testutils/dockerutils/client.go

221 lines
5.6 KiB
Go
Raw Normal View History

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
}