<commit_before>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package compose

import (
	"context"
	"fmt"
	"os"
	"path/filepath"

	"github.com/compose-spec/compose-go/types"
	"github.com/containerd/containerd/platforms"
	"github.com/docker/buildx/build"
	_ "github.com/docker/buildx/driver/docker" // required to get default driver registered
	"github.com/docker/buildx/util/buildflags"
	xprogress "github.com/docker/buildx/util/progress"
	"github.com/docker/cli/cli/command"
	bclient "github.com/moby/buildkit/client"
	"github.com/moby/buildkit/session"
	"github.com/moby/buildkit/session/auth/authprovider"
	specs "github.com/opencontainers/image-spec/specs-go/v1"

	"github.com/docker/compose/v2/pkg/api"
	"github.com/docker/compose/v2/pkg/progress"
	"github.com/docker/compose/v2/pkg/utils"
)

func (s *composeService) Build(ctx context.Context, project *types.Project, options api.BuildOptions) error {
	return progress.Run(ctx, func(ctx context.Context) error {
		return s.build(ctx, project, options)
	})
}

func (s *composeService) build(ctx context.Context, project *types.Project, options api.BuildOptions) error {
	opts := map[string]build.Options{}
	imagesToBuild := []string{}

	args := flatten(options.Args.Resolve(func(s string) (string, bool) {
		s, ok := project.Environment[s]
		return s, ok
	}))

	services, err := project.GetServices(options.Services...)
	if err != nil {
		return err
	}

	for _, service := range services {
		if service.Build != nil {
			imageName := getImageName(service, project.Name)
			imagesToBuild = append(imagesToBuild, imageName)
			buildOptions, err := s.toBuildOptions(project, service, imageName)
			if err != nil {
				return err
			}
			buildOptions.Pull = options.Pull
			buildOptions.BuildArgs = mergeArgs(buildOptions.BuildArgs, args)
			buildOptions.NoCache = options.NoCache
			buildOptions.CacheFrom, err = buildflags.ParseCacheEntry(service.Build.CacheFrom)
			if err != nil {
				return err
			}

			for _, image := range service.Build.CacheFrom {
				buildOptions.CacheFrom = append(buildOptions.CacheFrom, bclient.CacheOptionsEntry{
					Type:  "registry",
					Attrs: map[string]string{"ref": image},
				})
			}

			opts[imageName] = buildOptions
		}
	}

	_, err = s.doBuild(ctx, project, opts, options.Progress)
	if err == nil {
		if len(imagesToBuild) > 0 && !options.Quiet {
			utils.DisplayScanSuggestMsg()
		}
	}

	return err
}

func (s *composeService) ensureImagesExists(ctx context.Context, project *types.Project, quietPull bool) error {
	for _, service := range project.Services {
		if service.Image == "" && service.Build == nil {
			return fmt.Errorf("invalid service %q. Must specify either image or build", service.Name)
		}
	}

	images, err := s.getLocalImagesDigests(ctx, project)
	if err != nil {
		return err
	}

	err = s.pullRequiredImages(ctx, project, images, quietPull)
	if err != nil {
		return err
	}

	mode := xprogress.PrinterModeAuto
	if quietPull {
		mode = xprogress.PrinterModeQuiet
	}
	opts, err := s.getBuildOptions(project, images)
	if err != nil {
		return err
	}
	builtImages, err := s.doBuild(ctx, project, opts, mode)
	if err != nil {
		return err
	}

	if len(builtImages) > 0 {
		utils.DisplayScanSuggestMsg()
	}
	for name, digest := range builtImages {
		images[name] = digest
	}
	// set digest as com.docker.compose.image label so we can detect outdated containers
	for i, service := range project.Services {
		image := getImageName(service, project.Name)
		digest, ok := images[image]
		if ok {
			if project.Services[i].Labels == nil {
				project.Services[i].Labels = types.Labels{}
			}
			project.Services[i].Labels[api.ImageDigestLabel] = digest
			project.Services[i].Image = image
		}
	}
	return nil
}

func (s *composeService) getBuildOptions(project *types.Project, images map[string]string) (map[string]build.Options, error) {
	opts := map[string]build.Options{}
	for _, service := range project.Services {
		if service.Image == "" && service.Build == nil {
			return nil, fmt.Errorf("invalid service %q. Must specify either image or build", service.Name)
		}
		imageName := getImageName(service, project.Name)
		_, localImagePresent := images[imageName]

		if service.Build != nil {
			if localImagePresent && service.PullPolicy != types.PullPolicyBuild {
				continue
			}
			opt, err := s.toBuildOptions(project, service, imageName)
			if err != nil {
				return nil, err
			}
			opts[imageName] = opt
			continue
		}
	}
	return opts, nil

}

func (s *composeService) getLocalImagesDigests(ctx context.Context, project *types.Project) (map[string]string, error) {
	imageNames := []string{}
	for _, s := range project.Services {
		imgName := getImageName(s, project.Name)
		if !utils.StringContains(imageNames, imgName) {
			imageNames = append(imageNames, imgName)
		}
	}
	imgs, err := s.getImages(ctx, imageNames)
	if err != nil {
		return nil, err
	}
	images := map[string]string{}
	for name, info := range imgs {
		images[name] = info.ID
	}
	return images, nil
}

func (s *composeService) serverInfo(ctx context.Context) (command.ServerInfo, error) {
	ping, err := s.apiClient.Ping(ctx)
	if err != nil {
		return command.ServerInfo{}, err
	}
	serverInfo := command.ServerInfo{
		HasExperimental: ping.Experimental,
		OSType:          ping.OSType,
		BuildkitVersion: ping.BuilderVersion,
	}
	return serverInfo, err
}

func (s *composeService) doBuild(ctx context.Context, project *types.Project, opts map[string]build.Options, mode string) (map[string]string, error) {
	if len(opts) == 0 {
		return nil, nil
	}
	serverInfo, err := s.serverInfo(ctx)
	if err != nil {
		return nil, err
	}
	if buildkitEnabled, err := command.BuildKitEnabled(serverInfo); err != nil || !buildkitEnabled {
		return s.doBuildClassic(ctx, opts)
	}
	return s.doBuildBuildkit(ctx, project, opts, mode)
}

func (s *composeService) toBuildOptions(project *types.Project, service types.ServiceConfig, imageTag string) (build.Options, error) {
	var tags []string
	tags = append(tags, imageTag)

	buildArgs := flatten(service.Build.Args.Resolve(func(s string) (string, bool) {
		s, ok := project.Environment[s]
		return s, ok
	}))

	var plats []specs.Platform
	if platform, ok := project.Environment["DOCKER_DEFAULT_PLATFORM"]; ok {
		p, err := platforms.Parse(platform)
		if err != nil {
			return build.Options{}, err
		}
		plats = append(plats, p)
	}
	if service.Platform != "" {
		p, err := platforms.Parse(service.Platform)
		if err != nil {
			return build.Options{}, err
		}
		plats = append(plats, p)
	}

	return build.Options{
		Inputs: build.Inputs{
			ContextPath:    service.Build.Context,
			DockerfilePath: filepath.Join(service.Build.Context, service.Build.Dockerfile),
		},
		BuildArgs:   buildArgs,
		Tags:        tags,
		Target:      service.Build.Target,
		Exports:     []bclient.ExportEntry{{Type: "image", Attrs: map[string]string{}}},
		Platforms:   plats,
		Labels:      service.Build.Labels,
		NetworkMode: service.Build.Network,
		ExtraHosts:  service.Build.ExtraHosts,
		Session: []session.Attachable{
			authprovider.NewDockerAuthProvider(os.Stderr),
		},
	}, nil
}

func flatten(in types.MappingWithEquals) types.Mapping {
	if len(in) == 0 {
		return nil
	}
	out := types.Mapping{}
	for k, v := range in {
		if v == nil {
			continue
		}
		out[k] = *v
	}
	return out
}

func mergeArgs(m ...types.Mapping) types.Mapping {
	merged := types.Mapping{}
	for _, mapping := range m {
		for key, val := range mapping {
			merged[key] = val
		}
	}
	return merged
}
<commit_msg>use Dockerfile directly when path is absolute otherwise join it with Context path<commit_after>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package compose

import (
	"context"
	"fmt"
	"os"
	"path"
	"path/filepath"

	"github.com/compose-spec/compose-go/types"
	"github.com/containerd/containerd/platforms"
	"github.com/docker/buildx/build"
	_ "github.com/docker/buildx/driver/docker" // required to get default driver registered
	"github.com/docker/buildx/util/buildflags"
	xprogress "github.com/docker/buildx/util/progress"
	"github.com/docker/cli/cli/command"
	bclient "github.com/moby/buildkit/client"
	"github.com/moby/buildkit/session"
	"github.com/moby/buildkit/session/auth/authprovider"
	specs "github.com/opencontainers/image-spec/specs-go/v1"

	"github.com/docker/compose/v2/pkg/api"
	"github.com/docker/compose/v2/pkg/progress"
	"github.com/docker/compose/v2/pkg/utils"
)

func (s *composeService) Build(ctx context.Context, project *types.Project, options api.BuildOptions) error {
	return progress.Run(ctx, func(ctx context.Context) error {
		return s.build(ctx, project, options)
	})
}

func (s *composeService) build(ctx context.Context, project *types.Project, options api.BuildOptions) error {
	opts := map[string]build.Options{}
	imagesToBuild := []string{}

	args := flatten(options.Args.Resolve(func(s string) (string, bool) {
		s, ok := project.Environment[s]
		return s, ok
	}))

	services, err := project.GetServices(options.Services...)
	if err != nil {
		return err
	}

	for _, service := range services {
		if service.Build != nil {
			imageName := getImageName(service, project.Name)
			imagesToBuild = append(imagesToBuild, imageName)
			buildOptions, err := s.toBuildOptions(project, service, imageName)
			if err != nil {
				return err
			}
			buildOptions.Pull = options.Pull
			buildOptions.BuildArgs = mergeArgs(buildOptions.BuildArgs, args)
			buildOptions.NoCache = options.NoCache
			buildOptions.CacheFrom, err = buildflags.ParseCacheEntry(service.Build.CacheFrom)
			if err != nil {
				return err
			}

			for _, image := range service.Build.CacheFrom {
				buildOptions.CacheFrom = append(buildOptions.CacheFrom, bclient.CacheOptionsEntry{
					Type:  "registry",
					Attrs: map[string]string{"ref": image},
				})
			}

			opts[imageName] = buildOptions
		}
	}

	_, err = s.doBuild(ctx, project, opts, options.Progress)
	if err == nil {
		if len(imagesToBuild) > 0 && !options.Quiet {
			utils.DisplayScanSuggestMsg()
		}
	}

	return err
}

func (s *composeService) ensureImagesExists(ctx context.Context, project *types.Project, quietPull bool) error {
	for _, service := range project.Services {
		if service.Image == "" && service.Build == nil {
			return fmt.Errorf("invalid service %q. Must specify either image or build", service.Name)
		}
	}

	images, err := s.getLocalImagesDigests(ctx, project)
	if err != nil {
		return err
	}

	err = s.pullRequiredImages(ctx, project, images, quietPull)
	if err != nil {
		return err
	}

	mode := xprogress.PrinterModeAuto
	if quietPull {
		mode = xprogress.PrinterModeQuiet
	}
	opts, err := s.getBuildOptions(project, images)
	if err != nil {
		return err
	}
	builtImages, err := s.doBuild(ctx, project, opts, mode)
	if err != nil {
		return err
	}

	if len(builtImages) > 0 {
		utils.DisplayScanSuggestMsg()
	}
	for name, digest := range builtImages {
		images[name] = digest
	}
	// set digest as com.docker.compose.image label so we can detect outdated containers
	for i, service := range project.Services {
		image := getImageName(service, project.Name)
		digest, ok := images[image]
		if ok {
			if project.Services[i].Labels == nil {
				project.Services[i].Labels = types.Labels{}
			}
			project.Services[i].Labels[api.ImageDigestLabel] = digest
			project.Services[i].Image = image
		}
	}
	return nil
}

func (s *composeService) getBuildOptions(project *types.Project, images map[string]string) (map[string]build.Options, error) {
	opts := map[string]build.Options{}
	for _, service := range project.Services {
		if service.Image == "" && service.Build == nil {
			return nil, fmt.Errorf("invalid service %q. Must specify either image or build", service.Name)
		}
		imageName := getImageName(service, project.Name)
		_, localImagePresent := images[imageName]

		if service.Build != nil {
			if localImagePresent && service.PullPolicy != types.PullPolicyBuild {
				continue
			}
			opt, err := s.toBuildOptions(project, service, imageName)
			if err != nil {
				return nil, err
			}
			opts[imageName] = opt
			continue
		}
	}
	return opts, nil

}

func (s *composeService) getLocalImagesDigests(ctx context.Context, project *types.Project) (map[string]string, error) {
	imageNames := []string{}
	for _, s := range project.Services {
		imgName := getImageName(s, project.Name)
		if !utils.StringContains(imageNames, imgName) {
			imageNames = append(imageNames, imgName)
		}
	}
	imgs, err := s.getImages(ctx, imageNames)
	if err != nil {
		return nil, err
	}
	images := map[string]string{}
	for name, info := range imgs {
		images[name] = info.ID
	}
	return images, nil
}

func (s *composeService) serverInfo(ctx context.Context) (command.ServerInfo, error) {
	ping, err := s.apiClient.Ping(ctx)
	if err != nil {
		return command.ServerInfo{}, err
	}
	serverInfo := command.ServerInfo{
		HasExperimental: ping.Experimental,
		OSType:          ping.OSType,
		BuildkitVersion: ping.BuilderVersion,
	}
	return serverInfo, err
}

func (s *composeService) doBuild(ctx context.Context, project *types.Project, opts map[string]build.Options, mode string) (map[string]string, error) {
	if len(opts) == 0 {
		return nil, nil
	}
	serverInfo, err := s.serverInfo(ctx)
	if err != nil {
		return nil, err
	}
	if buildkitEnabled, err := command.BuildKitEnabled(serverInfo); err != nil || !buildkitEnabled {
		return s.doBuildClassic(ctx, opts)
	}
	return s.doBuildBuildkit(ctx, project, opts, mode)
}

func (s *composeService) toBuildOptions(project *types.Project, service types.ServiceConfig, imageTag string) (build.Options, error) {
	var tags []string
	tags = append(tags, imageTag)

	buildArgs := flatten(service.Build.Args.Resolve(func(s string) (string, bool) {
		s, ok := project.Environment[s]
		return s, ok
	}))

	var plats []specs.Platform
	if platform, ok := project.Environment["DOCKER_DEFAULT_PLATFORM"]; ok {
		p, err := platforms.Parse(platform)
		if err != nil {
			return build.Options{}, err
		}
		plats = append(plats, p)
	}
	if service.Platform != "" {
		p, err := platforms.Parse(service.Platform)
		if err != nil {
			return build.Options{}, err
		}
		plats = append(plats, p)
	}

	return build.Options{
		Inputs: build.Inputs{
			ContextPath:    service.Build.Context,
			DockerfilePath: dockerFilePath(service.Build.Context, service.Build.Dockerfile),
		},
		BuildArgs:   buildArgs,
		Tags:        tags,
		Target:      service.Build.Target,
		Exports:     []bclient.ExportEntry{{Type: "image", Attrs: map[string]string{}}},
		Platforms:   plats,
		Labels:      service.Build.Labels,
		NetworkMode: service.Build.Network,
		ExtraHosts:  service.Build.ExtraHosts,
		Session: []session.Attachable{
			authprovider.NewDockerAuthProvider(os.Stderr),
		},
	}, nil
}

func flatten(in types.MappingWithEquals) types.Mapping {
	if len(in) == 0 {
		return nil
	}
	out := types.Mapping{}
	for k, v := range in {
		if v == nil {
			continue
		}
		out[k] = *v
	}
	return out
}

func mergeArgs(m ...types.Mapping) types.Mapping {
	merged := types.Mapping{}
	for _, mapping := range m {
		for key, val := range mapping {
			merged[key] = val
		}
	}
	return merged
}

func dockerFilePath(context string, dockerfile string) string {
	if path.IsAbs(dockerfile) {
		return dockerfile
	}
	return filepath.Join(context, dockerfile)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package dns

import (
	"bytes"
	"encoding/json"
	skymsg "github.com/skynetservices/skydns/msg"
	"strings"
)

type TreeCache struct {
	ChildNodes map[string]*TreeCache
	Entries    map[string]interface{}
}

func NewTreeCache() *TreeCache {
	return &TreeCache{
		ChildNodes: make(map[string]*TreeCache),
		Entries:    make(map[string]interface{}),
	}
}

func (cache *TreeCache) Serialize() (string, error) {
	b, err := json.Marshal(cache)
	if err != nil {
		return "", err
	}

	var prettyJSON bytes.Buffer
	err = json.Indent(&prettyJSON, b, "", "\t")

	if err != nil {
		return "", err
	}
	return string(prettyJSON.Bytes()), nil
}

// setEntry creates the entire path if it doesn't already exist in the cache,
// then sets the given service record under the given key. The path this entry
// would have occupied in an etcd datastore is computed from the given fqdn and
// stored as the "Key" of the skydns service; this is only required because
// skydns expects the service record to contain a key in a specific format
// (presumably for legacy compatibility). Note that the fqnd string typically
// contains both the key and all elements in the path.
func (cache *TreeCache) setEntry(key string, val *skymsg.Service, fqdn string, path ...string) {
	// TODO: Consolidate setEntry and setSubCache into a single method with a
	// type switch.
	// TODO: Insted of passing the fqdn as an argument, we can reconstruct
	// it from the path, provided callers always pass the full path to the
	// object. This is currently *not* the case, since callers first create
	// a new, empty node, populate it, then parent it under the right path.
	// So we don't know the full key till the final parenting operation.
	node := cache.ensureChildNode(path...)

	// This key is used to construct the "target" for SRV record lookups.
	// For normal service/endpoint lookups, this will result in a key like:
	// /skydns/local/cluster/svc/svcNS/svcName/record-hash
	// but for headless services that govern pods requesting a specific
	// hostname (as used by petset), this will end up being:
	// /skydns/local/cluster/svc/svcNS/svcName/pod-hostname
	val.Key = skymsg.Path(fqdn)
	node.Entries[key] = val
}

func (cache *TreeCache) getSubCache(path ...string) *TreeCache {
	childCache := cache
	for _, subpath := range path {
		childCache = childCache.ChildNodes[subpath]
		if childCache == nil {
			return nil
		}
	}
	return childCache
}

// setSubCache inserts the given subtree under the given path:key. Usually the
// key is the name of a Kubernetes Service, and the path maps to the cluster
// subdomains matching the Service.
func (cache *TreeCache) setSubCache(key string, subCache *TreeCache, path ...string) {
	node := cache.ensureChildNode(path...)
	node.ChildNodes[key] = subCache
}

func (cache *TreeCache) getEntry(key string, path ...string) (interface{}, bool) {
	childNode := cache.getSubCache(path...)
	val, ok := childNode.Entries[key]
	return val, ok
}

func (cache *TreeCache) getValuesForPathWithWildcards(path ...string) []*skymsg.Service {
	retval := []*skymsg.Service{}
	nodesToExplore := []*TreeCache{cache}
	for idx, subpath := range path {
		nextNodesToExplore := []*TreeCache{}
		if idx == len(path)-1 {
			// if path ends on an entry, instead of a child node, add the entry
			for _, node := range nodesToExplore {
				if subpath == "*" {
					nextNodesToExplore = append(nextNodesToExplore, node)
				} else {
					if val, ok := node.Entries[subpath]; ok {
						retval = append(retval, val.(*skymsg.Service))
					} else {
						childNode := node.ChildNodes[subpath]
						if childNode != nil {
							nextNodesToExplore = append(nextNodesToExplore, childNode)
						}
					}
				}
			}
			nodesToExplore = nextNodesToExplore
			break
		}

		if subpath == "*" {
			for _, node := range nodesToExplore {
				for subkey, subnode := range node.ChildNodes {
					if !strings.HasPrefix(subkey, "_") {
						nextNodesToExplore = append(nextNodesToExplore, subnode)
					}
				}
			}
		} else {
			for _, node := range nodesToExplore {
				childNode := node.ChildNodes[subpath]
				if childNode != nil {
					nextNodesToExplore = append(nextNodesToExplore, childNode)
				}
			}
		}
		nodesToExplore = nextNodesToExplore
	}

	for _, node := range nodesToExplore {
		for _, val := range node.Entries {
			retval = append(retval, val.(*skymsg.Service))
		}
	}
	return retval
}

func (cache *TreeCache) deletePath(path ...string) bool {
	if len(path) == 0 {
		return false
	}
	if parentNode := cache.getSubCache(path[:len(path)-1]...); parentNode != nil {
		if _, ok := parentNode.ChildNodes[path[len(path)-1]]; ok {
			delete(parentNode.ChildNodes, path[len(path)-1])
			return true
		}
	}
	return false
}

func (cache *TreeCache) deleteEntry(key string, path ...string) bool {
	childNode := cache.getSubCache(path...)
	if childNode == nil {
		return false
	}
	if _, ok := childNode.Entries[key]; ok {
		delete(childNode.Entries, key)
		return true
	}
	return false
}

func (cache *TreeCache) appendValues(recursive bool, ref [][]interface{}) {
	for _, value := range cache.Entries {
		ref[0] = append(ref[0], value)
	}
	if recursive {
		for _, node := range cache.ChildNodes {
			node.appendValues(recursive, ref)
		}
	}
}

func (cache *TreeCache) ensureChildNode(path ...string) *TreeCache {
	childNode := cache
	for _, subpath := range path {
		newNode, ok := childNode.ChildNodes[subpath]
		if !ok {
			newNode = NewTreeCache()
			childNode.ChildNodes[subpath] = newNode
		}
		childNode = newNode
	}
	return childNode
}

// unused function. keeping it around in commented-fashion
// in the future, we might need some form of this function so that
// we can serialize to a file in a mounted empty dir..
//const (
//	dataFile = "data.dat"
//	crcFile  = "data.crc"
//)
//func (cache *TreeCache) Serialize(dir string) (string, error) {
//	cache.m.RLock()
//	defer cache.m.RUnlock()
//	b, err := json.Marshal(cache)
//	if err != nil {
//		return "", err
//	}
//
//	if err := ensureDir(dir, os.FileMode(0755)); err != nil {
//		return "", err
//	}
//	if err := ioutil.WriteFile(path.Join(dir, dataFile), b, 0644); err != nil {
//		return "", err
//	}
//	if err := ioutil.WriteFile(path.Join(dir, crcFile), getMD5(b), 0644); err != nil {
//		return "", err
//	}
//	return string(b), nil
//}

//func ensureDir(path string, perm os.FileMode) error {
//	s, err := os.Stat(path)
//	if err != nil || !s.IsDir() {
//		return os.Mkdir(path, perm)
//	}
//	return nil
//}

//func getMD5(b []byte) []byte {
//	h := md5.New()
//	h.Write(b)
//	return []byte(fmt.Sprintf("%x", h.Sum(nil)))
//}

// unused function. keeping it around in commented-fashion
// in the future, we might need some form of this function so that
// we can restart kube-dns, deserialize the tree and have a cache
// without having to wait for kube-dns to reach out to API server.
//func Deserialize(dir string) (*TreeCache, error) {
//	b, err := ioutil.ReadFile(path.Join(dir, dataFile))
//	if err != nil {
//		return nil, err
//	}
//
//	hash, err := ioutil.ReadFile(path.Join(dir, crcFile))
//	if err != nil {
//		return nil, err
//	}
//	if !reflect.DeepEqual(hash, getMD5(b)) {
//		return nil, fmt.Errorf("Checksum failed")
//	}
//
//	var cache TreeCache
//	err = json.Unmarshal(b, &cache)
//	if err != nil {
//		return nil, err
//	}
//	cache.m = &sync.RWMutex{}
//	return &cache, nil
//}
<commit_msg>Verify that the we get a non-nil subtree before consulting it.<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package dns

import (
	"bytes"
	"encoding/json"
	skymsg "github.com/skynetservices/skydns/msg"
	"strings"
)

type TreeCache struct {
	ChildNodes map[string]*TreeCache
	Entries    map[string]interface{}
}

func NewTreeCache() *TreeCache {
	return &TreeCache{
		ChildNodes: make(map[string]*TreeCache),
		Entries:    make(map[string]interface{}),
	}
}

func (cache *TreeCache) Serialize() (string, error) {
	b, err := json.Marshal(cache)
	if err != nil {
		return "", err
	}

	var prettyJSON bytes.Buffer
	err = json.Indent(&prettyJSON, b, "", "\t")

	if err != nil {
		return "", err
	}
	return string(prettyJSON.Bytes()), nil
}

// setEntry creates the entire path if it doesn't already exist in the cache,
// then sets the given service record under the given key. The path this entry
// would have occupied in an etcd datastore is computed from the given fqdn and
// stored as the "Key" of the skydns service; this is only required because
// skydns expects the service record to contain a key in a specific format
// (presumably for legacy compatibility). Note that the fqnd string typically
// contains both the key and all elements in the path.
func (cache *TreeCache) setEntry(key string, val *skymsg.Service, fqdn string, path ...string) {
	// TODO: Consolidate setEntry and setSubCache into a single method with a
	// type switch.
	// TODO: Insted of passing the fqdn as an argument, we can reconstruct
	// it from the path, provided callers always pass the full path to the
	// object. This is currently *not* the case, since callers first create
	// a new, empty node, populate it, then parent it under the right path.
	// So we don't know the full key till the final parenting operation.
	node := cache.ensureChildNode(path...)

	// This key is used to construct the "target" for SRV record lookups.
	// For normal service/endpoint lookups, this will result in a key like:
	// /skydns/local/cluster/svc/svcNS/svcName/record-hash
	// but for headless services that govern pods requesting a specific
	// hostname (as used by petset), this will end up being:
	// /skydns/local/cluster/svc/svcNS/svcName/pod-hostname
	val.Key = skymsg.Path(fqdn)
	node.Entries[key] = val
}

func (cache *TreeCache) getSubCache(path ...string) *TreeCache {
	childCache := cache
	for _, subpath := range path {
		childCache = childCache.ChildNodes[subpath]
		if childCache == nil {
			return nil
		}
	}
	return childCache
}

// setSubCache inserts the given subtree under the given path:key. Usually the
// key is the name of a Kubernetes Service, and the path maps to the cluster
// subdomains matching the Service.
func (cache *TreeCache) setSubCache(key string, subCache *TreeCache, path ...string) {
	node := cache.ensureChildNode(path...)
	node.ChildNodes[key] = subCache
}

func (cache *TreeCache) getEntry(key string, path ...string) (interface{}, bool) {
	childNode := cache.getSubCache(path...)
	if childNode == nil {
		return nil, false
	}
	val, ok := childNode.Entries[key]
	return val, ok
}

func (cache *TreeCache) getValuesForPathWithWildcards(path ...string) []*skymsg.Service {
	retval := []*skymsg.Service{}
	nodesToExplore := []*TreeCache{cache}
	for idx, subpath := range path {
		nextNodesToExplore := []*TreeCache{}
		if idx == len(path)-1 {
			// if path ends on an entry, instead of a child node, add the entry
			for _, node := range nodesToExplore {
				if subpath == "*" {
					nextNodesToExplore = append(nextNodesToExplore, node)
				} else {
					if val, ok := node.Entries[subpath]; ok {
						retval = append(retval, val.(*skymsg.Service))
					} else {
						childNode := node.ChildNodes[subpath]
						if childNode != nil {
							nextNodesToExplore = append(nextNodesToExplore, childNode)
						}
					}
				}
			}
			nodesToExplore = nextNodesToExplore
			break
		}

		if subpath == "*" {
			for _, node := range nodesToExplore {
				for subkey, subnode := range node.ChildNodes {
					if !strings.HasPrefix(subkey, "_") {
						nextNodesToExplore = append(nextNodesToExplore, subnode)
					}
				}
			}
		} else {
			for _, node := range nodesToExplore {
				childNode := node.ChildNodes[subpath]
				if childNode != nil {
					nextNodesToExplore = append(nextNodesToExplore, childNode)
				}
			}
		}
		nodesToExplore = nextNodesToExplore
	}

	for _, node := range nodesToExplore {
		for _, val := range node.Entries {
			retval = append(retval, val.(*skymsg.Service))
		}
	}
	return retval
}

func (cache *TreeCache) deletePath(path ...string) bool {
	if len(path) == 0 {
		return false
	}
	if parentNode := cache.getSubCache(path[:len(path)-1]...); parentNode != nil {
		if _, ok := parentNode.ChildNodes[path[len(path)-1]]; ok {
			delete(parentNode.ChildNodes, path[len(path)-1])
			return true
		}
	}
	return false
}

func (cache *TreeCache) deleteEntry(key string, path ...string) bool {
	childNode := cache.getSubCache(path...)
	if childNode == nil {
		return false
	}
	if _, ok := childNode.Entries[key]; ok {
		delete(childNode.Entries, key)
		return true
	}
	return false
}

func (cache *TreeCache) appendValues(recursive bool, ref [][]interface{}) {
	for _, value := range cache.Entries {
		ref[0] = append(ref[0], value)
	}
	if recursive {
		for _, node := range cache.ChildNodes {
			node.appendValues(recursive, ref)
		}
	}
}

func (cache *TreeCache) ensureChildNode(path ...string) *TreeCache {
	childNode := cache
	for _, subpath := range path {
		newNode, ok := childNode.ChildNodes[subpath]
		if !ok {
			newNode = NewTreeCache()
			childNode.ChildNodes[subpath] = newNode
		}
		childNode = newNode
	}
	return childNode
}

// unused function. keeping it around in commented-fashion
// in the future, we might need some form of this function so that
// we can serialize to a file in a mounted empty dir..
//const (
//	dataFile = "data.dat"
//	crcFile  = "data.crc"
//)
//func (cache *TreeCache) Serialize(dir string) (string, error) {
//	cache.m.RLock()
//	defer cache.m.RUnlock()
//	b, err := json.Marshal(cache)
//	if err != nil {
//		return "", err
//	}
//
//	if err := ensureDir(dir, os.FileMode(0755)); err != nil {
//		return "", err
//	}
//	if err := ioutil.WriteFile(path.Join(dir, dataFile), b, 0644); err != nil {
//		return "", err
//	}
//	if err := ioutil.WriteFile(path.Join(dir, crcFile), getMD5(b), 0644); err != nil {
//		return "", err
//	}
//	return string(b), nil
//}

//func ensureDir(path string, perm os.FileMode) error {
//	s, err := os.Stat(path)
//	if err != nil || !s.IsDir() {
//		return os.Mkdir(path, perm)
//	}
//	return nil
//}

//func getMD5(b []byte) []byte {
//	h := md5.New()
//	h.Write(b)
//	return []byte(fmt.Sprintf("%x", h.Sum(nil)))
//}

// unused function. keeping it around in commented-fashion
// in the future, we might need some form of this function so that
// we can restart kube-dns, deserialize the tree and have a cache
// without having to wait for kube-dns to reach out to API server.
//func Deserialize(dir string) (*TreeCache, error) {
//	b, err := ioutil.ReadFile(path.Join(dir, dataFile))
//	if err != nil {
//		return nil, err
//	}
//
//	hash, err := ioutil.ReadFile(path.Join(dir, crcFile))
//	if err != nil {
//		return nil, err
//	}
//	if !reflect.DeepEqual(hash, getMD5(b)) {
//		return nil, fmt.Errorf("Checksum failed")
//	}
//
//	var cache TreeCache
//	err = json.Unmarshal(b, &cache)
//	if err != nil {
//		return nil, err
//	}
//	cache.m = &sync.RWMutex{}
//	return &cache, nil
//}
<|endoftext|>


<commit_before>package docker

import (
	"bufio"
	"bytes"
	_docker "github.com/fsouza/go-dockerclient"
	"github.com/golang/glog"
	"io"
	"net"
	"path"
	"strings"
	"time"
)

type Docker struct {
	Endpoint string

	Cert string
	Key  string
	Ca   string

	docker *_docker.Client

	ContainerCreated func(*Container)
	ContainerStarted func(*Container)
}

type Port struct {
	ContainerPort int64  `json:"container_port"`
	HostPort      int64  `json:"host_port"`
	Type          string `json:"protocol"`
	AcceptIP      string `json:"accepts_ip"`
}

type Container struct {
	Id      string `json:"id"`
	Ip      string `json:"ip"`
	Image   string `json:"image"`
	ImageId string `json:"image_id"`

	Name    string `json:"name"`
	Command string `json:"command"`
	Ports   []Port `json:"ports"`
	Network _docker.NetworkSettings

	DockerData *_docker.Container `json:"docker_data"`

	docker *_docker.Client
}

type AuthIdentity struct {
	_docker.AuthConfiguration
}

type Image struct {
	Registry   string `json:"registry"`
	Repository string `json:"repository"`
	Tag        string `json:"tag"`
}

func (this Image) ImageString() string {
	s := this.Repository
	if this.Tag != "" {
		s = s + ":" + this.Tag
	}
	return s
}

func (this Image) Url() string {
	return path.Join(this.Registry, this.ImageString())
}

func ParseImageUrl(url string) Image {
	image := Image{}
	delim1 := strings.Index(url, "://")
	if delim1 < 0 {
		delim1 = 0
	} else {
		delim1 += 3
	}
	tag_index := strings.LastIndex(url[delim1:], ":")
	if tag_index > -1 {
		tag_index += delim1
		image.Tag = url[tag_index+1:]
	} else {
		tag_index = len(url)
	}
	project := path.Base(url[0:tag_index])
	account := path.Base(path.Dir(url[0:tag_index]))
	delim2 := strings.Index(url, account)
	image.Registry = url[0 : delim2-1]
	image.Repository = path.Join(account, project)
	return image
}

type ContainerControl struct {
	*_docker.Config

	// If false, the container starts up in daemon mode (as a service) - default
	RunOnce       bool                `json:"run_once,omitempty"`
	HostConfig    *_docker.HostConfig `json:"host_config"`
	ContainerName string              `json:"name,omitempty"`
}

// Endpoint and file paths
func NewTLSClient(endpoint string, cert, key, ca string) (c *Docker, err error) {
	c = &Docker{Endpoint: endpoint, Cert: cert, Ca: ca, Key: key}
	c.docker, err = _docker.NewTLSClient(endpoint, cert, key, ca)
	return c, err
}

func NewClient(endpoint string) (c *Docker, err error) {
	c = &Docker{Endpoint: endpoint}
	c.docker, err = _docker.NewClient(endpoint)
	return c, err
}

func (c *Docker) ListContainers() ([]*Container, error) {
	return c.FindContainers(nil)
}

func (c *Docker) FindContainersByName(name string) ([]*Container, error) {
	found := make([]*Container, 0)
	l, err := c.FindContainers(map[string][]string{
		"name": []string{name},
	})
	if err != nil {
		return nil, err
	}
	for _, cc := range l {
		err := cc.Inspect() // populates the Name, etc.
		glog.V(100).Infoln("Inspect container", *cc, "Err=", err)
		if err == nil && cc.Name == name {
			found = append(found, cc)
		}
	}
	return found, nil
}

func (c *Docker) FindContainers(filter map[string][]string) ([]*Container, error) {
	options := _docker.ListContainersOptions{
		All:  true,
		Size: true,
	}
	if filter != nil {
		options.Filters = filter
	}
	l, err := c.docker.ListContainers(options)
	if err != nil {
		return nil, err
	}
	out := []*Container{}
	for _, cc := range l {

		glog.V(100).Infoln("Matching", options, "Container==>", cc.Ports)
		c := &Container{
			Id:      cc.ID,
			Image:   cc.Image,
			Command: cc.Command,
			Ports:   get_ports(cc.Ports),
			docker:  c.docker,
		}
		c.Inspect()
		out = append(out, c)
	}
	return out, nil
}

func (c *Docker) PullImage(auth *AuthIdentity, image *Image) (<-chan error, error) {
	output_buff := bytes.NewBuffer(make([]byte, 1024*4))
	output := bufio.NewWriter(output_buff)

	err := c.docker.PullImage(_docker.PullImageOptions{
		Repository:   image.Repository,
		Registry:     image.Registry,
		Tag:          image.Tag,
		OutputStream: output,
	}, auth.AuthConfiguration)

	if err != nil {
		return nil, err
	}

	// Since the api doesn't have a channel, all we can do is read from the input
	// and then send a done signal when the input stream is exhausted.
	stopped := make(chan error)
	go func() {
		for {
			_, e := output_buff.ReadByte()
			if e == io.EOF {
				stopped <- nil
				return
			} else {
				stopped <- e
				return
			}
		}
	}()
	return stopped, err
}

func (c *Docker) StartContainer(auth *AuthIdentity, ct *ContainerControl) (*Container, error) {
	opts := _docker.CreateContainerOptions{
		Name:       ct.ContainerName,
		Config:     ct.Config,
		HostConfig: ct.HostConfig,
	}

	daemon := !ct.RunOnce
	// Detach mode (-d option in docker run)
	if daemon {
		opts.Config.AttachStdin = false
		opts.Config.AttachStdout = false
		opts.Config.AttachStderr = false
		opts.Config.StdinOnce = false
	}

	cc, err := c.docker.CreateContainer(opts)
	if err != nil {
		return nil, err
	}

	container := &Container{
		Id:     cc.ID,
		Image:  ct.Image,
		docker: c.docker,
	}

	if c.ContainerCreated != nil {
		c.ContainerCreated(container)
	}

	err = c.docker.StartContainer(cc.ID, ct.HostConfig)
	if err != nil {
		return nil, err
	}

	if c.ContainerStarted != nil {
		c.ContainerStarted(container)
	}

	err = container.Inspect()
	return container, err
}

func (c *Docker) StopContainer(auth *AuthIdentity, id string, timeout time.Duration) error {
	return c.docker.StopContainer(id, uint(timeout.Seconds()))
}

func (c *Docker) RemoveContainer(auth *AuthIdentity, id string, removeVolumes, force bool) error {
	return c.docker.RemoveContainer(_docker.RemoveContainerOptions{
		ID:            id,
		RemoveVolumes: removeVolumes,
		Force:         force,
	})
}

type Action int

const (
	Create Action = iota
	Start
	Stop
	Remove
	Die
)

// Docker event status are create -> start -> die -> stop for a container then destroy for docker -rm
var verbs map[string]Action = map[string]Action{
	"create":  Create,
	"start":   Start,
	"stop":    Stop,
	"destroy": Remove,
	"die":     Die,
}

func (c *Docker) WatchContainer(notify func(Action, *Container)) (chan<- bool, error) {
	return c.WatchContainerMatching(func(Action, *Container) bool { return true }, notify)
}

func (c *Docker) WatchContainerMatching(accept func(Action, *Container) bool, notify func(Action, *Container)) (chan<- bool, error) {
	stop := make(chan bool, 1)
	events := make(chan *_docker.APIEvents)
	err := c.docker.AddEventListener(events)
	if err != nil {
		return nil, err
	}
	go func() {
		for {
			select {
			case event := <-events:
				glog.V(100).Infoln("Docker event:", event)

				action, has := verbs[event.Status]
				if !has {
					continue
				}

				container := &Container{Id: event.ID, Image: event.From, docker: c.docker}
				if action != Remove {
					err := container.Inspect()
					if err != nil {
						glog.Warningln("Error inspecting container", event.ID)
						continue
					}
				}

				if watch != nil && accept(action, container) {
					notify(action, container)
				}

			case done := <-stop:
				if done {
					glog.Infoln("Watch terminated.")
					return
				}
			}
		}
	}()
	return stop, nil
}

func (c *Container) Inspect() error {
	cc, err := c.docker.InspectContainer(c.Id)
	if err != nil {
		return err
	}
	c.Name = cc.Name[1:] // there's this funny '/name' thing going on with how docker names containers
	c.ImageId = cc.Image
	c.Command = cc.Path + " " + strings.Join(cc.Args, " ")
	if cc.NetworkSettings != nil {
		c.Ip = cc.NetworkSettings.IPAddress
		c.Network = *cc.NetworkSettings
		c.Ports = get_ports(cc.NetworkSettings.PortMappingAPI())
	}
	c.DockerData = cc
	return nil
}

func get_ports(list []_docker.APIPort) []Port {
	out := make([]Port, len(list))
	for i, p := range list {
		out[i] = Port{
			ContainerPort: p.PrivatePort,
			HostPort:      p.PublicPort,
			Type:          p.Type,
			AcceptIP:      p.IP,
		}
	}
	return out
}

// Note that this depends on the context in which it is run.
// If this is run from the host (outside container), then it will return the address at eth0,
// but if it's run from inside a container, the eth0 interface is actually the docker0 interface
// on the host.
func GetEth0Ip() ([]string, error) {
	ips := []string{}
	intf, err := net.InterfaceByName("eth0")
	if err != nil {
		return ips, err
	}

	addrs, err := intf.Addrs()
	if err != nil {
		return ips, err
	}

	for _, a := range addrs {
		// parse the ip in CIDR form
		ip, _, err := net.ParseCIDR(a.String())
		if err != nil {
			return nil, err
		}
		ips = append(ips, ip.String())
	}
	return ips, nil
}
<commit_msg>Add RemoveImage method<commit_after>package docker

import (
	"bufio"
	"bytes"
	_docker "github.com/fsouza/go-dockerclient"
	"github.com/golang/glog"
	"io"
	"net"
	"path"
	"strings"
	"time"
)

type Docker struct {
	Endpoint string

	Cert string
	Key  string
	Ca   string

	docker *_docker.Client

	ContainerCreated func(*Container)
	ContainerStarted func(*Container)
}

type Port struct {
	ContainerPort int64  `json:"container_port"`
	HostPort      int64  `json:"host_port"`
	Type          string `json:"protocol"`
	AcceptIP      string `json:"accepts_ip"`
}

type Container struct {
	Id      string `json:"id"`
	Ip      string `json:"ip"`
	Image   string `json:"image"`
	ImageId string `json:"image_id"`

	Name    string `json:"name"`
	Command string `json:"command"`
	Ports   []Port `json:"ports"`
	Network _docker.NetworkSettings

	DockerData *_docker.Container `json:"docker_data"`

	docker *_docker.Client
}

type AuthIdentity struct {
	_docker.AuthConfiguration
}

type Image struct {
	Registry   string `json:"registry"`
	Repository string `json:"repository"`
	Tag        string `json:"tag"`
}

func (this Image) ImageString() string {
	s := this.Repository
	if this.Tag != "" {
		s = s + ":" + this.Tag
	}
	return s
}

func (this Image) Url() string {
	return path.Join(this.Registry, this.ImageString())
}

func ParseImageUrl(url string) Image {
	image := Image{}
	delim1 := strings.Index(url, "://")
	if delim1 < 0 {
		delim1 = 0
	} else {
		delim1 += 3
	}
	tag_index := strings.LastIndex(url[delim1:], ":")
	if tag_index > -1 {
		tag_index += delim1
		image.Tag = url[tag_index+1:]
	} else {
		tag_index = len(url)
	}
	project := path.Base(url[0:tag_index])
	account := path.Base(path.Dir(url[0:tag_index]))
	delim2 := strings.Index(url, account)
	image.Registry = url[0 : delim2-1]
	image.Repository = path.Join(account, project)
	return image
}

type ContainerControl struct {
	*_docker.Config

	// If false, the container starts up in daemon mode (as a service) - default
	RunOnce       bool                `json:"run_once,omitempty"`
	HostConfig    *_docker.HostConfig `json:"host_config"`
	ContainerName string              `json:"name,omitempty"`
}

// Endpoint and file paths
func NewTLSClient(endpoint string, cert, key, ca string) (c *Docker, err error) {
	c = &Docker{Endpoint: endpoint, Cert: cert, Ca: ca, Key: key}
	c.docker, err = _docker.NewTLSClient(endpoint, cert, key, ca)
	return c, err
}

func NewClient(endpoint string) (c *Docker, err error) {
	c = &Docker{Endpoint: endpoint}
	c.docker, err = _docker.NewClient(endpoint)
	return c, err
}

func (c *Docker) ListContainers() ([]*Container, error) {
	return c.FindContainers(nil)
}

func (c *Docker) FindContainersByName(name string) ([]*Container, error) {
	found := make([]*Container, 0)
	l, err := c.FindContainers(map[string][]string{
		"name": []string{name},
	})
	if err != nil {
		return nil, err
	}
	for _, cc := range l {
		err := cc.Inspect() // populates the Name, etc.
		glog.V(100).Infoln("Inspect container", *cc, "Err=", err)
		if err == nil && cc.Name == name {
			found = append(found, cc)
		}
	}
	return found, nil
}

func (c *Docker) FindContainers(filter map[string][]string) ([]*Container, error) {
	options := _docker.ListContainersOptions{
		All:  true,
		Size: true,
	}
	if filter != nil {
		options.Filters = filter
	}
	l, err := c.docker.ListContainers(options)
	if err != nil {
		return nil, err
	}
	out := []*Container{}
	for _, cc := range l {

		glog.V(100).Infoln("Matching", options, "Container==>", cc.Ports)
		c := &Container{
			Id:      cc.ID,
			Image:   cc.Image,
			Command: cc.Command,
			Ports:   get_ports(cc.Ports),
			docker:  c.docker,
		}
		c.Inspect()
		out = append(out, c)
	}
	return out, nil
}

func (c *Docker) PullImage(auth *AuthIdentity, image *Image) (<-chan error, error) {
	output_buff := bytes.NewBuffer(make([]byte, 1024*4))
	output := bufio.NewWriter(output_buff)

	err := c.docker.PullImage(_docker.PullImageOptions{
		Repository:   image.Repository,
		Registry:     image.Registry,
		Tag:          image.Tag,
		OutputStream: output,
	}, auth.AuthConfiguration)

	if err != nil {
		return nil, err
	}

	// Since the api doesn't have a channel, all we can do is read from the input
	// and then send a done signal when the input stream is exhausted.
	stopped := make(chan error)
	go func() {
		for {
			_, e := output_buff.ReadByte()
			if e == io.EOF {
				stopped <- nil
				return
			} else {
				stopped <- e
				return
			}
		}
	}()
	return stopped, err
}

func (c *Docker) StartContainer(auth *AuthIdentity, ct *ContainerControl) (*Container, error) {
	opts := _docker.CreateContainerOptions{
		Name:       ct.ContainerName,
		Config:     ct.Config,
		HostConfig: ct.HostConfig,
	}

	daemon := !ct.RunOnce
	// Detach mode (-d option in docker run)
	if daemon {
		opts.Config.AttachStdin = false
		opts.Config.AttachStdout = false
		opts.Config.AttachStderr = false
		opts.Config.StdinOnce = false
	}

	cc, err := c.docker.CreateContainer(opts)
	if err != nil {
		return nil, err
	}

	container := &Container{
		Id:     cc.ID,
		Image:  ct.Image,
		docker: c.docker,
	}

	if c.ContainerCreated != nil {
		c.ContainerCreated(container)
	}

	err = c.docker.StartContainer(cc.ID, ct.HostConfig)
	if err != nil {
		return nil, err
	}

	if c.ContainerStarted != nil {
		c.ContainerStarted(container)
	}

	err = container.Inspect()
	return container, err
}

func (c *Docker) StopContainer(auth *AuthIdentity, id string, timeout time.Duration) error {
	return c.docker.StopContainer(id, uint(timeout.Seconds()))
}

func (c *Docker) RemoveContainer(auth *AuthIdentity, id string, removeVolumes, force bool) error {
	return c.docker.RemoveContainer(_docker.RemoveContainerOptions{
		ID:            id,
		RemoveVolumes: removeVolumes,
		Force:         force,
	})
}

func (c *Docker) RemoveImage(image string, force, prune bool) error {
	return c.docker.RemoveImageExtended(image, _docker.RemoveImageOptions{
		Force:   force,
		NoPrune: !prune,
	})
}

type Action int

const (
	Create Action = iota
	Start
	Stop
	Remove
	Die
)

// Docker event status are create -> start -> die -> stop for a container then destroy for docker -rm
var verbs map[string]Action = map[string]Action{
	"create":  Create,
	"start":   Start,
	"stop":    Stop,
	"destroy": Remove,
	"die":     Die,
}

func (c *Docker) WatchContainer(notify func(Action, *Container)) (chan<- bool, error) {
	return c.WatchContainerMatching(func(Action, *Container) bool { return true }, notify)
}

func (c *Docker) WatchContainerMatching(accept func(Action, *Container) bool, notify func(Action, *Container)) (chan<- bool, error) {
	stop := make(chan bool, 1)
	events := make(chan *_docker.APIEvents)
	err := c.docker.AddEventListener(events)
	if err != nil {
		return nil, err
	}
	go func() {
		for {
			select {
			case event := <-events:
				glog.V(100).Infoln("Docker event:", event)

				action, has := verbs[event.Status]
				if !has {
					continue
				}

				container := &Container{Id: event.ID, Image: event.From, docker: c.docker}
				if action != Remove {
					err := container.Inspect()
					if err != nil {
						glog.Warningln("Error inspecting container", event.ID)
						continue
					}
				}

				if watch != nil && accept(action, container) {
					notify(action, container)
				}

			case done := <-stop:
				if done {
					glog.Infoln("Watch terminated.")
					return
				}
			}
		}
	}()
	return stop, nil
}

func (c *Container) Inspect() error {
	cc, err := c.docker.InspectContainer(c.Id)
	if err != nil {
		return err
	}
	c.Name = cc.Name[1:] // there's this funny '/name' thing going on with how docker names containers
	c.ImageId = cc.Image
	c.Command = cc.Path + " " + strings.Join(cc.Args, " ")
	if cc.NetworkSettings != nil {
		c.Ip = cc.NetworkSettings.IPAddress
		c.Network = *cc.NetworkSettings
		c.Ports = get_ports(cc.NetworkSettings.PortMappingAPI())
	}
	c.DockerData = cc
	return nil
}

func get_ports(list []_docker.APIPort) []Port {
	out := make([]Port, len(list))
	for i, p := range list {
		out[i] = Port{
			ContainerPort: p.PrivatePort,
			HostPort:      p.PublicPort,
			Type:          p.Type,
			AcceptIP:      p.IP,
		}
	}
	return out
}

// Note that this depends on the context in which it is run.
// If this is run from the host (outside container), then it will return the address at eth0,
// but if it's run from inside a container, the eth0 interface is actually the docker0 interface
// on the host.
func GetEth0Ip() ([]string, error) {
	ips := []string{}
	intf, err := net.InterfaceByName("eth0")
	if err != nil {
		return ips, err
	}

	addrs, err := intf.Addrs()
	if err != nil {
		return ips, err
	}

	for _, a := range addrs {
		// parse the ip in CIDR form
		ip, _, err := net.ParseCIDR(a.String())
		if err != nil {
			return nil, err
		}
		ips = append(ips, ip.String())
	}
	return ips, nil
}
<|endoftext|>


<commit_before>package pingdom

import (
	"encoding/json"
	"fmt"
)

// PingdomResponse represents a general response from the Pingdom API
type PingdomResponse struct {
	Message string `json:"message"`
}

// PingdomError represents an error response from the Pingdom API
type PingdomError struct {
	StatusCode int    `json:"statuscode"`
	StatusDesc string `json:"statusdesc"`
	Message    string `json:"errormessage"`
}

// CheckResponse represents the json response for a check from the Pingdom API
type CheckResponse struct {
	ID                       int                `json:"id"`
	Name                     string             `json:"name"`
	Resolution               int                `json:"resolution,omitempty"`
	SendNotificationWhenDown int                `json:"sendnotificationwhendown,omitempty"`
	NotifyAgainEvery         int                `json:"notifyagainevery,omitempty"`
	NotifyWhenBackup         bool               `json:"notifywhenbackup,omitempty"`
	Created                  int64              `json:"created,omitempty"`
	Hostname                 string             `json:"hostname,omitempty"`
	Status                   string             `json:"status,omitempty"`
	LastErrorTime            int64              `json:"lasterrortime,omitempty"`
	LastTestTime             int64              `json:"lasttesttime,omitempty"`
	LastResponseTime         int64              `json:"lastresponsetime,omitempty"`
	Paused                   bool               `json:"paused,omitempty"`
	IntegrationIds           []int              `json:"integrationids,omitempty"`
	Type                     CheckResponseType  `json:"type,omitempty"`
	Tags                     []CheckResponseTag `json:"tags,omitempty"`
	UserIds                  []int              `json:"userids,omitempty"`
	TeamIds                  []int              `json:"teamids,omitempty"`
}

type CheckResponseType struct {
	Name string                    `json:"-"`
	HTTP *CheckResponseHTTPDetails `json:"http,omitempty"`
	TCP  *CheckResponseTCPDetails  `json:"tcp,omitempty"`
}

type CheckResponseTag struct {
	Name  string      `json:"name"`
	Type  string      `json:"type"`
	Count interface{} `json:"count"`
}

// MaintenanceResponse represents the json response for a maintenance from the Pingdom API
type MaintenanceResponse struct {
	ID             int                      `json:"id"`
	Description    string                   `json:"description"`
	From           int64                    `json:"from"`
	To             int64                    `json:"to"`
	RecurrenceType string                   `json:"recurrencetype"`
	RepeatEvery    int                      `json:"repeatevery"`
	EffectiveTo    int64                    `json:"effectiveto"`
	Checks         MaintenanceCheckResponse `json:"checks"`
}

// MaintenanceCheckResponse represents Check reply in json MaintenanceResponse
type MaintenanceCheckResponse struct {
	Uptime []int `json:"uptime"`
	Tms    []int `json:"tms"`
}

// ProbeResponse represents the json response for probes from the PIngdom API
type ProbeResponse struct {
	ID         int    `json:"id"`
	Country    string `json:"country"`
	City       string `json:"city"`
	Name       string `json:"name"`
	Active     bool   `json:"active"`
	Hostname   string `json:"hostname"`
	IP         string `json:"ip"`
	IPv6       string `json:"ipv6"`
	CountryISO string `json:"countryiso"`
	Region     string `json:"region"`
}

// TeamResponse represents the json response for teams from the PIngdom API
type TeamResponse struct {
	ID    string `json:"id"`
	Name  string `json:"name"`
	Users []TeamUserResponse
}

// TeamUserResponse represents the json response for users in teams from the PIngdom API
type TeamUserResponse struct {
	ID    string `json:"id"`
	Email string `json:"email"`
	Name  string `json:"name"`
}

// TeamDeleteResponse represents the json response for delete team from the PIngdom API
type TeamDeleteResponse struct {
	Success bool `json:"success"`
}

func (c *CheckResponseType) UnmarshalJSON(b []byte) error {
	var raw interface{}

	err := json.Unmarshal(b, &raw)
	if err != nil {
		return err
	}

	switch v := raw.(type) {
	case string:
		c.Name = v
	case map[string]interface{}:
		if len(v) != 1 {
			return fmt.Errorf("Check detailed response `check.type` contains more than one object: %+v", v)
		}
		for k := range v {
			c.Name = k
		}

		// Allow continue use json.Unmarshall using a type != Unmarshaller
		// This avoid enter in a infinite loop
		type t CheckResponseType
		var rawCheckDetails t

		err := json.Unmarshal(b, &rawCheckDetails)
		if err != nil {
			return err
		}
		c.HTTP = rawCheckDetails.HTTP
		c.TCP = rawCheckDetails.TCP
	}
	return nil
}

// HttpCheck represents a Pingdom http check.
type CheckResponseHTTPDetails struct {
	Url              string            `json:"url,omitempty"`
	Encryption       bool              `json:"encryption,omitempty"`
	Port             int               `json:"port,omitempty"`
	Username         string            `json:"username,omitempty"`
	Password         string            `json:"password,omitempty"`
	ShouldContain    string            `json:"shouldcontain,omitempty"`
	ShouldNotContain string            `json:"shouldnotcontain,omitempty"`
	PostData         string            `json:"postdata,omitempty"`
	RequestHeaders   map[string]string `json:"requestheaders,omitempty"`
}

// HttpCheck represents a Pingdom http check.
type CheckResponseTCPDetails struct {
	Port           int    `json:"port,omitempty"`
	StringToSend   string `json:"stringtosend,omitempty"`
	StringToExpect string `json:"stringtoexpect,omitempty"`
}

// Return string representation of the PingdomError
func (r *PingdomError) Error() string {
	return fmt.Sprintf("%d %v: %v", r.StatusCode, r.StatusDesc, r.Message)
}

// private types used to unmarshall json responses from pingdom

type listChecksJsonResponse struct {
	Checks []CheckResponse `json:"checks"`
}

type listMaintenanceJsonResponse struct {
	Maintenances []MaintenanceResponse `json:"maintenance"`
}

type listProbesJsonResponse struct {
	Probes []ProbeResponse `json:"probes"`
}

type listTeamsJsonResponse struct {
	Teams []TeamResponse `json:"teams"`
}

type checkDetailsJsonResponse struct {
	Check *CheckResponse `json:"check"`
}

type maintenanceDetailsJsonResponse struct {
	Maintenance *MaintenanceResponse `json:"maintenance"`
}

type teamDetailsJsonResponse struct {
	Team *TeamResponse `json:"team"`
}

type errorJsonResponse struct {
	Error *PingdomError `json:"error"`
}
<commit_msg>Updated comments<commit_after>package pingdom

import (
	"encoding/json"
	"fmt"
)

// PingdomResponse represents a general response from the Pingdom API
type PingdomResponse struct {
	Message string `json:"message"`
}

// PingdomError represents an error response from the Pingdom API
type PingdomError struct {
	StatusCode int    `json:"statuscode"`
	StatusDesc string `json:"statusdesc"`
	Message    string `json:"errormessage"`
}

// CheckResponse represents the json response for a check from the Pingdom API
type CheckResponse struct {
	ID                       int                `json:"id"`
	Name                     string             `json:"name"`
	Resolution               int                `json:"resolution,omitempty"`
	SendNotificationWhenDown int                `json:"sendnotificationwhendown,omitempty"`
	NotifyAgainEvery         int                `json:"notifyagainevery,omitempty"`
	NotifyWhenBackup         bool               `json:"notifywhenbackup,omitempty"`
	Created                  int64              `json:"created,omitempty"`
	Hostname                 string             `json:"hostname,omitempty"`
	Status                   string             `json:"status,omitempty"`
	LastErrorTime            int64              `json:"lasterrortime,omitempty"`
	LastTestTime             int64              `json:"lasttesttime,omitempty"`
	LastResponseTime         int64              `json:"lastresponsetime,omitempty"`
	Paused                   bool               `json:"paused,omitempty"`
	IntegrationIds           []int              `json:"integrationids,omitempty"`
	Type                     CheckResponseType  `json:"type,omitempty"`
	Tags                     []CheckResponseTag `json:"tags,omitempty"`
	UserIds                  []int              `json:"userids,omitempty"`
	TeamIds                  []int              `json:"teamids,omitempty"`
}

type CheckResponseType struct {
	Name string                    `json:"-"`
	HTTP *CheckResponseHTTPDetails `json:"http,omitempty"`
	TCP  *CheckResponseTCPDetails  `json:"tcp,omitempty"`
}

type CheckResponseTag struct {
	Name  string      `json:"name"`
	Type  string      `json:"type"`
	Count interface{} `json:"count"`
}

// MaintenanceResponse represents the json response for a maintenance from the Pingdom API
type MaintenanceResponse struct {
	ID             int                      `json:"id"`
	Description    string                   `json:"description"`
	From           int64                    `json:"from"`
	To             int64                    `json:"to"`
	RecurrenceType string                   `json:"recurrencetype"`
	RepeatEvery    int                      `json:"repeatevery"`
	EffectiveTo    int64                    `json:"effectiveto"`
	Checks         MaintenanceCheckResponse `json:"checks"`
}

// MaintenanceCheckResponse represents Check reply in json MaintenanceResponse
type MaintenanceCheckResponse struct {
	Uptime []int `json:"uptime"`
	Tms    []int `json:"tms"`
}

// ProbeResponse represents the json response for probes from the PIngdom API
type ProbeResponse struct {
	ID         int    `json:"id"`
	Country    string `json:"country"`
	City       string `json:"city"`
	Name       string `json:"name"`
	Active     bool   `json:"active"`
	Hostname   string `json:"hostname"`
	IP         string `json:"ip"`
	IPv6       string `json:"ipv6"`
	CountryISO string `json:"countryiso"`
	Region     string `json:"region"`
}

// TeamResponse represents the json response for teams from the PIngdom API
type TeamResponse struct {
	ID    string `json:"id"`
	Name  string `json:"name"`
	Users []TeamUserResponse
}

// TeamUserResponse represents the json response for users in teams from the PIngdom API
type TeamUserResponse struct {
	ID    string `json:"id"`
	Email string `json:"email"`
	Name  string `json:"name"`
}

// TeamDeleteResponse represents the json response for delete team from the PIngdom API
type TeamDeleteResponse struct {
	Success bool `json:"success"`
}

func (c *CheckResponseType) UnmarshalJSON(b []byte) error {
	var raw interface{}

	err := json.Unmarshal(b, &raw)
	if err != nil {
		return err
	}

	switch v := raw.(type) {
	case string:
		c.Name = v
	case map[string]interface{}:
		if len(v) != 1 {
			return fmt.Errorf("Check detailed response `check.type` contains more than one object: %+v", v)
		}
		for k := range v {
			c.Name = k
		}

		// Allow continue use json.Unmarshall using a type != Unmarshaller
		// This avoid enter in a infinite loop
		type t CheckResponseType
		var rawCheckDetails t

		err := json.Unmarshal(b, &rawCheckDetails)
		if err != nil {
			return err
		}
		c.HTTP = rawCheckDetails.HTTP
		c.TCP = rawCheckDetails.TCP
	}
	return nil
}

// CheckResponseHTTPDetails represents the details specific to HTTP checks.
type CheckResponseHTTPDetails struct {
	Url              string            `json:"url,omitempty"`
	Encryption       bool              `json:"encryption,omitempty"`
	Port             int               `json:"port,omitempty"`
	Username         string            `json:"username,omitempty"`
	Password         string            `json:"password,omitempty"`
	ShouldContain    string            `json:"shouldcontain,omitempty"`
	ShouldNotContain string            `json:"shouldnotcontain,omitempty"`
	PostData         string            `json:"postdata,omitempty"`
	RequestHeaders   map[string]string `json:"requestheaders,omitempty"`
}

// CheckResponseTCPDetails represents the details specific to TCP checks.
type CheckResponseTCPDetails struct {
	Port           int    `json:"port,omitempty"`
	StringToSend   string `json:"stringtosend,omitempty"`
	StringToExpect string `json:"stringtoexpect,omitempty"`
}

// Return string representation of the PingdomError
func (r *PingdomError) Error() string {
	return fmt.Sprintf("%d %v: %v", r.StatusCode, r.StatusDesc, r.Message)
}

// private types used to unmarshall json responses from pingdom

type listChecksJsonResponse struct {
	Checks []CheckResponse `json:"checks"`
}

type listMaintenanceJsonResponse struct {
	Maintenances []MaintenanceResponse `json:"maintenance"`
}

type listProbesJsonResponse struct {
	Probes []ProbeResponse `json:"probes"`
}

type listTeamsJsonResponse struct {
	Teams []TeamResponse `json:"teams"`
}

type checkDetailsJsonResponse struct {
	Check *CheckResponse `json:"check"`
}

type maintenanceDetailsJsonResponse struct {
	Maintenance *MaintenanceResponse `json:"maintenance"`
}

type teamDetailsJsonResponse struct {
	Team *TeamResponse `json:"team"`
}

type errorJsonResponse struct {
	Error *PingdomError `json:"error"`
}
<|endoftext|>


<commit_before>package atlas

import (
	"bytes"
	"fmt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"io/ioutil"
	"net/http"
	"net/url"
	"testing"
)

const testURL = "http://localhost:10000"

func TestGetPageNum(t *testing.T) {
	url := "https://foo.example.com/"

	n := getPageNum(url)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
	url = "https://foo.example.com/?page=42"
	n = getPageNum(url)
	if n != "42" {
		t.Errorf("n=%s should be 42", n)
	}
	url = "https://foo.example.com/?country=fr&page=43"
	n = getPageNum(url)
	if n != "43" {
		t.Errorf("n=%s should be 43", n)
	}
	url = "https://foo.example.com/?country=fr&page=666&bar=1"
	n = getPageNum(url)
	if n != "666" {
		t.Errorf("n=%s should be 666", n)
	}

	url = ""
	n = getPageNum(url)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
}

func TestClienthandleAPIResponsese(t *testing.T) {
	var (
		r http.Response
		b bytes.Buffer
	)

	client, err := NewClient()
	err = client.handleAPIResponsese(nil)
	assert.Error(t, err, "should be in error")

	r = http.Response{StatusCode: 0}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	r = http.Response{StatusCode: 200}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	var jsonErr = `error:{status: 501, code: 500, detail: "test"}`

	fmt.Fprintf(&b, "%v", jsonErr)
	r.StatusCode = 300
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be in error")

	r.StatusCode = 500
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.Error(t, err, "should be in error")
}

func TestAddQueryParameters(t *testing.T) {
	p := AddQueryParameters("", map[string]string{})
	assert.Equal(t, "", p)
}

func TestAddQueryParameters_1(t *testing.T) {
	p := AddQueryParameters("", map[string]string{"": ""})
	assert.Equal(t, "?=", p)
}

func TestAddQueryParameters_2(t *testing.T) {
	p := AddQueryParameters("foo", map[string]string{"bar": "baz"})
	assert.Equal(t, "foo?bar=baz", p)
}

func TestClient_AddAPIKey(t *testing.T) {
	c, err := NewClient(Config{APIKey: "foo"})
	require.NoError(t, err)
	assert.NotNil(t, c)
	assert.NotEmpty(t, c)

	opts := map[string]string{}

	newk := c.addAPIKey(opts)
	assert.NotEmpty(t, c.config.APIKey)
	assert.Equal(t, 1, len(newk))
	assert.EqualValues(t, map[string]string{"key": "foo"}, newk)
}

func TestClient_PrepareRequest(t *testing.T) {
	c, err := NewClient(Config{endpoint: testURL})
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_PrepareRequest_2(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(apiEndpoint + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_PrepareRequest_3(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("FETCH", testURL +"/foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_MergeGlobalOptions(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{"foo": "bar"}
	c.opts = map[string]string{"baz": "xyz"}
	res := map[string]string{"foo": "bar", "baz": "xyz"}

	c.mergeGlobalOptions(opts)
	assert.EqualValues(t, res, opts)
}
<commit_msg>TesCfg is using the mock server now.<commit_after>package atlas

import (
	"bytes"
	"fmt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"io/ioutil"
	"net/http"
	"net/url"
	"testing"
)

const testURL = "http://localhost:10000"

func TestGetPageNum(t *testing.T) {
	nurl := "https://foo.example.com/"

	n := getPageNum(nurl)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
	nurl = "https://foo.example.com/?page=42"
	n = getPageNum(nurl)
	if n != "42" {
		t.Errorf("n=%s should be 42", n)
	}
	nurl = "https://foo.example.com/?country=fr&page=43"
	n = getPageNum(nurl)
	if n != "43" {
		t.Errorf("n=%s should be 43", n)
	}
	nurl = "https://foo.example.com/?country=fr&page=666&bar=1"
	n = getPageNum(nurl)
	if n != "666" {
		t.Errorf("n=%s should be 666", n)
	}

	nurl = ""
	n = getPageNum(nurl)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
}

func TestClienthandleAPIResponsese(t *testing.T) {
	var (
		r http.Response
		b bytes.Buffer
	)

	client, err := NewClient()
	err = client.handleAPIResponsese(nil)
	assert.Error(t, err, "should be in error")

	r = http.Response{StatusCode: 0}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	r = http.Response{StatusCode: 200}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	var jsonErr = `error:{status: 501, code: 500, detail: "test"}`

	fmt.Fprintf(&b, "%v", jsonErr)
	r.StatusCode = 300
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be in error")

	r.StatusCode = 500
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.Error(t, err, "should be in error")
}

func TestAddQueryParameters(t *testing.T) {
	p := AddQueryParameters("", map[string]string{})
	assert.Equal(t, "", p)
}

func TestAddQueryParameters_1(t *testing.T) {
	p := AddQueryParameters("", map[string]string{"": ""})
	assert.Equal(t, "?=", p)
}

func TestAddQueryParameters_2(t *testing.T) {
	p := AddQueryParameters("foo", map[string]string{"bar": "baz"})
	assert.Equal(t, "foo?bar=baz", p)
}

func TestClient_AddAPIKey(t *testing.T) {
	c, err := NewClient(Config{APIKey: "foo"})
	require.NoError(t, err)
	assert.NotNil(t, c)
	assert.NotEmpty(t, c)

	opts := map[string]string{}

	newk := c.addAPIKey(opts)
	assert.NotEmpty(t, c.config.APIKey)
	assert.Equal(t, 1, len(newk))
	assert.EqualValues(t, map[string]string{"key": "foo"}, newk)
}

func TestClient_PrepareRequest(t *testing.T) {
	c, err := NewClient(Config{endpoint: testURL})
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_PrepareRequest_2(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_PrepareRequest_3(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("FETCH", testURL +"/foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_MergeGlobalOptions(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{"foo": "bar"}
	c.opts = map[string]string{"baz": "xyz"}
	res := map[string]string{"foo": "bar", "baz": "xyz"}

	c.mergeGlobalOptions(opts)
	assert.EqualValues(t, res, opts)
}
<|endoftext|>


<commit_before>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package e2e

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"
	"gotest.tools/v3/assert"
	"gotest.tools/v3/icmd"
	"gotest.tools/v3/poll"

	"github.com/docker/compose/v2/cmd/compose"
)

var (
	// DockerExecutableName is the OS dependent Docker CLI binary name
	DockerExecutableName = "docker"

	// DockerComposeExecutableName is the OS dependent Docker CLI binary name
	DockerComposeExecutableName = "docker-" + compose.PluginName

	// DockerScanExecutableName is the OS dependent Docker CLI binary name
	DockerScanExecutableName = "docker-scan"
)

func init() {
	if runtime.GOOS == "windows" {
		DockerExecutableName = DockerExecutableName + ".exe"
		DockerComposeExecutableName = DockerComposeExecutableName + ".exe"
		DockerScanExecutableName = DockerScanExecutableName + ".exe"
	}
}

// CLI is used to wrap the CLI for end to end testing
type CLI struct {
	// ConfigDir for Docker configuration (set as DOCKER_CONFIG)
	ConfigDir string

	// HomeDir for tools that look for user files (set as HOME)
	HomeDir string

	// env overrides to apply to every invoked command
	//
	// To populate, use WithEnv when creating a CLI instance.
	env []string
}

// CLIOption to customize behavior for all commands for a CLI instance.
type CLIOption func(c *CLI)

// NewParallelCLI marks the parent test as parallel and returns a CLI instance
// suitable for usage across child tests.
func NewParallelCLI(t *testing.T, opts ...CLIOption) *CLI {
	t.Helper()
	t.Parallel()
	return NewCLI(t, opts...)
}

// NewCLI creates a CLI instance for running E2E tests.
func NewCLI(t testing.TB, opts ...CLIOption) *CLI {
	t.Helper()

	configDir := t.TempDir()
	initializePlugins(t, configDir)

	c := &CLI{
		ConfigDir: configDir,
		HomeDir:   t.TempDir(),
	}

	for _, opt := range opts {
		opt(c)
	}

	return c
}

// WithEnv sets environment variables that will be passed to commands.
func WithEnv(env ...string) CLIOption {
	return func(c *CLI) {
		c.env = append(c.env, env...)
	}
}

// initializePlugins copies the necessary plugin files to the temporary config
// directory for the test.
func initializePlugins(t testing.TB, d string) {
	t.Helper()

	t.Cleanup(func() {
		if t.Failed() {
			conf, _ := ioutil.ReadFile(filepath.Join(d, "config.json"))
			t.Errorf("Config: %s\n", string(conf))
			t.Error("Contents of config dir:")
			for _, p := range dirContents(d) {
				t.Errorf(p)
			}
		}
		_ = os.RemoveAll(d)
	})

	_ = os.MkdirAll(filepath.Join(d, "cli-plugins"), 0755)
	composePlugin, err := findExecutable(DockerComposeExecutableName, []string{"../../bin", "../../../bin"})
	if os.IsNotExist(err) {
		fmt.Println("WARNING: docker-compose cli-plugin not found")
	}
	if err == nil {
		err = CopyFile(composePlugin, filepath.Join(d, "cli-plugins", DockerComposeExecutableName))
		if err != nil {
			panic(err)
		}
		// We don't need a functional scan plugin, but a valid plugin binary
		err = CopyFile(composePlugin, filepath.Join(d, "cli-plugins", DockerScanExecutableName))
		if err != nil {
			panic(err)
		}
	}
}

func dirContents(dir string) []string {
	var res []string
	_ = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		res = append(res, path)
		return nil
	})
	return res
}

func findExecutable(executableName string, paths []string) (string, error) {
	for _, p := range paths {
		bin, err := filepath.Abs(path.Join(p, executableName))
		if err != nil {
			return "", err
		}

		if _, err := os.Stat(bin); os.IsNotExist(err) {
			continue
		}

		return bin, nil
	}

	return "", errors.Wrap(os.ErrNotExist, "executable not found")
}

// CopyFile copies a file from a sourceFile to a destinationFile setting permissions to 0755
func CopyFile(sourceFile string, destinationFile string) error {
	src, err := os.Open(sourceFile)
	if err != nil {
		return err
	}
	// nolint: errcheck
	defer src.Close()

	dst, err := os.OpenFile(destinationFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return err
	}
	// nolint: errcheck
	defer dst.Close()

	if _, err = io.Copy(dst, src); err != nil {
		return err
	}

	return err
}

// BaseEnvironment provides the minimal environment variables used across all
// Docker / Compose commands.
func (c *CLI) BaseEnvironment() []string {
	return []string{
		"HOME=" + c.HomeDir,
		"USER=" + os.Getenv("USER"),
		"DOCKER_CONFIG=" + c.ConfigDir,
		"KUBECONFIG=invalid",
	}
}

// NewCmd creates a cmd object configured with the test environment set
func (c *CLI) NewCmd(command string, args ...string) icmd.Cmd {
	return icmd.Cmd{
		Command: append([]string{command}, args...),
		Env:     append(c.BaseEnvironment(), c.env...),
	}
}

// NewCmdWithEnv creates a cmd object configured with the test environment set with additional env vars
func (c *CLI) NewCmdWithEnv(envvars []string, command string, args ...string) icmd.Cmd {
	// base env -> CLI overrides -> cmd overrides
	cmdEnv := append(c.BaseEnvironment(), c.env...)
	cmdEnv = append(cmdEnv, envvars...)
	return icmd.Cmd{
		Command: append([]string{command}, args...),
		Env:     cmdEnv,
	}
}

// MetricsSocket get the path where test metrics will be sent
func (c *CLI) MetricsSocket() string {
	return filepath.Join(c.ConfigDir, "./docker-cli.sock")
}

// NewDockerCmd creates a docker cmd without running it
func (c *CLI) NewDockerCmd(t testing.TB, args ...string) icmd.Cmd {
	for _, arg := range args {
		if arg == compose.PluginName {
			t.Fatal("This test called 'RunDockerCmd' for 'compose'. Please prefer 'RunDockerComposeCmd' to be able to test as a plugin and standalone")
		}
	}
	return c.NewCmd(DockerExecutableName, args...)
}

// RunDockerOrExitError runs a docker command and returns a result
func (c *CLI) RunDockerOrExitError(t testing.TB, args ...string) *icmd.Result {
	fmt.Printf("\t[%s] docker %s\n", t.Name(), strings.Join(args, " "))
	return icmd.RunCmd(c.NewDockerCmd(t, args...))
}

// RunCmd runs a command, expects no error and returns a result
func (c *CLI) RunCmd(t testing.TB, args ...string) *icmd.Result {
	fmt.Printf("\t[%s] %s\n", t.Name(), strings.Join(args, " "))
	assert.Assert(t, len(args) >= 1, "require at least one command in parameters")
	res := icmd.RunCmd(c.NewCmd(args[0], args[1:]...))
	res.Assert(t, icmd.Success)
	return res
}

// RunCmdInDir runs a command in a given dir, expects no error and returns a result
func (c *CLI) RunCmdInDir(t testing.TB, dir string, args ...string) *icmd.Result {
	fmt.Printf("\t[%s] %s\n", t.Name(), strings.Join(args, " "))
	assert.Assert(t, len(args) >= 1, "require at least one command in parameters")
	cmd := c.NewCmd(args[0], args[1:]...)
	cmd.Dir = dir
	res := icmd.RunCmd(cmd)
	res.Assert(t, icmd.Success)
	return res
}

// RunDockerCmd runs a docker command, expects no error and returns a result
func (c *CLI) RunDockerCmd(t testing.TB, args ...string) *icmd.Result {
	res := c.RunDockerOrExitError(t, args...)
	res.Assert(t, icmd.Success)
	return res
}

// RunDockerComposeCmd runs a docker compose command, expects no error and returns a result
func (c *CLI) RunDockerComposeCmd(t testing.TB, args ...string) *icmd.Result {
	res := c.RunDockerComposeCmdNoCheck(t, args...)
	res.Assert(t, icmd.Success)
	return res
}

// RunDockerComposeCmdNoCheck runs a docker compose command, don't presume of any expectation and returns a result
func (c *CLI) RunDockerComposeCmdNoCheck(t testing.TB, args ...string) *icmd.Result {
	return icmd.RunCmd(c.NewDockerComposeCmd(t, args...))
}

// NewDockerComposeCmd creates a command object for Compose, either in plugin
// or standalone mode (based on build tags).
func (c *CLI) NewDockerComposeCmd(t testing.TB, args ...string) icmd.Cmd {
	t.Helper()
	if composeStandaloneMode {
		return c.NewCmd(ComposeStandalonePath(t), args...)
	}
	args = append([]string{"compose"}, args...)
	return c.NewCmd(DockerExecutableName, args...)
}

// ComposeStandalonePath returns the path to the locally-built Compose
// standalone binary from the repo.
//
// This function will fail the test immediately if invoked when not running
// in standalone test mode.
func ComposeStandalonePath(t testing.TB) string {
	t.Helper()
	if !composeStandaloneMode {
		require.Fail(t, "Not running in standalone mode")
	}
	composeBinary, err := findExecutable(DockerComposeExecutableName, []string{"../../bin", "../../../bin"})
	require.NoError(t, err, "Could not find standalone Compose binary (%q)",
		DockerComposeExecutableName)
	return composeBinary
}

// StdoutContains returns a predicate on command result expecting a string in stdout
func StdoutContains(expected string) func(*icmd.Result) bool {
	return func(res *icmd.Result) bool {
		return strings.Contains(res.Stdout(), expected)
	}
}

// WaitForCmdResult try to execute a cmd until resulting output matches given predicate
func (c *CLI) WaitForCmdResult(t testing.TB, command icmd.Cmd, predicate func(*icmd.Result) bool, timeout time.Duration, delay time.Duration) {
	assert.Assert(t, timeout.Nanoseconds() > delay.Nanoseconds(), "timeout must be greater than delay")
	var res *icmd.Result
	checkStopped := func(logt poll.LogT) poll.Result {
		fmt.Printf("\t[%s] %s\n", t.Name(), strings.Join(command.Command, " "))
		res = icmd.RunCmd(command)
		if !predicate(res) {
			return poll.Continue("Cmd output did not match requirement: %q", res.Combined())
		}
		return poll.Success()
	}
	poll.WaitOn(t, checkStopped, poll.WithDelay(delay), poll.WithTimeout(timeout))
}

// WaitForCondition wait for predicate to execute to true
func (c *CLI) WaitForCondition(t testing.TB, predicate func() (bool, string), timeout time.Duration, delay time.Duration) {
	checkStopped := func(logt poll.LogT) poll.Result {
		pass, description := predicate()
		if !pass {
			return poll.Continue("Condition not met: %q", description)
		}
		return poll.Success()
	}
	poll.WaitOn(t, checkStopped, poll.WithDelay(delay), poll.WithTimeout(timeout))
}

// Lines split output into lines
func Lines(output string) []string {
	return strings.Split(strings.TrimSpace(output), "\n")
}

// HTTPGetWithRetry performs an HTTP GET on an `endpoint`, using retryDelay also as a request timeout.
// In the case of an error or the response status is not the expected one, it retries the same request,
// returning the response body as a string (empty if we could not reach it)
func HTTPGetWithRetry(t testing.TB, endpoint string, expectedStatus int, retryDelay time.Duration, timeout time.Duration) string {
	var (
		r   *http.Response
		err error
	)
	client := &http.Client{
		Timeout: retryDelay,
	}
	fmt.Printf("\t[%s] GET %s\n", t.Name(), endpoint)
	checkUp := func(t poll.LogT) poll.Result {
		r, err = client.Get(endpoint)
		if err != nil {
			return poll.Continue("reaching %q: Error %s", endpoint, err.Error())
		}
		if r.StatusCode == expectedStatus {
			return poll.Success()
		}
		return poll.Continue("reaching %q: %d != %d", endpoint, r.StatusCode, expectedStatus)
	}
	poll.WaitOn(t, checkUp, poll.WithDelay(retryDelay), poll.WithTimeout(timeout))
	if r != nil {
		b, err := ioutil.ReadAll(r.Body)
		assert.NilError(t, err)
		return string(b)
	}
	return ""
}
<commit_msg>e2e: improve test output on failures<commit_after>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package e2e

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"
	"gotest.tools/v3/assert"
	"gotest.tools/v3/icmd"
	"gotest.tools/v3/poll"

	"github.com/docker/compose/v2/cmd/compose"
)

var (
	// DockerExecutableName is the OS dependent Docker CLI binary name
	DockerExecutableName = "docker"

	// DockerComposeExecutableName is the OS dependent Docker CLI binary name
	DockerComposeExecutableName = "docker-" + compose.PluginName

	// DockerScanExecutableName is the OS dependent Docker CLI binary name
	DockerScanExecutableName = "docker-scan"
)

func init() {
	if runtime.GOOS == "windows" {
		DockerExecutableName = DockerExecutableName + ".exe"
		DockerComposeExecutableName = DockerComposeExecutableName + ".exe"
		DockerScanExecutableName = DockerScanExecutableName + ".exe"
	}
}

// CLI is used to wrap the CLI for end to end testing
type CLI struct {
	// ConfigDir for Docker configuration (set as DOCKER_CONFIG)
	ConfigDir string

	// HomeDir for tools that look for user files (set as HOME)
	HomeDir string

	// env overrides to apply to every invoked command
	//
	// To populate, use WithEnv when creating a CLI instance.
	env []string
}

// CLIOption to customize behavior for all commands for a CLI instance.
type CLIOption func(c *CLI)

// NewParallelCLI marks the parent test as parallel and returns a CLI instance
// suitable for usage across child tests.
func NewParallelCLI(t *testing.T, opts ...CLIOption) *CLI {
	t.Helper()
	t.Parallel()
	return NewCLI(t, opts...)
}

// NewCLI creates a CLI instance for running E2E tests.
func NewCLI(t testing.TB, opts ...CLIOption) *CLI {
	t.Helper()

	configDir := t.TempDir()
	initializePlugins(t, configDir)

	c := &CLI{
		ConfigDir: configDir,
		HomeDir:   t.TempDir(),
	}

	for _, opt := range opts {
		opt(c)
	}

	return c
}

// WithEnv sets environment variables that will be passed to commands.
func WithEnv(env ...string) CLIOption {
	return func(c *CLI) {
		c.env = append(c.env, env...)
	}
}

// initializePlugins copies the necessary plugin files to the temporary config
// directory for the test.
func initializePlugins(t testing.TB, configDir string) {
	t.Helper()

	t.Cleanup(func() {
		if t.Failed() {
			if conf, err := ioutil.ReadFile(filepath.Join(configDir, "config.json")); err == nil {
				t.Logf("Config: %s\n", string(conf))
			}
			t.Log("Contents of config dir:")
			for _, p := range dirContents(configDir) {
				t.Logf("  - %s", p)
			}
		}
	})

	require.NoError(t, os.MkdirAll(filepath.Join(configDir, "cli-plugins"), 0755),
		"Failed to create cli-plugins directory")
	composePlugin, err := findExecutable(DockerComposeExecutableName, []string{"../../bin", "../../../bin"})
	if os.IsNotExist(err) {
		t.Logf("WARNING: docker-compose cli-plugin not found")
	}
	if err == nil {
		CopyFile(t, composePlugin, filepath.Join(configDir, "cli-plugins", DockerComposeExecutableName))
		// We don't need a functional scan plugin, but a valid plugin binary
		CopyFile(t, composePlugin, filepath.Join(configDir, "cli-plugins", DockerScanExecutableName))
	}
}

func dirContents(dir string) []string {
	var res []string
	_ = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		res = append(res, path)
		return nil
	})
	return res
}

func findExecutable(executableName string, paths []string) (string, error) {
	for _, p := range paths {
		bin, err := filepath.Abs(path.Join(p, executableName))
		if err != nil {
			return "", err
		}

		if _, err := os.Stat(bin); os.IsNotExist(err) {
			continue
		}

		return bin, nil
	}

	return "", errors.Wrap(os.ErrNotExist, "executable not found")
}

// CopyFile copies a file from a sourceFile to a destinationFile setting permissions to 0755
func CopyFile(t testing.TB, sourceFile string, destinationFile string) {
	t.Helper()

	src, err := os.Open(sourceFile)
	require.NoError(t, err, "Failed to open source file: %s")
	// nolint: errcheck
	defer src.Close()

	dst, err := os.OpenFile(destinationFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
	require.NoError(t, err, "Failed to open destination file: %s", destinationFile)
	// nolint: errcheck
	defer dst.Close()

	_, err = io.Copy(dst, src)
	require.NoError(t, err, "Failed to copy file: %s", sourceFile)
}

// BaseEnvironment provides the minimal environment variables used across all
// Docker / Compose commands.
func (c *CLI) BaseEnvironment() []string {
	return []string{
		"HOME=" + c.HomeDir,
		"USER=" + os.Getenv("USER"),
		"DOCKER_CONFIG=" + c.ConfigDir,
		"KUBECONFIG=invalid",
	}
}

// NewCmd creates a cmd object configured with the test environment set
func (c *CLI) NewCmd(command string, args ...string) icmd.Cmd {
	return icmd.Cmd{
		Command: append([]string{command}, args...),
		Env:     append(c.BaseEnvironment(), c.env...),
	}
}

// NewCmdWithEnv creates a cmd object configured with the test environment set with additional env vars
func (c *CLI) NewCmdWithEnv(envvars []string, command string, args ...string) icmd.Cmd {
	// base env -> CLI overrides -> cmd overrides
	cmdEnv := append(c.BaseEnvironment(), c.env...)
	cmdEnv = append(cmdEnv, envvars...)
	return icmd.Cmd{
		Command: append([]string{command}, args...),
		Env:     cmdEnv,
	}
}

// MetricsSocket get the path where test metrics will be sent
func (c *CLI) MetricsSocket() string {
	return filepath.Join(c.ConfigDir, "./docker-cli.sock")
}

// NewDockerCmd creates a docker cmd without running it
func (c *CLI) NewDockerCmd(t testing.TB, args ...string) icmd.Cmd {
	t.Helper()
	for _, arg := range args {
		if arg == compose.PluginName {
			t.Fatal("This test called 'RunDockerCmd' for 'compose'. Please prefer 'RunDockerComposeCmd' to be able to test as a plugin and standalone")
		}
	}
	return c.NewCmd(DockerExecutableName, args...)
}

// RunDockerOrExitError runs a docker command and returns a result
func (c *CLI) RunDockerOrExitError(t testing.TB, args ...string) *icmd.Result {
	t.Helper()
	t.Logf("\t[%s] docker %s\n", t.Name(), strings.Join(args, " "))
	return icmd.RunCmd(c.NewDockerCmd(t, args...))
}

// RunCmd runs a command, expects no error and returns a result
func (c *CLI) RunCmd(t testing.TB, args ...string) *icmd.Result {
	t.Helper()
	t.Logf("\t[%s] %s\n", t.Name(), strings.Join(args, " "))
	assert.Assert(t, len(args) >= 1, "require at least one command in parameters")
	res := icmd.RunCmd(c.NewCmd(args[0], args[1:]...))
	res.Assert(t, icmd.Success)
	return res
}

// RunCmdInDir runs a command in a given dir, expects no error and returns a result
func (c *CLI) RunCmdInDir(t testing.TB, dir string, args ...string) *icmd.Result {
	t.Helper()
	t.Logf("\t[%s] %s\n", t.Name(), strings.Join(args, " "))
	assert.Assert(t, len(args) >= 1, "require at least one command in parameters")
	cmd := c.NewCmd(args[0], args[1:]...)
	cmd.Dir = dir
	res := icmd.RunCmd(cmd)
	res.Assert(t, icmd.Success)
	return res
}

// RunDockerCmd runs a docker command, expects no error and returns a result
func (c *CLI) RunDockerCmd(t testing.TB, args ...string) *icmd.Result {
	t.Helper()
	res := c.RunDockerOrExitError(t, args...)
	res.Assert(t, icmd.Success)
	return res
}

// RunDockerComposeCmd runs a docker compose command, expects no error and returns a result
func (c *CLI) RunDockerComposeCmd(t testing.TB, args ...string) *icmd.Result {
	t.Helper()
	res := c.RunDockerComposeCmdNoCheck(t, args...)
	res.Assert(t, icmd.Success)
	return res
}

// RunDockerComposeCmdNoCheck runs a docker compose command, don't presume of any expectation and returns a result
func (c *CLI) RunDockerComposeCmdNoCheck(t testing.TB, args ...string) *icmd.Result {
	t.Helper()
	return icmd.RunCmd(c.NewDockerComposeCmd(t, args...))
}

// NewDockerComposeCmd creates a command object for Compose, either in plugin
// or standalone mode (based on build tags).
func (c *CLI) NewDockerComposeCmd(t testing.TB, args ...string) icmd.Cmd {
	t.Helper()
	if composeStandaloneMode {
		return c.NewCmd(ComposeStandalonePath(t), args...)
	}
	args = append([]string{"compose"}, args...)
	return c.NewCmd(DockerExecutableName, args...)
}

// ComposeStandalonePath returns the path to the locally-built Compose
// standalone binary from the repo.
//
// This function will fail the test immediately if invoked when not running
// in standalone test mode.
func ComposeStandalonePath(t testing.TB) string {
	t.Helper()
	if !composeStandaloneMode {
		require.Fail(t, "Not running in standalone mode")
	}
	composeBinary, err := findExecutable(DockerComposeExecutableName, []string{"../../bin", "../../../bin"})
	require.NoError(t, err, "Could not find standalone Compose binary (%q)",
		DockerComposeExecutableName)
	return composeBinary
}

// StdoutContains returns a predicate on command result expecting a string in stdout
func StdoutContains(expected string) func(*icmd.Result) bool {
	return func(res *icmd.Result) bool {
		return strings.Contains(res.Stdout(), expected)
	}
}

// WaitForCmdResult try to execute a cmd until resulting output matches given predicate
func (c *CLI) WaitForCmdResult(
	t testing.TB,
	command icmd.Cmd,
	predicate func(*icmd.Result) bool,
	timeout time.Duration,
	delay time.Duration,
) {
	t.Helper()
	assert.Assert(t, timeout.Nanoseconds() > delay.Nanoseconds(), "timeout must be greater than delay")
	var res *icmd.Result
	checkStopped := func(logt poll.LogT) poll.Result {
		fmt.Printf("\t[%s] %s\n", t.Name(), strings.Join(command.Command, " "))
		res = icmd.RunCmd(command)
		if !predicate(res) {
			return poll.Continue("Cmd output did not match requirement: %q", res.Combined())
		}
		return poll.Success()
	}
	poll.WaitOn(t, checkStopped, poll.WithDelay(delay), poll.WithTimeout(timeout))
}

// WaitForCondition wait for predicate to execute to true
func (c *CLI) WaitForCondition(
	t testing.TB,
	predicate func() (bool, string),
	timeout time.Duration,
	delay time.Duration,
) {
	t.Helper()
	checkStopped := func(logt poll.LogT) poll.Result {
		pass, description := predicate()
		if !pass {
			return poll.Continue("Condition not met: %q", description)
		}
		return poll.Success()
	}
	poll.WaitOn(t, checkStopped, poll.WithDelay(delay), poll.WithTimeout(timeout))
}

// Lines split output into lines
func Lines(output string) []string {
	return strings.Split(strings.TrimSpace(output), "\n")
}

// HTTPGetWithRetry performs an HTTP GET on an `endpoint`, using retryDelay also as a request timeout.
// In the case of an error or the response status is not the expected one, it retries the same request,
// returning the response body as a string (empty if we could not reach it)
func HTTPGetWithRetry(
	t testing.TB,
	endpoint string,
	expectedStatus int,
	retryDelay time.Duration,
	timeout time.Duration,
) string {
	t.Helper()
	var (
		r   *http.Response
		err error
	)
	client := &http.Client{
		Timeout: retryDelay,
	}
	fmt.Printf("\t[%s] GET %s\n", t.Name(), endpoint)
	checkUp := func(t poll.LogT) poll.Result {
		r, err = client.Get(endpoint)
		if err != nil {
			return poll.Continue("reaching %q: Error %s", endpoint, err.Error())
		}
		if r.StatusCode == expectedStatus {
			return poll.Success()
		}
		return poll.Continue("reaching %q: %d != %d", endpoint, r.StatusCode, expectedStatus)
	}
	poll.WaitOn(t, checkUp, poll.WithDelay(retryDelay), poll.WithTimeout(timeout))
	if r != nil {
		b, err := ioutil.ReadAll(r.Body)
		assert.NilError(t, err)
		return string(b)
	}
	return ""
}
<|endoftext|>