PKGBUILD/pkg/testutils/dockerutils/builder.go

package dockerutils

import (
	"context"
	"fmt"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/go-connections/nat"
)

type Builder struct {
	client          *Client
	containerConfig *container.Config
	containerName   string
	hostConfig      *container.HostConfig
	networkConfig   *network.NetworkingConfig
	networks        map[string][]string
	ports           []string
	pull            bool
	waitForHealthy  bool
}

// AddEnv to the container
func (builder *Builder) AddEnv(key string, value string) *Builder {
	builder.containerConfig.Env = append(builder.containerConfig.Env, fmt.Sprintf("%v=%v", key, value))
	return builder
}

// AddLabel to the container
func (builder *Builder) AddLabel(key string, value string) *Builder {
	builder.containerConfig.Labels[key] = value
	return builder
}

// Env set environment variables to the container
func (builder *Builder) Env(env map[string]string) *Builder {
	builder.containerConfig.Labels = make(map[string]string)
	for key, value := range env {
		builder.containerConfig.Labels[key] = value
	}
	return builder
}

// Labels set labels to the container
func (builder *Builder) Labels(labels map[string]string) *Builder {
	builder.containerConfig.Labels = labels
	return builder
}

// Memory defines a memory limit for the container
func (builder *Builder) Memory(limit int64) *Builder {
	builder.hostConfig.Memory = limit
	return builder
}

// Mount a source volume or hostpath into the filesystem of the container
func (builder *Builder) Mount(source string, dest string) *Builder {
	builder.hostConfig.Binds = append(builder.hostConfig.Binds, fmt.Sprintf("%v:%v", source, dest))
	return builder
}

// Mounts a set of source volumes or hostpath into the filesystem of the
// container
func (builder *Builder) Mounts(mounts map[string]string) *Builder {
	for source, dest := range mounts {
		builder.Mount(source, dest)
	}
	return builder
}

// Port defines a port forwarding from the host machine to the container
// Examples:
// - 8080:8080
// - 10.6.231.10:8080:8080
// - 10.6.231.10:8080:8080/tcp
func (builder *Builder) Port(port string) *Builder {
	builder.ports = append(builder.ports, port)
	return builder
}

// Ports defines a set port forwarding rules from the host machine to the
// container
// Examples:
// - 8080:8080
// - 10.6.231.10:8080:8080
// - 10.6.231.10:8080:8080/tcp
func (builder *Builder) Ports(ports []string) *Builder {
	builder.ports = ports
	return builder
}

// Pull the image if absent
func (builder *Builder) Pull() *Builder {
	builder.pull = true
	return builder
}

// Start the container
func (builder *Builder) Start(ctx context.Context) (string, error) {

	// Pull container image if absent
	if builder.pull {
		err := builder.client.PullQuiet(ctx, builder.containerConfig.Image)
		if err != nil {
			return "", err
		}
	}

	// check if ports are available
	exposedPorts, portBindings, err := nat.ParsePortSpecs(builder.ports)
	if err != nil {
		return "", fmt.Errorf("Failed to parse ports %v: %v", builder.ports, err)
	}

	if len(portBindings) > 0 {
		builder.containerConfig.ExposedPorts = exposedPorts
		builder.hostConfig.PortBindings = portBindings
	}

	// add endpoint settings to existing networks
	networkExist := make(map[string]bool, 0)
	for networkName := range builder.networks {
		networkExist[networkName] = false
	}

	networks, err := builder.client.NetworkList(ctx, types.NetworkListOptions{})
	if err != nil {
		return "", fmt.Errorf("Failed to list networks: %v", err)
	}

	for _, nw := range networks {
		if aliases, present := builder.networks[nw.Name]; present {
			networkExist[nw.Name] = true
			builder.networkConfig.EndpointsConfig[nw.Name] = &network.EndpointSettings{
				Aliases:   aliases,
				NetworkID: nw.ID,
			}
		}
	}

	for nw, found := range networkExist {
		if !found {
			return "", fmt.Errorf("Failed to add endpoint settings for network %v. It does not exist", nw)
		}
	}

	// create container
	resp, err := builder.client.ContainerCreate(ctx, builder.containerConfig, builder.hostConfig, builder.networkConfig, builder.containerName)
	if err != nil {
		return "", fmt.Errorf("Failed to create container %v: %v", builder.containerName, err)
	}

	// start container
	err = builder.client.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{})
	if err != nil {
		stopError := builder.client.ContainerStopByIDs(ctx, time.Second, resp.ID)
		if stopError != nil {
			return "", fmt.Errorf("Failed to start container %v: %v\nUnable to remove container %v. Manual cleanup necessary", builder.containerName, err, builder.containerName)
		}
		return "", fmt.Errorf("Failed to start container %v: %v", builder.containerName, err)
	}

	// wait for healthy
	if builder.waitForHealthy {

		watcher := builder.client.GetWatcher()
		errorChannel := make(chan error, 1)
		doneChannel := make(chan struct{})
		err = watcher.AddListener(resp.ID, errorChannel, doneChannel)
		if err != nil {
			containerRemoveError := builder.client.ContainerRemove(ctx, resp.ID, types.ContainerRemoveOptions{Force: true})
			if containerRemoveError != nil {
				return "", fmt.Errorf("Failed while watching status for container %v: %v\nUnable to remove container %v: %v", resp.ID, err, resp.ID, containerRemoveError)
			}
			return "", fmt.Errorf("Failed while watching status for container %v: %v", resp.ID, err)
		}

		select {
		case err := <-errorChannel:
			if err != nil {
				containerRemoveError := builder.client.ContainerRemove(ctx, resp.ID, types.ContainerRemoveOptions{Force: true})
				if containerRemoveError != nil {
					return "", fmt.Errorf("Unable to remove container %v: %v", resp.ID, containerRemoveError)
				}
				return "", err
			}
		case <-doneChannel:
		}
	}

	return resp.ID, nil
}

func (builder *Builder) WaitForHealthy() *Builder {
	builder.waitForHealthy = true
	return builder
}