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
}

// Network add the container with aliasses to a specific network
func (builder *Builder) Network(networkName string, aliasses ...string) *Builder {
	builder.networks[networkName] = aliasses
	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
		}
	}

	// Network: portbinding Host->Container
	exposedPorts, portBindings, err := nat.ParsePortSpecs(builder.ports)
	if err != nil {
		return "", fmt.Errorf("unabel to parse ports: %v", err)
	}
	if len(portBindings) > 0 {
		time.Sleep(1 * time.Second)
		builder.containerConfig.ExposedPorts = exposedPorts
		builder.hostConfig.PortBindings = portBindings
	}

	// Network: Add container to container networks
	// Add the container to the first defined container network, if any one is
	// defined. If no one is defined, the docker API will add the container to
	// their default bridge docker0. The other networks will be added to the
	// container after the container start.
	var (
		networkNames = make([]string, 0)
		networks     = make([]types.NetworkResource, 0)
	)
	if len(builder.networks) > 0 {
		for networkName := range builder.networks {
			networkNames = append(networkNames, networkName)
		}
		var err error
		networks, err = builder.client.NetworkListByNames(ctx, networkNames...)
		if err != nil {
			return "", err
		}

		endpointSetting := &network.EndpointSettings{
			NetworkID: networks[0].ID,
			Aliases:   builder.networks[networkNames[0]],
		}

		builder.networkConfig.EndpointsConfig[networkNames[0]] = endpointSetting

		networkNames = networkNames[1:]
		networks = networks[1:]
	}

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

	err = builder.client.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{})
	if err != nil {

		shutdownErr := builder.client.ContainerStopByIDs(ctx, 1*time.Second, resp.ID)
		if shutdownErr != nil {
			return "", fmt.Errorf("Unable to start container: %v\nUnable to remove container %s: %v\nManual cleanup necessary", err, resp.ID, shutdownErr)
		}
		return "", fmt.Errorf("Unable to start container: %v", err)
	}

	// Network: Add more container networks
	for i, networkName := range networkNames {
		endpointSetting := &network.EndpointSettings{
			NetworkID: networks[i].ID,
			Aliases:   builder.networks[networkName],
		}
		err := builder.client.NetworkConnect(ctx, networks[i].ID, resp.ID, endpointSetting)
		if err != nil {
			return "", fmt.Errorf("Unable to append container endpoint to network %v", networkName)
		}
	}

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

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

		}
	}

	return resp.ID, err
}

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

// WithName set the name of the container
func (builder *Builder) WithName(containerName string) *Builder {
	builder.containerName = containerName
	return builder
}