<commit_before>// Copyright 2017-2018 Authors of Cilium
//
// 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 server

import (
	"fmt"
	"sync"
	"time"

	healthModels "github.com/cilium/cilium/api/v1/health/models"
	healthApi "github.com/cilium/cilium/api/v1/health/server"
	"github.com/cilium/cilium/api/v1/health/server/restapi"
	ciliumPkg "github.com/cilium/cilium/pkg/client"
	"github.com/cilium/cilium/pkg/health/defaults"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"

	"github.com/go-openapi/loads"
	"github.com/jessevdk/go-flags"
)

// AdminOption is an option for determining over which protocols the APIs are
// exposed.
type AdminOption string

const (
	// AdminOptionAny exposes every API over both Unix and HTTP sockets.
	AdminOptionAny AdminOption = "any"

	// AdminOptionUnix restricts most APIs to hosting over Unix sockets.
	AdminOptionUnix AdminOption = "unix"
)

var (
	log = logging.DefaultLogger.WithField(logfields.LogSubsys, "health-server")

	// PortToPaths is a convenience map for access to the ports and their
	// common string representations
	PortToPaths = map[int]string{
		defaults.HTTPPathPort: "Via L3",
	}

	// AdminOptions is the slice of all valid AdminOption values.
	AdminOptions = []AdminOption{
		AdminOptionAny,
		AdminOptionUnix,
	}
)

// Config stores the configuration data for a cilium-health server.
type Config struct {
	Debug         bool
	Passive       bool
	Admin         AdminOption
	CiliumURI     string
	ProbeInterval time.Duration
	ProbeDeadline time.Duration
}

// ipString is an IP address used as a more descriptive type name in maps.
type ipString string

// nodeMap maps IP addresses to healthNode objectss for convenient access to
// node information.
type nodeMap map[ipString]healthNode

// Server is the cilium-health daemon that is in charge of performing health
// and connectivity checks periodically, and serving the cilium-health API.
type Server struct {
	healthApi.Server  // Server to provide cilium-health API
	*ciliumPkg.Client // Client to "GET /healthz" on cilium daemon
	Config

	waitgroup  sync.WaitGroup      // Used to synchronize all goroutines
	tcpServers []*healthApi.Server // Servers for external pings
	startTime  time.Time

	// The lock protects against read and write access to the IP->Node map,
	// the list of statuses as most recently seen, and the last time a
	// probe was conducted.
	lock.RWMutex
	connectivity *healthReport
	localStatus  *healthModels.SelfStatus
}

// DumpUptime returns the time that this server has been running.
func (s *Server) DumpUptime() string {
	return time.Since(s.startTime).String()
}

// getNodes fetches the latest set of nodes in the cluster from the Cilium
// daemon, and updates the Server's 'nodes' map.
func (s *Server) getNodes() (nodeMap, error) {
	scopedLog := log
	if s.CiliumURI != "" {
		scopedLog = log.WithField("URI", s.CiliumURI)
	}
	scopedLog.Debug("Sending request for /healthz ...")

	resp, err := s.Daemon.GetHealthz(nil)
	if err != nil {
		return nil, fmt.Errorf("unable to get agent health: %s", err)
	}
	log.Debug("Got cilium /healthz")

	if resp == nil || resp.Payload == nil || resp.Payload.Cluster == nil {
		return nil, fmt.Errorf("received nil health response")
	}

	if resp.Payload.Cluster.Self != "" {
		s.localStatus = &healthModels.SelfStatus{
			Name: resp.Payload.Cluster.Self,
		}
	}

	nodes := make(nodeMap)
	for _, n := range resp.Payload.Cluster.Nodes {
		if n.PrimaryAddress != nil {
			if n.PrimaryAddress.IPV4 != nil {
				nodes[ipString(n.PrimaryAddress.IPV4.IP)] = NewHealthNode(n)
			}
			if n.PrimaryAddress.IPV6 != nil {
				nodes[ipString(n.PrimaryAddress.IPV6.IP)] = NewHealthNode(n)
			}
		}
		for _, addr := range n.SecondaryAddresses {
			nodes[ipString(addr.IP)] = NewHealthNode(n)
		}
		if n.HealthEndpointAddress != nil {
			if n.HealthEndpointAddress.IPV4 != nil {
				nodes[ipString(n.HealthEndpointAddress.IPV4.IP)] = NewHealthNode(n)
			}
			if n.HealthEndpointAddress.IPV6 != nil {
				nodes[ipString(n.HealthEndpointAddress.IPV6.IP)] = NewHealthNode(n)
			}
		}
	}
	return nodes, nil
}

// updateCluster makes the specified health report visible to the API.
//
// It only updates the server's API-visible health report if the provided
// report started after the current report.
func (s *Server) updateCluster(report *healthReport) {
	s.Lock()
	defer s.Unlock()

	if s.connectivity.startTime.Before(report.startTime) {
		s.connectivity = report
	}
}

// GetStatusResponse returns the most recent cluster connectivity status.
func (s *Server) GetStatusResponse() *healthModels.HealthStatusResponse {
	s.RLock()
	defer s.RUnlock()

	var name string
	// Check if localStatus is populated already. If not, the name is empty
	if s.localStatus != nil {
		name = s.localStatus.Name
	}

	return &healthModels.HealthStatusResponse{
		Local: &healthModels.SelfStatus{
			Name: name,
		},
		Nodes:     s.connectivity.nodes,
		Timestamp: s.connectivity.startTime.Format(time.RFC3339),
	}
}

// FetchStatusResponse updates the cluster with the latest set of nodes,
// runs a synchronous probe across the cluster, updates the connectivity cache
// and returns the results.
func (s *Server) FetchStatusResponse() (*healthModels.HealthStatusResponse, error) {
	nodes, err := s.getNodes()
	if err != nil {
		return nil, err
	}

	prober := newProber(s, nodes)
	if err := prober.Run(); err != nil {
		log.WithError(err).Info("Failed to run ping")
		return nil, err
	}
	log.Debug("Run complete")
	s.updateCluster(prober.getResults())

	return s.GetStatusResponse(), nil
}

// Run services that are not considered 'Passive': Actively probing other
// hosts and endpoints over ICMP and HTTP, and hosting a local Unix socket.
// Blocks indefinitely, or returns any errors that occur hosting the Unix
// socket API server.
func (s *Server) runActiveServices() error {
	// Run it once at the start so we get some initial status
	s.FetchStatusResponse()

	nodes, _ := s.getNodes()
	prober := newProber(s, nodes)
	prober.MaxRTT = s.ProbeInterval
	prober.OnIdle = func() {
		// Fetch results and update set of nodes to probe every
		// ProbeInterval
		s.updateCluster(prober.getResults())
		if nodes, err := s.getNodes(); err == nil {
			prober.setNodes(nodes)
		}
	}
	prober.RunLoop()
	defer prober.Stop()

	return s.Server.Serve()
}

// Serve spins up the following goroutines:
// * TCP API Server: Responders to the health API "/hello" message, one per path
//
// Also, if "Passive" is not set in s.Config:
// * Prober: Periodically run pings across the cluster at a configured interval
//   and update the server's connectivity status cache.
// * Unix API Server: Handle all health API requests over a unix socket.
func (s *Server) Serve() (err error) {
	for i := range s.tcpServers {
		s.waitgroup.Add(1)
		srv := s.tcpServers[i]
		go func() {
			defer s.waitgroup.Done()
			srv.Serve()
		}()
	}

	if !s.Config.Passive {
		err = s.runActiveServices()
	}
	s.waitgroup.Wait()

	return err
}

// Shutdown server and clean up resources
func (s *Server) Shutdown() {
	for i := range s.tcpServers {
		s.tcpServers[i].Shutdown()
	}
	if !s.Config.Passive {
		s.Server.Shutdown()
	}
}

func enableAPI(opt AdminOption, tcpPort int) bool {
	switch opt {
	case AdminOptionAny:
		return true
	case AdminOptionUnix:
		return tcpPort == 0
	default:
		return false
	}
}

// newServer instantiates a new instance of the API that serves the health
// API on the specified port. If tcpPort is 0, then a unix socket is opened
// which serves the entire API. If a tcpPort is specified, then it returns
// a server which only answers get requests for the root URL "/".
func (s *Server) newServer(spec *loads.Document, tcpPort int) *healthApi.Server {
	api := restapi.NewCiliumHealthAPI(spec)
	api.Logger = log.Printf

	// /hello
	api.GetHelloHandler = NewGetHelloHandler(s)

	if enableAPI(s.Config.Admin, tcpPort) {
		api.GetHealthzHandler = NewGetHealthzHandler(s)
		api.ConnectivityGetStatusHandler = NewGetStatusHandler(s)
		api.ConnectivityPutStatusProbeHandler = NewPutStatusProbeHandler(s)
	}

	srv := healthApi.NewServer(api)
	if tcpPort == 0 {
		srv.EnabledListeners = []string{"unix"}
		srv.SocketPath = flags.Filename(defaults.SockPath)
	} else {
		srv.EnabledListeners = []string{"http"}
		srv.Port = tcpPort
		srv.Host = "" // FIXME: Allow binding to specific IPs
	}
	srv.ConfigureAPI()

	return srv
}

// NewServer creates a server to handle health requests.
func NewServer(config Config) (*Server, error) {
	server := &Server{
		startTime:    time.Now(),
		Config:       config,
		tcpServers:   []*healthApi.Server{},
		connectivity: &healthReport{},
	}

	swaggerSpec, err := loads.Analyzed(healthApi.SwaggerJSON, "")
	if err != nil {
		return nil, err
	}

	if !config.Passive {
		cl, err := ciliumPkg.NewClient(config.CiliumURI)
		if err != nil {
			return nil, err
		}

		server.Client = cl
		server.Server = *server.newServer(swaggerSpec, 0)
	}
	for port := range PortToPaths {
		srv := server.newServer(swaggerSpec, port)
		server.tcpServers = append(server.tcpServers, srv)
	}

	return server, nil
}
<commit_msg>health: Exit server on first API error<commit_after>// Copyright 2017-2018 Authors of Cilium
//
// 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 server

import (
	"fmt"
	"time"

	healthModels "github.com/cilium/cilium/api/v1/health/models"
	healthApi "github.com/cilium/cilium/api/v1/health/server"
	"github.com/cilium/cilium/api/v1/health/server/restapi"
	ciliumPkg "github.com/cilium/cilium/pkg/client"
	"github.com/cilium/cilium/pkg/health/defaults"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"

	"github.com/go-openapi/loads"
	"github.com/jessevdk/go-flags"
)

// AdminOption is an option for determining over which protocols the APIs are
// exposed.
type AdminOption string

const (
	// AdminOptionAny exposes every API over both Unix and HTTP sockets.
	AdminOptionAny AdminOption = "any"

	// AdminOptionUnix restricts most APIs to hosting over Unix sockets.
	AdminOptionUnix AdminOption = "unix"
)

var (
	log = logging.DefaultLogger.WithField(logfields.LogSubsys, "health-server")

	// PortToPaths is a convenience map for access to the ports and their
	// common string representations
	PortToPaths = map[int]string{
		defaults.HTTPPathPort: "Via L3",
	}

	// AdminOptions is the slice of all valid AdminOption values.
	AdminOptions = []AdminOption{
		AdminOptionAny,
		AdminOptionUnix,
	}
)

// Config stores the configuration data for a cilium-health server.
type Config struct {
	Debug         bool
	Passive       bool
	Admin         AdminOption
	CiliumURI     string
	ProbeInterval time.Duration
	ProbeDeadline time.Duration
}

// ipString is an IP address used as a more descriptive type name in maps.
type ipString string

// nodeMap maps IP addresses to healthNode objectss for convenient access to
// node information.
type nodeMap map[ipString]healthNode

// Server is the cilium-health daemon that is in charge of performing health
// and connectivity checks periodically, and serving the cilium-health API.
type Server struct {
	healthApi.Server  // Server to provide cilium-health API
	*ciliumPkg.Client // Client to "GET /healthz" on cilium daemon
	Config

	tcpServers []*healthApi.Server // Servers for external pings
	startTime  time.Time

	// The lock protects against read and write access to the IP->Node map,
	// the list of statuses as most recently seen, and the last time a
	// probe was conducted.
	lock.RWMutex
	connectivity *healthReport
	localStatus  *healthModels.SelfStatus
}

// DumpUptime returns the time that this server has been running.
func (s *Server) DumpUptime() string {
	return time.Since(s.startTime).String()
}

// getNodes fetches the latest set of nodes in the cluster from the Cilium
// daemon, and updates the Server's 'nodes' map.
func (s *Server) getNodes() (nodeMap, error) {
	scopedLog := log
	if s.CiliumURI != "" {
		scopedLog = log.WithField("URI", s.CiliumURI)
	}
	scopedLog.Debug("Sending request for /healthz ...")

	resp, err := s.Daemon.GetHealthz(nil)
	if err != nil {
		return nil, fmt.Errorf("unable to get agent health: %s", err)
	}
	log.Debug("Got cilium /healthz")

	if resp == nil || resp.Payload == nil || resp.Payload.Cluster == nil {
		return nil, fmt.Errorf("received nil health response")
	}

	if resp.Payload.Cluster.Self != "" {
		s.localStatus = &healthModels.SelfStatus{
			Name: resp.Payload.Cluster.Self,
		}
	}

	nodes := make(nodeMap)
	for _, n := range resp.Payload.Cluster.Nodes {
		if n.PrimaryAddress != nil {
			if n.PrimaryAddress.IPV4 != nil {
				nodes[ipString(n.PrimaryAddress.IPV4.IP)] = NewHealthNode(n)
			}
			if n.PrimaryAddress.IPV6 != nil {
				nodes[ipString(n.PrimaryAddress.IPV6.IP)] = NewHealthNode(n)
			}
		}
		for _, addr := range n.SecondaryAddresses {
			nodes[ipString(addr.IP)] = NewHealthNode(n)
		}
		if n.HealthEndpointAddress != nil {
			if n.HealthEndpointAddress.IPV4 != nil {
				nodes[ipString(n.HealthEndpointAddress.IPV4.IP)] = NewHealthNode(n)
			}
			if n.HealthEndpointAddress.IPV6 != nil {
				nodes[ipString(n.HealthEndpointAddress.IPV6.IP)] = NewHealthNode(n)
			}
		}
	}
	return nodes, nil
}

// updateCluster makes the specified health report visible to the API.
//
// It only updates the server's API-visible health report if the provided
// report started after the current report.
func (s *Server) updateCluster(report *healthReport) {
	s.Lock()
	defer s.Unlock()

	if s.connectivity.startTime.Before(report.startTime) {
		s.connectivity = report
	}
}

// GetStatusResponse returns the most recent cluster connectivity status.
func (s *Server) GetStatusResponse() *healthModels.HealthStatusResponse {
	s.RLock()
	defer s.RUnlock()

	var name string
	// Check if localStatus is populated already. If not, the name is empty
	if s.localStatus != nil {
		name = s.localStatus.Name
	}

	return &healthModels.HealthStatusResponse{
		Local: &healthModels.SelfStatus{
			Name: name,
		},
		Nodes:     s.connectivity.nodes,
		Timestamp: s.connectivity.startTime.Format(time.RFC3339),
	}
}

// FetchStatusResponse updates the cluster with the latest set of nodes,
// runs a synchronous probe across the cluster, updates the connectivity cache
// and returns the results.
func (s *Server) FetchStatusResponse() (*healthModels.HealthStatusResponse, error) {
	nodes, err := s.getNodes()
	if err != nil {
		return nil, err
	}

	prober := newProber(s, nodes)
	if err := prober.Run(); err != nil {
		log.WithError(err).Info("Failed to run ping")
		return nil, err
	}
	log.Debug("Run complete")
	s.updateCluster(prober.getResults())

	return s.GetStatusResponse(), nil
}

// Run services that are not considered 'Passive': Actively probing other
// hosts and endpoints over ICMP and HTTP, and hosting a local Unix socket.
// Blocks indefinitely, or returns any errors that occur hosting the Unix
// socket API server.
func (s *Server) runActiveServices() error {
	// Run it once at the start so we get some initial status
	s.FetchStatusResponse()

	nodes, _ := s.getNodes()
	prober := newProber(s, nodes)
	prober.MaxRTT = s.ProbeInterval
	prober.OnIdle = func() {
		// Fetch results and update set of nodes to probe every
		// ProbeInterval
		s.updateCluster(prober.getResults())
		if nodes, err := s.getNodes(); err == nil {
			prober.setNodes(nodes)
		}
	}
	prober.RunLoop()
	defer prober.Stop()

	return s.Server.Serve()
}

// Serve spins up the following goroutines:
// * TCP API Server: Responders to the health API "/hello" message, one per path
//
// Also, if "Passive" is not set in s.Config:
// * Prober: Periodically run pings across the cluster at a configured interval
//   and update the server's connectivity status cache.
// * Unix API Server: Handle all health API requests over a unix socket.
//
// Callers should first defer the Server.Shutdown(), then call Serve().
func (s *Server) Serve() (err error) {
	errors := make(chan error)

	for i := range s.tcpServers {
		srv := s.tcpServers[i]
		go func() {
			errors <- srv.Serve()
		}()
	}

	if !s.Config.Passive {
		go func() {
			errors <- s.runActiveServices()
		}()
	}

	// Block for the first error, then return.
	err = <-errors
	return err
}

// Shutdown server and clean up resources
func (s *Server) Shutdown() {
	for i := range s.tcpServers {
		s.tcpServers[i].Shutdown()
	}
	if !s.Config.Passive {
		s.Server.Shutdown()
	}
}

func enableAPI(opt AdminOption, tcpPort int) bool {
	switch opt {
	case AdminOptionAny:
		return true
	case AdminOptionUnix:
		return tcpPort == 0
	default:
		return false
	}
}

// newServer instantiates a new instance of the API that serves the health
// API on the specified port. If tcpPort is 0, then a unix socket is opened
// which serves the entire API. If a tcpPort is specified, then it returns
// a server which only answers get requests for the root URL "/".
func (s *Server) newServer(spec *loads.Document, tcpPort int) *healthApi.Server {
	api := restapi.NewCiliumHealthAPI(spec)
	api.Logger = log.Printf

	// /hello
	api.GetHelloHandler = NewGetHelloHandler(s)

	if enableAPI(s.Config.Admin, tcpPort) {
		api.GetHealthzHandler = NewGetHealthzHandler(s)
		api.ConnectivityGetStatusHandler = NewGetStatusHandler(s)
		api.ConnectivityPutStatusProbeHandler = NewPutStatusProbeHandler(s)
	}

	srv := healthApi.NewServer(api)
	if tcpPort == 0 {
		srv.EnabledListeners = []string{"unix"}
		srv.SocketPath = flags.Filename(defaults.SockPath)
	} else {
		srv.EnabledListeners = []string{"http"}
		srv.Port = tcpPort
		srv.Host = "" // FIXME: Allow binding to specific IPs
	}
	srv.ConfigureAPI()

	return srv
}

// NewServer creates a server to handle health requests.
func NewServer(config Config) (*Server, error) {
	server := &Server{
		startTime:    time.Now(),
		Config:       config,
		tcpServers:   []*healthApi.Server{},
		connectivity: &healthReport{},
	}

	swaggerSpec, err := loads.Analyzed(healthApi.SwaggerJSON, "")
	if err != nil {
		return nil, err
	}

	if !config.Passive {
		cl, err := ciliumPkg.NewClient(config.CiliumURI)
		if err != nil {
			return nil, err
		}

		server.Client = cl
		server.Server = *server.newServer(swaggerSpec, 0)
	}
	for port := range PortToPaths {
		srv := server.newServer(swaggerSpec, port)
		server.tcpServers = append(server.tcpServers, srv)
	}

	return server, nil
}
<|endoftext|>


<commit_before>package redisconn

import (
	"bufio"
	"net"
	"time"

	"github.com/ngaut/deadline"
)

//not thread-safe
type Conn struct {
	addr       string
	nc         net.Conn
	closed     bool
	r          *bufio.Reader
	w          *bufio.Writer
	netTimeout int //second
}

const defaultBufSize = 4096

func NewConnection(addr string, netTimeout int) (*Conn, error) {
	return NewConnectionWithSize(addr, netTimeout, defaultBufSize, defaultBufSize)
}

func NewConnectionWithSize(addr string, netTimeout int, readSize int, writeSize int) (*Conn, error) {
	conn, err := net.DialTimeout("tcp", addr, time.Duration(netTimeout)*time.Second)
	if err != nil {
		return nil, err
	}

	return &Conn{
		addr:       addr,
		nc:         conn,
		closed:     false,
		r:          bufio.NewReaderSize(conn, readSize),
		w:          bufio.NewWriterSize(deadline.NewDeadlineWriter(conn, time.Duration(netTimeout)*time.Second), writeSize),
		netTimeout: netTimeout,
	}, nil
}

//require read to use bufio
func (c *Conn) Read(p []byte) (int, error) {
	panic("not allowed")
}

func (c *Conn) Flush() error {
	return c.w.Flush()
}

func (c *Conn) Write(p []byte) (int, error) {
	return c.w.Write(p)
}

func (c *Conn) BufioReader() *bufio.Reader {
	return c.r
}

func (c *Conn) SetWriteDeadline(t time.Time) error {
	return c.nc.SetWriteDeadline(t)
}

func (c *Conn) SetReadDeadline(t time.Time) error {
	return c.nc.SetReadDeadline(t)
}

func (c *Conn) SetDeadline(t time.Time) error {
	return c.nc.SetDeadline(t)
}

func (c *Conn) Close() {
	c.closed = true
	c.nc.Close()
}
<commit_msg>let out do write deadline itself<commit_after>package redisconn

import (
	"bufio"
	"net"
	"time"
)

//not thread-safe
type Conn struct {
	addr       string
	nc         net.Conn
	closed     bool
	r          *bufio.Reader
	w          *bufio.Writer
	netTimeout int //second
}

const defaultBufSize = 4096

func NewConnection(addr string, netTimeout int) (*Conn, error) {
	return NewConnectionWithSize(addr, netTimeout, defaultBufSize, defaultBufSize)
}

func NewConnectionWithSize(addr string, netTimeout int, readSize int, writeSize int) (*Conn, error) {
	conn, err := net.DialTimeout("tcp", addr, time.Duration(netTimeout)*time.Second)
	if err != nil {
		return nil, err
	}

	return &Conn{
		addr:       addr,
		nc:         conn,
		closed:     false,
		r:          bufio.NewReaderSize(conn, readSize),
		w:          bufio.NewWriterSize(conn, writeSize),
		netTimeout: netTimeout,
	}, nil
}

//require read to use bufio
func (c *Conn) Read(p []byte) (int, error) {
	panic("not allowed")
}

func (c *Conn) Flush() error {
	return c.w.Flush()
}

func (c *Conn) Write(p []byte) (int, error) {
	return c.w.Write(p)
}

func (c *Conn) BufioReader() *bufio.Reader {
	return c.r
}

func (c *Conn) SetWriteDeadline(t time.Time) error {
	return c.nc.SetWriteDeadline(t)
}

func (c *Conn) SetReadDeadline(t time.Time) error {
	return c.nc.SetReadDeadline(t)
}

func (c *Conn) SetDeadline(t time.Time) error {
	return c.nc.SetDeadline(t)
}

func (c *Conn) Close() {
	if c.closed {
		return
	}
	c.closed = true
	c.nc.Close()
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Skaffold 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 deploy

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	// k8syaml "k8s.io/apimachinery/pkg/util/yaml"
	// "k8s.io/client-go/kubernetes/scheme"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build/tag"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/v1alpha2"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	yaml "gopkg.in/yaml.v2"
)

type HelmDeployer struct {
	*v1alpha2.DeployConfig
	kubeContext string
	namespace   string
}

// NewHelmDeployer returns a new HelmDeployer for a DeployConfig filled
// with the needed configuration for `helm`
func NewHelmDeployer(cfg *v1alpha2.DeployConfig, kubeContext string, namespace string) *HelmDeployer {
	return &HelmDeployer{
		DeployConfig: cfg,
		kubeContext:  kubeContext,
		namespace:    namespace,
	}
}

func (h *HelmDeployer) Labels() map[string]string {
	return map[string]string{
		constants.Labels.Deployer: "helm",
	}
}

func (h *HelmDeployer) Deploy(ctx context.Context, out io.Writer, builds []build.Artifact) ([]Artifact, error) {
	deployResults := []Artifact{}
	for _, r := range h.HelmDeploy.Releases {
		results, err := h.deployRelease(out, r, builds)
		if err != nil {
			releaseName, _ := evaluateReleaseName(r.Name)
			return deployResults, errors.Wrapf(err, "deploying %s", releaseName)
		}
		deployResults = append(deployResults, results...)
	}
	return deployResults, nil
}

// Not implemented
func (h *HelmDeployer) Dependencies() ([]string, error) {
	return nil, nil
}

// Cleanup deletes what was deployed by calling Deploy.
func (h *HelmDeployer) Cleanup(ctx context.Context, out io.Writer) error {
	for _, r := range h.HelmDeploy.Releases {
		if err := h.deleteRelease(out, r); err != nil {
			releaseName, _ := evaluateReleaseName(r.Name)
			return errors.Wrapf(err, "deploying %s", releaseName)
		}
	}
	return nil
}

func (h *HelmDeployer) helm(out io.Writer, arg ...string) error {
	args := append([]string{"--kube-context", h.kubeContext}, arg...)

	cmd := exec.Command("helm", args...)
	cmd.Stdout = out
	cmd.Stderr = out

	return util.RunCmd(cmd)
}

func (h *HelmDeployer) deployRelease(out io.Writer, r v1alpha2.HelmRelease, builds []build.Artifact) ([]Artifact, error) {
	isInstalled := true

	releaseName, err := evaluateReleaseName(r.Name)
	if err != nil {
		return nil, errors.Wrap(err, "cannot parse the release name template")
	}
	if err := h.helm(out, "get", releaseName); err != nil {
		fmt.Fprintf(out, "Helm release %s not installed. Installing...\n", releaseName)
		isInstalled = false
	}
	params, err := JoinTagsToBuildResult(builds, r.Values)
	if err != nil {
		return nil, errors.Wrap(err, "matching build results to chart values")
	}

	var setOpts []string
	for k, v := range params {
		setOpts = append(setOpts, "--set")
		setOpts = append(setOpts, fmt.Sprintf("%s=%s", k, v.Tag))
	}

	// First build dependencies.
	logrus.Infof("Building helm dependencies...")
	if err := h.helm(out, "dep", "build", r.ChartPath); err != nil {
		return nil, errors.Wrap(err, "building helm dependencies")
	}

	var args []string
	if !isInstalled {
		args = append(args, "install", "--name", releaseName)
	} else {
		args = append(args, "upgrade", releaseName)
	}

	// There are 2 strategies:
	// 1) Deploy chart directly from filesystem path or from repository
	//    (like stable/kubernetes-dashboard). Version only applies to a
	//    chart from repository.
	// 2) Package chart into a .tgz archive with specific version and then deploy
	//    that packaged chart. This way user can apply any version and appVersion
	//    for the chart.
	if r.Packaged == nil {
		if r.Version != "" {
			args = append(args, "--version", r.Version)
		}
		args = append(args, r.ChartPath)
	} else {
		chartPath, err := h.packageChart(r)
		if err != nil {
			return nil, errors.WithMessage(err, "cannot package chart")
		}
		args = append(args, chartPath)
	}

	var ns string
	if h.namespace != "" {
		ns = h.namespace
	} else if r.Namespace != "" {
		ns = r.Namespace
	} else {
		ns = os.Getenv("SKAFFOLD_DEPLOY_NAMESPACE")
	}
	if ns != "" {
		args = append(args, "--namespace", ns)
	}
	if len(r.Overrides) != 0 {
		overrides, err := yaml.Marshal(r.Overrides)
		if err != nil {
			return nil, errors.Wrap(err, "cannot marshal overrides to create overrides values.yaml")
		}
		overridesFile, err := os.Create(constants.HelmOverridesFilename)
		if err != nil {
			return nil, errors.Wrapf(err, "cannot create file %s", constants.HelmOverridesFilename)
		}
		defer func() {
			overridesFile.Close()
			os.Remove(constants.HelmOverridesFilename)
		}()
		if _, err := overridesFile.WriteString(string(overrides)); err != nil {
			return nil, errors.Wrapf(err, "failed to write file %s", constants.HelmOverridesFilename)
		}
		args = append(args, "-f", constants.HelmOverridesFilename)
	}
	if r.ValuesFilePath != "" {
		args = append(args, "-f", r.ValuesFilePath)
	}

	setValues := r.SetValues
	if setValues == nil {
		setValues = map[string]string{}
	}
	if len(r.SetValueTemplates) != 0 {
		envMap := map[string]string{}
		for idx, b := range builds {
			suffix := ""
			if idx > 0 {
				suffix = strconv.Itoa(idx + 1)
			}
			m := tag.CreateEnvVarMap(b.ImageName, extractTag(b.Tag))
			for k, v := range m {
				envMap[k+suffix] = v
			}
			fmt.Printf("EnvVarMap: %#v\n", envMap)
		}
		for k, v := range r.SetValueTemplates {
			t, err := util.ParseEnvTemplate(v)
			if err != nil {
				return nil, errors.Wrapf(err, "failed to parse setValueTemplates")
			}
			result, err := util.ExecuteEnvTemplate(t, envMap)
			if err != nil {
				return nil, errors.Wrapf(err, "failed to generate setValueTemplates")
			} else {
				setValues[k] = result
			}
		}
	}
	for k, v := range setValues {
		setOpts = append(setOpts, "--set")
		setOpts = append(setOpts, fmt.Sprintf("%s=%s", k, v))
	}
	if r.Wait {
		args = append(args, "--wait")
	}
	args = append(args, setOpts...)

	helmErr := h.helm(out, args...)
	return h.getDeployResults(ns, r.Name), helmErr
}

// imageName if the given string includes a fully qualified docker image name then lets trim just the tag part out
func extractTag(imageName string) string {
	idx := strings.LastIndex(imageName, "/")
	if idx < 0 {
		return imageName
	}
	tag := imageName[idx+1:]
	idx = strings.Index(tag, ":")
	if idx > 0 {
		return tag[idx+1:]
	}
	return tag
}

// packageChart packages the chart and returns path to the chart archive file.
// If this function returns an error, it will always be wrapped.
func (h *HelmDeployer) packageChart(r v1alpha2.HelmRelease) (string, error) {
	tmp := os.TempDir()
	packageArgs := []string{"package", r.ChartPath, "--destination", tmp}
	if r.Packaged.Version != "" {
		v, err := concretize(r.Packaged.Version)
		if err != nil {
			return "", errors.Wrap(err, `concretize "packaged.version" template`)
		}
		packageArgs = append(packageArgs, "--version", v)
	}
	if r.Packaged.AppVersion != "" {
		av, err := concretize(r.Packaged.AppVersion)
		if err != nil {
			return "", errors.Wrap(err, `concretize "packaged.appVersion" template`)
		}
		packageArgs = append(packageArgs, "--app-version", av)
	}

	buf := &bytes.Buffer{}
	err := h.helm(buf, packageArgs...)
	output := strings.TrimSpace(buf.String())
	if err != nil {
		return "", errors.Wrapf(err, "package chart into a .tgz archive (%s)", output)
	}

	fpath, err := extractChartFilename(output, tmp)
	if err != nil {
		return "", err
	}

	return filepath.Join(tmp, fpath), nil
}

func (h *HelmDeployer) getReleaseInfo(release string) (*bufio.Reader, error) {
	var releaseInfo bytes.Buffer
	if err := h.helm(&releaseInfo, "get", release); err != nil {
		return nil, fmt.Errorf("error retrieving helm deployment info: %s", releaseInfo.String())
	}
	return bufio.NewReader(&releaseInfo), nil
}

// Retrieve info about all releases using helm get
// Skaffold labels will be applied to each deployed k8s object
// Since helm isn't always consistent with retrieving results, don't return errors here
func (h *HelmDeployer) getDeployResults(namespace string, release string) []Artifact {
	b, err := h.getReleaseInfo(release)
	if err != nil {
		logrus.Warnf(err.Error())
		return nil
	}
	return parseReleaseInfo(namespace, b)
}

func (h *HelmDeployer) deleteRelease(out io.Writer, r v1alpha2.HelmRelease) error {
	releaseName, err := evaluateReleaseName(r.Name)
	if err != nil {
		return errors.Wrap(err, "cannot parse the release name template")
	}

	if err := h.helm(out, "delete", releaseName, "--purge"); err != nil {
		logrus.Debugf("deleting release %s: %v\n", releaseName, err)
	}

	return nil
}

func evaluateReleaseName(nameTemplate string) (string, error) {
	tmpl, err := util.ParseEnvTemplate(nameTemplate)
	if err != nil {
		return "", errors.Wrap(err, "parsing template")
	}

	return util.ExecuteEnvTemplate(tmpl, nil)
}

// concretize parses and executes template s with OS environment variables.
// If s is not a template but a simple string, returns unchanged s.
func concretize(s string) (string, error) {
	tmpl, err := util.ParseEnvTemplate(s)
	if err != nil {
		return "", errors.Wrap(err, "parsing template")
	}

	tmpl.Option("missingkey=error")
	return util.ExecuteEnvTemplate(tmpl, nil)
}

func extractChartFilename(s, tmp string) (string, error) {
	s = strings.TrimSpace(s)
	idx := strings.Index(s, tmp)
	if idx == -1 {
		return "", errors.New("cannot locate packaged chart archive")
	}

	return s[idx+len(tmp):], nil
}
<commit_msg>fix: CI<commit_after>/*
Copyright 2018 The Skaffold 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 deploy

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	// k8syaml "k8s.io/apimachinery/pkg/util/yaml"
	// "k8s.io/client-go/kubernetes/scheme"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build/tag"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/v1alpha2"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	yaml "gopkg.in/yaml.v2"
)

type HelmDeployer struct {
	*v1alpha2.DeployConfig
	kubeContext string
	namespace   string
}

// NewHelmDeployer returns a new HelmDeployer for a DeployConfig filled
// with the needed configuration for `helm`
func NewHelmDeployer(cfg *v1alpha2.DeployConfig, kubeContext string, namespace string) *HelmDeployer {
	return &HelmDeployer{
		DeployConfig: cfg,
		kubeContext:  kubeContext,
		namespace:    namespace,
	}
}

func (h *HelmDeployer) Labels() map[string]string {
	return map[string]string{
		constants.Labels.Deployer: "helm",
	}
}

func (h *HelmDeployer) Deploy(ctx context.Context, out io.Writer, builds []build.Artifact) ([]Artifact, error) {
	deployResults := []Artifact{}
	for _, r := range h.HelmDeploy.Releases {
		results, err := h.deployRelease(out, r, builds)
		if err != nil {
			releaseName, _ := evaluateReleaseName(r.Name)
			return deployResults, errors.Wrapf(err, "deploying %s", releaseName)
		}
		deployResults = append(deployResults, results...)
	}
	return deployResults, nil
}

// Not implemented
func (h *HelmDeployer) Dependencies() ([]string, error) {
	return nil, nil
}

// Cleanup deletes what was deployed by calling Deploy.
func (h *HelmDeployer) Cleanup(ctx context.Context, out io.Writer) error {
	for _, r := range h.HelmDeploy.Releases {
		if err := h.deleteRelease(out, r); err != nil {
			releaseName, _ := evaluateReleaseName(r.Name)
			return errors.Wrapf(err, "deploying %s", releaseName)
		}
	}
	return nil
}

func (h *HelmDeployer) helm(out io.Writer, arg ...string) error {
	args := append([]string{"--kube-context", h.kubeContext}, arg...)

	cmd := exec.Command("helm", args...)
	cmd.Stdout = out
	cmd.Stderr = out

	return util.RunCmd(cmd)
}

func (h *HelmDeployer) deployRelease(out io.Writer, r v1alpha2.HelmRelease, builds []build.Artifact) ([]Artifact, error) {
	isInstalled := true

	releaseName, err := evaluateReleaseName(r.Name)
	if err != nil {
		return nil, errors.Wrap(err, "cannot parse the release name template")
	}
	if err := h.helm(out, "get", releaseName); err != nil {
		fmt.Fprintf(out, "Helm release %s not installed. Installing...\n", releaseName)
		isInstalled = false
	}
	params, err := JoinTagsToBuildResult(builds, r.Values)
	if err != nil {
		return nil, errors.Wrap(err, "matching build results to chart values")
	}

	var setOpts []string
	for k, v := range params {
		setOpts = append(setOpts, "--set")
		setOpts = append(setOpts, fmt.Sprintf("%s=%s", k, v.Tag))
	}

	// First build dependencies.
	logrus.Infof("Building helm dependencies...")
	if err := h.helm(out, "dep", "build", r.ChartPath); err != nil {
		return nil, errors.Wrap(err, "building helm dependencies")
	}

	var args []string
	if !isInstalled {
		args = append(args, "install", "--name", releaseName)
	} else {
		args = append(args, "upgrade", releaseName)
	}

	// There are 2 strategies:
	// 1) Deploy chart directly from filesystem path or from repository
	//    (like stable/kubernetes-dashboard). Version only applies to a
	//    chart from repository.
	// 2) Package chart into a .tgz archive with specific version and then deploy
	//    that packaged chart. This way user can apply any version and appVersion
	//    for the chart.
	if r.Packaged == nil {
		if r.Version != "" {
			args = append(args, "--version", r.Version)
		}
		args = append(args, r.ChartPath)
	} else {
		chartPath, err := h.packageChart(r)
		if err != nil {
			return nil, errors.WithMessage(err, "cannot package chart")
		}
		args = append(args, chartPath)
	}

	var ns string
	if h.namespace != "" {
		ns = h.namespace
	} else if r.Namespace != "" {
		ns = r.Namespace
	} else {
		ns = os.Getenv("SKAFFOLD_DEPLOY_NAMESPACE")
	}
	if ns != "" {
		args = append(args, "--namespace", ns)
	}
	if len(r.Overrides) != 0 {
		overrides, err := yaml.Marshal(r.Overrides)
		if err != nil {
			return nil, errors.Wrap(err, "cannot marshal overrides to create overrides values.yaml")
		}
		overridesFile, err := os.Create(constants.HelmOverridesFilename)
		if err != nil {
			return nil, errors.Wrapf(err, "cannot create file %s", constants.HelmOverridesFilename)
		}
		defer func() {
			overridesFile.Close()
			os.Remove(constants.HelmOverridesFilename)
		}()
		if _, err := overridesFile.WriteString(string(overrides)); err != nil {
			return nil, errors.Wrapf(err, "failed to write file %s", constants.HelmOverridesFilename)
		}
		args = append(args, "-f", constants.HelmOverridesFilename)
	}
	if r.ValuesFilePath != "" {
		args = append(args, "-f", r.ValuesFilePath)
	}

	setValues := r.SetValues
	if setValues == nil {
		setValues = map[string]string{}
	}
	if len(r.SetValueTemplates) != 0 {
		envMap := map[string]string{}
		for idx, b := range builds {
			suffix := ""
			if idx > 0 {
				suffix = strconv.Itoa(idx + 1)
			}
			m := tag.CreateEnvVarMap(b.ImageName, extractTag(b.Tag))
			for k, v := range m {
				envMap[k+suffix] = v
			}
			fmt.Printf("EnvVarMap: %#v\n", envMap)
		}
		for k, v := range r.SetValueTemplates {
			t, err := util.ParseEnvTemplate(v)
			if err != nil {
				return nil, errors.Wrapf(err, "failed to parse setValueTemplates")
			}
			result, err := util.ExecuteEnvTemplate(t, envMap)
			if err != nil {
				return nil, errors.Wrapf(err, "failed to generate setValueTemplates")
			}
			setValues[k] = result
		}
	}
	for k, v := range setValues {
		setOpts = append(setOpts, "--set")
		setOpts = append(setOpts, fmt.Sprintf("%s=%s", k, v))
	}
	if r.Wait {
		args = append(args, "--wait")
	}
	args = append(args, setOpts...)

	helmErr := h.helm(out, args...)
	return h.getDeployResults(ns, r.Name), helmErr
}

// imageName if the given string includes a fully qualified docker image name then lets trim just the tag part out
func extractTag(imageName string) string {
	idx := strings.LastIndex(imageName, "/")
	if idx < 0 {
		return imageName
	}
	tag := imageName[idx+1:]
	idx = strings.Index(tag, ":")
	if idx > 0 {
		return tag[idx+1:]
	}
	return tag
}

// packageChart packages the chart and returns path to the chart archive file.
// If this function returns an error, it will always be wrapped.
func (h *HelmDeployer) packageChart(r v1alpha2.HelmRelease) (string, error) {
	tmp := os.TempDir()
	packageArgs := []string{"package", r.ChartPath, "--destination", tmp}
	if r.Packaged.Version != "" {
		v, err := concretize(r.Packaged.Version)
		if err != nil {
			return "", errors.Wrap(err, `concretize "packaged.version" template`)
		}
		packageArgs = append(packageArgs, "--version", v)
	}
	if r.Packaged.AppVersion != "" {
		av, err := concretize(r.Packaged.AppVersion)
		if err != nil {
			return "", errors.Wrap(err, `concretize "packaged.appVersion" template`)
		}
		packageArgs = append(packageArgs, "--app-version", av)
	}

	buf := &bytes.Buffer{}
	err := h.helm(buf, packageArgs...)
	output := strings.TrimSpace(buf.String())
	if err != nil {
		return "", errors.Wrapf(err, "package chart into a .tgz archive (%s)", output)
	}

	fpath, err := extractChartFilename(output, tmp)
	if err != nil {
		return "", err
	}

	return filepath.Join(tmp, fpath), nil
}

func (h *HelmDeployer) getReleaseInfo(release string) (*bufio.Reader, error) {
	var releaseInfo bytes.Buffer
	if err := h.helm(&releaseInfo, "get", release); err != nil {
		return nil, fmt.Errorf("error retrieving helm deployment info: %s", releaseInfo.String())
	}
	return bufio.NewReader(&releaseInfo), nil
}

// Retrieve info about all releases using helm get
// Skaffold labels will be applied to each deployed k8s object
// Since helm isn't always consistent with retrieving results, don't return errors here
func (h *HelmDeployer) getDeployResults(namespace string, release string) []Artifact {
	b, err := h.getReleaseInfo(release)
	if err != nil {
		logrus.Warnf(err.Error())
		return nil
	}
	return parseReleaseInfo(namespace, b)
}

func (h *HelmDeployer) deleteRelease(out io.Writer, r v1alpha2.HelmRelease) error {
	releaseName, err := evaluateReleaseName(r.Name)
	if err != nil {
		return errors.Wrap(err, "cannot parse the release name template")
	}

	if err := h.helm(out, "delete", releaseName, "--purge"); err != nil {
		logrus.Debugf("deleting release %s: %v\n", releaseName, err)
	}

	return nil
}

func evaluateReleaseName(nameTemplate string) (string, error) {
	tmpl, err := util.ParseEnvTemplate(nameTemplate)
	if err != nil {
		return "", errors.Wrap(err, "parsing template")
	}

	return util.ExecuteEnvTemplate(tmpl, nil)
}

// concretize parses and executes template s with OS environment variables.
// If s is not a template but a simple string, returns unchanged s.
func concretize(s string) (string, error) {
	tmpl, err := util.ParseEnvTemplate(s)
	if err != nil {
		return "", errors.Wrap(err, "parsing template")
	}

	tmpl.Option("missingkey=error")
	return util.ExecuteEnvTemplate(tmpl, nil)
}

func extractChartFilename(s, tmp string) (string, error) {
	s = strings.TrimSpace(s)
	idx := strings.Index(s, tmp)
	if idx == -1 {
		return "", errors.New("cannot locate packaged chart archive")
	}

	return s[idx+len(tmp):], nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Skaffold 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 util

import (
	"context"
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/spf13/afero"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/output/log"
)

var stdin []byte
var Fs = afero.NewOsFs()

// ReadConfiguration reads a `skaffold.yaml` configuration and
// returns its content.
func ReadConfiguration(filename string) ([]byte, error) {
	switch {
	case filename == "":
		return nil, errors.New("filename not specified")
	case filename == "-":
		if len(stdin) == 0 {
			var err error
			stdin, err = ioutil.ReadAll(os.Stdin)
			if err != nil {
				return []byte{}, err
			}
		}
		return stdin, nil
	case IsURL(filename):
		return Download(filename)
	default:
		fp := filename
		if !filepath.IsAbs(fp) {
			dir, err := os.Getwd()
			if err != nil {
				return []byte{}, err
			}
			fp = filepath.Join(dir, fp)
		}
		contents, err := afero.ReadFile(Fs, fp)
		if err != nil {
			// If the config file is the default `skaffold.yaml`,
			// then we also try to read `skaffold.yml`.
			if filename == "skaffold.yaml" {
				log.Entry(context.TODO()).Infof("Could not open skaffold.yaml: \"%s\"", err)
				log.Entry(context.TODO()).Info("Trying to read from skaffold.yml instead")
				contents, errIgnored := afero.ReadFile(Fs, filepath.Join(filepath.Dir(fp), "skaffold.yml"))
				if errIgnored != nil {
					// Return original error because it's the one that matters
					return nil, err
				}

				return contents, nil
			}
		}

		return contents, err
	}
}

func ReadFile(filename string) ([]byte, error) {
	if !filepath.IsAbs(filename) {
		dir, err := os.Getwd()
		if err != nil {
			return []byte{}, err
		}
		filename = filepath.Join(dir, filename)
	}
	return afero.ReadFile(Fs, filename)
}
<commit_msg>fix: add godoc to Fs var in config.go (#7004)<commit_after>/*
Copyright 2019 The Skaffold 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 util

import (
	"context"
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/spf13/afero"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/output/log"
)

// Fs is the underlying filesystem to use for reading skaffold project files & configuration.  OS FS by default
var Fs = afero.NewOsFs()

var stdin []byte

// ReadConfiguration reads a `skaffold.yaml` configuration and
// returns its content.
func ReadConfiguration(filename string) ([]byte, error) {
	switch {
	case filename == "":
		return nil, errors.New("filename not specified")
	case filename == "-":
		if len(stdin) == 0 {
			var err error
			stdin, err = ioutil.ReadAll(os.Stdin)
			if err != nil {
				return []byte{}, err
			}
		}
		return stdin, nil
	case IsURL(filename):
		return Download(filename)
	default:
		fp := filename
		if !filepath.IsAbs(fp) {
			dir, err := os.Getwd()
			if err != nil {
				return []byte{}, err
			}
			fp = filepath.Join(dir, fp)
		}
		contents, err := afero.ReadFile(Fs, fp)
		if err != nil {
			// If the config file is the default `skaffold.yaml`,
			// then we also try to read `skaffold.yml`.
			if filename == "skaffold.yaml" {
				log.Entry(context.TODO()).Infof("Could not open skaffold.yaml: \"%s\"", err)
				log.Entry(context.TODO()).Info("Trying to read from skaffold.yml instead")
				contents, errIgnored := afero.ReadFile(Fs, filepath.Join(filepath.Dir(fp), "skaffold.yml"))
				if errIgnored != nil {
					// Return original error because it's the one that matters
					return nil, err
				}

				return contents, nil
			}
		}

		return contents, err
	}
}

func ReadFile(filename string) ([]byte, error) {
	if !filepath.IsAbs(filename) {
		dir, err := os.Getwd()
		if err != nil {
			return []byte{}, err
		}
		filename = filepath.Join(dir, filename)
	}
	return afero.ReadFile(Fs, filename)
}
<|endoftext|>


<commit_before>// Copyright 2020 The SQLFlow 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 sql

import (
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
	"sqlflow.org/sqlflow/pkg/database"
	"sqlflow.org/sqlflow/pkg/pipe"
)

const (
	testTrainSelectWithLimit = `
SELECT * FROM iris.train
liMIT 10
TO TRAIN xgboost.gbtree
WITH
    objective="multi:softprob",
    train.num_boost_round = 30,
    eta = 0.4,
    num_class = 3
COLUMN sepal_length, sepal_width, petal_length, petal_width
LABEL class
INTO sqlflow_models.my_xgboost_model;
`
	testTrainSelectIris = testSelectIris + `
TO TRAIN DNNClassifier
WITH
  model.n_classes = 3,
  model.hidden_units = [10, 20]
COLUMN sepal_length, sepal_width, petal_length, petal_width
LABEL class
INTO sqlflow_models.my_dnn_model;
`
	testPredictSelectIris = `
SELECT *
FROM iris.test
TO PREDICT iris.predict.class
USING sqlflow_models.my_dnn_model;
`
	testClusteringTrain = `SELECT sepal_length, sepal_width, petal_length, petal_width
FROM iris.train
TO TRAIN sqlflow_models.DeepEmbeddingClusterModel
WITH
  model.pretrain_dims = [10,10,3],
  model.n_clusters = 3,
  model.pretrain_epochs=5,
  train.batch_size=10,
  train.verbose=1
INTO sqlflow_models.my_clustering_model;
`
	testClusteringPredict = `
SELECT sepal_length, sepal_width, petal_length, petal_width
FROM iris.test
TO PREDICT iris.predict.class
USING sqlflow_models.my_clustering_model;
`
	testXGBoostTrainSelectIris = `
SELECT *
FROM iris.train
TO TRAIN xgboost.gbtree
WITH
    objective="multi:softprob",
    train.num_boost_round = 30,
    eta = 0.4,
    num_class = 3
COLUMN sepal_length, sepal_width, petal_length, petal_width
LABEL class
INTO sqlflow_models.my_xgboost_model;
`
	testExplainTreeModelSelectIris = `
SELECT * FROM iris.train
TO EXPLAIN sqlflow_models.my_xgboost_model
USING TreeExplainer;
`
	testXGBoostPredictIris = `
SELECT *
FROM iris.test
TO PREDICT iris.predict.class
USING sqlflow_models.my_xgboost_model;
`
	testXGBoostTrainSelectHousing = `
SELECT *
FROM housing.train
TO TRAIN xgboost.gbtree
WITH
	objective="reg:squarederror",
	train.num_boost_round = 30,
	validation.select="SELECT * FROM housing.train LIMIT 20"
COLUMN f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13
LABEL target
INTO sqlflow_models.my_xgb_regression_model;
`
	testXGBoostPredictHousing = `
SELECT *
FROM housing.test
TO PREDICT housing.xgb_predict.target
USING sqlflow_models.my_xgb_regression_model;
`
)

func TestRunSQLProgram(t *testing.T) {
	a := assert.New(t)
	modelDir := ""
	a.NotPanics(func() {
		stream := RunSQLProgram(`
SELECT sepal_length as sl, sepal_width as sw, class FROM iris.train
TO TRAIN xgboost.gbtree
WITH
    objective="multi:softprob",
    train.num_boost_round = 30,
    eta = 0.4,
    num_class = 3
LABEL class
INTO sqlflow_models.my_xgboost_model_by_program;

SELECT sepal_length as sl, sepal_width as sw FROM iris.test
TO PREDICT iris.predict.class
USING sqlflow_models.my_xgboost_model_by_program;

SELECT sepal_length as sl, sepal_width as sw, class FROM iris.train
TO EXPLAIN sqlflow_models.my_xgboost_model_by_program
USING TreeExplainer;
`, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecuteXGBoostClassifier(t *testing.T) {
	a := assert.New(t)
	modelDir := ""
	a.NotPanics(func() {
		stream := RunSQLProgram(testTrainSelectWithLimit, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testXGBoostPredictIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
	a.NotPanics(func() {
		stream := RunSQLProgram(testXGBoostTrainSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testExplainTreeModelSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testXGBoostPredictIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecuteXGBoostRegression(t *testing.T) {
	a := assert.New(t)
	modelDir := ""
	a.NotPanics(func() {
		stream := RunSQLProgram(testXGBoostTrainSelectHousing, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testExplainTreeModelSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testXGBoostPredictHousing, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecutorTrainAndPredictDNN(t *testing.T) {
	a := assert.New(t)
	modelDir := ""
	a.NotPanics(func() {
		stream := RunSQLProgram(testTrainSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testPredictSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecutorTrainAndPredictClusteringLocalFS(t *testing.T) {
	a := assert.New(t)
	modelDir, e := ioutil.TempDir("/tmp", "sqlflow_models")
	a.Nil(e)
	defer os.RemoveAll(modelDir)
	a.NotPanics(func() {
		stream := RunSQLProgram(testClusteringTrain, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testClusteringPredict, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecutorTrainAndPredictDNNLocalFS(t *testing.T) {
	a := assert.New(t)
	modelDir, e := ioutil.TempDir("/tmp", "sqlflow_models")
	a.Nil(e)
	defer os.RemoveAll(modelDir)
	a.NotPanics(func() {
		stream := RunSQLProgram(testTrainSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testPredictSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecutorTrainAndPredictionDNNClassifierDENSE(t *testing.T) {
	if getEnv("SQLFLOW_TEST_DB", "mysql") == "hive" {
		t.Skip(fmt.Sprintf("%s: skip Hive test", getEnv("SQLFLOW_TEST_DB", "mysql")))
	}
	a := assert.New(t)
	a.NotPanics(func() {
		trainSQL := `SELECT * FROM iris.train_dense
TO TRAIN DNNClassifier
WITH
model.n_classes = 3,
model.hidden_units = [10, 20],
train.epoch = 200,
train.batch_size = 10,
train.verbose = 1
COLUMN NUMERIC(dense, 4)
LABEL class
INTO sqlflow_models.my_dense_dnn_model;`
		stream := RunSQLProgram(trainSQL, "", database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))

		predSQL := `SELECT * FROM iris.test_dense
TO PREDICT iris.predict_dense.class
USING sqlflow_models.my_dense_dnn_model
;`
		stream = RunSQLProgram(predSQL, "", database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestLogChanWriter_Write(t *testing.T) {
	a := assert.New(t)
	rd, wr := pipe.Pipe()
	go func() {
		defer wr.Close()
		cw := &logChanWriter{wr: wr}
		cw.Write([]byte("hello\n世界"))
		cw.Write([]byte("hello\n世界"))
		cw.Write([]byte("\n"))
		cw.Write([]byte("世界\n世界\n世界\n"))
	}()

	c := rd.ReadAll()

	a.Equal("hello\n", <-c)
	a.Equal("世界hello\n", <-c)
	a.Equal("世界\n", <-c)
	a.Equal("世界\n", <-c)
	a.Equal("世界\n", <-c)
	a.Equal("世界\n", <-c)
	_, more := <-c
	a.False(more)
}
<commit_msg>fix dec test loss got nan (#2175)<commit_after>// Copyright 2020 The SQLFlow 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 sql

import (
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
	"sqlflow.org/sqlflow/pkg/database"
	"sqlflow.org/sqlflow/pkg/pipe"
)

const (
	testTrainSelectWithLimit = `
SELECT * FROM iris.train
liMIT 10
TO TRAIN xgboost.gbtree
WITH
    objective="multi:softprob",
    train.num_boost_round = 30,
    eta = 0.4,
    num_class = 3
COLUMN sepal_length, sepal_width, petal_length, petal_width
LABEL class
INTO sqlflow_models.my_xgboost_model;
`
	testTrainSelectIris = testSelectIris + `
TO TRAIN DNNClassifier
WITH
  model.n_classes = 3,
  model.hidden_units = [10, 20]
COLUMN sepal_length, sepal_width, petal_length, petal_width
LABEL class
INTO sqlflow_models.my_dnn_model;
`
	testPredictSelectIris = `
SELECT *
FROM iris.test
TO PREDICT iris.predict.class
USING sqlflow_models.my_dnn_model;
`
	testClusteringTrain = `SELECT (sepal_length - 4.4) / 3.5 as sepal_length, (sepal_width - 2.0) / 2.2 as sepal_width, (petal_length - 1) / 5.9 as petal_length, (petal_width - 0.1) / 2.4 as petal_width
FROM iris.train
TO TRAIN sqlflow_models.DeepEmbeddingClusterModel
WITH
  model.pretrain_dims = [10,10,3],
  model.n_clusters = 3,
  model.pretrain_epochs=5,
  train.batch_size=10,
  train.verbose=1
INTO sqlflow_models.my_clustering_model;
`
	testClusteringPredict = `
SELECT (sepal_length - 4.4) / 3.5 as sepal_length, (sepal_width - 2.0) / 2.2 as sepal_width, (petal_length - 1) / 5.9 as petal_length, (petal_width - 0.1) / 2.4 as petal_width
FROM iris.test
TO PREDICT iris.predict.class
USING sqlflow_models.my_clustering_model;
`
	testXGBoostTrainSelectIris = `
SELECT *
FROM iris.train
TO TRAIN xgboost.gbtree
WITH
    objective="multi:softprob",
    train.num_boost_round = 30,
    eta = 0.4,
    num_class = 3
COLUMN sepal_length, sepal_width, petal_length, petal_width
LABEL class
INTO sqlflow_models.my_xgboost_model;
`
	testExplainTreeModelSelectIris = `
SELECT * FROM iris.train
TO EXPLAIN sqlflow_models.my_xgboost_model
USING TreeExplainer;
`
	testXGBoostPredictIris = `
SELECT *
FROM iris.test
TO PREDICT iris.predict.class
USING sqlflow_models.my_xgboost_model;
`
	testXGBoostTrainSelectHousing = `
SELECT *
FROM housing.train
TO TRAIN xgboost.gbtree
WITH
	objective="reg:squarederror",
	train.num_boost_round = 30,
	validation.select="SELECT * FROM housing.train LIMIT 20"
COLUMN f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13
LABEL target
INTO sqlflow_models.my_xgb_regression_model;
`
	testXGBoostPredictHousing = `
SELECT *
FROM housing.test
TO PREDICT housing.xgb_predict.target
USING sqlflow_models.my_xgb_regression_model;
`
)

func TestRunSQLProgram(t *testing.T) {
	a := assert.New(t)
	modelDir := ""
	a.NotPanics(func() {
		stream := RunSQLProgram(`
SELECT sepal_length as sl, sepal_width as sw, class FROM iris.train
TO TRAIN xgboost.gbtree
WITH
    objective="multi:softprob",
    train.num_boost_round = 30,
    eta = 0.4,
    num_class = 3
LABEL class
INTO sqlflow_models.my_xgboost_model_by_program;

SELECT sepal_length as sl, sepal_width as sw FROM iris.test
TO PREDICT iris.predict.class
USING sqlflow_models.my_xgboost_model_by_program;

SELECT sepal_length as sl, sepal_width as sw, class FROM iris.train
TO EXPLAIN sqlflow_models.my_xgboost_model_by_program
USING TreeExplainer;
`, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecuteXGBoostClassifier(t *testing.T) {
	a := assert.New(t)
	modelDir := ""
	a.NotPanics(func() {
		stream := RunSQLProgram(testTrainSelectWithLimit, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testXGBoostPredictIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
	a.NotPanics(func() {
		stream := RunSQLProgram(testXGBoostTrainSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testExplainTreeModelSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testXGBoostPredictIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecuteXGBoostRegression(t *testing.T) {
	a := assert.New(t)
	modelDir := ""
	a.NotPanics(func() {
		stream := RunSQLProgram(testXGBoostTrainSelectHousing, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testExplainTreeModelSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testXGBoostPredictHousing, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecutorTrainAndPredictDNN(t *testing.T) {
	a := assert.New(t)
	modelDir := ""
	a.NotPanics(func() {
		stream := RunSQLProgram(testTrainSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testPredictSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecutorTrainAndPredictClusteringLocalFS(t *testing.T) {
	a := assert.New(t)
	modelDir, e := ioutil.TempDir("/tmp", "sqlflow_models")
	a.Nil(e)
	defer os.RemoveAll(modelDir)
	a.NotPanics(func() {
		stream := RunSQLProgram(testClusteringTrain, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testClusteringPredict, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecutorTrainAndPredictDNNLocalFS(t *testing.T) {
	a := assert.New(t)
	modelDir, e := ioutil.TempDir("/tmp", "sqlflow_models")
	a.Nil(e)
	defer os.RemoveAll(modelDir)
	a.NotPanics(func() {
		stream := RunSQLProgram(testTrainSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
		stream = RunSQLProgram(testPredictSelectIris, modelDir, database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestExecutorTrainAndPredictionDNNClassifierDENSE(t *testing.T) {
	if getEnv("SQLFLOW_TEST_DB", "mysql") == "hive" {
		t.Skip(fmt.Sprintf("%s: skip Hive test", getEnv("SQLFLOW_TEST_DB", "mysql")))
	}
	a := assert.New(t)
	a.NotPanics(func() {
		trainSQL := `SELECT * FROM iris.train_dense
TO TRAIN DNNClassifier
WITH
model.n_classes = 3,
model.hidden_units = [10, 20],
train.epoch = 200,
train.batch_size = 10,
train.verbose = 1
COLUMN NUMERIC(dense, 4)
LABEL class
INTO sqlflow_models.my_dense_dnn_model;`
		stream := RunSQLProgram(trainSQL, "", database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))

		predSQL := `SELECT * FROM iris.test_dense
TO PREDICT iris.predict_dense.class
USING sqlflow_models.my_dense_dnn_model
;`
		stream = RunSQLProgram(predSQL, "", database.GetSessionFromTestingDB())
		a.True(goodStream(stream.ReadAll()))
	})
}

func TestLogChanWriter_Write(t *testing.T) {
	a := assert.New(t)
	rd, wr := pipe.Pipe()
	go func() {
		defer wr.Close()
		cw := &logChanWriter{wr: wr}
		cw.Write([]byte("hello\n世界"))
		cw.Write([]byte("hello\n世界"))
		cw.Write([]byte("\n"))
		cw.Write([]byte("世界\n世界\n世界\n"))
	}()

	c := rd.ReadAll()

	a.Equal("hello\n", <-c)
	a.Equal("世界hello\n", <-c)
	a.Equal("世界\n", <-c)
	a.Equal("世界\n", <-c)
	a.Equal("世界\n", <-c)
	a.Equal("世界\n", <-c)
	_, more := <-c
	a.False(more)
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the LGPLv3, see LICENCE file for details.

package charmstore

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sort"
	"time"

	"github.com/juju/errgo"
	jc "github.com/juju/testing/checkers"
	"gopkg.in/juju/charm.v3"
	"gopkg.in/juju/charm.v3/testing"
	gc "launchpad.net/gocheck"

	"github.com/juju/charmstore/internal/blobstore"
	"github.com/juju/charmstore/internal/mongodoc"
	"github.com/juju/charmstore/internal/storetesting"
	"github.com/juju/charmstore/params"
)

type StoreSuite struct {
	storetesting.IsolatedMgoSuite
}

var _ = gc.Suite(&StoreSuite{})

func (s *StoreSuite) checkAddCharm(c *gc.C, ch charm.Charm) {
	store, err := NewStore(s.Session.DB("foo"))
	c.Assert(err, gc.IsNil)
	url := mustParseReference("cs:precise/wordpress-23")

	// Add the charm to the store.
	beforeAdding := time.Now()
	err = store.AddCharmWithArchive(url, ch)
	c.Assert(err, gc.IsNil)
	afterAdding := time.Now()

	var doc mongodoc.Entity
	err = store.DB.Entities().FindId("cs:precise/wordpress-23").One(&doc)
	c.Assert(err, gc.IsNil)

	// The entity doc has been correctly added to the mongo collection.
	size, hash := mustGetSizeAndHash(ch)
	sort.Strings(doc.CharmProvidedInterfaces)
	sort.Strings(doc.CharmRequiredInterfaces)

	// Check the upload time and then reset it to its zero value
	// so that we can test the deterministic parts later.
	c.Assert(doc.UploadTime, jc.TimeBetween(beforeAdding, afterAdding))

	doc.UploadTime = time.Time{}

	blobName := doc.BlobName
	c.Assert(blobName, gc.Matches, "[0-9a-z]+")
	doc.BlobName = ""
	c.Assert(doc, jc.DeepEquals, mongodoc.Entity{
		URL:     url,
		BaseURL: mustParseReference("cs:wordpress"),
		BlobHash:                hash,
		Size:                    size,
		CharmMeta:               ch.Meta(),
		CharmActions:            ch.Actions(),
		CharmConfig:             ch.Config(),
		CharmProvidedInterfaces: []string{"http", "logging", "monitoring"},
		CharmRequiredInterfaces: []string{"mysql", "varnish"},
	})

	// The charm archive has been properly added to the blob store.
	r, obtainedSize, err := store.BlobStore.Open(blobName)
	c.Assert(err, gc.IsNil)
	c.Assert(obtainedSize, gc.Equals, size)
	data, err := ioutil.ReadAll(r)
	c.Assert(err, gc.IsNil)
	charmArchive, err := charm.ReadCharmArchiveBytes(data)
	c.Assert(err, gc.IsNil)
	c.Assert(charmArchive.Meta(), jc.DeepEquals, ch.Meta())
	c.Assert(charmArchive.Config(), jc.DeepEquals, ch.Config())
	c.Assert(charmArchive.Actions(), jc.DeepEquals, ch.Actions())
	c.Assert(charmArchive.Revision(), jc.DeepEquals, ch.Revision())

	// Try inserting the charm again - it should fail because the charm is
	// already there.
	err = store.AddCharmWithArchive(url, ch)
	c.Assert(errgo.Cause(err), gc.Equals, params.ErrDuplicateUpload)
}

func (s *StoreSuite) checkAddBundle(c *gc.C, bundle charm.Bundle) {
	store, err := NewStore(s.Session.DB("foo"))
	c.Assert(err, gc.IsNil)
	url := mustParseReference("cs:bundle/wordpress-simple-42")

	// Add the bundle to the store.
	beforeAdding := time.Now()
	err = store.AddBundleWithArchive(url, bundle)
	c.Assert(err, gc.IsNil)
	afterAdding := time.Now()

	var doc mongodoc.Entity
	err = store.DB.Entities().FindId("cs:bundle/wordpress-simple-42").One(&doc)
	c.Assert(err, gc.IsNil)
	sort.Sort(orderedURLs(doc.BundleCharms))

	// Check the upload time and then reset it to its zero value
	// so that we can test the deterministic parts later.
	c.Assert(doc.UploadTime, jc.TimeBetween(beforeAdding, afterAdding))
	doc.UploadTime = time.Time{}

	// The blob name is random, but we check that it's
	// in the correct format, and non-empty.
	blobName := doc.BlobName
	c.Assert(blobName, gc.Matches, "[0-9a-z]+")
	doc.BlobName = ""

	// The entity doc has been correctly added to the mongo collection.
	size, hash := mustGetSizeAndHash(bundle)
	c.Assert(doc, jc.DeepEquals, mongodoc.Entity{
		URL:          url,
		BaseURL:      mustParseReference("cs:wordpress-simple"),
		BlobHash:     hash,
		Size:         size,
		BundleData:   bundle.Data(),
		BundleReadMe: bundle.ReadMe(),
		BundleCharms: []*charm.Reference{
			mustParseReference("mysql"),
			mustParseReference("wordpress"),
		},
	})

	// The bundle archive has been properly added to the blob store.
	r, obtainedSize, err := store.BlobStore.Open(blobName)
	c.Assert(err, gc.IsNil)
	c.Assert(obtainedSize, gc.Equals, size)
	data, err := ioutil.ReadAll(r)
	c.Assert(err, gc.IsNil)
	bundleArchive, err := charm.ReadBundleArchiveBytes(data)
	c.Assert(err, gc.IsNil)
	c.Assert(bundleArchive.Data(), jc.DeepEquals, bundle.Data())
	c.Assert(bundleArchive.ReadMe(), jc.DeepEquals, bundle.ReadMe())

	// Try inserting the bundle again - it should fail because the bundle is
	// already there.
	err = store.AddBundleWithArchive(url, bundle)
	c.Assert(err, gc.Equals, params.ErrDuplicateUpload)
}

type orderedURLs []*charm.Reference

func (o orderedURLs) Less(i, j int) bool {
	return o[i].String() < o[j].String()
}

func (o orderedURLs) Swap(i, j int) {
	o[i], o[j] = o[j], o[i]
}

func (o orderedURLs) Len() int {
	return len(o)
}

var expandURLTests = []struct {
	inStore []string
	expand  string
	expect  []string
}{{
	inStore: []string{"cs:precise/wordpress-23"},
	expand:  "wordpress",
	expect:  []string{"cs:precise/wordpress-23"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:precise/wordpress-24"},
	expand:  "wordpress",
	expect:  []string{"cs:precise/wordpress-23", "cs:precise/wordpress-24"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24"},
	expand:  "precise/wordpress",
	expect:  []string{"cs:precise/wordpress-23"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24", "cs:foo/bar-434"},
	expand:  "wordpress",
	expect:  []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-23", "cs:trusty/wordpress-24"},
	expand:  "wordpress-23",
	expect:  []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-23"},
}, {
	inStore: []string{"cs:~user/precise/wordpress-23", "cs:~user/trusty/wordpress-23"},
	expand:  "~user/precise/wordpress",
	expect:  []string{"cs:~user/precise/wordpress-23"},
}, {
	inStore: []string{"cs:~user/precise/wordpress-23", "cs:~user/trusty/wordpress-23"},
	expand:  "~user/wordpress",
	expect:  []string{"cs:~user/precise/wordpress-23", "cs:~user/trusty/wordpress-23"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24", "cs:foo/bar-434"},
	expand:  "precise/wordpress-23",
	expect:  []string{"cs:precise/wordpress-23"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24", "cs:foo/bar-434"},
	expand:  "arble",
	expect:  []string{},
}}

func (s *StoreSuite) TestExpandURL(c *gc.C) {
	wordpress := testing.Charms.CharmDir("wordpress")
	for i, test := range expandURLTests {
		c.Logf("test %d: %q from %q", i, test.expand, test.inStore)
		store, err := NewStore(s.Session.DB("foo"))
		c.Assert(err, gc.IsNil)
		_, err = store.DB.Entities().RemoveAll(nil)
		c.Assert(err, gc.IsNil)
		urls := mustParseReferences(test.inStore)
		for _, url := range urls {
			err := store.AddCharmWithArchive(url, wordpress)
			c.Assert(err, gc.IsNil)
		}
		gotURLs, err := store.ExpandURL((*charm.Reference)(mustParseReference(test.expand)))
		c.Assert(err, gc.IsNil)

		gotURLStrs := urlStrings(gotURLs)
		sort.Strings(gotURLStrs)
		c.Assert(gotURLStrs, jc.DeepEquals, test.expect)
	}
}

func urlStrings(urls []*charm.Reference) []string {
	urlStrs := make([]string, len(urls))
	for i, url := range urls {
		urlStrs[i] = url.String()
	}
	return urlStrs
}

func mustParseReferences(urlStrs []string) []*charm.Reference {
	urls := make([]*charm.Reference, len(urlStrs))
	for i, u := range urlStrs {
		urls[i] = mustParseReference(u)
	}
	return urls
}

func (s *StoreSuite) TestAddCharmDir(c *gc.C) {
	charmDir := testing.Charms.CharmDir("wordpress")
	s.checkAddCharm(c, charmDir)
}

func (s *StoreSuite) TestAddCharmArchive(c *gc.C) {
	charmArchive := testing.Charms.CharmArchive(c.MkDir(), "wordpress")
	s.checkAddCharm(c, charmArchive)
}

func (s *StoreSuite) TestAddBundleDir(c *gc.C) {
	bundleDir := testing.Charms.BundleDir("wordpress")
	s.checkAddBundle(c, bundleDir)
}

func (s *StoreSuite) TestAddBundleArchive(c *gc.C) {
	bundleArchive, err := charm.ReadBundleArchive(
		testing.Charms.BundleArchivePath(c.MkDir(), "wordpress"),
	)
	c.Assert(err, gc.IsNil)
	s.checkAddBundle(c, bundleArchive)
}

func mustGetSizeAndHash(c interface{}) (int64, string) {
	var r io.ReadWriter
	var err error
	switch c := c.(type) {
	case archiverTo:
		r = new(bytes.Buffer)
		err = c.ArchiveTo(r)
	case *charm.BundleArchive:
		r, err = os.Open(c.Path)
	case *charm.CharmArchive:
		r, err = os.Open(c.Path)
	default:
		panic(fmt.Sprintf("unable to get size and hash for type %T", c))
	}
	if err != nil {
		panic(err)
	}
	hash := blobstore.NewHash()
	size, err := io.Copy(hash, r)
	if err != nil {
		panic(err)
	}
	return size, fmt.Sprintf("%x", hash.Sum(nil))
}

func mustParseReference(url string) *charm.Reference {
	ref, err := charm.ParseReference(url)
	if err != nil {
		panic(err)
	}
	return ref
}
<commit_msg>gofmt<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the LGPLv3, see LICENCE file for details.

package charmstore

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sort"
	"time"

	"github.com/juju/errgo"
	jc "github.com/juju/testing/checkers"
	"gopkg.in/juju/charm.v3"
	"gopkg.in/juju/charm.v3/testing"
	gc "launchpad.net/gocheck"

	"github.com/juju/charmstore/internal/blobstore"
	"github.com/juju/charmstore/internal/mongodoc"
	"github.com/juju/charmstore/internal/storetesting"
	"github.com/juju/charmstore/params"
)

type StoreSuite struct {
	storetesting.IsolatedMgoSuite
}

var _ = gc.Suite(&StoreSuite{})

func (s *StoreSuite) checkAddCharm(c *gc.C, ch charm.Charm) {
	store, err := NewStore(s.Session.DB("foo"))
	c.Assert(err, gc.IsNil)
	url := mustParseReference("cs:precise/wordpress-23")

	// Add the charm to the store.
	beforeAdding := time.Now()
	err = store.AddCharmWithArchive(url, ch)
	c.Assert(err, gc.IsNil)
	afterAdding := time.Now()

	var doc mongodoc.Entity
	err = store.DB.Entities().FindId("cs:precise/wordpress-23").One(&doc)
	c.Assert(err, gc.IsNil)

	// The entity doc has been correctly added to the mongo collection.
	size, hash := mustGetSizeAndHash(ch)
	sort.Strings(doc.CharmProvidedInterfaces)
	sort.Strings(doc.CharmRequiredInterfaces)

	// Check the upload time and then reset it to its zero value
	// so that we can test the deterministic parts later.
	c.Assert(doc.UploadTime, jc.TimeBetween(beforeAdding, afterAdding))

	doc.UploadTime = time.Time{}

	blobName := doc.BlobName
	c.Assert(blobName, gc.Matches, "[0-9a-z]+")
	doc.BlobName = ""
	c.Assert(doc, jc.DeepEquals, mongodoc.Entity{
		URL:                     url,
		BaseURL:                 mustParseReference("cs:wordpress"),
		BlobHash:                hash,
		Size:                    size,
		CharmMeta:               ch.Meta(),
		CharmActions:            ch.Actions(),
		CharmConfig:             ch.Config(),
		CharmProvidedInterfaces: []string{"http", "logging", "monitoring"},
		CharmRequiredInterfaces: []string{"mysql", "varnish"},
	})

	// The charm archive has been properly added to the blob store.
	r, obtainedSize, err := store.BlobStore.Open(blobName)
	c.Assert(err, gc.IsNil)
	c.Assert(obtainedSize, gc.Equals, size)
	data, err := ioutil.ReadAll(r)
	c.Assert(err, gc.IsNil)
	charmArchive, err := charm.ReadCharmArchiveBytes(data)
	c.Assert(err, gc.IsNil)
	c.Assert(charmArchive.Meta(), jc.DeepEquals, ch.Meta())
	c.Assert(charmArchive.Config(), jc.DeepEquals, ch.Config())
	c.Assert(charmArchive.Actions(), jc.DeepEquals, ch.Actions())
	c.Assert(charmArchive.Revision(), jc.DeepEquals, ch.Revision())

	// Try inserting the charm again - it should fail because the charm is
	// already there.
	err = store.AddCharmWithArchive(url, ch)
	c.Assert(errgo.Cause(err), gc.Equals, params.ErrDuplicateUpload)
}

func (s *StoreSuite) checkAddBundle(c *gc.C, bundle charm.Bundle) {
	store, err := NewStore(s.Session.DB("foo"))
	c.Assert(err, gc.IsNil)
	url := mustParseReference("cs:bundle/wordpress-simple-42")

	// Add the bundle to the store.
	beforeAdding := time.Now()
	err = store.AddBundleWithArchive(url, bundle)
	c.Assert(err, gc.IsNil)
	afterAdding := time.Now()

	var doc mongodoc.Entity
	err = store.DB.Entities().FindId("cs:bundle/wordpress-simple-42").One(&doc)
	c.Assert(err, gc.IsNil)
	sort.Sort(orderedURLs(doc.BundleCharms))

	// Check the upload time and then reset it to its zero value
	// so that we can test the deterministic parts later.
	c.Assert(doc.UploadTime, jc.TimeBetween(beforeAdding, afterAdding))
	doc.UploadTime = time.Time{}

	// The blob name is random, but we check that it's
	// in the correct format, and non-empty.
	blobName := doc.BlobName
	c.Assert(blobName, gc.Matches, "[0-9a-z]+")
	doc.BlobName = ""

	// The entity doc has been correctly added to the mongo collection.
	size, hash := mustGetSizeAndHash(bundle)
	c.Assert(doc, jc.DeepEquals, mongodoc.Entity{
		URL:          url,
		BaseURL:      mustParseReference("cs:wordpress-simple"),
		BlobHash:     hash,
		Size:         size,
		BundleData:   bundle.Data(),
		BundleReadMe: bundle.ReadMe(),
		BundleCharms: []*charm.Reference{
			mustParseReference("mysql"),
			mustParseReference("wordpress"),
		},
	})

	// The bundle archive has been properly added to the blob store.
	r, obtainedSize, err := store.BlobStore.Open(blobName)
	c.Assert(err, gc.IsNil)
	c.Assert(obtainedSize, gc.Equals, size)
	data, err := ioutil.ReadAll(r)
	c.Assert(err, gc.IsNil)
	bundleArchive, err := charm.ReadBundleArchiveBytes(data)
	c.Assert(err, gc.IsNil)
	c.Assert(bundleArchive.Data(), jc.DeepEquals, bundle.Data())
	c.Assert(bundleArchive.ReadMe(), jc.DeepEquals, bundle.ReadMe())

	// Try inserting the bundle again - it should fail because the bundle is
	// already there.
	err = store.AddBundleWithArchive(url, bundle)
	c.Assert(err, gc.Equals, params.ErrDuplicateUpload)
}

type orderedURLs []*charm.Reference

func (o orderedURLs) Less(i, j int) bool {
	return o[i].String() < o[j].String()
}

func (o orderedURLs) Swap(i, j int) {
	o[i], o[j] = o[j], o[i]
}

func (o orderedURLs) Len() int {
	return len(o)
}

var expandURLTests = []struct {
	inStore []string
	expand  string
	expect  []string
}{{
	inStore: []string{"cs:precise/wordpress-23"},
	expand:  "wordpress",
	expect:  []string{"cs:precise/wordpress-23"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:precise/wordpress-24"},
	expand:  "wordpress",
	expect:  []string{"cs:precise/wordpress-23", "cs:precise/wordpress-24"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24"},
	expand:  "precise/wordpress",
	expect:  []string{"cs:precise/wordpress-23"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24", "cs:foo/bar-434"},
	expand:  "wordpress",
	expect:  []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-23", "cs:trusty/wordpress-24"},
	expand:  "wordpress-23",
	expect:  []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-23"},
}, {
	inStore: []string{"cs:~user/precise/wordpress-23", "cs:~user/trusty/wordpress-23"},
	expand:  "~user/precise/wordpress",
	expect:  []string{"cs:~user/precise/wordpress-23"},
}, {
	inStore: []string{"cs:~user/precise/wordpress-23", "cs:~user/trusty/wordpress-23"},
	expand:  "~user/wordpress",
	expect:  []string{"cs:~user/precise/wordpress-23", "cs:~user/trusty/wordpress-23"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24", "cs:foo/bar-434"},
	expand:  "precise/wordpress-23",
	expect:  []string{"cs:precise/wordpress-23"},
}, {
	inStore: []string{"cs:precise/wordpress-23", "cs:trusty/wordpress-24", "cs:foo/bar-434"},
	expand:  "arble",
	expect:  []string{},
}}

func (s *StoreSuite) TestExpandURL(c *gc.C) {
	wordpress := testing.Charms.CharmDir("wordpress")
	for i, test := range expandURLTests {
		c.Logf("test %d: %q from %q", i, test.expand, test.inStore)
		store, err := NewStore(s.Session.DB("foo"))
		c.Assert(err, gc.IsNil)
		_, err = store.DB.Entities().RemoveAll(nil)
		c.Assert(err, gc.IsNil)
		urls := mustParseReferences(test.inStore)
		for _, url := range urls {
			err := store.AddCharmWithArchive(url, wordpress)
			c.Assert(err, gc.IsNil)
		}
		gotURLs, err := store.ExpandURL((*charm.Reference)(mustParseReference(test.expand)))
		c.Assert(err, gc.IsNil)

		gotURLStrs := urlStrings(gotURLs)
		sort.Strings(gotURLStrs)
		c.Assert(gotURLStrs, jc.DeepEquals, test.expect)
	}
}

func urlStrings(urls []*charm.Reference) []string {
	urlStrs := make([]string, len(urls))
	for i, url := range urls {
		urlStrs[i] = url.String()
	}
	return urlStrs
}

func mustParseReferences(urlStrs []string) []*charm.Reference {
	urls := make([]*charm.Reference, len(urlStrs))
	for i, u := range urlStrs {
		urls[i] = mustParseReference(u)
	}
	return urls
}

func (s *StoreSuite) TestAddCharmDir(c *gc.C) {
	charmDir := testing.Charms.CharmDir("wordpress")
	s.checkAddCharm(c, charmDir)
}

func (s *StoreSuite) TestAddCharmArchive(c *gc.C) {
	charmArchive := testing.Charms.CharmArchive(c.MkDir(), "wordpress")
	s.checkAddCharm(c, charmArchive)
}

func (s *StoreSuite) TestAddBundleDir(c *gc.C) {
	bundleDir := testing.Charms.BundleDir("wordpress")
	s.checkAddBundle(c, bundleDir)
}

func (s *StoreSuite) TestAddBundleArchive(c *gc.C) {
	bundleArchive, err := charm.ReadBundleArchive(
		testing.Charms.BundleArchivePath(c.MkDir(), "wordpress"),
	)
	c.Assert(err, gc.IsNil)
	s.checkAddBundle(c, bundleArchive)
}

func mustGetSizeAndHash(c interface{}) (int64, string) {
	var r io.ReadWriter
	var err error
	switch c := c.(type) {
	case archiverTo:
		r = new(bytes.Buffer)
		err = c.ArchiveTo(r)
	case *charm.BundleArchive:
		r, err = os.Open(c.Path)
	case *charm.CharmArchive:
		r, err = os.Open(c.Path)
	default:
		panic(fmt.Sprintf("unable to get size and hash for type %T", c))
	}
	if err != nil {
		panic(err)
	}
	hash := blobstore.NewHash()
	size, err := io.Copy(hash, r)
	if err != nil {
		panic(err)
	}
	return size, fmt.Sprintf("%x", hash.Sum(nil))
}

func mustParseReference(url string) *charm.Reference {
	ref, err := charm.ParseReference(url)
	if err != nil {
		panic(err)
	}
	return ref
}
<|endoftext|>


<commit_before>// Copyright The prometheus-operator 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 versionutil_test

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"testing"

	"github.com/prometheus/common/version"
	"github.com/stretchr/testify/assert"

	"github.com/prometheus-operator/prometheus-operator/pkg/versionutil"
)

func TestShouldPrintVersion(t *testing.T) {
	const program = "foo"

	restore := setAllVersionFieldsTo("test-value")
	defer restore()

	tests := map[string]struct {
		flag        string
		expOutput   string
		program     string
		shouldPrint bool
	}{
		"Should print only version": {
			flag:      "--short-version",
			expOutput: "test-value",
		},
		"Should print full version": {
			flag: "--version",
			expOutput: fmt.Sprintf(
				"%s, version test-value (branch: test-value, revision: test-value)\n"+
					"  build user:       test-value\n"+
					"  build date:       test-value\n"+
					"  go version:       test-value", program),
		},
	}
	for tn, tc := range tests {
		t.Run(tn, func(t *testing.T) {
			restore := snapshotOSArgsAndFlags()
			defer restore()

			// given
			os.Args = []string{program, tc.flag}
			flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)

			// when
			versionutil.RegisterParseFlags()
			// then
			assert.Equal(t, true, versionutil.ShouldPrintVersion())

			// when
			var buf bytes.Buffer
			versionutil.Print(&buf, program)
			// then
			assert.Equal(t, tc.expOutput, buf.String())
		})
	}
}

func TestShouldNotPrintVersion(t *testing.T) {
	// given
	restore := snapshotOSArgsAndFlags()
	defer restore()

	// no flags set
	os.Args = []string{"cmd"}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	// when
	versionutil.RegisterParseFlags()

	// then
	assert.False(t, versionutil.ShouldPrintVersion())
}

// snapshotOSArgsAndFlags the os.Args and flag.CommandLine and allows you to restore them.
// inspired by: https://golang.org/src/flag/flag_test.go#L318
func snapshotOSArgsAndFlags() func() {
	oldArgs := os.Args
	oldFlags := flag.CommandLine

	return func() {
		os.Args = oldArgs
		flag.CommandLine = oldFlags
	}
}

// setAllVersionFieldsToTestFixture sets all version fields to a given value.
// Simplifies test cases by ensuring that values are predictable.
func setAllVersionFieldsTo(fixture string) func() {
	oldVersion := version.Version
	oldRevision := version.Revision
	oldBranch := version.Branch
	oldBuildUser := version.BuildUser
	oldBuildDate := version.BuildDate
	oldGoVersion := version.GoVersion

	version.Version = fixture
	version.Revision = fixture
	version.Branch = fixture
	version.BuildUser = fixture
	version.BuildDate = fixture
	version.GoVersion = fixture

	return func() {
		version.Version = oldVersion
		version.Revision = oldRevision
		version.Branch = oldBranch
		version.BuildUser = oldBuildUser
		version.BuildDate = oldBuildDate
		version.GoVersion = oldGoVersion
	}
}
<commit_msg>pkg/versionutil: fix tests (#3614)<commit_after>// Copyright The prometheus-operator 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 versionutil_test

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"testing"

	"github.com/prometheus/common/version"
	"github.com/stretchr/testify/assert"

	"github.com/prometheus-operator/prometheus-operator/pkg/versionutil"
)

func TestShouldPrintVersion(t *testing.T) {
	const program = "foo"

	restore := setAllVersionFieldsTo("test-value")
	defer restore()

	tests := map[string]struct {
		flag        string
		expOutput   string
		program     string
		shouldPrint bool
	}{
		"Should print only version": {
			flag:      "--short-version",
			expOutput: "test-value",
		},
		"Should print full version": {
			flag: "--version",
			expOutput: fmt.Sprintf(
				"%s, version test-value (branch: test-value, revision: test-value)\n"+
					"  build user:       test-value\n"+
					"  build date:       test-value\n"+
					"  go version:       test-value\n"+
					"  platform:         linux/amd64", program),
		},
	}
	for tn, tc := range tests {
		t.Run(tn, func(t *testing.T) {
			restore := snapshotOSArgsAndFlags()
			defer restore()

			// given
			os.Args = []string{program, tc.flag}
			flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)

			// when
			versionutil.RegisterParseFlags()
			// then
			assert.Equal(t, true, versionutil.ShouldPrintVersion())

			// when
			var buf bytes.Buffer
			versionutil.Print(&buf, program)
			// then
			assert.Equal(t, tc.expOutput, buf.String())
		})
	}
}

func TestShouldNotPrintVersion(t *testing.T) {
	// given
	restore := snapshotOSArgsAndFlags()
	defer restore()

	// no flags set
	os.Args = []string{"cmd"}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	// when
	versionutil.RegisterParseFlags()

	// then
	assert.False(t, versionutil.ShouldPrintVersion())
}

// snapshotOSArgsAndFlags the os.Args and flag.CommandLine and allows you to restore them.
// inspired by: https://golang.org/src/flag/flag_test.go#L318
func snapshotOSArgsAndFlags() func() {
	oldArgs := os.Args
	oldFlags := flag.CommandLine

	return func() {
		os.Args = oldArgs
		flag.CommandLine = oldFlags
	}
}

// setAllVersionFieldsToTestFixture sets all version fields to a given value.
// Simplifies test cases by ensuring that values are predictable.
func setAllVersionFieldsTo(fixture string) func() {
	oldVersion := version.Version
	oldRevision := version.Revision
	oldBranch := version.Branch
	oldBuildUser := version.BuildUser
	oldBuildDate := version.BuildDate
	oldGoVersion := version.GoVersion

	version.Version = fixture
	version.Revision = fixture
	version.Branch = fixture
	version.BuildUser = fixture
	version.BuildDate = fixture
	version.GoVersion = fixture

	return func() {
		version.Version = oldVersion
		version.Revision = oldRevision
		version.Branch = oldBranch
		version.BuildUser = oldBuildUser
		version.BuildDate = oldBuildDate
		version.GoVersion = oldGoVersion
	}
}
<|endoftext|>


<commit_before>package linode

import (
	"context"
	"fmt"

	"github.com/TimothyYe/godns/internal/utils"
	"github.com/linode/linodego"
	log "github.com/sirupsen/logrus"
)

type DNSClient struct {
	linodeClient *linodego.Client
}

func CreateLinodeDNSClient(linodeClient *linodego.Client) IDNSClient {
	dnsClient := DNSClient{
		linodeClient: linodeClient,
	}
	return &dnsClient
}

func (dnsClient *DNSClient) UpdateDNSRecordIP(domain string, subdomain string, ip string) error {
	if subdomain == utils.RootDomain {
		subdomain = ""
	}

	domainID, err := dnsClient.getDomainID(domain)
	if err != nil {
		return err
	}

	recordExists, recordID, err := dnsClient.getDomainRecordID(domainID, subdomain)
	if err != nil {
		return err
	}
	if !recordExists {
		recordID, _ = dnsClient.createDomainRecord(domainID, subdomain)
	}

	err = dnsClient.updateDomainRecord(domainID, recordID, ip)
	if err != nil {
		return err
	}

	return nil
}

func (dnsClient *DNSClient) getDomainID(name string) (int, error) {
	f := linodego.Filter{}
	f.AddField(linodego.Eq, "domain", name)
	fStr, err := f.MarshalJSON()
	if err != nil {
		log.Fatal(err)
	}

	opts := linodego.NewListOptions(0, string(fStr))
	res, err := dnsClient.linodeClient.ListDomains(context.Background(), opts)
	if err != nil {
		return 0, err
	}
	if len(res) == 0 {
		return 0, fmt.Errorf("No domains found for name %s", name)
	}
	return res[0].ID, nil
}

func (dnsClient *DNSClient) getDomainRecordID(domainID int, name string) (bool, int, error) {
	f := linodego.Filter{}
	f.AddField(linodego.Eq, "name", name)
	fStr, err := f.MarshalJSON()
	if err != nil {
		log.Fatal(err)
	}
	opts := linodego.NewListOptions(0, string(fStr))
	res, err := dnsClient.linodeClient.ListDomainRecords(context.Background(), domainID, opts)
	if err != nil {
		return false, 0, err
	}
	if len(res) == 0 {
		return false, 0, nil
	}
	return true, res[0].ID, nil
}

func (dnsClient *DNSClient) createDomainRecord(domainID int, name string) (int, error) {
	opts := &linodego.DomainRecordCreateOptions{
		Type:   "A",
		Name:   name,
		Target: "127.0.0.1",
		TTLSec: 30,
	}
	record, err := dnsClient.linodeClient.CreateDomainRecord(context.Background(), domainID, *opts)
	if err != nil {
		return 0, err
	}
	return record.ID, nil
}

func (dnsClient *DNSClient) updateDomainRecord(domainID int, id int, ip string) error {
	opts := &linodego.DomainRecordUpdateOptions{Target: ip}
	_, err := dnsClient.linodeClient.UpdateDomainRecord(context.Background(), domainID, id, *opts)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Work around subdomain filtering bug in linode API (#154)<commit_after>package linode

import (
	"context"
	"fmt"

	"github.com/TimothyYe/godns/internal/utils"
	"github.com/linode/linodego"
	log "github.com/sirupsen/logrus"
)

type DNSClient struct {
	linodeClient *linodego.Client
}

func CreateLinodeDNSClient(linodeClient *linodego.Client) IDNSClient {
	dnsClient := DNSClient{
		linodeClient: linodeClient,
	}
	return &dnsClient
}

func (dnsClient *DNSClient) UpdateDNSRecordIP(domain string, subdomain string, ip string) error {
	if subdomain == utils.RootDomain {
		subdomain = ""
	}

	domainID, err := dnsClient.getDomainID(domain)
	if err != nil {
		return err
	}

	recordExists, recordID, err := dnsClient.getDomainRecordID(domainID, subdomain)
	if err != nil {
		return err
	}
	if !recordExists {
		recordID, _ = dnsClient.createDomainRecord(domainID, subdomain)
	}

	err = dnsClient.updateDomainRecord(domainID, recordID, ip)
	if err != nil {
		return err
	}

	return nil
}

func (dnsClient *DNSClient) getDomainID(name string) (int, error) {
	f := linodego.Filter{}
	f.AddField(linodego.Eq, "domain", name)
	fStr, err := f.MarshalJSON()
	if err != nil {
		log.Fatal(err)
	}

	opts := linodego.NewListOptions(0, string(fStr))
	res, err := dnsClient.linodeClient.ListDomains(context.Background(), opts)
	if err != nil {
		return 0, err
	}
	if len(res) == 0 {
		return 0, fmt.Errorf("No domains found for name %s", name)
	}
	return res[0].ID, nil
}

func (dnsClient *DNSClient) getDomainRecordID(domainID int, name string) (bool, int, error) {
	res, err := dnsClient.linodeClient.ListDomainRecords(context.Background(), domainID, nil)
	if err != nil {
		return false, 0, err
	}
	if len(res) == 0 {
		return false, 0, nil
	}
	for _, record := range res {
		if record.Name == name {
			return true, record.ID, nil
		}
	}
	return false, 0, nil
}

func (dnsClient *DNSClient) createDomainRecord(domainID int, name string) (int, error) {
	opts := &linodego.DomainRecordCreateOptions{
		Type:   "A",
		Name:   name,
		Target: "127.0.0.1",
		TTLSec: 30,
	}
	record, err := dnsClient.linodeClient.CreateDomainRecord(context.Background(), domainID, *opts)
	if err != nil {
		return 0, err
	}
	return record.ID, nil
}

func (dnsClient *DNSClient) updateDomainRecord(domainID int, id int, ip string) error {
	opts := &linodego.DomainRecordUpdateOptions{Target: ip}
	_, err := dnsClient.linodeClient.UpdateDomainRecord(context.Background(), domainID, id, *opts)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>