<commit_before>package monitor

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/sivel/overseer/status"
)

type HTTPStatusConfig struct {
	Name                 string
	URL                  *url.URL
	Codes                []int
	NotificationInterval time.Duration
	Verify               bool
	Timeout              time.Duration
	Method               string
}

type HTTPStatus struct {
	config *HTTPStatusConfig
	status *status.Status
}

func NewHTTPStatus(conf map[string]interface{}) Monitor {
	var err error
	monitor := new(HTTPStatus)

	var pURL *url.URL
	if urlInterface, ok := conf["url"]; ok {
		pURL, err = url.Parse(urlInterface.(string))
		if err != nil {
			log.Fatalf("Invalid URL: %s", conf["url"].(string))
		} else if !ok {
			log.Fatalf("No URL provided")
		}
	}

	var name string = pURL.String()
	if nameInterface, ok := conf["name"]; ok {
		name = nameInterface.(string)
	}

	var codes []int = []int{200}
	if codesInterface, ok := conf["codes"]; ok {
		for _, code := range codesInterface.([]interface{}) {
			codes = append(codes, int(code.(float64)))
		}
	}

	var notificationInterval time.Duration = time.Second * 60
	if ni, ok := conf["notification_interval"]; ok {
		notificationInterval, err = time.ParseDuration(ni.(string))
	}

	var verify bool = false
	if verifyInterface, ok := conf["verify"]; ok {
		verify = verifyInterface.(bool)
	}

	var timeout time.Duration = time.Second * 2
	if timeoutInterface, ok := conf["timeout"]; ok {
		timeout, _ = time.ParseDuration(timeoutInterface.(string))
	}

	var method string = "HEAD"
	if methodInterface, ok := conf["method"]; ok {
		method = strings.ToUpper(methodInterface.(string))
	}

	monitor.config = &HTTPStatusConfig{
		Name:                 name,
		URL:                  pURL,
		Codes:                codes,
		NotificationInterval: notificationInterval,
		Verify:               verify,
		Timeout:              timeout,
		Method:               method,
	}
	monitor.status = status.NewStatus(
		name,
		status.UNKNOWN,
		status.UNKNOWN,
		notificationInterval,
		time.Now(),
		time.Now(),
		0,
		"",
	)
	return monitor
}

func (m *HTTPStatus) Watch(statusChan chan *status.Status) {
	for {
		m.Check()
		statusChan <- m.status
		time.Sleep(time.Second * 10)
	}
}

func isValidCode(code int, codes []int) bool {
	var valid bool = false
	for _, c := range codes {
		if c == code {
			valid = true
		}
	}
	return valid
}

func checkChanged(current int, last int, startOfLastStatus time.Time) (bool, time.Time) {
	var start time.Time = startOfLastStatus
	var changed bool = false
	if current != last {
		changed = true
		start = time.Now()
	}
	return changed, start
}

func (m *HTTPStatus) Check() {
	fmt.Println("HTTPStatus Check Running for " + m.config.URL.String())

	transport := http.Transport{
		Dial: func(network, addr string) (net.Conn, error) {
			return net.DialTimeout(network, addr, m.config.Timeout)
		},
	}

	if m.config.URL.Scheme == "https" {
		transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: m.config.Verify}
	}

	client := http.Client{
		Transport: &transport,
	}

	requestStart := time.Now()
	resp, err := client.Do(&http.Request{Method: "HEAD", URL: m.config.URL})
	duration := time.Now().UnixNano() - requestStart.UnixNano()

	var current int = status.UP
	var message string = "OK"
	if err != nil {
		current = status.DOWN
		message = err.Error()
	} else {
		defer resp.Body.Close()

		if !isValidCode(resp.StatusCode, m.config.Codes) {
			current = status.DOWN
			message = fmt.Sprintf("Invalid response code: %d", resp.StatusCode)
		}
	}

	_, start := checkChanged(current, m.status.Current, m.status.StartOfCurrentStatus)

	m.status = status.NewStatus(
		m.config.Name,
		current,
		m.status.Current,
		m.config.NotificationInterval,
		start,
		m.status.LastNotification,
		duration,
		message,
	)
}
<commit_msg>Use supplied check_interval<commit_after>package monitor

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/sivel/overseer/status"
)

type HTTPStatusConfig struct {
	Name                 string
	URL                  *url.URL
	Codes                []int
	CheckInterval        time.Duration
	NotificationInterval time.Duration
	Verify               bool
	Timeout              time.Duration
	Method               string
}

type HTTPStatus struct {
	config *HTTPStatusConfig
	status *status.Status
}

func NewHTTPStatus(conf map[string]interface{}) Monitor {
	var err error
	monitor := new(HTTPStatus)

	var pURL *url.URL
	if urlInterface, ok := conf["url"]; ok {
		pURL, err = url.Parse(urlInterface.(string))
		if err != nil {
			log.Fatalf("Invalid URL: %s", conf["url"].(string))
		} else if !ok {
			log.Fatalf("No URL provided")
		}
	}

	var name string = pURL.String()
	if nameInterface, ok := conf["name"]; ok {
		name = nameInterface.(string)
	}

	var codes []int = []int{200}
	if codesInterface, ok := conf["codes"]; ok {
		for _, code := range codesInterface.([]interface{}) {
			codes = append(codes, int(code.(float64)))
		}
	}

	var checkInterval time.Duration = time.Second * 10
	if ci, ok := conf["check_interval"]; ok {
		checkInterval, err = time.ParseDuration(ci.(string))
	}

	var notificationInterval time.Duration = time.Second * 60
	if ni, ok := conf["notification_interval"]; ok {
		notificationInterval, err = time.ParseDuration(ni.(string))
	}

	var verify bool = false
	if verifyInterface, ok := conf["verify"]; ok {
		verify = verifyInterface.(bool)
	}

	var timeout time.Duration = time.Second * 2
	if timeoutInterface, ok := conf["timeout"]; ok {
		timeout, _ = time.ParseDuration(timeoutInterface.(string))
	}

	var method string = "HEAD"
	if methodInterface, ok := conf["method"]; ok {
		method = strings.ToUpper(methodInterface.(string))
	}

	monitor.config = &HTTPStatusConfig{
		Name:                 name,
		URL:                  pURL,
		Codes:                codes,
		CheckInterval:        checkInterval,
		NotificationInterval: notificationInterval,
		Verify:               verify,
		Timeout:              timeout,
		Method:               method,
	}
	monitor.status = status.NewStatus(
		name,
		status.UNKNOWN,
		status.UNKNOWN,
		notificationInterval,
		time.Now(),
		time.Now(),
		0,
		"",
	)
	return monitor
}

func (m *HTTPStatus) Watch(statusChan chan *status.Status) {
	for {
		m.Check()
		statusChan <- m.status
		time.Sleep(m.config.CheckInterval)
	}
}

func isValidCode(code int, codes []int) bool {
	var valid bool = false
	for _, c := range codes {
		if c == code {
			valid = true
		}
	}
	return valid
}

func checkChanged(current int, last int, startOfLastStatus time.Time) (bool, time.Time) {
	var start time.Time = startOfLastStatus
	var changed bool = false
	if current != last {
		changed = true
		start = time.Now()
	}
	return changed, start
}

func (m *HTTPStatus) Check() {
	fmt.Println("HTTPStatus Check Running for " + m.config.URL.String())

	transport := http.Transport{
		Dial: func(network, addr string) (net.Conn, error) {
			return net.DialTimeout(network, addr, m.config.Timeout)
		},
	}

	if m.config.URL.Scheme == "https" {
		transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: m.config.Verify}
	}

	client := http.Client{
		Transport: &transport,
	}

	requestStart := time.Now()
	resp, err := client.Do(&http.Request{Method: "HEAD", URL: m.config.URL})
	duration := time.Now().UnixNano() - requestStart.UnixNano()

	var current int = status.UP
	var message string = "OK"
	if err != nil {
		current = status.DOWN
		message = err.Error()
	} else {
		defer resp.Body.Close()

		if !isValidCode(resp.StatusCode, m.config.Codes) {
			current = status.DOWN
			message = fmt.Sprintf("Invalid response code: %d", resp.StatusCode)
		}
	}

	_, start := checkChanged(current, m.status.Current, m.status.StartOfCurrentStatus)

	m.status = status.NewStatus(
		m.config.Name,
		current,
		m.status.Current,
		m.config.NotificationInterval,
		start,
		m.status.LastNotification,
		duration,
		message,
	)
}
<|endoftext|>


<commit_before>package udp

import (
	"bytes"
	"net"
	"testing"
	"time"

	"github.com/CotaPreco/Horus/message"
	"github.com/stretchr/testify/assert"
)

const (
	RHOST = "0.0.0.0"
	RPORT = 10000
)

type ReceiveStrategy struct {
	Ch chan string
}

func (s *ReceiveStrategy) CanReceive(message []byte) bool {
	s.Ch <- "CanReceive:"
	return true
}

func (s *ReceiveStrategy) Receive(message []byte) message.MessageInterface {
	s.Ch <- " " + string(bytes.Trim(message, "\x00"))
	return nil
}

func TestReceive(t *testing.T) {
	ch := make(chan string)

	strategy := &ReceiveStrategy{
		Ch: ch,
	}

	receiver := NewUdpReceiver(RHOST, RPORT, strategy)
	go receiver.Receive()

	client, _ := net.DialUDP(
		"udp",
		&net.UDPAddr{
			IP:   net.IPv4zero,
			Port: 0,
		},
		&net.UDPAddr{
			IP:   net.ParseIP(RHOST),
			Port: RPORT,
		},
	)

	time.Sleep(time.Millisecond + 100)

	client.Write([]byte("A"))
	client.Write([]byte("B"))
	client.Write([]byte("C"))
	client.Write([]byte("D"))

	client.Close()

	time.Sleep(time.Millisecond + 100)

	assert.Equal(t, <-ch+<-ch, "CanReceive: A")
	assert.Equal(t, <-ch+<-ch, "CanReceive: B")
	assert.Equal(t, <-ch+<-ch, "CanReceive: C")
	assert.Equal(t, <-ch+<-ch, "CanReceive: D")
}
<commit_msg>Packet size must be specified<commit_after>package udp

import (
	"bytes"
	"net"
	"testing"
	"time"

	"github.com/CotaPreco/Horus/message"
	"github.com/stretchr/testify/assert"
)

const (
	RHOST = "0.0.0.0"
	RPORT = 10000
)

type ReceiveStrategy struct {
	Ch chan string
}

func (s *ReceiveStrategy) CanReceive(message []byte) bool {
	s.Ch <- "CanReceive:"
	return true
}

func (s *ReceiveStrategy) Receive(message []byte) message.MessageInterface {
	s.Ch <- " " + string(bytes.Trim(message, "\x00"))
	return nil
}

func TestReceive(t *testing.T) {
	ch := make(chan string)

	strategy := &ReceiveStrategy{
		Ch: ch,
	}

	receiver := NewUdpReceiver(RHOST, RPORT, 1, strategy)
	go receiver.Receive()

	client, _ := net.DialUDP(
		"udp",
		&net.UDPAddr{
			IP:   net.IPv4zero,
			Port: 0,
		},
		&net.UDPAddr{
			IP:   net.ParseIP(RHOST),
			Port: RPORT,
		},
	)

	time.Sleep(time.Millisecond + 100)

	client.Write([]byte("A"))
	client.Write([]byte("B"))
	client.Write([]byte("C"))
	client.Write([]byte("D"))

	client.Close()

	time.Sleep(time.Millisecond + 100)

	assert.Equal(t, <-ch+<-ch, "CanReceive: A")
	assert.Equal(t, <-ch+<-ch, "CanReceive: B")
	assert.Equal(t, <-ch+<-ch, "CanReceive: C")
	assert.Equal(t, <-ch+<-ch, "CanReceive: D")
}
<|endoftext|>


<commit_before>package types

import (
	"fmt"
	"strings"
)

// VerifyRequest contains the verification params.
type VerifyRequest struct {
	// URL to hit during provider verification.
	ProviderBaseURL string

	// Local/HTTP paths to Pact files.
	PactURLs []string

	// Pact Broker URL for broker-based verification
	BrokerURL string

	// Tags to find in Broker for matrix-based testing
	Tags []string

	// URL to retrieve valid Provider States.
	ProviderStatesURL string

	// URL to post currentp provider state to on the Provider API.
	ProviderStatesSetupURL string

	// Username when authenticating to a Pact Broker.
	BrokerUsername string

	// Password when authenticating to a Pact Broker.
	BrokerPassword string

	// Arguments to the VerificationProvider
	// Deprecated: This will be deleted after the native library replaces Ruby deps.
	Args []string
}

// Validate checks that the minimum fields are provided.
// Deprecated: This map be deleted after the native library replaces Ruby deps,
// and should not be used outside of this library.
func (v *VerifyRequest) Validate() error {
	v.Args = []string{}
	if v.ProviderBaseURL != "" {
		v.Args = append(v.Args, fmt.Sprintf("--provider-base-url %s", v.ProviderBaseURL))
	} else {
		return fmt.Errorf("ProviderBaseURL is mandatory.")
	}

	if len(v.PactURLs) != 0 {
		v.Args = append(v.Args, fmt.Sprintf("--pact-urls %s", strings.Join(v.PactURLs[:], ",")))
	} else {
		return fmt.Errorf("PactURLs is mandatory.")
	}

	if v.ProviderStatesSetupURL != "" {
		v.Args = append(v.Args, fmt.Sprintf("--provider-states-setup-url %s", v.ProviderStatesSetupURL))
	}

	if v.ProviderStatesURL != "" {
		v.Args = append(v.Args, fmt.Sprintf("--provider-states-url %s", v.ProviderStatesURL))
	}

	if v.BrokerUsername != "" {
		v.Args = append(v.Args, fmt.Sprintf("--broker-username %s", v.BrokerUsername))
	}

	if v.BrokerPassword != "" {
		v.Args = append(v.Args, fmt.Sprintf("--broker-password %s", v.BrokerPassword))
	}
	return nil
}
<commit_msg>fix(windows): split pact verification CLI arg tuples into individual elements #9<commit_after>package types

import (
	"fmt"
	"strings"
)

// VerifyRequest contains the verification params.
type VerifyRequest struct {
	// URL to hit during provider verification.
	ProviderBaseURL string

	// Local/HTTP paths to Pact files.
	PactURLs []string

	// Pact Broker URL for broker-based verification
	BrokerURL string

	// Tags to find in Broker for matrix-based testing
	Tags []string

	// URL to retrieve valid Provider States.
	ProviderStatesURL string

	// URL to post currentp provider state to on the Provider API.
	ProviderStatesSetupURL string

	// Username when authenticating to a Pact Broker.
	BrokerUsername string

	// Password when authenticating to a Pact Broker.
	BrokerPassword string

	// Arguments to the VerificationProvider
	// Deprecated: This will be deleted after the native library replaces Ruby deps.
	Args []string
}

// Validate checks that the minimum fields are provided.
// Deprecated: This map be deleted after the native library replaces Ruby deps,
// and should not be used outside of this library.
func (v *VerifyRequest) Validate() error {
	v.Args = []string{}
	if v.ProviderBaseURL != "" {
		v.Args = append(v.Args, "--provider-base-url")
		v.Args = append(v.Args, v.ProviderBaseURL)
	} else {
		return fmt.Errorf("ProviderBaseURL is mandatory.")
	}

	if len(v.PactURLs) != 0 {
		v.Args = append(v.Args, "--pact-urls")
		v.Args = append(v.Args, strings.Join(v.PactURLs[:], ","))
	} else {
		return fmt.Errorf("PactURLs is mandatory.")
	}

	if v.ProviderStatesSetupURL != "" {
		v.Args = append(v.Args, "--provider-states-setup-url")
		v.Args = append(v.Args, v.ProviderStatesSetupURL)
	}

	if v.ProviderStatesURL != "" {
		v.Args = append(v.Args, "--provider-states-url")
		v.Args = append(v.Args, v.ProviderStatesURL)
	}

	if v.BrokerUsername != "" {
		v.Args = append(v.Args, "--broker-username")
		v.Args = append(v.Args, v.BrokerUsername)
	}

	if v.BrokerPassword != "" {
		v.Args = append(v.Args, "--broker-password")
		v.Args = append(v.Args, v.BrokerPassword)
	}
	return nil
}
<|endoftext|>


<commit_before>// Package geoJSON converts geoJSON TO s2 and back to GeoJSON
// this is mostly an endpoint to visualize the simplifications
package geoJSON

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"time"

	log "github.com/Sirupsen/logrus"
)

// Endpoint is  the name of the geojson handler endpoint
const Endpoint = "/tos2/geojson/"

//Handler handles a request for a geojsonPoint
func Handler(w http.ResponseWriter, r *http.Request) {
	t := time.Now()
	// parse form
	var precision int
	var err error
	values := r.URL.Query()
	if p, ok := values["precision"]; ok {
		precision, err = strconv.Atoi(p[0])
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
		}
	} else {
		// set max precision
		precision = 30
	}
	log.Debugf("Request with precision: %v", precision)
	// request
	resp, err := Matcher(r)
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	geoj, err := resp.ToGeoJSON(precision)
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	// response
	encoder := json.NewEncoder(w)
	w.Header().Set("Content-Type", "application/json")
	err = encoder.Encode(geoj)
	log.Debug(time.Since(t))
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// Matcher exctract from the url witch geoJSON object we want
func Matcher(r *http.Request) (p GeoJSON, err error) {
	objectType := r.URL.Path[len(Endpoint):]
	dec := json.NewDecoder(r.Body)
	switch objectType {
	case "point":
		// TODO this is ugly
		pp := Point{}
		err = dec.Decode(&pp)
		p = pp
	case "polygon":
		pp := Polygon{}
		err = dec.Decode(&pp)
		p = pp
	case "multipolygon":
		pp := MultiPolygon{}
		err = dec.Decode(&pp)
		p = pp
	default:
		err = fmt.Errorf("Bad geoJSON object type")
	}
	return p, err
}
<commit_msg>Enforce endpoint <-> geojson type check.<commit_after>// Package geoJSON converts geoJSON TO s2 and back to GeoJSON
// this is mostly an endpoint to visualize the simplifications
package geoJSON

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
)

// Endpoint is  the name of the geojson handler endpoint
const Endpoint = "/tos2/geojson/"

//Handler handles a request for a geojsonPoint
func Handler(w http.ResponseWriter, r *http.Request) {
	t := time.Now()
	// parse form
	var precision int
	var err error
	values := r.URL.Query()
	if p, ok := values["precision"]; ok {
		precision, err = strconv.Atoi(p[0])
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
		}
	} else {
		// set max precision
		precision = 30
	}
	log.Debugf("Request with precision: %v", precision)
	// request
	resp, err := Matcher(r)
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	geoj, err := resp.ToGeoJSON(precision)
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	// response
	encoder := json.NewEncoder(w)
	w.Header().Set("Content-Type", "application/json")
	err = encoder.Encode(geoj)
	log.Debug(time.Since(t))
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// Matcher exctract from the url witch geoJSON object we want
func Matcher(r *http.Request) (p GeoJSON, err error) {
	objectType := r.URL.Path[len(Endpoint):]
	dec := json.NewDecoder(r.Body)
	switch objectType {
	case "point":
		// TODO this is ugly
		pp := Point{}
		err = dec.Decode(&pp)
		p = pp
		if strings.ToLower(pp.Type) != "point" {
			err = fmt.Errorf("%v not  a geoJSON point", pp.Type)
		}
	case "polygon":
		pp := Polygon{}
		err = dec.Decode(&pp)
		p = pp
		if strings.ToLower(pp.Type) != "polygon" {
			err = fmt.Errorf("%v not  a geoJSON polygon", pp.Type)
		}
	case "multipolygon":
		pp := MultiPolygon{}
		err = dec.Decode(&pp)
		p = pp
		if strings.ToLower(pp.Type) != "multipolygon" {
			err = fmt.Errorf("%v not a geoJSON multipolygon", pp.Type)
		}
	default:
		err = fmt.Errorf("Bad geoJSON object type")
	}
	return p, err
}
<|endoftext|>


<commit_before>//: ----------------------------------------------------------------------------
//: Copyright (C) 2018 Alan Willis.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    segmentio.go
//: details: vflow kafka producer plugin
//: author:  Alan Willis
//: date:    12/05/2018
//:
//: 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.
//: ----------------------------------------------------------------------------
// +build kafkav2

package producer

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"log"
	"net"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/segmentio/kafka-go"
	"github.com/segmentio/kafka-go/gzip"
	"github.com/segmentio/kafka-go/lz4"
	"github.com/segmentio/kafka-go/snappy"
	"gopkg.in/yaml.v2"
)

// Kafka represents kafka producer
type Kafka struct {
	producer *kafka.Writer
	config   KafkaConfig
	logger   *log.Logger
}

// Config represents kafka configuration
type KafkaConfig struct {
	run             kafka.WriterConfig
	Brokers         []string `yaml:"brokers" env:"BROKERS"`
	BootstrapServer string   `yaml:"bootstrap-server" env:"BOOTSTRAP_SERVER"`
	ClientID        string   `yaml:"client-id" env:"CLIENT_ID"`
	Compression     string   `yaml:"compression" env:"COMPRESSION"`
	MaxAttempts     int      `yaml:"max-attempts" env:"MAX_ATTEMPTS"`
	QueueSize       int      `yaml:"queue-size" env:"QUEUE_SIZE"`
	BatchSize       int      `yaml:"batch-size" env:"BATCH_SIZE"`
	KeepAlive       int      `yaml:"keepalive" env:"KEEPALIVE"`
	IOTimeout       int      `yaml:"connect-timeout" env:"CONNECT_TIMEOUT"`
	RequiredAcks    int      `yaml:"required-acks" env:"REQUIRED_ACKS"`
	PeriodicFlush   int      `yaml:"pflush" env:"PERIODIC_FLUSH"`
	TLSCertFile     string   `yaml:"tls-cert" env:"TLS_CERT"`
	TLSKeyFile      string   `yaml:"tls-key" env:"TLS_KEY"`
	CAFile          string   `yaml:"ca-file" env:"CA_FILE"`
	VerifySSL       bool     `yaml:"verify-ssl" env:"VERIFY_SSL"`
}

func (k *Kafka) setup(configFile string, logger *log.Logger) error {
	var err error

	// set default values
	k.config = KafkaConfig{
		Brokers:       []string{"localhost:9092"},
		ClientID:      "vFlow.Kafka",
		MaxAttempts:   10,
		QueueSize:     1024,
		BatchSize:     256,
		KeepAlive:     180,
		IOTimeout:     10,
		RequiredAcks:  -1,
		PeriodicFlush: 20,
		VerifySSL:     true,
	}

	// setup logger
	k.logger = logger

	// load configuration file if available
	if err = k.load(configFile); err != nil {
		logger.Println(err)
	}

	// get env config
	k.loadEnv("VFLOW_KAFKA")

	// lookup bootstrap server
	if k.config.BootstrapServer != "" {
		brokers, err := bootstrapLookup(k.config.BootstrapServer)
		if err != nil {
			k.logger.Printf("error getting bootstrap servers: %v", err)
		} else {
			k.config.Brokers = brokers
		}
	}

	// init kafka configuration
	k.config.run = kafka.WriterConfig{
		Brokers: k.config.Brokers,
		Dialer: &kafka.Dialer{
			ClientID:  k.config.ClientID,
			Timeout:   time.Second * time.Duration(k.config.IOTimeout),
			KeepAlive: time.Second * time.Duration(k.config.KeepAlive),
			DualStack: true,
		},
		Balancer:      &kafka.Hash{},
		MaxAttempts:   k.config.MaxAttempts,
		QueueCapacity: k.config.QueueSize,
		BatchSize:     k.config.BatchSize,
		ReadTimeout:   time.Second * time.Duration(k.config.IOTimeout),
		WriteTimeout:  time.Second * time.Duration(k.config.IOTimeout),
		RequiredAcks:  k.config.RequiredAcks,
		Async:         false,
	}

	if tlsConfig := k.tlsConfig(); tlsConfig != nil {
		k.config.run.Dialer.TLS = tlsConfig
		k.logger.Println("Kafka client TLS enabled")
	}

	switch k.config.Compression {
	case "gzip":
		k.config.run.CompressionCodec = gzip.NewCompressionCodec()
	case "lz4":
		k.config.run.CompressionCodec = lz4.NewCompressionCodec()
	case "snappy":
		k.config.run.CompressionCodec = snappy.NewCompressionCodec()
	}

	return err
}

func (k *Kafka) inputMsg(topic string, mCh chan []byte, ec *uint64) {

	k.config.run.Topic = topic
	k.logger.Printf("start producer: Kafka, brokers: %+v, topic: %s\n",
		k.config.run.Brokers, k.config.run.Topic)
	k.producer = kafka.NewWriter(k.config.run)

	batch := make([]kafka.Message, 0, k.config.BatchSize)

	var shutdown = false
	var pflush = false
	var pftimer = time.NewTimer(time.Second * time.Duration(k.config.PeriodicFlush))

	for {
		select {
		case message, ok := <-mCh:
			if ok {
				batch = append(batch, kafka.Message{Value: message})
			} else {
				shutdown = true
			}
		case <-pftimer.C:
			pflush = true
		}

		if len(batch) == k.config.BatchSize || shutdown || pflush {

			if !pftimer.Stop() {
				pflush = false
			}

			err := k.producer.WriteMessages(context.Background(), batch...)

			if err != nil {
				k.logger.Printf("error writing to kafka: %v", err)
				*ec++
			}

			if shutdown {
				k.logger.Printf("shutting down kafka writer, flushed %d records", len(batch))
				break
			}

			pftimer.Reset(time.Second * time.Duration(k.config.PeriodicFlush))
			batch = nil
		}
	}

	k.producer.Close()
}

func (k *Kafka) load(f string) error {
	b, err := ioutil.ReadFile(f)
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(b, &k.config)
	if err != nil {
		return err
	}

	return nil
}

func (k Kafka) tlsConfig() *tls.Config {
	var t *tls.Config

	if k.config.TLSCertFile != "" && k.config.TLSKeyFile != "" && k.config.CAFile != "" {
		cert, err := tls.LoadX509KeyPair(k.config.TLSCertFile, k.config.TLSKeyFile)
		if err != nil {
			k.logger.Fatal("Kafka TLS error: ", err)
		}

		caCert, err := ioutil.ReadFile(k.config.CAFile)
		if err != nil {
			k.logger.Fatal("Kafka TLS error: ", err)
		}

		caCertPool := x509.NewCertPool()
		caCertPool.AppendCertsFromPEM(caCert)

		t = &tls.Config{
			Certificates:       []tls.Certificate{cert},
			RootCAs:            caCertPool,
			InsecureSkipVerify: !k.config.VerifySSL,
		}
	}

	return t
}

func (k *Kafka) loadEnv(prefix string) {
	v := reflect.ValueOf(&k.config).Elem()
	t := v.Type()
	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		env := t.Field(i).Tag.Get("env")
		if env == "" {
			continue
		}

		val, ok := os.LookupEnv(prefix + "_" + env)
		if !ok {
			continue
		}

		switch f.Kind() {
		case reflect.Int:
			valInt, err := strconv.Atoi(val)
			if err != nil {
				k.logger.Println(err)
				continue
			}
			f.SetInt(int64(valInt))
		case reflect.String:
			f.SetString(val)
		case reflect.Slice:
			for _, elm := range strings.Split(val, ";") {
				f.Index(0).SetString(elm)
			}
		case reflect.Bool:
			valBool, err := strconv.ParseBool(val)
			if err != nil {
				k.logger.Println(err)
				continue
			}
			f.SetBool(valBool)
		}
	}
}

func bootstrapLookup(endpoint string) ([]string, error) {

	var err error
	var brokers []string

	host, port, err := net.SplitHostPort(endpoint)
	if err != nil {
		return brokers, err
	}

	addrs, err := net.LookupHost(host)

	if err != nil {
		return brokers, err
	}

	for _, ip := range addrs {
		brokers = append(brokers, strings.Join([]string{ip, port}, ":"))
	}

	return brokers, err
}
<commit_msg>fix unhandled error<commit_after>//: ----------------------------------------------------------------------------
//: Copyright (C) 2018 Alan Willis.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    segmentio.go
//: details: vflow kafka producer plugin
//: author:  Alan Willis
//: date:    12/05/2018
//:
//: 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.
//: ----------------------------------------------------------------------------
// +build kafkav2

package producer

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"log"
	"net"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/segmentio/kafka-go"
	"github.com/segmentio/kafka-go/gzip"
	"github.com/segmentio/kafka-go/lz4"
	"github.com/segmentio/kafka-go/snappy"
	"gopkg.in/yaml.v2"
)

// Kafka represents kafka producer
type Kafka struct {
	producer *kafka.Writer
	config   KafkaConfig
	logger   *log.Logger
}

// Config represents kafka configuration
type KafkaConfig struct {
	run             kafka.WriterConfig
	Brokers         []string `yaml:"brokers" env:"BROKERS"`
	BootstrapServer string   `yaml:"bootstrap-server" env:"BOOTSTRAP_SERVER"`
	ClientID        string   `yaml:"client-id" env:"CLIENT_ID"`
	Compression     string   `yaml:"compression" env:"COMPRESSION"`
	MaxAttempts     int      `yaml:"max-attempts" env:"MAX_ATTEMPTS"`
	QueueSize       int      `yaml:"queue-size" env:"QUEUE_SIZE"`
	BatchSize       int      `yaml:"batch-size" env:"BATCH_SIZE"`
	KeepAlive       int      `yaml:"keepalive" env:"KEEPALIVE"`
	IOTimeout       int      `yaml:"connect-timeout" env:"CONNECT_TIMEOUT"`
	RequiredAcks    int      `yaml:"required-acks" env:"REQUIRED_ACKS"`
	PeriodicFlush   int      `yaml:"pflush" env:"PERIODIC_FLUSH"`
	TLSCertFile     string   `yaml:"tls-cert" env:"TLS_CERT"`
	TLSKeyFile      string   `yaml:"tls-key" env:"TLS_KEY"`
	CAFile          string   `yaml:"ca-file" env:"CA_FILE"`
	VerifySSL       bool     `yaml:"verify-ssl" env:"VERIFY_SSL"`
}

func (k *Kafka) setup(configFile string, logger *log.Logger) error {
	var err error

	// set default values
	k.config = KafkaConfig{
		Brokers:       []string{"localhost:9092"},
		ClientID:      "vFlow.Kafka",
		MaxAttempts:   10,
		QueueSize:     1024,
		BatchSize:     256,
		KeepAlive:     180,
		IOTimeout:     10,
		RequiredAcks:  -1,
		PeriodicFlush: 20,
		VerifySSL:     true,
	}

	// setup logger
	k.logger = logger

	// load configuration file if available
	if err = k.load(configFile); err != nil {
		logger.Println(err)
	}

	// get env config
	k.loadEnv("VFLOW_KAFKA")

	// lookup bootstrap server
	if k.config.BootstrapServer != "" {
		brokers, err := bootstrapLookup(k.config.BootstrapServer)
		if err != nil {
			k.logger.Printf("error getting bootstrap servers: %v", err)
		} else {
			k.config.Brokers = brokers
		}
	}

	// init kafka configuration
	k.config.run = kafka.WriterConfig{
		Brokers: k.config.Brokers,
		Dialer: &kafka.Dialer{
			ClientID:  k.config.ClientID,
			Timeout:   time.Second * time.Duration(k.config.IOTimeout),
			KeepAlive: time.Second * time.Duration(k.config.KeepAlive),
			DualStack: true,
		},
		Balancer:      &kafka.Hash{},
		MaxAttempts:   k.config.MaxAttempts,
		QueueCapacity: k.config.QueueSize,
		BatchSize:     k.config.BatchSize,
		ReadTimeout:   time.Second * time.Duration(k.config.IOTimeout),
		WriteTimeout:  time.Second * time.Duration(k.config.IOTimeout),
		RequiredAcks:  k.config.RequiredAcks,
		Async:         false,
	}

	if tlsConfig := k.tlsConfig(); tlsConfig != nil {
		k.config.run.Dialer.TLS = tlsConfig
		k.logger.Println("Kafka client TLS enabled")
	}

	switch k.config.Compression {
	case "gzip":
		k.config.run.CompressionCodec = gzip.NewCompressionCodec()
	case "lz4":
		k.config.run.CompressionCodec = lz4.NewCompressionCodec()
	case "snappy":
		k.config.run.CompressionCodec = snappy.NewCompressionCodec()
	}

	return err
}

func (k *Kafka) inputMsg(topic string, mCh chan []byte, ec *uint64) {

	k.config.run.Topic = topic
	k.logger.Printf("start producer: Kafka, brokers: %+v, topic: %s\n",
		k.config.run.Brokers, k.config.run.Topic)
	k.producer = kafka.NewWriter(k.config.run)

	batch := make([]kafka.Message, 0, k.config.BatchSize)

	var shutdown = false
	var pflush = false
	var pftimer = time.NewTimer(time.Second * time.Duration(k.config.PeriodicFlush))

	for {
		select {
		case message, ok := <-mCh:
			if ok {
				batch = append(batch, kafka.Message{Value: message})
			} else {
				shutdown = true
			}
		case <-pftimer.C:
			pflush = true
		}

		if len(batch) == k.config.BatchSize || shutdown || pflush {

			if !pftimer.Stop() {
				pflush = false
			}

			err := k.producer.WriteMessages(context.Background(), batch...)

			if err != nil {
				k.logger.Printf("error writing to kafka: %v", err)
				*ec++
			}

			if shutdown {
				k.logger.Printf("shutting down kafka writer, flushed %d records", len(batch))
				break
			}

			pftimer.Reset(time.Second * time.Duration(k.config.PeriodicFlush))
			batch = nil
		}
	}

	if err := k.producer.Close(); err != nil {
		k.logger.Printf("error shutting down kafka writer: %v", err)
	}
}

func (k *Kafka) load(f string) error {
	b, err := ioutil.ReadFile(f)
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(b, &k.config)
	if err != nil {
		return err
	}

	return nil
}

func (k Kafka) tlsConfig() *tls.Config {
	var t *tls.Config

	if k.config.TLSCertFile != "" && k.config.TLSKeyFile != "" && k.config.CAFile != "" {
		cert, err := tls.LoadX509KeyPair(k.config.TLSCertFile, k.config.TLSKeyFile)
		if err != nil {
			k.logger.Fatal("Kafka TLS error: ", err)
		}

		caCert, err := ioutil.ReadFile(k.config.CAFile)
		if err != nil {
			k.logger.Fatal("Kafka TLS error: ", err)
		}

		caCertPool := x509.NewCertPool()
		caCertPool.AppendCertsFromPEM(caCert)

		t = &tls.Config{
			Certificates:       []tls.Certificate{cert},
			RootCAs:            caCertPool,
			InsecureSkipVerify: !k.config.VerifySSL,
		}
	}

	return t
}

func (k *Kafka) loadEnv(prefix string) {
	v := reflect.ValueOf(&k.config).Elem()
	t := v.Type()
	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		env := t.Field(i).Tag.Get("env")
		if env == "" {
			continue
		}

		val, ok := os.LookupEnv(prefix + "_" + env)
		if !ok {
			continue
		}

		switch f.Kind() {
		case reflect.Int:
			valInt, err := strconv.Atoi(val)
			if err != nil {
				k.logger.Println(err)
				continue
			}
			f.SetInt(int64(valInt))
		case reflect.String:
			f.SetString(val)
		case reflect.Slice:
			for _, elm := range strings.Split(val, ";") {
				f.Index(0).SetString(elm)
			}
		case reflect.Bool:
			valBool, err := strconv.ParseBool(val)
			if err != nil {
				k.logger.Println(err)
				continue
			}
			f.SetBool(valBool)
		}
	}
}

func bootstrapLookup(endpoint string) ([]string, error) {

	var err error
	var brokers []string

	host, port, err := net.SplitHostPort(endpoint)
	if err != nil {
		return brokers, err
	}

	addrs, err := net.LookupHost(host)

	if err != nil {
		return brokers, err
	}

	for _, ip := range addrs {
		brokers = append(brokers, strings.Join([]string{ip, port}, ":"))
	}

	return brokers, err
}
<|endoftext|>


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

package state_test

import (
	"time"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
	"gopkg.in/mgo.v2/bson"

	"github.com/juju/juju/state"
	"github.com/juju/juju/testing"
	"github.com/juju/juju/testing/factory"
)

type MeterStateSuite struct {
	ConnSuite
	unit           *state.Unit
	factory        *factory.Factory
	metricsManager *state.MetricsManager
}

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

func (s *MeterStateSuite) SetUpTest(c *gc.C) {
	s.ConnSuite.SetUpTest(c)
	s.factory = factory.NewFactory(s.State)
	s.unit = s.factory.MakeUnit(c, nil)
	c.Assert(s.unit.Series(), gc.Equals, "quantal")
	var err error
	s.metricsManager, err = s.State.MetricsManager()
	c.Assert(err, jc.ErrorIsNil)
}

func (s *MeterStateSuite) TestMeterStatus(c *gc.C) {
	status, err := s.unit.GetMeterStatus()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(status.Code, gc.Equals, state.MeterNotSet)
	err = s.unit.SetMeterStatus("GREEN", "Additional information.")
	c.Assert(err, jc.ErrorIsNil)
	status, err = s.unit.GetMeterStatus()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(status.Code, gc.Equals, state.MeterGreen)
}

func (s *MeterStateSuite) TestMeterStatusIncludesEnvUUID(c *gc.C) {
	jujuDB := s.MgoSuite.Session.DB("juju")
	meterStatus := jujuDB.C("meterStatus")
	var docs []bson.M
	err := meterStatus.Find(nil).All(&docs)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(docs, gc.HasLen, 1)
	c.Assert(docs[0]["env-uuid"], gc.Equals, s.State.EnvironUUID())
}

func (s *MeterStateSuite) TestSetMeterStatusIncorrect(c *gc.C) {
	err := s.unit.SetMeterStatus("NOT SET", "Additional information.")
	c.Assert(err, gc.ErrorMatches, `invalid meter status "NOT SET"`)
	status, err := s.unit.GetMeterStatus()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(status.Code, gc.Equals, state.MeterNotSet)

	err = s.unit.SetMeterStatus("this-is-not-a-valid-status", "Additional information.")
	c.Assert(err, gc.ErrorMatches, `invalid meter status "NOT AVAILABLE"`)
	status, err = s.unit.GetMeterStatus()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(status.Code, gc.Equals, state.MeterNotSet)
}

func (s *MeterStateSuite) TestSetMeterStatusWhenDying(c *gc.C) {
	preventUnitDestroyRemove(c, s.unit)
	testWhenDying(c, s.unit, contentionErr, contentionErr, func() error {
		err := s.unit.SetMeterStatus("GREEN", "Additional information.")
		if err != nil {
			return err
		}
		status, err := s.unit.GetMeterStatus()
		c.Assert(err, jc.ErrorIsNil)
		c.Assert(status.Code, gc.Equals, state.MeterNotSet)
		return nil
	})
}

func (s *MeterStateSuite) TestMeterStatusRemovedWithUnit(c *gc.C) {
	err := s.unit.SetMeterStatus("GREEN", "Information.")
	c.Assert(err, jc.ErrorIsNil)
	err = s.unit.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)
	err = s.unit.Remove()
	c.Assert(err, jc.ErrorIsNil)
	status, err := s.unit.GetMeterStatus()
	c.Assert(err, gc.ErrorMatches, "cannot retrieve meter status for unit .*: not found")
	c.Assert(status.Code, gc.Equals, state.MeterNotAvailable)
}

func (s *MeterStateSuite) TestMeterStatusWatcherRespondstoMeterStatus(c *gc.C) {
	watcher := s.unit.WatchMeterStatus()
	err := s.unit.SetMeterStatus("GREEN", "Information.")
	c.Assert(err, jc.ErrorIsNil)
	assertMeterStatusChanged(c, watcher)
}

func (s *MeterStateSuite) TestMeterStatusWatcherRespondsToMetricsManager(c *gc.C) {
	mm, err := s.State.MetricsManager()
	c.Assert(err, jc.ErrorIsNil)
	watcher := s.unit.WatchMeterStatus()
	err = mm.SetLastSuccessfulSend(time.Now())
	c.Assert(err, jc.ErrorIsNil)
	for i := 0; i < 3; i++ {
		err := mm.IncrementConsecutiveErrors()
		c.Assert(err, jc.ErrorIsNil)
	}
	status := mm.MeterStatus()
	c.Assert(status.Code, gc.Equals, state.MeterAmber)
	assertMeterStatusChanged(c, watcher)
}

func assertMeterStatusChanged(c *gc.C, w state.NotifyWatcher) {
	select {
	case <-w.Changes():
	case <-time.After(testing.LongWait):
		c.Fatalf("expected event from watcher by now")
	}
}

func assertMetricsManagerAmberState(c *gc.C, metricsManager *state.MetricsManager) {
	err := metricsManager.SetLastSuccessfulSend(time.Now())
	c.Assert(err, jc.ErrorIsNil)
	for i := 0; i < 3; i++ {
		err := metricsManager.IncrementConsecutiveErrors()
		c.Assert(err, jc.ErrorIsNil)
	}
	status := metricsManager.MeterStatus()
	c.Assert(status.Code, gc.Equals, state.MeterAmber)
}

// TODO (mattyw) This function could be moved into a metricsmanager testing package.
func assertMetricsManagerRedState(c *gc.C, metricsManager *state.MetricsManager) {
	// To enter the red state we need to set a last successful send as over 1 week ago
	err := metricsManager.SetLastSuccessfulSend(time.Now().Add(-8 * 24 * time.Hour))
	c.Assert(err, jc.ErrorIsNil)
	for i := 0; i < 3; i++ {
		err := metricsManager.IncrementConsecutiveErrors()
		c.Assert(err, jc.ErrorIsNil)
	}
	status := metricsManager.MeterStatus()
	c.Assert(status.Code, gc.Equals, state.MeterRed)
}

// TestMeterStatusMetricsManagerCombinations test every possible combination
// of meter status from the unit and the metrics manager.
func (s *MeterStateSuite) TestMeterStatusMetricsManagerCombinations(c *gc.C) {
	greenMetricsMangager := func() {}
	amberMetricsManager := func() {
		assertMetricsManagerAmberState(c, s.metricsManager)
	}
	redMetricsManager := func() {
		assertMetricsManagerRedState(c, s.metricsManager)
	}
	greenUnit := func() {
		err := s.unit.SetMeterStatus("GREEN", "Unit")
		c.Assert(err, jc.ErrorIsNil)
	}
	amberUnit := func() {
		err := s.unit.SetMeterStatus("AMBER", "Unit")
		c.Assert(err, jc.ErrorIsNil)
	}
	redUnit := func() {
		err := s.unit.SetMeterStatus("RED", "Unit")
		c.Assert(err, jc.ErrorIsNil)
	}

	tests := []struct {
		about          string
		metricsManager func()
		unit           func()
		expectedCode   state.MeterStatusCode
		expectedInfo   string
	}{{
		"green metrics manager and green unit returns green overall",
		greenMetricsMangager,
		greenUnit,
		state.MeterGreen,
		"Unit",
	}, {
		"amber metrics manager and amber unit returns amber overall",
		amberMetricsManager,
		amberUnit,
		state.MeterAmber,
		"Unit",
	}, {
		"red metrics manager and red unit returns red overall",
		redMetricsManager,
		redUnit,
		state.MeterRed,
		"failed to send metrics, exceeded grace period",
	}, {

		"red metrics manager and amber unit returns red overall",
		redMetricsManager,
		amberUnit,
		state.MeterRed,
		"failed to send metrics, exceeded grace period",
	}, {
		"red metrics manager and green unit returns red overall",
		redMetricsManager,
		greenUnit,
		state.MeterRed,
		"failed to send metrics, exceeded grace period",
	}, {
		"amber metrics manager and red unit returns red overall",
		amberMetricsManager,
		redUnit,
		state.MeterRed,
		"Unit",
	}, {
		"amber metrics manager and green unit returns amber overall",
		amberMetricsManager,
		greenUnit,
		state.MeterAmber,
		"failed to send metrics",
	}, {
		"green metrics manager and red unit returns red overall",
		greenMetricsMangager,
		redUnit,
		state.MeterRed,
		"Unit",
	}, {
		"green metrics manager and amber unit returns amber overall",
		greenMetricsMangager,
		amberUnit,
		state.MeterAmber,
		"Unit",
	}}

	for i, test := range tests {
		c.Logf("running test %d %s", i, test.about)
		test.metricsManager()
		test.unit()
		status, err := s.unit.GetMeterStatus()
		c.Assert(err, jc.ErrorIsNil)
		c.Check(status.Code, gc.Equals, test.expectedCode)
	}
}

func (s *MeterStateSuite) TestMeterStatusCombination(c *gc.C) {
	var (
		Red          = state.MeterStatus{state.MeterRed, ""}
		Amber        = state.MeterStatus{state.MeterAmber, ""}
		Green        = state.MeterStatus{state.MeterGreen, ""}
		NotSet       = state.MeterStatus{state.MeterNotSet, ""}
		NotAvailable = state.MeterStatus{state.MeterNotAvailable, ""}
	)
	c.Assert(state.CombineMeterStatus(Red, Red).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Red, Amber).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Red, Green).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Red, NotSet).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Red, NotAvailable).Code, gc.Equals, NotAvailable.Code)

	c.Assert(state.CombineMeterStatus(Amber, Red).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Amber, Amber).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(Amber, Green).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(Amber, NotSet).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(Amber, NotAvailable).Code, gc.Equals, NotAvailable.Code)

	c.Assert(state.CombineMeterStatus(Green, Red).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Green, Amber).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(Green, Green).Code, gc.Equals, Green.Code)
	c.Assert(state.CombineMeterStatus(Green, NotSet).Code, gc.Equals, NotSet.Code)
	c.Assert(state.CombineMeterStatus(Green, NotAvailable).Code, gc.Equals, NotAvailable.Code)

	c.Assert(state.CombineMeterStatus(NotSet, Red).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(NotSet, Amber).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(NotSet, Green).Code, gc.Equals, NotSet.Code)
	c.Assert(state.CombineMeterStatus(NotSet, NotSet).Code, gc.Equals, NotSet.Code)
	c.Assert(state.CombineMeterStatus(NotSet, NotAvailable).Code, gc.Equals, NotAvailable.Code)

	c.Assert(state.CombineMeterStatus(NotAvailable, Red).Code, gc.Equals, NotAvailable.Code)
	c.Assert(state.CombineMeterStatus(NotAvailable, Amber).Code, gc.Equals, NotAvailable.Code)
	c.Assert(state.CombineMeterStatus(NotAvailable, Green).Code, gc.Equals, NotAvailable.Code)
	c.Assert(state.CombineMeterStatus(NotAvailable, NotSet).Code, gc.Equals, NotAvailable.Code)
	c.Assert(state.CombineMeterStatus(NotAvailable, NotAvailable).Code, gc.Equals, NotAvailable.Code)
}
<commit_msg>meterstatesuite: Added combined test<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state_test

import (
	"time"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
	"gopkg.in/mgo.v2/bson"

	"github.com/juju/juju/state"
	"github.com/juju/juju/testing"
	"github.com/juju/juju/testing/factory"
)

type MeterStateSuite struct {
	ConnSuite
	unit           *state.Unit
	factory        *factory.Factory
	metricsManager *state.MetricsManager
}

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

func (s *MeterStateSuite) SetUpTest(c *gc.C) {
	s.ConnSuite.SetUpTest(c)
	s.factory = factory.NewFactory(s.State)
	s.unit = s.factory.MakeUnit(c, nil)
	c.Assert(s.unit.Series(), gc.Equals, "quantal")
	var err error
	s.metricsManager, err = s.State.MetricsManager()
	c.Assert(err, jc.ErrorIsNil)
}

func (s *MeterStateSuite) TestMeterStatus(c *gc.C) {
	status, err := s.unit.GetMeterStatus()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(status.Code, gc.Equals, state.MeterNotSet)
	err = s.unit.SetMeterStatus("GREEN", "Additional information.")
	c.Assert(err, jc.ErrorIsNil)
	status, err = s.unit.GetMeterStatus()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(status.Code, gc.Equals, state.MeterGreen)
}

func (s *MeterStateSuite) TestMeterStatusIncludesEnvUUID(c *gc.C) {
	jujuDB := s.MgoSuite.Session.DB("juju")
	meterStatus := jujuDB.C("meterStatus")
	var docs []bson.M
	err := meterStatus.Find(nil).All(&docs)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(docs, gc.HasLen, 1)
	c.Assert(docs[0]["env-uuid"], gc.Equals, s.State.EnvironUUID())
}

func (s *MeterStateSuite) TestSetMeterStatusIncorrect(c *gc.C) {
	err := s.unit.SetMeterStatus("NOT SET", "Additional information.")
	c.Assert(err, gc.ErrorMatches, `invalid meter status "NOT SET"`)
	status, err := s.unit.GetMeterStatus()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(status.Code, gc.Equals, state.MeterNotSet)

	err = s.unit.SetMeterStatus("this-is-not-a-valid-status", "Additional information.")
	c.Assert(err, gc.ErrorMatches, `invalid meter status "NOT AVAILABLE"`)
	status, err = s.unit.GetMeterStatus()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(status.Code, gc.Equals, state.MeterNotSet)
}

func (s *MeterStateSuite) TestSetMeterStatusWhenDying(c *gc.C) {
	preventUnitDestroyRemove(c, s.unit)
	testWhenDying(c, s.unit, contentionErr, contentionErr, func() error {
		err := s.unit.SetMeterStatus("GREEN", "Additional information.")
		if err != nil {
			return err
		}
		status, err := s.unit.GetMeterStatus()
		c.Assert(err, jc.ErrorIsNil)
		c.Assert(status.Code, gc.Equals, state.MeterNotSet)
		return nil
	})
}

func (s *MeterStateSuite) TestMeterStatusRemovedWithUnit(c *gc.C) {
	err := s.unit.SetMeterStatus("GREEN", "Information.")
	c.Assert(err, jc.ErrorIsNil)
	err = s.unit.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)
	err = s.unit.Remove()
	c.Assert(err, jc.ErrorIsNil)
	status, err := s.unit.GetMeterStatus()
	c.Assert(err, gc.ErrorMatches, "cannot retrieve meter status for unit .*: not found")
	c.Assert(status.Code, gc.Equals, state.MeterNotAvailable)
}

func (s *MeterStateSuite) TestMeterStatusWatcherRespondstoMeterStatus(c *gc.C) {
	watcher := s.unit.WatchMeterStatus()
	assertMeterStatusChanged(c, watcher)
	err := s.unit.SetMeterStatus("GREEN", "Information.")
	c.Assert(err, jc.ErrorIsNil)
	assertMeterStatusChanged(c, watcher)
}

func (s *MeterStateSuite) TestMeterStatusWatcherRespondsToMetricsManager(c *gc.C) {
	mm, err := s.State.MetricsManager()
	c.Assert(err, jc.ErrorIsNil)
	watcher := s.unit.WatchMeterStatus()
	assertMeterStatusChanged(c, watcher)
	err = mm.SetLastSuccessfulSend(time.Now())
	c.Assert(err, jc.ErrorIsNil)
	for i := 0; i < 3; i++ {
		err := mm.IncrementConsecutiveErrors()
		c.Assert(err, jc.ErrorIsNil)
	}
	status := mm.MeterStatus()
	c.Assert(status.Code, gc.Equals, state.MeterAmber)
	assertMeterStatusChanged(c, watcher)
}

func (s *MeterStateSuite) TestMeterStatusWatcherRespondsToMetricsManagerAndStatus(c *gc.C) {
	mm, err := s.State.MetricsManager()
	c.Assert(err, jc.ErrorIsNil)
	watcher := s.unit.WatchMeterStatus()
	assertMeterStatusChanged(c, watcher)
	err = mm.SetLastSuccessfulSend(time.Now())
	c.Assert(err, jc.ErrorIsNil)
	for i := 0; i < 3; i++ {
		err := mm.IncrementConsecutiveErrors()
		c.Assert(err, jc.ErrorIsNil)
	}
	status := mm.MeterStatus()
	c.Assert(status.Code, gc.Equals, state.MeterAmber)
	err = s.unit.SetMeterStatus("GREEN", "Information.")
	c.Assert(err, jc.ErrorIsNil)
	assertMeterStatusChanged(c, watcher)
	select {
	case <-watcher.Changes():
		c.Fatalf("unexpected event from watcher")
	case <-time.After(testing.ShortWait):
	}
}

func assertMeterStatusChanged(c *gc.C, w state.NotifyWatcher) {
	select {
	case <-w.Changes():
	case <-time.After(testing.LongWait):
		c.Fatalf("expected event from watcher by now")
	}
}

func assertMetricsManagerAmberState(c *gc.C, metricsManager *state.MetricsManager) {
	err := metricsManager.SetLastSuccessfulSend(time.Now())
	c.Assert(err, jc.ErrorIsNil)
	for i := 0; i < 3; i++ {
		err := metricsManager.IncrementConsecutiveErrors()
		c.Assert(err, jc.ErrorIsNil)
	}
	status := metricsManager.MeterStatus()
	c.Assert(status.Code, gc.Equals, state.MeterAmber)
}

// TODO (mattyw) This function could be moved into a metricsmanager testing package.
func assertMetricsManagerRedState(c *gc.C, metricsManager *state.MetricsManager) {
	// To enter the red state we need to set a last successful send as over 1 week ago
	err := metricsManager.SetLastSuccessfulSend(time.Now().Add(-8 * 24 * time.Hour))
	c.Assert(err, jc.ErrorIsNil)
	for i := 0; i < 3; i++ {
		err := metricsManager.IncrementConsecutiveErrors()
		c.Assert(err, jc.ErrorIsNil)
	}
	status := metricsManager.MeterStatus()
	c.Assert(status.Code, gc.Equals, state.MeterRed)
}

// TestMeterStatusMetricsManagerCombinations test every possible combination
// of meter status from the unit and the metrics manager.
func (s *MeterStateSuite) TestMeterStatusMetricsManagerCombinations(c *gc.C) {
	greenMetricsMangager := func() {}
	amberMetricsManager := func() {
		assertMetricsManagerAmberState(c, s.metricsManager)
	}
	redMetricsManager := func() {
		assertMetricsManagerRedState(c, s.metricsManager)
	}
	greenUnit := func() {
		err := s.unit.SetMeterStatus("GREEN", "Unit")
		c.Assert(err, jc.ErrorIsNil)
	}
	amberUnit := func() {
		err := s.unit.SetMeterStatus("AMBER", "Unit")
		c.Assert(err, jc.ErrorIsNil)
	}
	redUnit := func() {
		err := s.unit.SetMeterStatus("RED", "Unit")
		c.Assert(err, jc.ErrorIsNil)
	}

	tests := []struct {
		about          string
		metricsManager func()
		unit           func()
		expectedCode   state.MeterStatusCode
		expectedInfo   string
	}{{
		"green metrics manager and green unit returns green overall",
		greenMetricsMangager,
		greenUnit,
		state.MeterGreen,
		"Unit",
	}, {
		"amber metrics manager and amber unit returns amber overall",
		amberMetricsManager,
		amberUnit,
		state.MeterAmber,
		"Unit",
	}, {
		"red metrics manager and red unit returns red overall",
		redMetricsManager,
		redUnit,
		state.MeterRed,
		"failed to send metrics, exceeded grace period",
	}, {

		"red metrics manager and amber unit returns red overall",
		redMetricsManager,
		amberUnit,
		state.MeterRed,
		"failed to send metrics, exceeded grace period",
	}, {
		"red metrics manager and green unit returns red overall",
		redMetricsManager,
		greenUnit,
		state.MeterRed,
		"failed to send metrics, exceeded grace period",
	}, {
		"amber metrics manager and red unit returns red overall",
		amberMetricsManager,
		redUnit,
		state.MeterRed,
		"Unit",
	}, {
		"amber metrics manager and green unit returns amber overall",
		amberMetricsManager,
		greenUnit,
		state.MeterAmber,
		"failed to send metrics",
	}, {
		"green metrics manager and red unit returns red overall",
		greenMetricsMangager,
		redUnit,
		state.MeterRed,
		"Unit",
	}, {
		"green metrics manager and amber unit returns amber overall",
		greenMetricsMangager,
		amberUnit,
		state.MeterAmber,
		"Unit",
	}}

	for i, test := range tests {
		c.Logf("running test %d %s", i, test.about)
		test.metricsManager()
		test.unit()
		status, err := s.unit.GetMeterStatus()
		c.Assert(err, jc.ErrorIsNil)
		c.Check(status.Code, gc.Equals, test.expectedCode)
	}
}

func (s *MeterStateSuite) TestMeterStatusCombination(c *gc.C) {
	var (
		Red          = state.MeterStatus{state.MeterRed, ""}
		Amber        = state.MeterStatus{state.MeterAmber, ""}
		Green        = state.MeterStatus{state.MeterGreen, ""}
		NotSet       = state.MeterStatus{state.MeterNotSet, ""}
		NotAvailable = state.MeterStatus{state.MeterNotAvailable, ""}
	)
	c.Assert(state.CombineMeterStatus(Red, Red).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Red, Amber).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Red, Green).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Red, NotSet).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Red, NotAvailable).Code, gc.Equals, NotAvailable.Code)

	c.Assert(state.CombineMeterStatus(Amber, Red).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Amber, Amber).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(Amber, Green).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(Amber, NotSet).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(Amber, NotAvailable).Code, gc.Equals, NotAvailable.Code)

	c.Assert(state.CombineMeterStatus(Green, Red).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(Green, Amber).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(Green, Green).Code, gc.Equals, Green.Code)
	c.Assert(state.CombineMeterStatus(Green, NotSet).Code, gc.Equals, NotSet.Code)
	c.Assert(state.CombineMeterStatus(Green, NotAvailable).Code, gc.Equals, NotAvailable.Code)

	c.Assert(state.CombineMeterStatus(NotSet, Red).Code, gc.Equals, Red.Code)
	c.Assert(state.CombineMeterStatus(NotSet, Amber).Code, gc.Equals, Amber.Code)
	c.Assert(state.CombineMeterStatus(NotSet, Green).Code, gc.Equals, NotSet.Code)
	c.Assert(state.CombineMeterStatus(NotSet, NotSet).Code, gc.Equals, NotSet.Code)
	c.Assert(state.CombineMeterStatus(NotSet, NotAvailable).Code, gc.Equals, NotAvailable.Code)

	c.Assert(state.CombineMeterStatus(NotAvailable, Red).Code, gc.Equals, NotAvailable.Code)
	c.Assert(state.CombineMeterStatus(NotAvailable, Amber).Code, gc.Equals, NotAvailable.Code)
	c.Assert(state.CombineMeterStatus(NotAvailable, Green).Code, gc.Equals, NotAvailable.Code)
	c.Assert(state.CombineMeterStatus(NotAvailable, NotSet).Code, gc.Equals, NotAvailable.Code)
	c.Assert(state.CombineMeterStatus(NotAvailable, NotAvailable).Code, gc.Equals, NotAvailable.Code)
}
<|endoftext|>


<commit_before>package qb

import (
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

type unknownType struct{}

type mapperTestUser struct {
	ID          string `qb:"constraints:primary_key"`
	FacebookID  int64  `qb:"constraints:ref(facebook.id)"`
	ProfileID   int64  `qb:"constraints:ref(profile.id)"`
	ProfileName string `qb:"constraints:ref(profile.name)"`
	Email       string `qb:"type:varchar(255); constraints:unique,notnull"`
	FullName    string `qb:"constraints:notnull,default"`
	Password    string `qb:"type:text"`
	UserType    string `qb:"constraints:default(guest)"`
	Premium     bool
	CreatedAt   time.Time  `qb:"constraints:notnull"`
	DeletedAt   *time.Time `qb:"constraints:null"`
	Level       int
	Money       float32
	Score       float64
	Unknown     unknownType
}

type MapperTestSqliteAutoIncrementUser struct {
	ID int64 `qb:"constraints:auto_increment"`
}

func TestMapper(t *testing.T) {

	mapper := NewMapper("mysql")

	userTable, err := mapper.ToTable(mapperTestUser{})

	assert.Nil(t, err)
	fmt.Println(userTable.SQL())
}

func TestMapperSqliteAutoIncrement(t *testing.T) {

	mapper := NewMapper("sqlite3")
	sqliteAutoIncrementUserTable, err := mapper.ToTable(MapperTestSqliteAutoIncrementUser{})

	assert.Nil(t, err)
	fmt.Println(sqliteAutoIncrementUserTable.SQL())

}

type MapperTestUserErr struct {
	ID    string `qb:"type:varchar(255);tag_should_raise_err:val;"`
	Email string `qb:"wrongtag:"`
}

func TestMapperError(t *testing.T) {

	mapper := NewMapper("postgres")

	userErrTable, err := mapper.ToTable(MapperTestUserErr{})

	assert.NotNil(t, err)
	assert.Empty(t, userErrTable)
}

type InvalidConstraint struct {
	ID string `qb:"constraints:invalid_constraint"`
}

func TestMapperInvalidConstraint(t *testing.T) {

	mapper := NewMapper("mysql")

	invalidConstraintTable, err := mapper.ToTable(InvalidConstraint{})

	assert.Nil(t, invalidConstraintTable)
	assert.NotNil(t, err)
}

func TestMapperUtilFuncs(t *testing.T) {

	mapper := NewMapper("mysql")

	assert.Equal(t, mapper.ColName("CreatedAt"), "created_at")

	kv := mapper.ToMap(MapperTestUserErr{})
	assert.Equal(t, kv, map[string]interface{}{})
}
<commit_msg>fix #34<commit_after>package qb

import (
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

func TestMapper(t *testing.T) {

	type UnknownType struct{}

	type User struct {
		ID          string `qb:"constraints:primary_key"`
		FacebookID  int64  `qb:"constraints:ref(facebook.id)"`
		ProfileID   int64  `qb:"constraints:ref(profile.id)"`
		ProfileName string `qb:"constraints:ref(profile.name)"`
		Email       string `qb:"type:varchar(255); constraints:unique,notnull"`
		FullName    string `qb:"constraints:notnull,default"`
		Password    string `qb:"type:text"`
		UserType    string `qb:"constraints:default(guest)"`
		Premium     bool
		CreatedAt   time.Time  `qb:"constraints:notnull"`
		DeletedAt   *time.Time `qb:"constraints:null"`
		Level       int
		Money       float32
		Score       float64
		Unknown     UnknownType
	}

	mapper := NewMapper("mysql")

	userTable, err := mapper.ToTable(User{})

	assert.Nil(t, err)
	fmt.Println(userTable.SQL())
}

func TestMapperSqliteAutoIncrement(t *testing.T) {

	type User struct {
		ID int64 `qb:"constraints:auto_increment"`
	}

	mapper := NewMapper("sqlite3")
	table, err := mapper.ToTable(User{})

	assert.Nil(t, err)
	fmt.Println(table.SQL())
}

func TestMapperError(t *testing.T) {

	type UserErr struct {
		ID    string `qb:"type:varchar(255);tag_should_raise_err:val;"`
		Email string `qb:"wrongtag:"`
	}

	mapper := NewMapper("postgres")

	userErrTable, err := mapper.ToTable(UserErr{})

	assert.NotNil(t, err)
	assert.Empty(t, userErrTable)
}

type InvalidConstraint struct {
	ID string `qb:"constraints:invalid_constraint"`
}

func TestMapperInvalidConstraint(t *testing.T) {

	mapper := NewMapper("mysql")

	invalidConstraintTable, err := mapper.ToTable(InvalidConstraint{})

	assert.Nil(t, invalidConstraintTable)
	assert.NotNil(t, err)
}

func TestMapperUtilFuncs(t *testing.T) {

	type UserErr struct {
		ID    string `qb:"type:varchar(255);tag_should_raise_err:val;"`
		Email string `qb:"wrongtag:"`
	}

	mapper := NewMapper("mysql")

	assert.Equal(t, mapper.ColName("CreatedAt"), "created_at")

	kv := mapper.ToMap(UserErr{})
	assert.Equal(t, kv, map[string]interface{}{})
}
<|endoftext|>


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

package discoverspaces_test

import (
	"errors"

	"github.com/juju/names"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/apiserver/common"
	"github.com/juju/juju/apiserver/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	apiservertesting "github.com/juju/juju/apiserver/testing"
	coretesting "github.com/juju/juju/testing"
)

type DiscoverSpacesSuite struct {
	coretesting.BaseSuite
	apiservertesting.StubNetwork

	resources  *common.Resources
	authorizer apiservertesting.FakeAuthorizer
	facade     *discoverspaces.DiscoverSpacesAPI
}

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

func (s *DiscoverSpacesSuite) SetUpSuite(c *gc.C) {
	s.StubNetwork.SetUpSuite(c)
	s.BaseSuite.SetUpSuite(c)
}

func (s *DiscoverSpacesSuite) TearDownSuite(c *gc.C) {
	s.BaseSuite.TearDownSuite(c)
}

func (s *DiscoverSpacesSuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)
	apiservertesting.BackingInstance.SetUp(
		c,
		apiservertesting.StubZonedEnvironName,
		apiservertesting.WithZones,
		apiservertesting.WithSpaces,
		apiservertesting.WithSubnets)

	s.resources = common.NewResources()
	s.authorizer = apiservertesting.FakeAuthorizer{
		Tag:            names.NewUserTag("admin"),
		EnvironManager: true,
	}

	var err error
	s.facade, err = discoverspaces.NewDiscoverSpacesAPIWithBacking(
		apiservertesting.BackingInstance, s.resources, s.authorizer,
	)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(s.facade, gc.NotNil)
}

func (s *DiscoverSpacesSuite) TearDownTest(c *gc.C) {
	if s.resources != nil {
		s.resources.StopAll()
	}
	s.BaseSuite.TearDownTest(c)
}

func (s *DiscoverSpacesSuite) TestEnvironConfigFailure(c *gc.C) {
	apiservertesting.BackingInstance.SetErrors(errors.New("boom"))

	result, err := s.facade.EnvironConfig()
	c.Assert(err, gc.ErrorMatches, "boom")
	c.Assert(result, jc.DeepEquals, params.EnvironConfigResult{})

	apiservertesting.BackingInstance.CheckCallNames(c, "EnvironConfig")
}

func (s *DiscoverSpacesSuite) TestEnvironConfigSuccess(c *gc.C) {
	result, err := s.facade.EnvironConfig()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, jc.DeepEquals, params.EnvironConfigResult{
		Config: apiservertesting.BackingInstance.EnvConfig.AllAttrs(),
	})

	apiservertesting.BackingInstance.CheckCallNames(c, "EnvironConfig")
}

func (s *DiscoverSpacesSuite) TestListSpaces(c *gc.C) {
	result, err := s.facade.ListSpaces()
	c.Assert(err, jc.ErrorIsNil)

	expectedResult := []params.ProviderSpace{
		{Name: "default",
			Subnets: []params.Subnet{
				{CIDR: "192.168.0.0/24",
					ProviderId: "provider-192.168.0.0/24",
					SpaceTag:   "space-default",
					Zones:      []string{"foo"},
					Status:     "in-use"},
				{CIDR: "192.168.3.0/24",
					ProviderId: "provider-192.168.3.0/24",
					VLANTag:    23,
					SpaceTag:   "space-default",
					Zones:      []string{"bar", "bam"}}}},
		{Name: "dmz",
			Subnets: []params.Subnet{
				{CIDR: "192.168.1.0/24",
					ProviderId: "provider-192.168.1.0/24",
					VLANTag:    23,
					SpaceTag:   "space-dmz",
					Zones:      []string{"bar", "bam"}}}},
		{Name: "private",
			Subnets: []params.Subnet{
				{CIDR: "192.168.2.0/24",
					ProviderId: "provider-192.168.2.0/24",
					SpaceTag:   "space-private",
					Zones:      []string{"foo"},
					Status:     "in-use"}}}}
	c.Assert(result.Results, jc.DeepEquals, expectedResult)
	apiservertesting.BackingInstance.CheckCallNames(c, "AllSpaces")
}

func (s *DiscoverSpacesSuite) TestListSpacesFailure(c *gc.C) {
	apiservertesting.BackingInstance.SetErrors(errors.New("boom"))

	result, err := s.facade.ListSpaces()
	c.Assert(err, gc.ErrorMatches, "boom")
	c.Assert(result, jc.DeepEquals, params.DiscoverSpacesResults{})

	apiservertesting.BackingInstance.CheckCallNames(c, "AllSpaces")
}
<commit_msg>Indentation arguments<commit_after>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package discoverspaces_test

import (
	"errors"

	"github.com/juju/names"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/apiserver/common"
	"github.com/juju/juju/apiserver/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	apiservertesting "github.com/juju/juju/apiserver/testing"
	coretesting "github.com/juju/juju/testing"
)

type DiscoverSpacesSuite struct {
	coretesting.BaseSuite
	apiservertesting.StubNetwork

	resources  *common.Resources
	authorizer apiservertesting.FakeAuthorizer
	facade     *discoverspaces.DiscoverSpacesAPI
}

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

func (s *DiscoverSpacesSuite) SetUpSuite(c *gc.C) {
	s.StubNetwork.SetUpSuite(c)
	s.BaseSuite.SetUpSuite(c)
}

func (s *DiscoverSpacesSuite) TearDownSuite(c *gc.C) {
	s.BaseSuite.TearDownSuite(c)
}

func (s *DiscoverSpacesSuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)
	apiservertesting.BackingInstance.SetUp(
		c,
		apiservertesting.StubZonedEnvironName,
		apiservertesting.WithZones,
		apiservertesting.WithSpaces,
		apiservertesting.WithSubnets)

	s.resources = common.NewResources()
	s.authorizer = apiservertesting.FakeAuthorizer{
		Tag:            names.NewUserTag("admin"),
		EnvironManager: true,
	}

	var err error
	s.facade, err = discoverspaces.NewDiscoverSpacesAPIWithBacking(
		apiservertesting.BackingInstance, s.resources, s.authorizer,
	)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(s.facade, gc.NotNil)
}

func (s *DiscoverSpacesSuite) TearDownTest(c *gc.C) {
	if s.resources != nil {
		s.resources.StopAll()
	}
	s.BaseSuite.TearDownTest(c)
}

func (s *DiscoverSpacesSuite) TestEnvironConfigFailure(c *gc.C) {
	apiservertesting.BackingInstance.SetErrors(errors.New("boom"))

	result, err := s.facade.EnvironConfig()
	c.Assert(err, gc.ErrorMatches, "boom")
	c.Assert(result, jc.DeepEquals, params.EnvironConfigResult{})

	apiservertesting.BackingInstance.CheckCallNames(c, "EnvironConfig")
}

func (s *DiscoverSpacesSuite) TestEnvironConfigSuccess(c *gc.C) {
	result, err := s.facade.EnvironConfig()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, jc.DeepEquals, params.EnvironConfigResult{
		Config: apiservertesting.BackingInstance.EnvConfig.AllAttrs(),
	})

	apiservertesting.BackingInstance.CheckCallNames(c, "EnvironConfig")
}

func (s *DiscoverSpacesSuite) TestListSpaces(c *gc.C) {
	result, err := s.facade.ListSpaces()
	c.Assert(err, jc.ErrorIsNil)

	expectedResult := []params.ProviderSpace{{
		Name: "default",
		Subnets: []params.Subnet{
			{CIDR: "192.168.0.0/24",
				ProviderId: "provider-192.168.0.0/24",
				SpaceTag:   "space-default",
				Zones:      []string{"foo"},
				Status:     "in-use"},
			{CIDR: "192.168.3.0/24",
				ProviderId: "provider-192.168.3.0/24",
				VLANTag:    23,
				SpaceTag:   "space-default",
				Zones:      []string{"bar", "bam"}}}}, {
		Name: "dmz",
		Subnets: []params.Subnet{
			{CIDR: "192.168.1.0/24",
				ProviderId: "provider-192.168.1.0/24",
				VLANTag:    23,
				SpaceTag:   "space-dmz",
				Zones:      []string{"bar", "bam"}}}}, {
		Name: "private",
		Subnets: []params.Subnet{
			{CIDR: "192.168.2.0/24",
				ProviderId: "provider-192.168.2.0/24",
				SpaceTag:   "space-private",
				Zones:      []string{"foo"},
				Status:     "in-use"}},
	}}
	c.Assert(result.Results, jc.DeepEquals, expectedResult)
	apiservertesting.BackingInstance.CheckCallNames(c, "AllSpaces")
}

func (s *DiscoverSpacesSuite) TestListSpacesFailure(c *gc.C) {
	apiservertesting.BackingInstance.SetErrors(errors.New("boom"))

	result, err := s.facade.ListSpaces()
	c.Assert(err, gc.ErrorMatches, "boom")
	c.Assert(result, jc.DeepEquals, params.DiscoverSpacesResults{})

	apiservertesting.BackingInstance.CheckCallNames(c, "AllSpaces")
}
<|endoftext|>


<commit_before>// Copyright 2014 The Serviced 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 facade

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/docker/docker/pkg/parsers"
	"github.com/zenoss/glog"
	dockerclient "github.com/zenoss/go-dockerclient"

	"github.com/control-center/serviced/commons"
	"github.com/control-center/serviced/commons/docker"
	"github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/domain/servicedefinition"
	"github.com/control-center/serviced/domain/servicetemplate"
	"github.com/control-center/serviced/isvcs"
)

type reloadLogstashContainer func(ctx datastore.Context, f *Facade) error

var LogstashContainerReloader reloadLogstashContainer = reloadLogstashContainerImpl

var getDockerClient = func() (*dockerclient.Client, error) { return dockerclient.NewClient("unix:///var/run/docker.sock") }

//AddServiceTemplate  adds a service template to the system. Returns the id of the template added
func (f *Facade) AddServiceTemplate(ctx datastore.Context, serviceTemplate servicetemplate.ServiceTemplate) (string, error) {
	hash, err := serviceTemplate.Hash()
	if err != nil {
		return "", err
	}
	serviceTemplate.ID = hash

	if st, _ := f.templateStore.Get(ctx, hash); st != nil {
		// This id already exists in the system
		glog.Infof("Not replacing existing template %s", hash)
		return hash, nil
	}

	if err = f.templateStore.Put(ctx, serviceTemplate); err != nil {
		return "", err
	}

	// this takes a while so don't block the main thread
	go LogstashContainerReloader(ctx, f)
	return hash, err
}

//UpdateServiceTemplate updates a service template
func (f *Facade) UpdateServiceTemplate(ctx datastore.Context, template servicetemplate.ServiceTemplate) error {
	if err := f.templateStore.Put(ctx, template); err != nil {
		return err
	}
	go LogstashContainerReloader(ctx, f) // don't block the main thread
	return nil
}

//RemoveServiceTemplate removes the service template from the system
func (f *Facade) RemoveServiceTemplate(ctx datastore.Context, id string) error {
	if _, err := f.templateStore.Get(ctx, id); err != nil {
		return fmt.Errorf("Unable to find template: %s", id)
	}

	glog.V(2).Infof("Facade.RemoveServiceTemplate: %s", id)
	if err := f.templateStore.Delete(ctx, id); err != nil {
		return err
	}

	go LogstashContainerReloader(ctx, f)
	return nil
}

func (f *Facade) GetServiceTemplates(ctx datastore.Context) (map[string]servicetemplate.ServiceTemplate, error) {
	glog.V(2).Infof("Facade.GetServiceTemplates")
	results, err := f.templateStore.GetServiceTemplates(ctx)
	templateMap := make(map[string]servicetemplate.ServiceTemplate)
	if err != nil {
		glog.V(2).Infof("Facade.GetServiceTemplates: err=%s", err)
		return templateMap, err
	}
	for _, st := range results {
		templateMap[st.ID] = *st
	}
	return templateMap, nil
}

func getImageIDs(sds ...servicedefinition.ServiceDefinition) []string {
	set := map[string]struct{}{}
	for _, sd := range sds {
		for _, img := range getImageIDs(sd.Services...) {
			set[img] = struct{}{}
		}
		if sd.ImageID != "" {
			set[sd.ImageID] = struct{}{}
		}
	}
	result := []string{}
	for img, _ := range set {
		result = append(result, img)
	}
	return result
}

func pullTemplateImages(template *servicetemplate.ServiceTemplate) error {
	for _, img := range getImageIDs(template.Services...) {
		imageID, err := commons.ParseImageID(img)
		if err != nil {
			return err
		}
		tag := imageID.Tag
		if tag == "" {
			tag = "latest"
		}
		image := fmt.Sprintf("%s:%s", imageID.BaseName(), tag)
		glog.Infof("Pulling image %s", image)
		if err := docker.PullImage(image); err != nil {
			glog.Warningf("Unable to pull image %s", image)
		}
	}
	return nil
}

var deployments = make(map[string]map[string]string)

// UpdateDeployTemplateStatus updates the deployment status of the service being deployed
func UpdateDeployTemplateStatus(deploymentID string, status string) {
	if _, ok := deployments[deploymentID]; !ok {
		deployments[deploymentID] = make(map[string]string)
	}

	deployments[deploymentID]["lastStatus"] = deployments[deploymentID]["status"]
	deployments[deploymentID]["status"] = status
}

// gather a list of all active DeploymentIDs
func (f *Facade) DeployTemplateActive(active *[]map[string]string) error {
	// we initialize the data container to something here in case it has not been initialized yet
	*active = make([]map[string]string, 0)
	for _, v := range deployments {
		*active = append(*active, v)
	}

	return nil
}

// DeployTemplateStatus sets the status of a deployed service or template
func (f *Facade) DeployTemplateStatus(deploymentID string, status *string) error {
	if _, ok := deployments[deploymentID]; ok {
		if deployments[deploymentID]["lastStatus"] != deployments[deploymentID]["status"] {
			deployments[deploymentID]["lastStatus"] = deployments[deploymentID]["status"]
			*status = deployments[deploymentID]["status"]
		} else if deployments[deploymentID]["status"] != "" {
			time.Sleep(100 * time.Millisecond)
			f.DeployTemplateStatus(deploymentID, status)
		}
	} else {
		*status = ""
	}

	return nil
}

//DeployTemplate creates and deployes a service to the pool and returns the tenant id of the newly deployed service
func (f *Facade) DeployTemplate(ctx datastore.Context, poolID string, templateID string, deploymentID string) (string, error) {
	// add an entry for reporting status
	deployments[deploymentID] = map[string]string{
		"TemplateID":   templateID,
		"DeploymentID": deploymentID,
		"PoolID":       poolID,
		"status":       "Starting",
		"lastStatus":   "",
	}
	defer delete(deployments, deploymentID)

	UpdateDeployTemplateStatus(deploymentID, "deploy_loading_template|"+templateID)
	template, err := f.templateStore.Get(ctx, templateID)
	if err != nil {
		glog.Errorf("unable to load template: %s", templateID)
		return "", err
	}

	//check that deployment id does not already exist
	svcs, err := f.serviceStore.GetServicesByDeployment(ctx, deploymentID)
	if err != nil {
		glog.Errorf("unable to validate deploymentID %v while deploying %v", deploymentID, templateID)
		return "", err
	}
	for _, svc := range svcs {
		if svc.DeploymentID == deploymentID {
			return "", fmt.Errorf("deployment ID %v is already in use", deploymentID)
		}
	}

	//now that we know the template name, set it in the status
	deployments[deploymentID]["templateName"] = template.Name

	UpdateDeployTemplateStatus(deploymentID, "deploy_loading_resource_pool|"+poolID)
	pool, err := f.GetResourcePool(ctx, poolID)
	if err != nil {
		glog.Errorf("Unable to load resource pool: %s", poolID)
		return "", err
	}
	if pool == nil {
		return "", fmt.Errorf("poolid %s not found", poolID)
	}

	UpdateDeployTemplateStatus(deploymentID, "deploy_pulling_images")
	if err := pullTemplateImages(template); err != nil {
		glog.Errorf("Unable to pull one or more images")
		return "", err
	}

	volumes := make(map[string]string)
	var tenantID string
	err = f.deployServiceDefinitions(ctx, template.Services, poolID, "", volumes, deploymentID, &tenantID)

	return tenantID, err
}

func (f *Facade) DeployService(ctx datastore.Context, parentID string, sd servicedefinition.ServiceDefinition) (string, error) {
	parent, err := service.NewStore().Get(ctx, parentID)
	if err != nil {
		return "", fmt.Errorf("could not get parent '%s': %s", parentID, err)
	}

	tenantId, err := f.GetTenantID(ctx, parentID)
	if err != nil {
		return "", fmt.Errorf("getting tenant id: %s", err)
	}

	volumes := make(map[string]string)
	return f.deployServiceDefinition(ctx, sd, parent.PoolID, parentID, volumes, parent.DeploymentID, &tenantId)
}

func (f *Facade) deployServiceDefinition(ctx datastore.Context, sd servicedefinition.ServiceDefinition, pool string, parentServiceID string, volumes map[string]string, deploymentId string, tenantId *string) (string, error) {
	// Always deploy in stopped state, starting is a separate step
	ds := int(service.SVCStop)

	exportedVolumes := make(map[string]string)
	for k, v := range volumes {
		exportedVolumes[k] = v
	}
	svc, err := service.BuildService(sd, parentServiceID, pool, ds, deploymentId)
	if err != nil {
		return "", err
	}

	UpdateDeployTemplateStatus(deploymentId, "deploy_loading_service|"+svc.Name)
	getSvc := func(svcID string) (service.Service, error) {
		svc, err := f.GetService(ctx, svcID)
		return *svc, err
	}
	findChild := func(svcID, childName string) (service.Service, error) {
		svc, err := f.FindChildService(ctx, svcID, childName)
		return *svc, err
	}

	//for each endpoint, evaluate its Application
	if err = svc.EvaluateEndpointTemplates(getSvc, findChild); err != nil {
		return "", err
	}

	//for each endpoint, evaluate its Application
	if err = svc.EvaluateEndpointTemplates(getSvc, findChild); err != nil {
		return "", err
	}

	if parentServiceID == "" {
		*tenantId = svc.ID
	}

	// Using the tenant id, tag the base image with the tenantID
	if svc.ImageID != "" {
		UpdateDeployTemplateStatus(deploymentId, "deploy_renaming_image|"+svc.Name)
		name, err := renameImageID(f.dockerRegistry, svc.ImageID, *tenantId)
		if err != nil {
			glog.Errorf("malformed imageId: %s", svc.ImageID)
			return "", err
		}

		_, err = docker.FindImage(name, false)
		if err != nil {
			if err != docker.ErrNoSuchImage && !strings.HasPrefix(err.Error(), "No such id:") {
				glog.Error(err)
				return "", err
			}
			UpdateDeployTemplateStatus(deploymentId, "deploy_loading_image|"+name)
			image, err := docker.FindImage(svc.ImageID, false)
			if err != nil {
				msg := fmt.Errorf("could not look up image %s: %s. Check your docker login and retry application deployment.", svc.ImageID, err)
				glog.Error(err.Error())
				return "", msg
			}
			UpdateDeployTemplateStatus(deploymentId, "deploy_tagging_image|"+name)
			if _, err := image.Tag(name); err != nil {
				glog.Errorf("could not tag image: %s (%v)", image.ID, err)
				return "", err
			}
		}
		svc.ImageID = name
	}

	err = f.AddService(ctx, *svc)
	if err != nil {
		return "", err
	}

	return svc.ID, f.deployServiceDefinitions(ctx, sd.Services, pool, svc.ID, exportedVolumes, deploymentId, tenantId)
}

func (f *Facade) deployServiceDefinitions(ctx datastore.Context, sds []servicedefinition.ServiceDefinition, pool string, parentServiceID string, volumes map[string]string, deploymentId string, tenantId *string) error {
	// ensure that all images in the templates exist
	imageIds := make(map[string]struct{})
	for _, svc := range sds {
		getSubServiceImageIDs(imageIds, svc)
	}

	for imageId, _ := range imageIds {
		_, err := docker.FindImage(imageId, false)
		if err != nil {
			msg := fmt.Errorf("could not look up image %s: %s. Check your docker login and retry service deployment.", imageId, err)
			glog.Error(err.Error())
			return msg
		}
	}

	for _, sd := range sds {
		if _, err := f.deployServiceDefinition(ctx, sd, pool, parentServiceID, volumes, deploymentId, tenantId); err != nil {
			return err
		}
	}
	return nil
}

func getSubServiceImageIDs(ids map[string]struct{}, svc servicedefinition.ServiceDefinition) {
	found := struct{}{}

	if len(svc.ImageID) != 0 {
		ids[svc.ImageID] = found
	}
	for _, s := range svc.Services {
		getSubServiceImageIDs(ids, s)
	}
}

func renameImageID(dockerRegistry, imageId, tenantId string) (string, error) {

	repo, _ := parsers.ParseRepositoryTag(imageId)
	re := regexp.MustCompile("/?([^/]+)\\z")
	matches := re.FindStringSubmatch(repo)
	if matches == nil {
		return "", errors.New("malformed imageid")
	}
	name := matches[1]
	return fmt.Sprintf("%s/%s/%s", dockerRegistry, tenantId, name), nil
}

// writeLogstashConfiguration takes all the available
// services and writes out the filters section for logstash.
// This is required before logstash startsup
func writeLogstashConfiguration(templates map[string]servicetemplate.ServiceTemplate) error {
	// FIXME: eventually this file should live in the DFS or the config should
	// live in zookeeper to allow the agents to get to this
	if err := dao.WriteConfigurationFile(templates); err != nil {
		return err
	}
	return nil
}

// Anytime the available service definitions are modified
// we need to restart the logstash container so it can write out
// its new filter set.
// This method depends on the elasticsearch container being up and running.
func reloadLogstashContainerImpl(ctx datastore.Context, f *Facade) error {
	templates, err := f.GetServiceTemplates(ctx)
	if err != nil {
		glog.Fatalf("Could not write logstash configuration: %s", err)
	}

	if err := writeLogstashConfiguration(templates); err != nil {
		glog.Fatalf("Could not write logstash configuration: %s", err)
		return err
	}
	glog.V(2).Info("Starting logstash container")
	if err := isvcs.Mgr.Notify("restart logstash"); err != nil {
		glog.Fatalf("Could not start logstash container: %s", err)
		return err
	}
	return nil
}
<commit_msg>CC-839: Fix nil pointer dereference<commit_after>// Copyright 2014 The Serviced 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 facade

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/docker/docker/pkg/parsers"
	"github.com/zenoss/glog"
	dockerclient "github.com/zenoss/go-dockerclient"

	"github.com/control-center/serviced/commons"
	"github.com/control-center/serviced/commons/docker"
	"github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/domain/servicedefinition"
	"github.com/control-center/serviced/domain/servicetemplate"
	"github.com/control-center/serviced/isvcs"
)

type reloadLogstashContainer func(ctx datastore.Context, f *Facade) error

var LogstashContainerReloader reloadLogstashContainer = reloadLogstashContainerImpl

var getDockerClient = func() (*dockerclient.Client, error) { return dockerclient.NewClient("unix:///var/run/docker.sock") }

//AddServiceTemplate  adds a service template to the system. Returns the id of the template added
func (f *Facade) AddServiceTemplate(ctx datastore.Context, serviceTemplate servicetemplate.ServiceTemplate) (string, error) {
	hash, err := serviceTemplate.Hash()
	if err != nil {
		return "", err
	}
	serviceTemplate.ID = hash

	if st, _ := f.templateStore.Get(ctx, hash); st != nil {
		// This id already exists in the system
		glog.Infof("Not replacing existing template %s", hash)
		return hash, nil
	}

	if err = f.templateStore.Put(ctx, serviceTemplate); err != nil {
		return "", err
	}

	// this takes a while so don't block the main thread
	go LogstashContainerReloader(ctx, f)
	return hash, err
}

//UpdateServiceTemplate updates a service template
func (f *Facade) UpdateServiceTemplate(ctx datastore.Context, template servicetemplate.ServiceTemplate) error {
	if err := f.templateStore.Put(ctx, template); err != nil {
		return err
	}
	go LogstashContainerReloader(ctx, f) // don't block the main thread
	return nil
}

//RemoveServiceTemplate removes the service template from the system
func (f *Facade) RemoveServiceTemplate(ctx datastore.Context, id string) error {
	if _, err := f.templateStore.Get(ctx, id); err != nil {
		return fmt.Errorf("Unable to find template: %s", id)
	}

	glog.V(2).Infof("Facade.RemoveServiceTemplate: %s", id)
	if err := f.templateStore.Delete(ctx, id); err != nil {
		return err
	}

	go LogstashContainerReloader(ctx, f)
	return nil
}

func (f *Facade) GetServiceTemplates(ctx datastore.Context) (map[string]servicetemplate.ServiceTemplate, error) {
	glog.V(2).Infof("Facade.GetServiceTemplates")
	results, err := f.templateStore.GetServiceTemplates(ctx)
	templateMap := make(map[string]servicetemplate.ServiceTemplate)
	if err != nil {
		glog.V(2).Infof("Facade.GetServiceTemplates: err=%s", err)
		return templateMap, err
	}
	for _, st := range results {
		templateMap[st.ID] = *st
	}
	return templateMap, nil
}

func getImageIDs(sds ...servicedefinition.ServiceDefinition) []string {
	set := map[string]struct{}{}
	for _, sd := range sds {
		for _, img := range getImageIDs(sd.Services...) {
			set[img] = struct{}{}
		}
		if sd.ImageID != "" {
			set[sd.ImageID] = struct{}{}
		}
	}
	result := []string{}
	for img, _ := range set {
		result = append(result, img)
	}
	return result
}

func pullTemplateImages(template *servicetemplate.ServiceTemplate) error {
	for _, img := range getImageIDs(template.Services...) {
		imageID, err := commons.ParseImageID(img)
		if err != nil {
			return err
		}
		tag := imageID.Tag
		if tag == "" {
			tag = "latest"
		}
		image := fmt.Sprintf("%s:%s", imageID.BaseName(), tag)
		glog.Infof("Pulling image %s", image)
		if err := docker.PullImage(image); err != nil {
			glog.Warningf("Unable to pull image %s", image)
		}
	}
	return nil
}

var deployments = make(map[string]map[string]string)

// UpdateDeployTemplateStatus updates the deployment status of the service being deployed
func UpdateDeployTemplateStatus(deploymentID string, status string) {
	if _, ok := deployments[deploymentID]; !ok {
		deployments[deploymentID] = make(map[string]string)
	}

	deployments[deploymentID]["lastStatus"] = deployments[deploymentID]["status"]
	deployments[deploymentID]["status"] = status
}

// gather a list of all active DeploymentIDs
func (f *Facade) DeployTemplateActive(active *[]map[string]string) error {
	// we initialize the data container to something here in case it has not been initialized yet
	*active = make([]map[string]string, 0)
	for _, v := range deployments {
		*active = append(*active, v)
	}

	return nil
}

// DeployTemplateStatus sets the status of a deployed service or template
func (f *Facade) DeployTemplateStatus(deploymentID string, status *string) error {
	if _, ok := deployments[deploymentID]; ok {
		if deployments[deploymentID]["lastStatus"] != deployments[deploymentID]["status"] {
			deployments[deploymentID]["lastStatus"] = deployments[deploymentID]["status"]
			*status = deployments[deploymentID]["status"]
		} else if deployments[deploymentID]["status"] != "" {
			time.Sleep(100 * time.Millisecond)
			f.DeployTemplateStatus(deploymentID, status)
		}
	} else {
		*status = ""
	}

	return nil
}

//DeployTemplate creates and deployes a service to the pool and returns the tenant id of the newly deployed service
func (f *Facade) DeployTemplate(ctx datastore.Context, poolID string, templateID string, deploymentID string) (string, error) {
	// add an entry for reporting status
	deployments[deploymentID] = map[string]string{
		"TemplateID":   templateID,
		"DeploymentID": deploymentID,
		"PoolID":       poolID,
		"status":       "Starting",
		"lastStatus":   "",
	}
	defer delete(deployments, deploymentID)

	UpdateDeployTemplateStatus(deploymentID, "deploy_loading_template|"+templateID)
	template, err := f.templateStore.Get(ctx, templateID)
	if err != nil {
		glog.Errorf("unable to load template: %s", templateID)
		return "", err
	}

	//check that deployment id does not already exist
	svcs, err := f.serviceStore.GetServicesByDeployment(ctx, deploymentID)
	if err != nil {
		glog.Errorf("unable to validate deploymentID %v while deploying %v", deploymentID, templateID)
		return "", err
	}
	for _, svc := range svcs {
		if svc.DeploymentID == deploymentID {
			return "", fmt.Errorf("deployment ID %v is already in use", deploymentID)
		}
	}

	//now that we know the template name, set it in the status
	deployments[deploymentID]["templateName"] = template.Name

	UpdateDeployTemplateStatus(deploymentID, "deploy_loading_resource_pool|"+poolID)
	pool, err := f.GetResourcePool(ctx, poolID)
	if err != nil {
		glog.Errorf("Unable to load resource pool: %s", poolID)
		return "", err
	}
	if pool == nil {
		return "", fmt.Errorf("poolid %s not found", poolID)
	}

	UpdateDeployTemplateStatus(deploymentID, "deploy_pulling_images")
	if err := pullTemplateImages(template); err != nil {
		glog.Errorf("Unable to pull one or more images")
		return "", err
	}

	volumes := make(map[string]string)
	var tenantID string
	err = f.deployServiceDefinitions(ctx, template.Services, poolID, "", volumes, deploymentID, &tenantID)

	return tenantID, err
}

func (f *Facade) DeployService(ctx datastore.Context, parentID string, sd servicedefinition.ServiceDefinition) (string, error) {
	parent, err := service.NewStore().Get(ctx, parentID)
	if err != nil {
		return "", fmt.Errorf("could not get parent '%s': %s", parentID, err)
	}

	tenantId, err := f.GetTenantID(ctx, parentID)
	if err != nil {
		return "", fmt.Errorf("getting tenant id: %s", err)
	}

	volumes := make(map[string]string)
	return f.deployServiceDefinition(ctx, sd, parent.PoolID, parentID, volumes, parent.DeploymentID, &tenantId)
}

func (f *Facade) deployServiceDefinition(ctx datastore.Context, sd servicedefinition.ServiceDefinition, pool string, parentServiceID string, volumes map[string]string, deploymentId string, tenantId *string) (string, error) {
	// Always deploy in stopped state, starting is a separate step
	ds := int(service.SVCStop)

	exportedVolumes := make(map[string]string)
	for k, v := range volumes {
		exportedVolumes[k] = v
	}
	svc, err := service.BuildService(sd, parentServiceID, pool, ds, deploymentId)
	if err != nil {
		return "", err
	}

	UpdateDeployTemplateStatus(deploymentId, "deploy_loading_service|"+svc.Name)
	getSvc := func(svcID string) (service.Service, error) {
		svc, err := f.GetService(ctx, svcID)
		if err != nil {
			return service.Service{}, err
		}
		return *svc, err
	}
	findChild := func(svcID, childName string) (service.Service, error) {
		svc, err := f.FindChildService(ctx, svcID, childName)
		if err != nil {
			return service.Service{}, err
		}
		return *svc, err
	}

	//for each endpoint, evaluate its Application
	if err = svc.EvaluateEndpointTemplates(getSvc, findChild); err != nil {
		return "", err
	}

	//for each endpoint, evaluate its Application
	if err = svc.EvaluateEndpointTemplates(getSvc, findChild); err != nil {
		return "", err
	}

	if parentServiceID == "" {
		*tenantId = svc.ID
	}

	// Using the tenant id, tag the base image with the tenantID
	if svc.ImageID != "" {
		UpdateDeployTemplateStatus(deploymentId, "deploy_renaming_image|"+svc.Name)
		name, err := renameImageID(f.dockerRegistry, svc.ImageID, *tenantId)
		if err != nil {
			glog.Errorf("malformed imageId: %s", svc.ImageID)
			return "", err
		}

		_, err = docker.FindImage(name, false)
		if err != nil {
			if err != docker.ErrNoSuchImage && !strings.HasPrefix(err.Error(), "No such id:") {
				glog.Error(err)
				return "", err
			}
			UpdateDeployTemplateStatus(deploymentId, "deploy_loading_image|"+name)
			image, err := docker.FindImage(svc.ImageID, false)
			if err != nil {
				msg := fmt.Errorf("could not look up image %s: %s. Check your docker login and retry application deployment.", svc.ImageID, err)
				glog.Error(err.Error())
				return "", msg
			}
			UpdateDeployTemplateStatus(deploymentId, "deploy_tagging_image|"+name)
			if _, err := image.Tag(name); err != nil {
				glog.Errorf("could not tag image: %s (%v)", image.ID, err)
				return "", err
			}
		}
		svc.ImageID = name
	}

	err = f.AddService(ctx, *svc)
	if err != nil {
		return "", err
	}

	return svc.ID, f.deployServiceDefinitions(ctx, sd.Services, pool, svc.ID, exportedVolumes, deploymentId, tenantId)
}

func (f *Facade) deployServiceDefinitions(ctx datastore.Context, sds []servicedefinition.ServiceDefinition, pool string, parentServiceID string, volumes map[string]string, deploymentId string, tenantId *string) error {
	// ensure that all images in the templates exist
	imageIds := make(map[string]struct{})
	for _, svc := range sds {
		getSubServiceImageIDs(imageIds, svc)
	}

	for imageId, _ := range imageIds {
		_, err := docker.FindImage(imageId, false)
		if err != nil {
			msg := fmt.Errorf("could not look up image %s: %s. Check your docker login and retry service deployment.", imageId, err)
			glog.Error(err.Error())
			return msg
		}
	}

	for _, sd := range sds {
		if _, err := f.deployServiceDefinition(ctx, sd, pool, parentServiceID, volumes, deploymentId, tenantId); err != nil {
			return err
		}
	}
	return nil
}

func getSubServiceImageIDs(ids map[string]struct{}, svc servicedefinition.ServiceDefinition) {
	found := struct{}{}

	if len(svc.ImageID) != 0 {
		ids[svc.ImageID] = found
	}
	for _, s := range svc.Services {
		getSubServiceImageIDs(ids, s)
	}
}

func renameImageID(dockerRegistry, imageId, tenantId string) (string, error) {

	repo, _ := parsers.ParseRepositoryTag(imageId)
	re := regexp.MustCompile("/?([^/]+)\\z")
	matches := re.FindStringSubmatch(repo)
	if matches == nil {
		return "", errors.New("malformed imageid")
	}
	name := matches[1]
	return fmt.Sprintf("%s/%s/%s", dockerRegistry, tenantId, name), nil
}

// writeLogstashConfiguration takes all the available
// services and writes out the filters section for logstash.
// This is required before logstash startsup
func writeLogstashConfiguration(templates map[string]servicetemplate.ServiceTemplate) error {
	// FIXME: eventually this file should live in the DFS or the config should
	// live in zookeeper to allow the agents to get to this
	if err := dao.WriteConfigurationFile(templates); err != nil {
		return err
	}
	return nil
}

// Anytime the available service definitions are modified
// we need to restart the logstash container so it can write out
// its new filter set.
// This method depends on the elasticsearch container being up and running.
func reloadLogstashContainerImpl(ctx datastore.Context, f *Facade) error {
	templates, err := f.GetServiceTemplates(ctx)
	if err != nil {
		glog.Fatalf("Could not write logstash configuration: %s", err)
	}

	if err := writeLogstashConfiguration(templates); err != nil {
		glog.Fatalf("Could not write logstash configuration: %s", err)
		return err
	}
	glog.V(2).Info("Starting logstash container")
	if err := isvcs.Mgr.Notify("restart logstash"); err != nil {
		glog.Fatalf("Could not start logstash container: %s", err)
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package arel

type NamedFunctionNode struct{}
<commit_msg>NamedFunctionNode should include a FunctionNode<commit_after>package arel

type NamedFunctionNode struct {
	FunctionNode
}
<|endoftext|>


<commit_before>// Copyright 2015 Square Inc.
//
// 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 query

import (
	"testing"
	"time"

	"github.com/square/metrics/testing_support/assert"
)

func Test_parseRelativeTime(t *testing.T) {
	now := time.Unix(1413321866, 0).UTC()

	// Valid relative timestamps
	timestampTests := []struct {
		timeString        string
		expectedTimestamp int64
		expectSuccess     bool
	}{
		// Valid relative timestamps
		{"-2s", 1413321864000, true},
		{"-3m", 1413321686000, true},
		{"-4h", 1413307466000, true},
		{"-5d", 1412889866000, true},
		{"-3w", 1411507466000, true},
		{"-1M", 1410729866000, true},
		{"-1y", 1381785866000, true},
		{"1s", 1413321867000, true},
		{"+1s", 1413321867000, true},
		{"5d", 1413322298000, true},
		// Bad relative timestamps
		{"5dd", -1, false},
		{"-5dd", -1, false},
		{"-5z", -1, false},
	}

	for _, c := range timestampTests {
		ts, err := parseDate(c.timeString, now)
		if err != nil && c.expectSuccess {
			t.Fatal("Received unexpected error from parseRelativeTime: ", err)
		}

		if ts != c.expectedTimestamp {
			t.Fatalf("Expected %d but received %d", c.expectedTimestamp, ts)
		}
	}
}

func TestUnescapeLiteral(t *testing.T) {
	a := assert.New(t)
	a.EqString(unescapeLiteral("'foo'"), "foo")
	a.EqString(unescapeLiteral("foo"), "foo")
	a.EqString(unescapeLiteral("nodes.cpu.io"), "nodes.cpu.io")
	a.EqString(unescapeLiteral(`"hello"`), `hello`)
	a.EqString(unescapeLiteral(`"\"hello\""`), `"hello"`)
	a.EqString(unescapeLiteral(`'\"hello\"'`), `"hello"`)
	a.EqString(unescapeLiteral("\"\\`\""), "`")
}

func testFunction1() (string, string) {
	return functionName(0), functionName(1)
}

func TestFunctionName(t *testing.T) {
	a := assert.New(t)
	a.EqString(functionName(0), "TestFunctionName")
	first, second := testFunction1()
	a.EqString(first, "testFunction1")
	a.EqString(second, "TestFunctionName")
}
<commit_msg>fix relative time in test<commit_after>// Copyright 2015 Square Inc.
//
// 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 query

import (
	"testing"
	"time"

	"github.com/square/metrics/testing_support/assert"
)

func Test_parseRelativeTime(t *testing.T) {
	now := time.Unix(1413321866, 0).UTC()

	// Valid relative timestamps
	timestampTests := []struct {
		timeString        string
		expectedTimestamp int64
		expectSuccess     bool
	}{
		// Valid relative timestamps
		{"-2s", 1413321864000, true},
		{"-3m", 1413321686000, true},
		{"-4h", 1413307466000, true},
		{"-5d", 1412889866000, true},
		{"-3w", 1411507466000, true},
		{"-1M", 1410729866000, true},
		{"-1y", 1381785866000, true},
		{"1s", 1413321867000, true},
		{"+1s", 1413321867000, true},
		{"5d", 1413753866000, true},
		// Bad relative timestamps
		{"5dd", -1, false},
		{"-5dd", -1, false},
		{"-5z", -1, false},
	}

	for _, c := range timestampTests {
		ts, err := parseDate(c.timeString, now)
		if err != nil && c.expectSuccess {
			t.Fatal("Received unexpected error from parseRelativeTime: ", err)
		}

		if ts != c.expectedTimestamp {
			t.Fatalf("Expected %d but received %d", c.expectedTimestamp, ts)
		}
	}
}

func TestUnescapeLiteral(t *testing.T) {
	a := assert.New(t)
	a.EqString(unescapeLiteral("'foo'"), "foo")
	a.EqString(unescapeLiteral("foo"), "foo")
	a.EqString(unescapeLiteral("nodes.cpu.io"), "nodes.cpu.io")
	a.EqString(unescapeLiteral(`"hello"`), `hello`)
	a.EqString(unescapeLiteral(`"\"hello\""`), `"hello"`)
	a.EqString(unescapeLiteral(`'\"hello\"'`), `"hello"`)
	a.EqString(unescapeLiteral("\"\\`\""), "`")
}

func testFunction1() (string, string) {
	return functionName(0), functionName(1)
}

func TestFunctionName(t *testing.T) {
	a := assert.New(t)
	a.EqString(functionName(0), "TestFunctionName")
	first, second := testFunction1()
	a.EqString(first, "testFunction1")
	a.EqString(second, "TestFunctionName")
}
<|endoftext|>