<commit_before>package api

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"

	"github.com/grafana/grafana-plugin-sdk-go/data"
	"github.com/grafana/grafana/pkg/api/response"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/services/datasourceproxy"
	"github.com/grafana/grafana/pkg/services/datasources"
	apimodels "github.com/grafana/grafana/pkg/services/ngalert/api/tooling/definitions"
	"github.com/grafana/grafana/pkg/services/ngalert/eval"
	ngmodels "github.com/grafana/grafana/pkg/services/ngalert/models"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/tsdb"
	"github.com/grafana/grafana/pkg/util"
	"github.com/pkg/errors"
	"gopkg.in/macaron.v1"
	"gopkg.in/yaml.v3"
)

var searchRegex = regexp.MustCompile(`\{(\w+)\}`)

var NotImplementedResp = ErrResp(http.StatusNotImplemented, errors.New("endpoint not implemented"), "")

func toMacaronPath(path string) string {
	return string(searchRegex.ReplaceAllFunc([]byte(path), func(s []byte) []byte {
		m := string(s[1 : len(s)-1])
		return []byte(fmt.Sprintf(":%s", m))
	}))
}

func backendType(ctx *models.ReqContext, cache datasources.CacheService) (apimodels.Backend, error) {
	recipient := ctx.Params("Recipient")
	if recipient == apimodels.GrafanaBackend.String() {
		return apimodels.GrafanaBackend, nil
	}
	if datasourceID, err := strconv.ParseInt(recipient, 10, 64); err == nil {
		if ds, err := cache.GetDatasource(datasourceID, ctx.SignedInUser, ctx.SkipCache); err == nil {
			switch ds.Type {
			case "loki", "prometheus":
				return apimodels.LoTexRulerBackend, nil
			case "alertmanager":
				return apimodels.AlertmanagerBackend, nil
			default:
				return 0, fmt.Errorf("unexpected backend type (%v)", ds.Type)
			}
		}
	}
	return 0, fmt.Errorf("unexpected backend type (%v)", recipient)
}

// macaron unsafely asserts the http.ResponseWriter is an http.CloseNotifier, which will panic.
// Here we impl it, which will ensure this no longer happens, but neither will we take
// advantage cancelling upstream requests when the downstream has closed.
// NB: http.CloseNotifier is a deprecated ifc from before the context pkg.
type safeMacaronWrapper struct {
	http.ResponseWriter
}

func (w *safeMacaronWrapper) CloseNotify() <-chan bool {
	return make(chan bool)
}

// replacedResponseWriter overwrites the underlying responsewriter used by a *models.ReqContext.
// It's ugly because it needs to replace a value behind a few nested pointers.
func replacedResponseWriter(ctx *models.ReqContext) (*models.ReqContext, *response.NormalResponse) {
	resp := response.CreateNormalResponse(make(http.Header), nil, 0)
	cpy := *ctx
	cpyMCtx := *cpy.Context
	cpyMCtx.Resp = macaron.NewResponseWriter(ctx.Req.Method, &safeMacaronWrapper{resp})
	cpy.Context = &cpyMCtx
	return &cpy, resp
}

type AlertingProxy struct {
	DataProxy *datasourceproxy.DatasourceProxyService
}

// withReq proxies a different request
func (p *AlertingProxy) withReq(
	ctx *models.ReqContext,
	method string,
	u *url.URL,
	body io.Reader,
	extractor func(*response.NormalResponse) (interface{}, error),
	headers map[string]string,
) response.Response {
	req, err := http.NewRequest(method, u.String(), body)
	if err != nil {
		return ErrResp(http.StatusBadRequest, err, "")
	}
	for h, v := range headers {
		req.Header.Add(h, v)
	}
	newCtx, resp := replacedResponseWriter(ctx)
	newCtx.Req.Request = req
	p.DataProxy.ProxyDatasourceRequestWithID(newCtx, ctx.ParamsInt64("Recipient"))

	status := resp.Status()
	if status >= 400 {
		errMessage := string(resp.Body())
		// if Content-Type is application/json
		// and it is successfully decoded and contains a message
		// return this as response error message
		if strings.HasPrefix(resp.Header().Get("Content-Type"), "application/json") {
			var m map[string]interface{}
			if err := json.Unmarshal(resp.Body(), &m); err == nil {
				if message, ok := m["message"]; ok {
					errMessage = message.(string)
				}
			}
		}
		return ErrResp(status, errors.New(errMessage), "")
	}

	t, err := extractor(resp)
	if err != nil {
		return ErrResp(http.StatusInternalServerError, err, "")
	}

	b, err := json.Marshal(t)
	if err != nil {
		return ErrResp(http.StatusInternalServerError, err, "")
	}

	return response.JSON(status, b)
}

func yamlExtractor(v interface{}) func(*response.NormalResponse) (interface{}, error) {
	return func(resp *response.NormalResponse) (interface{}, error) {
		contentType := resp.Header().Get("Content-Type")
		if !strings.Contains(contentType, "yaml") {
			return nil, fmt.Errorf("unexpected content type from upstream. expected YAML, got %v", contentType)
		}
		decoder := yaml.NewDecoder(bytes.NewReader(resp.Body()))
		decoder.KnownFields(true)

		err := decoder.Decode(v)

		return v, err
	}
}

func jsonExtractor(v interface{}) func(*response.NormalResponse) (interface{}, error) {
	if v == nil {
		// json unmarshal expects a pointer
		v = &map[string]interface{}{}
	}
	return func(resp *response.NormalResponse) (interface{}, error) {
		contentType := resp.Header().Get("Content-Type")
		if !strings.Contains(contentType, "json") {
			return nil, fmt.Errorf("unexpected content type from upstream. expected JSON, got %v", contentType)
		}
		return v, json.Unmarshal(resp.Body(), v)
	}
}

func messageExtractor(resp *response.NormalResponse) (interface{}, error) {
	return map[string]string{"message": string(resp.Body())}, nil
}

func validateCondition(c ngmodels.Condition, user *models.SignedInUser, skipCache bool, datasourceCache datasources.CacheService) error {
	if len(c.Data) == 0 {
		return nil
	}

	refIDs, err := validateQueriesAndExpressions(c.Data, user, skipCache, datasourceCache)
	if err != nil {
		return err
	}

	t := make([]string, 0, len(refIDs))
	for refID := range refIDs {
		t = append(t, refID)
	}
	if _, ok := refIDs[c.Condition]; !ok {
		return fmt.Errorf("condition %s not found in any query or expression: it should be one of: [%s]", c.Condition, strings.Join(t, ","))
	}
	return nil
}

func validateQueriesAndExpressions(data []ngmodels.AlertQuery, user *models.SignedInUser, skipCache bool, datasourceCache datasources.CacheService) (map[string]struct{}, error) {
	refIDs := make(map[string]struct{})
	if len(data) == 0 {
		return nil, nil
	}

	for _, query := range data {
		datasourceUID, err := query.GetDatasource()
		if err != nil {
			return nil, err
		}

		isExpression, err := query.IsExpression()
		if err != nil {
			return nil, err
		}
		if isExpression {
			refIDs[query.RefID] = struct{}{}
			continue
		}

		_, err = datasourceCache.GetDatasourceByUID(datasourceUID, user, skipCache)
		if err != nil {
			return nil, fmt.Errorf("invalid query %s: %w: %s", query.RefID, err, datasourceUID)
		}
		refIDs[query.RefID] = struct{}{}
	}
	return refIDs, nil
}

func conditionEval(c *models.ReqContext, cmd ngmodels.EvalAlertConditionCommand, datasourceCache datasources.CacheService, dataService *tsdb.Service, cfg *setting.Cfg) response.Response {
	evalCond := ngmodels.Condition{
		Condition: cmd.Condition,
		OrgID:     c.SignedInUser.OrgId,
		Data:      cmd.Data,
	}
	if err := validateCondition(evalCond, c.SignedInUser, c.SkipCache, datasourceCache); err != nil {
		return ErrResp(http.StatusBadRequest, err, "invalid condition")
	}

	now := cmd.Now
	if now.IsZero() {
		now = timeNow()
	}

	evaluator := eval.Evaluator{Cfg: cfg}
	evalResults, err := evaluator.ConditionEval(&evalCond, now, dataService)
	if err != nil {
		return ErrResp(http.StatusBadRequest, err, "Failed to evaluate conditions")
	}

	frame := evalResults.AsDataFrame()
	return response.JSONStreaming(http.StatusOK, util.DynMap{
		"instances": []*data.Frame{&frame},
	})
}

// ErrorResp creates a response with a visible error
func ErrResp(status int, err error, msg string, args ...interface{}) *response.NormalResponse {
	if msg != "" {
		err = errors.WithMessagef(err, msg, args...)
	}
	return response.Error(status, err.Error(), nil)
}
<commit_msg>Alerting: Recact html responses (#35277)<commit_after>package api

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"

	"github.com/grafana/grafana-plugin-sdk-go/data"
	"github.com/grafana/grafana/pkg/api/response"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/services/datasourceproxy"
	"github.com/grafana/grafana/pkg/services/datasources"
	apimodels "github.com/grafana/grafana/pkg/services/ngalert/api/tooling/definitions"
	"github.com/grafana/grafana/pkg/services/ngalert/eval"
	ngmodels "github.com/grafana/grafana/pkg/services/ngalert/models"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/tsdb"
	"github.com/grafana/grafana/pkg/util"
	"github.com/pkg/errors"
	"gopkg.in/macaron.v1"
	"gopkg.in/yaml.v3"
)

var searchRegex = regexp.MustCompile(`\{(\w+)\}`)

var NotImplementedResp = ErrResp(http.StatusNotImplemented, errors.New("endpoint not implemented"), "")

func toMacaronPath(path string) string {
	return string(searchRegex.ReplaceAllFunc([]byte(path), func(s []byte) []byte {
		m := string(s[1 : len(s)-1])
		return []byte(fmt.Sprintf(":%s", m))
	}))
}

func backendType(ctx *models.ReqContext, cache datasources.CacheService) (apimodels.Backend, error) {
	recipient := ctx.Params("Recipient")
	if recipient == apimodels.GrafanaBackend.String() {
		return apimodels.GrafanaBackend, nil
	}
	if datasourceID, err := strconv.ParseInt(recipient, 10, 64); err == nil {
		if ds, err := cache.GetDatasource(datasourceID, ctx.SignedInUser, ctx.SkipCache); err == nil {
			switch ds.Type {
			case "loki", "prometheus":
				return apimodels.LoTexRulerBackend, nil
			case "alertmanager":
				return apimodels.AlertmanagerBackend, nil
			default:
				return 0, fmt.Errorf("unexpected backend type (%v)", ds.Type)
			}
		}
	}
	return 0, fmt.Errorf("unexpected backend type (%v)", recipient)
}

// macaron unsafely asserts the http.ResponseWriter is an http.CloseNotifier, which will panic.
// Here we impl it, which will ensure this no longer happens, but neither will we take
// advantage cancelling upstream requests when the downstream has closed.
// NB: http.CloseNotifier is a deprecated ifc from before the context pkg.
type safeMacaronWrapper struct {
	http.ResponseWriter
}

func (w *safeMacaronWrapper) CloseNotify() <-chan bool {
	return make(chan bool)
}

// replacedResponseWriter overwrites the underlying responsewriter used by a *models.ReqContext.
// It's ugly because it needs to replace a value behind a few nested pointers.
func replacedResponseWriter(ctx *models.ReqContext) (*models.ReqContext, *response.NormalResponse) {
	resp := response.CreateNormalResponse(make(http.Header), nil, 0)
	cpy := *ctx
	cpyMCtx := *cpy.Context
	cpyMCtx.Resp = macaron.NewResponseWriter(ctx.Req.Method, &safeMacaronWrapper{resp})
	cpy.Context = &cpyMCtx
	return &cpy, resp
}

type AlertingProxy struct {
	DataProxy *datasourceproxy.DatasourceProxyService
}

// withReq proxies a different request
func (p *AlertingProxy) withReq(
	ctx *models.ReqContext,
	method string,
	u *url.URL,
	body io.Reader,
	extractor func(*response.NormalResponse) (interface{}, error),
	headers map[string]string,
) response.Response {
	req, err := http.NewRequest(method, u.String(), body)
	if err != nil {
		return ErrResp(http.StatusBadRequest, err, "")
	}
	for h, v := range headers {
		req.Header.Add(h, v)
	}
	newCtx, resp := replacedResponseWriter(ctx)
	newCtx.Req.Request = req
	p.DataProxy.ProxyDatasourceRequestWithID(newCtx, ctx.ParamsInt64("Recipient"))

	status := resp.Status()
	if status >= 400 {
		errMessage := string(resp.Body())
		// if Content-Type is application/json
		// and it is successfully decoded and contains a message
		// return this as response error message
		if strings.HasPrefix(resp.Header().Get("Content-Type"), "application/json") {
			var m map[string]interface{}
			if err := json.Unmarshal(resp.Body(), &m); err == nil {
				if message, ok := m["message"]; ok {
					errMessage = message.(string)
				}
			}
		} else if strings.HasPrefix(resp.Header().Get("Content-Type"), "text/html") {
			// if Content-Type is text/html
			// do not return the body
			errMessage = "redacted html"
		}
		return ErrResp(status, errors.New(errMessage), "")
	}

	t, err := extractor(resp)
	if err != nil {
		return ErrResp(http.StatusInternalServerError, err, "")
	}

	b, err := json.Marshal(t)
	if err != nil {
		return ErrResp(http.StatusInternalServerError, err, "")
	}

	return response.JSON(status, b)
}

func yamlExtractor(v interface{}) func(*response.NormalResponse) (interface{}, error) {
	return func(resp *response.NormalResponse) (interface{}, error) {
		contentType := resp.Header().Get("Content-Type")
		if !strings.Contains(contentType, "yaml") {
			return nil, fmt.Errorf("unexpected content type from upstream. expected YAML, got %v", contentType)
		}
		decoder := yaml.NewDecoder(bytes.NewReader(resp.Body()))
		decoder.KnownFields(true)

		err := decoder.Decode(v)

		return v, err
	}
}

func jsonExtractor(v interface{}) func(*response.NormalResponse) (interface{}, error) {
	if v == nil {
		// json unmarshal expects a pointer
		v = &map[string]interface{}{}
	}
	return func(resp *response.NormalResponse) (interface{}, error) {
		contentType := resp.Header().Get("Content-Type")
		if !strings.Contains(contentType, "json") {
			return nil, fmt.Errorf("unexpected content type from upstream. expected JSON, got %v", contentType)
		}
		return v, json.Unmarshal(resp.Body(), v)
	}
}

func messageExtractor(resp *response.NormalResponse) (interface{}, error) {
	return map[string]string{"message": string(resp.Body())}, nil
}

func validateCondition(c ngmodels.Condition, user *models.SignedInUser, skipCache bool, datasourceCache datasources.CacheService) error {
	if len(c.Data) == 0 {
		return nil
	}

	refIDs, err := validateQueriesAndExpressions(c.Data, user, skipCache, datasourceCache)
	if err != nil {
		return err
	}

	t := make([]string, 0, len(refIDs))
	for refID := range refIDs {
		t = append(t, refID)
	}
	if _, ok := refIDs[c.Condition]; !ok {
		return fmt.Errorf("condition %s not found in any query or expression: it should be one of: [%s]", c.Condition, strings.Join(t, ","))
	}
	return nil
}

func validateQueriesAndExpressions(data []ngmodels.AlertQuery, user *models.SignedInUser, skipCache bool, datasourceCache datasources.CacheService) (map[string]struct{}, error) {
	refIDs := make(map[string]struct{})
	if len(data) == 0 {
		return nil, nil
	}

	for _, query := range data {
		datasourceUID, err := query.GetDatasource()
		if err != nil {
			return nil, err
		}

		isExpression, err := query.IsExpression()
		if err != nil {
			return nil, err
		}
		if isExpression {
			refIDs[query.RefID] = struct{}{}
			continue
		}

		_, err = datasourceCache.GetDatasourceByUID(datasourceUID, user, skipCache)
		if err != nil {
			return nil, fmt.Errorf("invalid query %s: %w: %s", query.RefID, err, datasourceUID)
		}
		refIDs[query.RefID] = struct{}{}
	}
	return refIDs, nil
}

func conditionEval(c *models.ReqContext, cmd ngmodels.EvalAlertConditionCommand, datasourceCache datasources.CacheService, dataService *tsdb.Service, cfg *setting.Cfg) response.Response {
	evalCond := ngmodels.Condition{
		Condition: cmd.Condition,
		OrgID:     c.SignedInUser.OrgId,
		Data:      cmd.Data,
	}
	if err := validateCondition(evalCond, c.SignedInUser, c.SkipCache, datasourceCache); err != nil {
		return ErrResp(http.StatusBadRequest, err, "invalid condition")
	}

	now := cmd.Now
	if now.IsZero() {
		now = timeNow()
	}

	evaluator := eval.Evaluator{Cfg: cfg}
	evalResults, err := evaluator.ConditionEval(&evalCond, now, dataService)
	if err != nil {
		return ErrResp(http.StatusBadRequest, err, "Failed to evaluate conditions")
	}

	frame := evalResults.AsDataFrame()
	return response.JSONStreaming(http.StatusOK, util.DynMap{
		"instances": []*data.Frame{&frame},
	})
}

// ErrorResp creates a response with a visible error
func ErrResp(status int, err error, msg string, args ...interface{}) *response.NormalResponse {
	if msg != "" {
		err = errors.WithMessagef(err, msg, args...)
	}
	return response.Error(status, err.Error(), nil)
}
<|endoftext|>


<commit_before>package elasticsearch

import (
	"github.com/grafana/grafana/pkg/components/simplejson"
)

// Query represents the time series query model of the datasource
type Query struct {
	TimeField     string       `json:"timeField"`
	RawQuery      string       `json:"query"`
	BucketAggs    []*BucketAgg `json:"bucketAggs"`
	Metrics       []*MetricAgg `json:"metrics"`
	Alias         string       `json:"alias"`
	Interval      string
	IntervalMs    int64
	RefID         string
	MaxDataPoints int64
}

// BucketAgg represents a bucket aggregation of the time series query model of the datasource
type BucketAgg struct {
	Field    string           `json:"field"`
	ID       string           `json:"id"`
	Settings *simplejson.Json `json:"settings"`
	Type     string           `jsons:"type"`
}

// MetricAgg represents a metric aggregation of the time series query model of the datasource
type MetricAgg struct {
	Field             string            `json:"field"`
	Hide              bool              `json:"hide"`
	ID                string            `json:"id"`
	PipelineAggregate string            `json:"pipelineAgg"`
	PipelineVariables map[string]string `json:"pipelineVariables"`
	Settings          *simplejson.Json  `json:"settings"`
	Meta              *simplejson.Json  `json:"meta"`
	Type              string            `json:"type"`
}

var metricAggType = map[string]string{
	"count":          "Count",
	"avg":            "Average",
	"sum":            "Sum",
	"max":            "Max",
	"min":            "Min",
	"extended_stats": "Extended Stats",
	"percentiles":    "Percentiles",
	"top_metrics":    "Top Metrics",
	"cardinality":    "Unique Count",
	"moving_avg":     "Moving Average",
	"moving_fn":      "Moving Function",
	"cumulative_sum": "Cumulative Sum",
	"derivative":     "Derivative",
	"serial_diff":    "Serial Difference",
	"bucket_script":  "Bucket Script",
	"raw_document":   "Raw Document",
	"rate":           "Rate",
}

var extendedStats = map[string]string{
	"avg":                        "Avg",
	"min":                        "Min",
	"max":                        "Max",
	"sum":                        "Sum",
	"count":                      "Count",
	"std_deviation":              "Std Dev",
	"std_deviation_bounds_upper": "Std Dev Upper",
	"std_deviation_bounds_lower": "Std Dev Lower",
}

var pipelineAggType = map[string]string{
	"moving_avg":     "moving_avg",
	"moving_fn":      "moving_fn",
	"cumulative_sum": "cumulative_sum",
	"derivative":     "derivative",
	"serial_diff":    "serial_diff",
	"bucket_script":  "bucket_script",
}

var scriptableAggType = map[string]string{
	"avg":            "avg",
	"sum":            "sum",
	"max":            "max",
	"min":            "min",
	"extended_stats": "extended_stats",
	"percentiles":    "percentiles",
	"bucket_script":  "bucket_script",
}

var pipelineAggWithMultipleBucketPathsType = map[string]string{
	"bucket_script": "bucket_script",
}

func isPipelineAgg(metricType string) bool {
	if _, ok := pipelineAggType[metricType]; ok {
		return true
	}
	return false
}

func isMetricAggregationWithInlineScriptSupport(metricType string) bool {
	if _, ok := scriptableAggType[metricType]; ok {
		return true
	}
	return false
}

func isPipelineAggWithMultipleBucketPaths(metricType string) bool {
	if _, ok := pipelineAggWithMultipleBucketPathsType[metricType]; ok {
		return true
	}
	return false
}

func describeMetric(metricType, field string) string {
	text := metricAggType[metricType]
	if metricType == countType {
		return text
	}
	return text + " " + field
}
<commit_msg>Elasticsearch: fix typo (#58519)<commit_after>package elasticsearch

import (
	"github.com/grafana/grafana/pkg/components/simplejson"
)

// Query represents the time series query model of the datasource
type Query struct {
	TimeField     string       `json:"timeField"`
	RawQuery      string       `json:"query"`
	BucketAggs    []*BucketAgg `json:"bucketAggs"`
	Metrics       []*MetricAgg `json:"metrics"`
	Alias         string       `json:"alias"`
	Interval      string
	IntervalMs    int64
	RefID         string
	MaxDataPoints int64
}

// BucketAgg represents a bucket aggregation of the time series query model of the datasource
type BucketAgg struct {
	Field    string           `json:"field"`
	ID       string           `json:"id"`
	Settings *simplejson.Json `json:"settings"`
	Type     string           `json:"type"`
}

// MetricAgg represents a metric aggregation of the time series query model of the datasource
type MetricAgg struct {
	Field             string            `json:"field"`
	Hide              bool              `json:"hide"`
	ID                string            `json:"id"`
	PipelineAggregate string            `json:"pipelineAgg"`
	PipelineVariables map[string]string `json:"pipelineVariables"`
	Settings          *simplejson.Json  `json:"settings"`
	Meta              *simplejson.Json  `json:"meta"`
	Type              string            `json:"type"`
}

var metricAggType = map[string]string{
	"count":          "Count",
	"avg":            "Average",
	"sum":            "Sum",
	"max":            "Max",
	"min":            "Min",
	"extended_stats": "Extended Stats",
	"percentiles":    "Percentiles",
	"top_metrics":    "Top Metrics",
	"cardinality":    "Unique Count",
	"moving_avg":     "Moving Average",
	"moving_fn":      "Moving Function",
	"cumulative_sum": "Cumulative Sum",
	"derivative":     "Derivative",
	"serial_diff":    "Serial Difference",
	"bucket_script":  "Bucket Script",
	"raw_document":   "Raw Document",
	"rate":           "Rate",
}

var extendedStats = map[string]string{
	"avg":                        "Avg",
	"min":                        "Min",
	"max":                        "Max",
	"sum":                        "Sum",
	"count":                      "Count",
	"std_deviation":              "Std Dev",
	"std_deviation_bounds_upper": "Std Dev Upper",
	"std_deviation_bounds_lower": "Std Dev Lower",
}

var pipelineAggType = map[string]string{
	"moving_avg":     "moving_avg",
	"moving_fn":      "moving_fn",
	"cumulative_sum": "cumulative_sum",
	"derivative":     "derivative",
	"serial_diff":    "serial_diff",
	"bucket_script":  "bucket_script",
}

var scriptableAggType = map[string]string{
	"avg":            "avg",
	"sum":            "sum",
	"max":            "max",
	"min":            "min",
	"extended_stats": "extended_stats",
	"percentiles":    "percentiles",
	"bucket_script":  "bucket_script",
}

var pipelineAggWithMultipleBucketPathsType = map[string]string{
	"bucket_script": "bucket_script",
}

func isPipelineAgg(metricType string) bool {
	if _, ok := pipelineAggType[metricType]; ok {
		return true
	}
	return false
}

func isMetricAggregationWithInlineScriptSupport(metricType string) bool {
	if _, ok := scriptableAggType[metricType]; ok {
		return true
	}
	return false
}

func isPipelineAggWithMultipleBucketPaths(metricType string) bool {
	if _, ok := pipelineAggWithMultipleBucketPathsType[metricType]; ok {
		return true
	}
	return false
}

func describeMetric(metricType, field string) string {
	text := metricAggType[metricType]
	if metricType == countType {
		return text
	}
	return text + " " + field
}
<|endoftext|>


<commit_before>package erede

import (
	"io"
)

type Webservice struct {
	user             string
	password         string
	environment      int
	verbose          bool
	XMLRequestLogger io.Writer
	XMLResultLogger  io.Writer
}

func NewWebservice(user, pw string, env int) *Webservice {
	return &Webservice{user, pw, env, false, nil, nil}
}

func (ws *Webservice) URL() string {
	if ws.environment == ProductionEnv {
		return URLPROD
	}
	return URLDEV
}

func (ws *Webservice) Verbose(tf bool) {
	ws.verbose = tf
}

func (ws *Webservice) NewCreditCardTransaction(cardNumber, cvc2, expiresMonth, expiresYear, orderID string, moneyAmount float64) *Transaction {
	if len(expiresYear) > 2 {
		expiresYear = expiresYear[2:] // will become an issue after year 9999
	}
	t := &Transaction{
		CardNumber:       cardNumber,
		CardExpiryDate:   CardExpDate{expiresMonth, expiresYear},
		CVC2:             cvc2,
		OrderID:          orderID,
		OrderMoneyAmount: moneyAmount,
		OrderType:        "credit",
	}
	t.ws = ws
	t.location = "Sao Paulo"
	return t
}

func (ws *Webservice) NewDebitCardTransaction(cardNumber, cvc2, expiresMonth, expiresYear, orderID string, moneyAmount float64, debitDescription, debitStoreURL string) *Transaction {
	if len(expiresYear) > 2 {
		expiresYear = expiresYear[2:] // will become an issue after year 9999
	}
	t := &Transaction{
		CardNumber:       cardNumber,
		CardExpiryDate:   CardExpDate{expiresMonth, expiresYear},
		CVC2:             cvc2,
		OrderID:          orderID,
		OrderMoneyAmount: moneyAmount,
		OrderType:        "debit",
	}
	t.location = "Sao Paulo"
	t.DebitInfo.DateTime = DTNow()
	t.DebitInfo.Description = debitDescription
	t.DebitInfo.URL = debitStoreURL
	t.skipRisk = true
	t.ws = ws
	return t
}

func (ws *Webservice) NewBoletoTransaction(orderID string, moneyAmount float64, details TransactionBoletoInfo) *Transaction {
	t := &Transaction{
		OrderID:          orderID,
		OrderMoneyAmount: moneyAmount,
		OrderType:        "boleto",
		BoletoInfo:       details,
	}
	t.ws = ws
	return t
}

func (ws *Webservice) NewQueryTransaction(gatewayref string) *QueryTx {
	qtx := &QueryTx{
		GatewayReference: gatewayref,
	}
	qtx.ws = ws
	return qtx
}

//TODO: proper fulfill and pre methods
<commit_msg>consultbolete Webservice<commit_after>package erede

import (
	"bytes"
	"encoding/xml"
	"io"
	"log"
	"net/http"
)

type Webservice struct {
	user             string
	password         string
	environment      int
	verbose          bool
	XMLRequestLogger io.Writer
	XMLResultLogger  io.Writer
}

func NewWebservice(user, pw string, env int) *Webservice {
	return &Webservice{user, pw, env, false, nil, nil}
}

func (ws *Webservice) URL() string {
	if ws.environment == ProductionEnv {
		return URLPROD
	}
	return URLDEV
}

func (ws *Webservice) Verbose(tf bool) {
	ws.verbose = tf
}

func (ws *Webservice) NewCreditCardTransaction(cardNumber, cvc2, expiresMonth, expiresYear, orderID string, moneyAmount float64) *Transaction {
	if len(expiresYear) > 2 {
		expiresYear = expiresYear[2:] // will become an issue after year 9999
	}
	t := &Transaction{
		CardNumber:       cardNumber,
		CardExpiryDate:   CardExpDate{expiresMonth, expiresYear},
		CVC2:             cvc2,
		OrderID:          orderID,
		OrderMoneyAmount: moneyAmount,
		OrderType:        "credit",
	}
	t.ws = ws
	t.location = "Sao Paulo"
	return t
}

func (ws *Webservice) NewDebitCardTransaction(cardNumber, cvc2, expiresMonth, expiresYear, orderID string, moneyAmount float64, debitDescription, debitStoreURL string) *Transaction {
	if len(expiresYear) > 2 {
		expiresYear = expiresYear[2:] // will become an issue after year 9999
	}
	t := &Transaction{
		CardNumber:       cardNumber,
		CardExpiryDate:   CardExpDate{expiresMonth, expiresYear},
		CVC2:             cvc2,
		OrderID:          orderID,
		OrderMoneyAmount: moneyAmount,
		OrderType:        "debit",
	}
	t.location = "Sao Paulo"
	t.DebitInfo.DateTime = DTNow()
	t.DebitInfo.Description = debitDescription
	t.DebitInfo.URL = debitStoreURL
	t.skipRisk = true
	t.ws = ws
	return t
}

func (ws *Webservice) NewBoletoTransaction(orderID string, moneyAmount float64, details TransactionBoletoInfo) *Transaction {
	t := &Transaction{
		OrderID:          orderID,
		OrderMoneyAmount: moneyAmount,
		OrderType:        "boleto",
		BoletoInfo:       details,
	}
	t.ws = ws
	return t
}

func (ws *Webservice) ConsultBoleto(gatewayRef string, reqccbuf, respccbuf io.Writer) (*QueryResponse, error) {
	tpl := new(bytes.Buffer)
	tpl.WriteString(`<Request version="2">`)
	tpl.WriteString("<Authentication><AcquirerCode><rdcd_pv>")
	tpl.WriteString(ws.user)
	tpl.WriteString("</rdcd_pv></AcquirerCode><password>")
	tpl.WriteString(ws.password)
	tpl.WriteString("</password></Authentication>")
	tpl.WriteString("<Transaction><HistoricTxn><reference>")
	tpl.WriteString(gatewayRef)
	tpl.WriteString("</reference>")
	tpl.WriteString("<method>query</method></HistoricTxn></Transaction></Request>")
	if Verbose {
		log.Println("Sending ConsultBoleto: ", tpl.String())
	}
	if reqccbuf != nil {
		reqccbuf.Write(tpl.Bytes())
	}
	resp, err := http.Post(URL(), "application/xml", tpl)
	if err != nil {
		return nil, err
	}
	tpl.Reset()
	defer resp.Body.Close()
	io.Copy(tpl, resp.Body)
	if Verbose {
		log.Println("XML RESPONSE", tpl.String())
	}
	if respccbuf != nil {
		respccbuf.Write(tpl.Bytes())
	}
	tr := &QueryResponse{}
	err = xml.Unmarshal(tpl.Bytes(), tr)
	return tr, err
}

func (ws *Webservice) NewQueryTransaction(gatewayref string) *QueryTx {
	qtx := &QueryTx{
		GatewayReference: gatewayref,
	}
	qtx.ws = ws
	return qtx
}

//TODO: proper fulfill and pre methods
<|endoftext|>


<commit_before>package executor

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
)

func TestMain(m *testing.M) {
	switch os.Getenv("TEST_MAIN") {
	case "app":
		appMain()
	default:
		os.Exit(m.Run())
	}
}

func appMain() {
	if len(os.Args) < 2 {
		fmt.Fprintln(os.Stderr, "no command provided")
		os.Exit(1)
	}
	args := os.Args[1:]
	shift := func() string {
		s := args[0]
		args = args[1:]
		return s
	}
	for len(args) > 0 {
		switch cmd := shift(); cmd {
		case "sleep":
			if len(args) < 1 {
				fmt.Fprintln(os.Stderr, "expected arg for sleep")
				os.Exit(1)
			}
			dur, err := time.ParseDuration(shift())
			if err != nil {
				fmt.Fprintf(os.Stderr, "could not parse sleep time: %v", err)
				os.Exit(1)
			}
			time.Sleep(dur)
		case "echo":
			fmt.Println(strings.Join(args, " "))
			args = args[:0]
		case "write":
			if len(args) < 2 {
				fmt.Fprintln(os.Stderr, "expected two args for write")
				os.Exit(1)
			}
			msg := shift()
			file := shift()
			ioutil.WriteFile(file, []byte(msg), 0666)
		default:
			fmt.Fprintln(os.Stderr, "unknown command:", cmd)
			os.Exit(1)
		}
	}
}

var (
	constraint = &structs.Resources{
		CPU:      250,
		MemoryMB: 256,
		Networks: []*structs.NetworkResource{
			&structs.NetworkResource{
				MBits:        50,
				DynamicPorts: []structs.Port{{Label: "http"}},
			},
		},
	}
)

func mockAllocDir(t *testing.T) (string, *allocdir.AllocDir) {
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]

	allocDir := allocdir.NewAllocDir(filepath.Join(os.TempDir(), alloc.ID))
	if err := allocDir.Build([]*structs.Task{task}); err != nil {
		log.Panicf("allocDir.Build() failed: %v", err)
	}

	return task.Name, allocDir
}

func testExecutor(t *testing.T, buildExecutor func() Executor, compatible func(*testing.T)) {
	if compatible != nil {
		compatible(t)
	}

	command := func(name string, args ...string) Executor {
		e := buildExecutor()
		SetCommand(e, name, args)
		e.Command().Env = append(os.Environ(), "TEST_MAIN=app")
		return e
	}

	Executor_Start_Invalid(t, command)
	Executor_Start_Wait_Failure_Code(t, command)
	Executor_Start_Wait(t, command)
	Executor_Start_Kill(t, command)
	Executor_Open(t, command, buildExecutor)
}

type buildExecCommand func(name string, args ...string) Executor

func Executor_Start_Invalid(t *testing.T, command buildExecCommand) {
	invalid := "/bin/foobar"
	e := command(invalid, "1")

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()
	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err == nil {
		log.Panicf("Start(%v) should have failed", invalid)
	}
}

func Executor_Start_Wait_Failure_Code(t *testing.T, command buildExecCommand) {
	e := command(os.Args[0], "fail")

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()
	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	if err := e.Wait(); err == nil {
		log.Panicf("Wait() should have failed")
	}
}

func Executor_Start_Wait(t *testing.T, command buildExecCommand) {
	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()

	taskDir, ok := alloc.TaskDirs[task]
	if !ok {
		log.Panicf("No task directory found for task %v", task)
	}

	expected := "hello world"
	file := filepath.Join(allocdir.TaskLocal, "output.txt")
	absFilePath := filepath.Join(taskDir, file)
	e := command(os.Args[0], "sleep", "1s", "write", expected, file)

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	if res := e.Wait(); !res.Successful() {
		log.Panicf("Wait() failed: %v", res)
	}

	output, err := ioutil.ReadFile(absFilePath)
	if err != nil {
		log.Panicf("Couldn't read file %v", absFilePath)
	}

	act := string(output)
	if act != expected {
		log.Panicf("Command output incorrectly: want %v; got %v", expected, act)
	}
}

func Executor_Start_Kill(t *testing.T, command buildExecCommand) {
	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()

	taskDir, ok := alloc.TaskDirs[task]
	if !ok {
		log.Panicf("No task directory found for task %v", task)
	}

	filePath := filepath.Join(taskDir, "output")
	e := command(os.Args[0], "sleep", "1s", "write", "failure", filePath)

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	if err := e.Shutdown(); err != nil {
		log.Panicf("Shutdown() failed: %v", err)
	}

	time.Sleep(1500 * time.Millisecond)

	// Check that the file doesn't exist.
	if _, err := os.Stat(filePath); err == nil {
		log.Panicf("Stat(%v) should have failed: task not killed", filePath)
	}
}

func Executor_Open(t *testing.T, command buildExecCommand, newExecutor func() Executor) {
	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()

	taskDir, ok := alloc.TaskDirs[task]
	if !ok {
		log.Panicf("No task directory found for task %v", task)
	}

	expected := "hello world"
	file := filepath.Join(allocdir.TaskLocal, "output.txt")
	absFilePath := filepath.Join(taskDir, file)
	e := command(os.Args[0], "sleep", "1s", "write", expected, file)

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	id, err := e.ID()
	if err != nil {
		log.Panicf("ID() failed: %v", err)
	}

	e2 := newExecutor()
	if err := e2.Open(id); err != nil {
		log.Panicf("Open(%v) failed: %v", id, err)
	}

	if res := e2.Wait(); !res.Successful() {
		log.Panicf("Wait() failed: %v", res)
	}

	output, err := ioutil.ReadFile(absFilePath)
	if err != nil {
		log.Panicf("Couldn't read file %v", absFilePath)
	}

	act := string(output)
	if act != expected {
		log.Panicf("Command output incorrectly: want %v; got %v", expected, act)
	}
}

func Executor_Open_Invalid(t *testing.T, command buildExecCommand, newExecutor func() Executor) {
	task, alloc := mockAllocDir(t)
	e := command("echo", "foo")

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	id, err := e.ID()
	if err != nil {
		log.Panicf("ID() failed: %v", err)
	}

	// Destroy the allocdir which removes the exit code.
	alloc.Destroy()

	e2 := newExecutor()
	if err := e2.Open(id); err == nil {
		log.Panicf("Open(%v) should have failed", id)
	}
}
<commit_msg>Improve names and documentation in test code.<commit_after>package executor

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
)

// testBinary is the path to the running test binary
var testBinary = os.Args[0]

func TestMain(m *testing.M) {
	// The tests in this package recursively execute the test binary produced
	// by go test. The TEST_MAIN environment variable controls the recursive
	// execution.
	switch tm := os.Getenv("TEST_MAIN"); tm {
	case "":
		os.Exit(m.Run())
	case "app":
		appMain()
	default:
		fmt.Fprintf(os.Stderr, "unexpected value for TEST_MAIN, \"%s\"\n", tm)
		os.Exit(1)
	}
}

// setTestAppEnv sets the environement of cmd for a recursive call into
// TestMain.
func setTestAppEnv(cmd *exec.Cmd) {
	cmd.Env = append(os.Environ(), "TEST_MAIN=app")
}

func appMain() {
	if len(os.Args) < 2 {
		fmt.Fprintln(os.Stderr, "no command provided")
		os.Exit(1)
	}

	args := os.Args[1:]

	// popArg removes the first argument from args and returns it.
	popArg := func() string {
		s := args[0]
		args = args[1:]
		return s
	}

	// execute a sequence of operations from args
	for len(args) > 0 {
		switch cmd := popArg(); cmd {

		case "sleep":
			// sleep <dur>: sleep for a duration indicated by the first
			// argument
			if len(args) < 1 {
				fmt.Fprintln(os.Stderr, "expected arg for sleep")
				os.Exit(1)
			}
			dur, err := time.ParseDuration(popArg())
			if err != nil {
				fmt.Fprintf(os.Stderr, "could not parse sleep time: %v", err)
				os.Exit(1)
			}
			time.Sleep(dur)

		case "echo":
			// echo <msg ...>: write the remaining arguments to stdout each
			// separated by a single space and followed by a newline.
			fmt.Println(strings.Join(args, " "))
			args = args[:0]

		case "write":
			// write <msg> <file>: write a message to a file. The first
			// argument is the msg. The second argument is the path to the
			// target file.
			if len(args) < 2 {
				fmt.Fprintln(os.Stderr, "expected two args for write")
				os.Exit(1)
			}
			msg := popArg()
			file := popArg()
			ioutil.WriteFile(file, []byte(msg), 0666)

		default:
			fmt.Fprintln(os.Stderr, "unknown command:", cmd)
			os.Exit(1)
		}
	}
}

var (
	constraint = &structs.Resources{
		CPU:      250,
		MemoryMB: 256,
		Networks: []*structs.NetworkResource{
			&structs.NetworkResource{
				MBits:        50,
				DynamicPorts: []structs.Port{{Label: "http"}},
			},
		},
	}
)

func mockAllocDir(t *testing.T) (string, *allocdir.AllocDir) {
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]

	allocDir := allocdir.NewAllocDir(filepath.Join(os.TempDir(), alloc.ID))
	if err := allocDir.Build([]*structs.Task{task}); err != nil {
		log.Panicf("allocDir.Build() failed: %v", err)
	}

	return task.Name, allocDir
}

func testExecutor(t *testing.T, buildExecutor func() Executor, compatible func(*testing.T)) {
	if compatible != nil {
		compatible(t)
	}

	command := func(name string, args ...string) Executor {
		e := buildExecutor()
		SetCommand(e, name, args)
		setTestAppEnv(e.Command())
		return e
	}

	Executor_Start_Invalid(t, command)
	Executor_Start_Wait_Failure_Code(t, command)
	Executor_Start_Wait(t, command)
	Executor_Start_Kill(t, command)
	Executor_Open(t, command, buildExecutor)
}

type buildExecCommand func(name string, args ...string) Executor

func Executor_Start_Invalid(t *testing.T, command buildExecCommand) {
	invalid := "/bin/foobar"
	e := command(invalid, "1")

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()
	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err == nil {
		log.Panicf("Start(%v) should have failed", invalid)
	}
}

func Executor_Start_Wait_Failure_Code(t *testing.T, command buildExecCommand) {
	e := command(testBinary, "fail")

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()
	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	if err := e.Wait(); err == nil {
		log.Panicf("Wait() should have failed")
	}
}

func Executor_Start_Wait(t *testing.T, command buildExecCommand) {
	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()

	taskDir, ok := alloc.TaskDirs[task]
	if !ok {
		log.Panicf("No task directory found for task %v", task)
	}

	expected := "hello world"
	file := filepath.Join(allocdir.TaskLocal, "output.txt")
	absFilePath := filepath.Join(taskDir, file)
	e := command(testBinary, "sleep", "1s", "write", expected, file)

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	if res := e.Wait(); !res.Successful() {
		log.Panicf("Wait() failed: %v", res)
	}

	output, err := ioutil.ReadFile(absFilePath)
	if err != nil {
		log.Panicf("Couldn't read file %v", absFilePath)
	}

	act := string(output)
	if act != expected {
		log.Panicf("Command output incorrectly: want %v; got %v", expected, act)
	}
}

func Executor_Start_Kill(t *testing.T, command buildExecCommand) {
	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()

	taskDir, ok := alloc.TaskDirs[task]
	if !ok {
		log.Panicf("No task directory found for task %v", task)
	}

	filePath := filepath.Join(taskDir, "output")
	e := command(testBinary, "sleep", "1s", "write", "failure", filePath)

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	if err := e.Shutdown(); err != nil {
		log.Panicf("Shutdown() failed: %v", err)
	}

	time.Sleep(1500 * time.Millisecond)

	// Check that the file doesn't exist.
	if _, err := os.Stat(filePath); err == nil {
		log.Panicf("Stat(%v) should have failed: task not killed", filePath)
	}
}

func Executor_Open(t *testing.T, command buildExecCommand, newExecutor func() Executor) {
	task, alloc := mockAllocDir(t)
	defer alloc.Destroy()

	taskDir, ok := alloc.TaskDirs[task]
	if !ok {
		log.Panicf("No task directory found for task %v", task)
	}

	expected := "hello world"
	file := filepath.Join(allocdir.TaskLocal, "output.txt")
	absFilePath := filepath.Join(taskDir, file)
	e := command(testBinary, "sleep", "1s", "write", expected, file)

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	id, err := e.ID()
	if err != nil {
		log.Panicf("ID() failed: %v", err)
	}

	e2 := newExecutor()
	if err := e2.Open(id); err != nil {
		log.Panicf("Open(%v) failed: %v", id, err)
	}

	if res := e2.Wait(); !res.Successful() {
		log.Panicf("Wait() failed: %v", res)
	}

	output, err := ioutil.ReadFile(absFilePath)
	if err != nil {
		log.Panicf("Couldn't read file %v", absFilePath)
	}

	act := string(output)
	if act != expected {
		log.Panicf("Command output incorrectly: want %v; got %v", expected, act)
	}
}

func Executor_Open_Invalid(t *testing.T, command buildExecCommand, newExecutor func() Executor) {
	task, alloc := mockAllocDir(t)
	e := command("echo", "foo")

	if err := e.Limit(constraint); err != nil {
		log.Panicf("Limit() failed: %v", err)
	}

	if err := e.ConfigureTaskDir(task, alloc); err != nil {
		log.Panicf("ConfigureTaskDir(%v, %v) failed: %v", task, alloc, err)
	}

	if err := e.Start(); err != nil {
		log.Panicf("Start() failed: %v", err)
	}

	id, err := e.ID()
	if err != nil {
		log.Panicf("ID() failed: %v", err)
	}

	// Destroy the allocdir which removes the exit code.
	alloc.Destroy()

	e2 := newExecutor()
	if err := e2.Open(id); err == nil {
		log.Panicf("Open(%v) should have failed", id)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"image"
	"io"
	"log"
	"os"

	"github.com/lucasb-eyer/go-colorful"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/config"
	"github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/sphere-go-led-controller/ui"
	"github.com/ninjasphere/sphere-go-led-controller/util"
	"github.com/tarm/goserial"
)

type LedController struct {
	controlEnabled bool
	controlLayout  *ui.PaneLayout
	pairingLayout  *ui.PairingLayout
	conn           *ninja.Connection
	serial         io.ReadWriteCloser
	waiting        chan bool
}

func NewLedController(conn *ninja.Connection) (*LedController, error) {

	c := &serial.Config{Name: "/dev/tty.ledmatrix", Baud: 115200}
	s, err := serial.OpenPort(c)
	if err != nil {
		return nil, err
	}

	// Send a blank image to the led matrix
	util.WriteLEDMatrix(image.NewRGBA(image.Rect(0, 0, 16, 16)), s)

	controller := &LedController{
		conn:          conn,
		pairingLayout: ui.NewPairingLayout(conn),
		serial:        s,
		waiting:       make(chan bool),
	}

	conn.MustExportService(controller, "$node/"+config.Serial()+"/led-controller", &model.ServiceAnnouncement{
		Schema: "/service/led-controller",
	})

	return controller, nil
}

func (c *LedController) start(enableControl bool) {
	c.controlEnabled = enableControl

	go func() {
		for {
			if c.controlEnabled {

				if c.controlLayout == nil {

					log.Println("Enabling layout... clearing LED")

					util.WriteLEDMatrix(image.NewRGBA(image.Rect(0, 0, 16, 16)), c.serial)

					c.controlLayout = getPaneLayout(c.conn)
					log.Println("Finished control layout")
				}

				image, wake, err := c.controlLayout.Render()
				if err != nil {
					log.Fatal(err)
				}

				util.WriteLEDMatrix(image, c.serial)

				if wake != nil {
					log.Println("Waiting as the UI is asleep")
					select {
					case <-wake:
						log.Println("UI woke up!")
					case <-c.waiting:
						log.Println("Got a command from rpc...")
					}
				}

			} else {

				image, err := c.pairingLayout.Render()
				if err != nil {
					log.Fatal(err)
				}
				util.WriteLEDMatrix(image, c.serial)

			}
		}
	}()
}

func (c *LedController) EnableControl() error {
	c.controlEnabled = true
	c.gotCommand()
	return nil
}

func (c *LedController) DisableControl() error {
	c.controlEnabled = false
	c.gotCommand()
	return nil
}

type PairingCodeRequest struct {
	Code        string `json:"code"`
	DisplayTime int    `json:"displayTime"`
}

func (c *LedController) DisplayPairingCode(req *PairingCodeRequest) error {
	c.controlEnabled = false
	c.pairingLayout.ShowCode(req.Code)
	c.gotCommand()
	return nil
}

type ColorRequest struct {
	Color       string `json:"color"`
	DisplayTime int    `json:"displayTime"`
}

func (c *LedController) DisplayColor(req *ColorRequest) error {
	col, err := colorful.Hex(req.Color)

	if err != nil {
		return err
	}

	c.controlEnabled = false
	c.pairingLayout.ShowColor(col)
	c.gotCommand()
	return nil
}

type IconRequest struct {
	Icon        string `json:"icon"`
	DisplayTime int    `json:"displayTime"`
}

func (c *LedController) DisplayIcon(req *IconRequest) error {
	c.controlEnabled = false
	c.pairingLayout.ShowIcon(req.Icon)
	c.gotCommand()
	return nil
}

func (c *LedController) gotCommand() {
	select {
	case c.waiting <- true:
	default:
	}
}

// Load from a config file instead...
func getPaneLayout(conn *ninja.Connection) *ui.PaneLayout {
	layout, wake := ui.NewPaneLayout(false)

	mediaPane := ui.NewMediaPane(&ui.MediaPaneImages{
		Volume: "images/media-volume-speaker.gif",
		Mute:   "images/media-volume-mute.png",
		Play:   "images/media-play.png",
		Pause:  "images/media-pause.png",
		Stop:   "images/media-stop.png",
		Next:   "images/media-next.png",
	}, conn)
	layout.AddPane(mediaPane)

	if len(os.Getenv("CERTIFICATION")) > 0 {
		layout.AddPane(ui.NewCertPane(conn.GetMqttClient()))
	} else {
		//layout.AddPane(ui.NewTextScrollPane("Exit Music (For A Film)"))

		heaterPane := ui.NewOnOffPane("images/heater-off.png", "images/heater-on.gif", func(state bool) {
			log.Printf("Heater state: %t", state)
		}, conn, "heater")
		layout.AddPane(heaterPane)
	}

	lightPane := ui.NewLightPane("images/light-off.png", "images/light-on.png", func(state bool) {
		log.Printf("Light on-off state: %t", state)
	}, func(state float64) {
		log.Printf("Light color state: %f", state)
	}, conn)
	layout.AddPane(lightPane)

	fanPane := ui.NewOnOffPane("images/fan-off.png", "images/fan-on.gif", func(state bool) {
		log.Printf("Fan state: %t", state)
	}, conn, "fan")

	layout.AddPane(fanPane)

	go func() {
		<-wake
	}()

	go layout.Wake()

	return layout
}
<commit_msg>If we timeout (10sec) writing a frame to the led matrix, reboot!<commit_after>package main

import (
	"image"
	"io"
	"log"
	"os"
	"os/exec"
	"time"

	"github.com/lucasb-eyer/go-colorful"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/config"
	"github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/sphere-go-led-controller/ui"
	"github.com/ninjasphere/sphere-go-led-controller/util"
	"github.com/tarm/goserial"
)

type LedController struct {
	controlEnabled bool
	controlLayout  *ui.PaneLayout
	pairingLayout  *ui.PairingLayout
	conn           *ninja.Connection
	serial         io.ReadWriteCloser
	waiting        chan bool
}

func NewLedController(conn *ninja.Connection) (*LedController, error) {

	c := &serial.Config{Name: "/dev/tty.ledmatrix", Baud: 115200}
	s, err := serial.OpenPort(c)
	if err != nil {
		return nil, err
	}

	// Send a blank image to the led matrix
	util.WriteLEDMatrix(image.NewRGBA(image.Rect(0, 0, 16, 16)), s)

	controller := &LedController{
		conn:          conn,
		pairingLayout: ui.NewPairingLayout(conn),
		serial:        s,
		waiting:       make(chan bool),
	}

	conn.MustExportService(controller, "$node/"+config.Serial()+"/led-controller", &model.ServiceAnnouncement{
		Schema: "/service/led-controller",
	})

	return controller, nil
}

func (c *LedController) start(enableControl bool) {
	c.controlEnabled = enableControl

	frameWritten := make(chan bool)

	go func() {
		for {
			if c.controlEnabled {

				if c.controlLayout == nil {

					log.Println("Enabling layout... clearing LED")

					util.WriteLEDMatrix(image.NewRGBA(image.Rect(0, 0, 16, 16)), c.serial)

					c.controlLayout = getPaneLayout(c.conn)
					log.Println("Finished control layout")
				}

				image, wake, err := c.controlLayout.Render()
				if err != nil {
					log.Fatal(err)
				}

				go func() {
					util.WriteLEDMatrix(image, c.serial)
					frameWritten <- true
				}()

				select {
				case <-frameWritten:
					// All good.
				case <-time.After(10 * time.Second):
					log.Println("Timeout writing to LED matrix. REBOOTING!")
					// Timed out writing to the led matrix. For now. Boot!
					cmd := exec.Command("reboot")
					output, err := cmd.Output()

					log.Printf("Output from reboot: %s err: %s", output, err)
				}

				if wake != nil {
					log.Println("Waiting as the UI is asleep")
					select {
					case <-wake:
						log.Println("UI woke up!")
					case <-c.waiting:
						log.Println("Got a command from rpc...")
					}
				}

			} else {

				image, err := c.pairingLayout.Render()
				if err != nil {
					log.Fatal(err)
				}
				util.WriteLEDMatrix(image, c.serial)

			}
		}
	}()
}

func (c *LedController) EnableControl() error {
	c.controlEnabled = true
	c.gotCommand()
	return nil
}

func (c *LedController) DisableControl() error {
	c.controlEnabled = false
	c.gotCommand()
	return nil
}

type PairingCodeRequest struct {
	Code        string `json:"code"`
	DisplayTime int    `json:"displayTime"`
}

func (c *LedController) DisplayPairingCode(req *PairingCodeRequest) error {
	c.controlEnabled = false
	c.pairingLayout.ShowCode(req.Code)
	c.gotCommand()
	return nil
}

type ColorRequest struct {
	Color       string `json:"color"`
	DisplayTime int    `json:"displayTime"`
}

func (c *LedController) DisplayColor(req *ColorRequest) error {
	col, err := colorful.Hex(req.Color)

	if err != nil {
		return err
	}

	c.controlEnabled = false
	c.pairingLayout.ShowColor(col)
	c.gotCommand()
	return nil
}

type IconRequest struct {
	Icon        string `json:"icon"`
	DisplayTime int    `json:"displayTime"`
}

func (c *LedController) DisplayIcon(req *IconRequest) error {
	c.controlEnabled = false
	c.pairingLayout.ShowIcon(req.Icon)
	c.gotCommand()
	return nil
}

func (c *LedController) gotCommand() {
	select {
	case c.waiting <- true:
	default:
	}
}

// Load from a config file instead...
func getPaneLayout(conn *ninja.Connection) *ui.PaneLayout {
	layout, wake := ui.NewPaneLayout(false)

	mediaPane := ui.NewMediaPane(&ui.MediaPaneImages{
		Volume: "images/media-volume-speaker.gif",
		Mute:   "images/media-volume-mute.png",
		Play:   "images/media-play.png",
		Pause:  "images/media-pause.png",
		Stop:   "images/media-stop.png",
		Next:   "images/media-next.png",
	}, conn)
	layout.AddPane(mediaPane)

	if len(os.Getenv("CERTIFICATION")) > 0 {
		layout.AddPane(ui.NewCertPane(conn.GetMqttClient()))
	} else {
		//layout.AddPane(ui.NewTextScrollPane("Exit Music (For A Film)"))

		heaterPane := ui.NewOnOffPane("images/heater-off.png", "images/heater-on.gif", func(state bool) {
			log.Printf("Heater state: %t", state)
		}, conn, "heater")
		layout.AddPane(heaterPane)
	}

	lightPane := ui.NewLightPane("images/light-off.png", "images/light-on.png", func(state bool) {
		log.Printf("Light on-off state: %t", state)
	}, func(state float64) {
		log.Printf("Light color state: %f", state)
	}, conn)
	layout.AddPane(lightPane)

	fanPane := ui.NewOnOffPane("images/fan-off.png", "images/fan-on.gif", func(state bool) {
		log.Printf("Fan state: %t", state)
	}, conn, "fan")

	layout.AddPane(fanPane)

	go func() {
		<-wake
	}()

	go layout.Wake()

	return layout
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package encdb defines an encrypted database used within Mute.
Such an encrypted database consists of two file for a given database file with
name "dbname":

  dbname.db
  dbname.key

The file "dbname.db" is an AES-256 encrypted sqlite3 file managed by the
package "mute/go-sqlcipher". The file named "dbname.key" is an
AES-256 encrypted text file which contains the (randomly generated) raw
encryption key for "dbname.db". To decrypt the key file the key derivation
function PBKDF2 is applied to a supplied passphrase (with a configurable number
of iterations) and the derived key is used as the AES-256 key for "dbname.key".

This design allows a very cheap rekey of the database, because only the key
file needs to be changed and the database file itself doesn't have to be
modified for a rekey operation.
*/
package encdb

import (
	"database/sql"
	"encoding/hex"
	"fmt"
	"os"

	_ "github.com/mutecomm/go-sqlcipher" // encDB is tightly coupled to go-sqlcipher
	"github.com/mutecomm/mute/log"
	"github.com/mutecomm/mute/util"
)

const (
	dbsuffix  = ".db"
	keysuffix = ".key"
)

func createTables(db *sql.DB, createStmts []string) error {
	for _, stmt := range createStmts {
		_, err := db.Exec(stmt)
		if err != nil {
			return log.Errorf("encdb: %q: %s", err, stmt)
		}
	}
	return nil
}

// Create tries to create an encrypted database with the given passphrase and
// iter many KDF iterations. Thereby, dbname is the prefix of the following
// two database files which will be created and must not exist already:
//
//  dbname.db
//  dbname.key
//
// The SQL database is initialized with the statements given in createStmts.
// In case of error (for example, the database files do exist already or
// cannot be created) an error is returned.
func Create(dbname string, passphrase []byte, iter int, createStmts []string) error {
	dbfile := dbname + dbsuffix
	keyfile := dbname + keysuffix
	// make sure files do not exist already
	if _, err := os.Stat(dbfile); err == nil {
		return log.Errorf("encdb: dbfile '%s' exists already", dbfile)
	}
	if _, err := os.Stat(keyfile); err == nil {
		return log.Errorf("encdb: dbfile '%s' exists already", keyfile)
	}
	// create keyfile
	key, err := generateKeyfile(keyfile, passphrase, iter)
	if err != nil {
		return err
	}
	// create DB
	db, err := sql.Open("sqlite3", dbfile)
	if err != nil {
		return log.Error(err)
	}
	defer db.Close()
	_, err = db.Exec(fmt.Sprintf("PRAGMA key = \"x'%s'\";", hex.EncodeToString(key)))
	if err != nil {
		return log.Error(err)
	}
	_, err = db.Exec("PRAGMA cipher_page_size = 4096;")
	if err != nil {
		return log.Error(err)
	}
	_, err = db.Exec("PRAGMA auto_vacuum = FULL;")
	if err != nil {
		return log.Error(err)
	}
	return createTables(db, createStmts)
}

// Open tries to open an encrypted database with the given passphrase.
// Thereby, dbname is the prefix of the following two database files (which
// must already exist):
//
//  dbname.db
//  dbname.key
//
// In case of error (for example, the database files do not exist or the
// passphrase is wrong) an error is returned.
func Open(dbname string, passphrase []byte) (*sql.DB, error) {
	dbfile := dbname + dbsuffix
	keyfile := dbname + keysuffix
	// make sure files exists
	if _, err := os.Stat(dbfile); err != nil {
		return nil, log.Error(err)
	}
	if _, err := os.Stat(keyfile); err != nil {
		return nil, log.Error(err)
	}
	// get key from keyfile
	key, err := readKeyfile(keyfile, passphrase)
	if err != nil {
		return nil, err
	}
	// open DB
	db, err := sql.Open("sqlite3", dbfile)
	if err != nil {
		return nil, log.Error(err)
	}
	_, err = db.Exec(fmt.Sprintf("PRAGMA key = \"x'%s'\";", hex.EncodeToString(key)))
	if err != nil {
		return nil, log.Error(err)
	}
	_, err = db.Exec("PRAGMA cipher_page_size = 4096;")
	if err != nil {
		return nil, log.Error(err)
	}
	// test key
	_, err = db.Exec("SELECT count(*) FROM sqlite_master;")
	if err != nil {
		return nil, log.Error(err)
	}
	// enable foreign key support
	if _, err := db.Exec("PRAGMA foreign_keys = ON;"); err != nil {
		return nil, log.Error(err)
	}
	return db, nil
}

// Rekey tries to rekey an encrypted database with the given newPassphrase and
// newIter many KDF iterations. The correct oldPassphrase must be supplied.
// Thereby, dbname is the prefix of the following two database files (which must
// already exist):
//
//  dbname.db
//  dbname.key
//
// Rekey replaces the dbname.key file and leaves the dbname.db file unmodified,
// allowing for very fast rekey operations. In case of error (for example, the
// database files do not exist or the oldPassphrase is wrong) an error is
// returned.
func Rekey(dbname string, oldPassphrase, newPassphrase []byte, newIter int) error {
	encdb, err := Open(dbname, oldPassphrase)
	if err != nil {
		return err
	}
	defer encdb.Close()
	keyfile := dbname + keysuffix
	return replaceKeyfile(keyfile, oldPassphrase, newPassphrase, newIter)
}

var autoVacuumModes = []string{
	"NONE",
	"FULL",
	"INCREMENTAL",
}

// Status returns the autoVacuum and freelistCount of db.
func Status(db *sql.DB) (autoVacuum string, freelistCount int64, err error) {
	var av int64
	err = db.QueryRow("PRAGMA auto_vacuum;").Scan(&av)
	if err != nil {
		return "", 0, err
	}
	autoVacuum = autoVacuumModes[av]
	err = db.QueryRow("PRAGMA freelist_count;").Scan(&freelistCount)
	if err != nil {
		return "", 0, err
	}
	return
}

// Vacuum executes VACUUM command in db. If autoVacuumMode is not nil and
// different from the current one, the auto_vacuum mode is changed before
// VACUUM is executed.
func Vacuum(db *sql.DB, autoVacuumMode string) error {
	if autoVacuumMode != "" {
		if !util.ContainsString(autoVacuumModes, autoVacuumMode) {
			return log.Errorf("encdb: unknown auto_vacuum mode: %s", autoVacuumMode)
		}
		var av int64
		err := db.QueryRow("PRAGMA auto_vacuum;").Scan(&av)
		if err != nil {
			return err
		}
		currentMode := autoVacuumModes[av]
		if currentMode != autoVacuumMode {
			_, err = db.Exec(fmt.Sprintf("PRAGMA auto_vacuum = %s;", autoVacuumMode))
			if err != nil {
				return err
			}
		}
	}
	_, err := db.Exec("VACUUM;")
	if err != nil {
		return log.Error(err)
	}
	return nil
}

// Incremental executes incremental_vacuum to free up to pages many pages. If
// pages is 0, all pages are freed. If the current auto_vacuum mode is not
// INCREMENTAL, an error is returned.
func Incremental(db *sql.DB, pages int64) error {
	var av int64
	err := db.QueryRow("PRAGMA auto_vacuum;").Scan(&av)
	if err != nil {
		return err
	}
	if autoVacuumModes[av] != "INCREMENTAL" {
		return log.Error("encdb: current auto_vacuum mode is not INCREMENTAL")
	}
	_, err = db.Exec(fmt.Sprintf("PRAGMA incremental_vacuum(%d);", pages))
	if err != nil {
		log.Error(err)
	}
	return nil
}
<commit_msg>encdb: use new go-sqlcipher API<commit_after>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package encdb defines an encrypted database used within Mute.
Such an encrypted database consists of two file for a given database file with
name "dbname":

  dbname.db
  dbname.key

The file "dbname.db" is an AES-256 encrypted sqlite3 file managed by the
package "mute/go-sqlcipher". The file named "dbname.key" is an
AES-256 encrypted text file which contains the (randomly generated) raw
encryption key for "dbname.db". To decrypt the key file the key derivation
function PBKDF2 is applied to a supplied passphrase (with a configurable number
of iterations) and the derived key is used as the AES-256 key for "dbname.key".

This design allows a very cheap rekey of the database, because only the key
file needs to be changed and the database file itself doesn't have to be
modified for a rekey operation.
*/
package encdb

import (
	"database/sql"
	"encoding/hex"
	"fmt"
	"os"

	_ "github.com/mutecomm/go-sqlcipher" // encDB is tightly coupled to go-sqlcipher
	"github.com/mutecomm/mute/log"
	"github.com/mutecomm/mute/util"
)

const (
	dbsuffix  = ".db"
	keysuffix = ".key"
)

func createTables(db *sql.DB, createStmts []string) error {
	for _, stmt := range createStmts {
		_, err := db.Exec(stmt)
		if err != nil {
			return log.Errorf("encdb: %q: %s", err, stmt)
		}
	}
	return nil
}

// Create tries to create an encrypted database with the given passphrase and
// iter many KDF iterations. Thereby, dbname is the prefix of the following
// two database files which will be created and must not exist already:
//
//  dbname.db
//  dbname.key
//
// The SQL database is initialized with the statements given in createStmts.
// In case of error (for example, the database files do exist already or
// cannot be created) an error is returned.
func Create(dbname string, passphrase []byte, iter int, createStmts []string) error {
	dbfile := dbname + dbsuffix
	keyfile := dbname + keysuffix
	// make sure files do not exist already
	if _, err := os.Stat(dbfile); err == nil {
		return log.Errorf("encdb: dbfile '%s' exists already", dbfile)
	}
	if _, err := os.Stat(keyfile); err == nil {
		return log.Errorf("encdb: dbfile '%s' exists already", keyfile)
	}
	// create keyfile
	key, err := generateKeyfile(keyfile, passphrase, iter)
	if err != nil {
		return err
	}
	// create DB
	dbfile += fmt.Sprintf("?_pragma_key=%s&_pragma_cipher_page_size=4096",
		hex.EncodeToString(key))
	db, err := sql.Open("sqlite3", dbfile)
	if err != nil {
		return log.Error(err)
	}
	defer db.Close()
	_, err = db.Exec("PRAGMA auto_vacuum = FULL;")
	if err != nil {
		return log.Error(err)
	}
	return createTables(db, createStmts)
}

// Open tries to open an encrypted database with the given passphrase.
// Thereby, dbname is the prefix of the following two database files (which
// must already exist):
//
//  dbname.db
//  dbname.key
//
// In case of error (for example, the database files do not exist or the
// passphrase is wrong) an error is returned.
func Open(dbname string, passphrase []byte) (*sql.DB, error) {
	dbfile := dbname + dbsuffix
	keyfile := dbname + keysuffix
	// make sure files exists
	if _, err := os.Stat(dbfile); err != nil {
		return nil, log.Error(err)
	}
	if _, err := os.Stat(keyfile); err != nil {
		return nil, log.Error(err)
	}
	// get key from keyfile
	key, err := readKeyfile(keyfile, passphrase)
	if err != nil {
		return nil, err
	}
	// open DB
	dbfile += fmt.Sprintf("?_pragma_key=%s&_pragma_cipher_page_size=4096",
		hex.EncodeToString(key))
	db, err := sql.Open("sqlite3", dbfile)
	if err != nil {
		return nil, log.Error(err)
	}
	// test key
	_, err = db.Exec("SELECT count(*) FROM sqlite_master;")
	if err != nil {
		return nil, log.Error(err)
	}
	// enable foreign key support
	if _, err := db.Exec("PRAGMA foreign_keys = ON;"); err != nil {
		return nil, log.Error(err)
	}
	return db, nil
}

// Rekey tries to rekey an encrypted database with the given newPassphrase and
// newIter many KDF iterations. The correct oldPassphrase must be supplied.
// Thereby, dbname is the prefix of the following two database files (which must
// already exist):
//
//  dbname.db
//  dbname.key
//
// Rekey replaces the dbname.key file and leaves the dbname.db file unmodified,
// allowing for very fast rekey operations. In case of error (for example, the
// database files do not exist or the oldPassphrase is wrong) an error is
// returned.
func Rekey(dbname string, oldPassphrase, newPassphrase []byte, newIter int) error {
	encdb, err := Open(dbname, oldPassphrase)
	if err != nil {
		return err
	}
	defer encdb.Close()
	keyfile := dbname + keysuffix
	return replaceKeyfile(keyfile, oldPassphrase, newPassphrase, newIter)
}

var autoVacuumModes = []string{
	"NONE",
	"FULL",
	"INCREMENTAL",
}

// Status returns the autoVacuum and freelistCount of db.
func Status(db *sql.DB) (autoVacuum string, freelistCount int64, err error) {
	var av int64
	err = db.QueryRow("PRAGMA auto_vacuum;").Scan(&av)
	if err != nil {
		return "", 0, err
	}
	autoVacuum = autoVacuumModes[av]
	err = db.QueryRow("PRAGMA freelist_count;").Scan(&freelistCount)
	if err != nil {
		return "", 0, err
	}
	return
}

// Vacuum executes VACUUM command in db. If autoVacuumMode is not nil and
// different from the current one, the auto_vacuum mode is changed before
// VACUUM is executed.
func Vacuum(db *sql.DB, autoVacuumMode string) error {
	if autoVacuumMode != "" {
		if !util.ContainsString(autoVacuumModes, autoVacuumMode) {
			return log.Errorf("encdb: unknown auto_vacuum mode: %s", autoVacuumMode)
		}
		var av int64
		err := db.QueryRow("PRAGMA auto_vacuum;").Scan(&av)
		if err != nil {
			return err
		}
		currentMode := autoVacuumModes[av]
		if currentMode != autoVacuumMode {
			_, err = db.Exec(fmt.Sprintf("PRAGMA auto_vacuum = %s;", autoVacuumMode))
			if err != nil {
				return err
			}
		}
	}
	_, err := db.Exec("VACUUM;")
	if err != nil {
		return log.Error(err)
	}
	return nil
}

// Incremental executes incremental_vacuum to free up to pages many pages. If
// pages is 0, all pages are freed. If the current auto_vacuum mode is not
// INCREMENTAL, an error is returned.
func Incremental(db *sql.DB, pages int64) error {
	var av int64
	err := db.QueryRow("PRAGMA auto_vacuum;").Scan(&av)
	if err != nil {
		return err
	}
	if autoVacuumModes[av] != "INCREMENTAL" {
		return log.Error("encdb: current auto_vacuum mode is not INCREMENTAL")
	}
	_, err = db.Exec(fmt.Sprintf("PRAGMA incremental_vacuum(%d);", pages))
	if err != nil {
		log.Error(err)
	}
	return nil
}
<|endoftext|>


<commit_before>package qb

import (
	_ "github.com/lib/pq"
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestEngine(t *testing.T) {

	engine, err := NewEngine("postgres", "user=root dbname=pqtest")

	assert.Equal(t, err, nil)
	assert.Equal(t, engine.Driver(), "postgres")
	assert.Equal(t, engine.Ping(), engine.DB().Ping())
	assert.Equal(t, engine.Dsn(), "user=root dbname=pqtest")
}
<commit_msg>add invalid & exec tests of engine, increase code coverage of engine tests<commit_after>package qb

import (
	_ "github.com/lib/pq"
	"github.com/stretchr/testify/assert"
	"testing"
	"fmt"
)

func TestEngine(t *testing.T) {

	engine, err := NewEngine("postgres", "user=root dbname=pqtest")

	assert.Equal(t, err, nil)
	assert.Equal(t, engine.Driver(), "postgres")
	assert.Equal(t, engine.Ping(), engine.DB().Ping())
	assert.Equal(t, engine.Dsn(), "user=root dbname=pqtest")
}

func TestInvalidEngine(t *testing.T) {

	engine, err := NewEngine("invalid", "")
	assert.NotEqual(t, err, nil)
	assert.Equal(t, engine, (*Engine)(nil))
}

func TestEngineExec(t *testing.T) {

	engine, err := NewEngine("postgres", "user=root dbname=pqtest")

	query, err := NewBuilder().Insert("user", "full_name").Values("Aras Can Akin").Build()
	fmt.Println(query)
	assert.Equal(t, err, nil)

	res, err := engine.Exec(query)
	assert.Equal(t, res, nil)
	assert.NotEqual(t, err, nil)
}
<|endoftext|>


<commit_before>package sudoku

import (
	"sort"
	"testing"
)

func TestBasicCellList(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	grid.Load(SOLVED_TEST_GRID)
	row := CellList(grid.Row(2))
	if !row.SameRow() {
		t.Log("The items of a row were not all of the same row.")
		t.Fail()
	}

	if row.SameCol() {
		t.Log("For some reason we thought all the cells in a row were in the same col")
		t.Fail()
	}

	var refs []cellRef

	for i := 0; i < DIM; i++ {
		refs = append(refs, cellRef{2, i})
	}

	if !row.sameAsRefs(refs) {
		t.Error("sameAsRefs didn't match values for row 2")
	}

	col := CellList(grid.Col(2))
	if !col.SameCol() {
		t.Log("The items in the col were not int he same col.")
		t.Fail()
	}

	if col.SameRow() {
		t.Log("For some reason we thought all the cells in a col were in the same row")
		t.Fail()
	}

	block := CellList(grid.Block(2))
	if !block.SameBlock() {
		t.Log("The items in the block were not int he same block.")
		t.Fail()
	}

	if block.SameRow() {
		t.Log("For some reason we thought all the cells in a col were in the same row")
		t.Fail()
	}

	if block.SameCol() {
		t.Log("For some reason we thought all the cells in a block were in the same col")
		t.Fail()
	}

	nums := row.CollectNums(func(cell *Cell) int {
		return cell.Row
	})

	if !nums.Same() {
		t.Log("Collecting rows gave us different numbers/.")
		t.Fail()
	}

	isZeroRow := func(cell *Cell) bool {
		return cell.Row == 0
	}

	cells := grid.Block(0).Filter(isZeroRow)

	if len(cells) != BLOCK_DIM {
		t.Log("We got back the wrong number of cells when filtering")
		t.Fail()
	}

	if !cells.SameRow() {
		t.Log("We got back cells not inthe same row.")
		t.Fail()
	}

	description := cells.Description()

	if description != "(0,0), (0,1), and (0,2)" {
		t.Log("Got wrong description of cellList: ", description)
		t.Fail()
	}

	unsortedList := CellList{grid.Cell(0, 1), grid.Cell(1, 2), grid.Cell(0, 0)}
	unsortedList.Sort()
	if unsortedList[0].Row != 0 || unsortedList[0].Col != 0 ||
		unsortedList[1].Row != 0 || unsortedList[1].Col != 1 ||
		unsortedList[2].Row != 1 || unsortedList[2].Col != 2 {
		t.Error("Cell List didn't get sorted: ", unsortedList)
	}

}

type chainTestConfiguration struct {
	name string
	one  []cellRef
	two  []cellRef
}

type chainTestResult struct {
	name          string
	originalIndex int
	value         float64
}

type chainTestResults []chainTestResult

func (self chainTestResults) Len() int {
	return len(self)
}

func (self chainTestResults) Less(i, j int) bool {
	return self[i].value < self[j].value
}

func (self chainTestResults) Swap(i, j int) {
	self[i], self[j] = self[j], self[i]
}

func TestChainDissimilarity(t *testing.T) {

	//The first bit is where we configure the tests.
	//We should add cases here in the order of similar to dissimilar. The test will then verify
	//they come out in that order.

	tests := []chainTestConfiguration{
		{
			"same row same block",
			[]cellRef{{0, 0}},
			[]cellRef{{0, 1}},
		},
		//These next two should be the same difficulty.
		//TODO: might need to generalize the test to allow me to say
		//that two can be equivalent.
		{
			"same row different blocks",
			[]cellRef{{0, 0}, {0, 1}},
			[]cellRef{{0, 3}, {0, 4}},
		},
		{
			"same col different blocks",
			[]cellRef{{0, 0}, {1, 0}},
			[]cellRef{{3, 0}, {4, 0}},
		},
		{
			"single cell opposite corners",
			[]cellRef{{0, 0}},
			[]cellRef{{8, 8}},
		},
	}

	//Now run the tests

	grid := NewGrid()

	var results chainTestResults

	for i, test := range tests {
		var listOne CellList
		var listTwo CellList
		for _, ref := range test.one {
			listOne = append(listOne, ref.Cell(grid))
		}
		for _, ref := range test.two {
			listTwo = append(listTwo, ref.Cell(grid))
		}
		dissimilarity := listOne.ChainDissimilarity(listTwo)
		if dissimilarity < 0.0 {
			t.Fatal(test.name, "failed with a dissimilarity less than 0.0: ", dissimilarity)
		}
		if dissimilarity > 1.0 {
			t.Fatal(test.name, "failed with a dissimilarity great than 1.0:", dissimilarity)
		}
		result := chainTestResult{test.name, i, dissimilarity}
		results = append(results, result)
	}

	//TODO: sort them and see if their originalIndexes are now now in order.
	sort.Sort(results)

	for i, result := range results {
		if result.originalIndex != i {
			t.Error(result.name, "was in position", i, " but it was supposed to be in position", result.originalIndex, ". Value:", result.value)
		}
	}

}

func TestFilledNums(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	if !grid.LoadFromFile(puzzlePath("nakedpairblock1.sdk")) {
		t.Fatal("Couldn't load file")
	}

	filledNums := grid.Row(0).FilledNums()

	if !filledNums.SameContentAs(IntSlice{3, 7, 8, 9}) {
		t.Error("Filled nums had wrong nums", filledNums)
	}

}

func TestIntList(t *testing.T) {
	numArr := [...]int{1, 1, 1}
	if !IntSlice(numArr[:]).Same() {
		t.Log("We didn't think that a num list with all of the same ints was the same.")
		t.Fail()
	}
	differentNumArr := [...]int{1, 2, 1}
	if IntSlice(differentNumArr[:]).Same() {
		t.Log("We thought a list of different ints were the same")
		t.Fail()
	}
	description := IntSlice(numArr[:]).Description()
	if description != "1, 1, and 1" {
		t.Log("Didn't get right description: ", description)
		t.Fail()
	}

	unsortedList := IntSlice{3, 2, 1}
	unsortedList.Sort()
	if !unsortedList.SameAs(IntSlice{1, 2, 3}) {
		t.Error("IntSlice.Sort did not sort the list.")
	}

	oneList := IntSlice{1}

	description = oneList.Description()

	if description != "1" {
		t.Log("Didn't get the right description for a short intlist: ", description)
		t.Fail()
	}

	twoList := IntSlice{1, 1}

	description = twoList.Description()

	if description != "1 and 1" {
		t.Log("Did'get the the right description for a two-item intList: ", description)
		t.Fail()
	}
}

func TestInverseSubset(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	cells := grid.Row(0)

	indexes := IntSlice([]int{4, 6, 2})

	subset := cells.InverseSubset(indexes)

	if len(subset) != DIM-3 {
		t.Error("Inverse subset gave wrong number of results")
	}

	for _, cell := range subset {
		if cell.Col == 2 || cell.Col == 4 || cell.Col == 6 {
			t.Error("Inverse subset included cells it shouldn't have.")
		}
	}

}

func TestIntSliceIntersection(t *testing.T) {
	one := IntSlice([]int{1, 3, 2, 5})
	two := IntSlice([]int{2, 7, 6, 5})

	result := one.Intersection(two)

	if len(result) != 2 {
		t.Error("Intersection had wrong number of items")
	}

	sort.Ints(result)

	if result[0] != 2 || result[1] != 5 {
		t.Error("Intersection result was wrong.")
	}
}

func TestIntSliceDifference(t *testing.T) {
	one := IntSlice([]int{1, 2, 3, 4, 5, 6})
	two := IntSlice([]int{3, 4, 7})

	result := one.Difference(two)

	if !result.SameContentAs(IntSlice([]int{1, 2, 5, 6})) {
		t.Error("Int slice difference gave wrong result: ", result)
	}
}

func TestSameContentAs(t *testing.T) {
	one := IntSlice([]int{2, 3, 1})
	two := IntSlice([]int{2, 1, 3})

	if !one.SameContentAs(two) {
		t.Log("Didn't think two equivalent slices were the same.")
		t.Fail()
	}

	if !one.SameAs([]int{2, 3, 1}) {
		t.Log("We mutated one")
		t.Fail()
	}

	if !two.SameAs([]int{2, 1, 3}) {
		t.Log("We mutated two")
		t.Fail()
	}

	onePair := IntSlice([]int{3, 2})
	twoPair := IntSlice([]int{2, 3})

	if !onePair.SameContentAs(twoPair) {
		t.Log("Didn't think two equivalent pairs were the same.")
		t.Fail()
	}

}
<commit_msg>Another test of chainDissimilarity<commit_after>package sudoku

import (
	"sort"
	"testing"
)

func TestBasicCellList(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	grid.Load(SOLVED_TEST_GRID)
	row := CellList(grid.Row(2))
	if !row.SameRow() {
		t.Log("The items of a row were not all of the same row.")
		t.Fail()
	}

	if row.SameCol() {
		t.Log("For some reason we thought all the cells in a row were in the same col")
		t.Fail()
	}

	var refs []cellRef

	for i := 0; i < DIM; i++ {
		refs = append(refs, cellRef{2, i})
	}

	if !row.sameAsRefs(refs) {
		t.Error("sameAsRefs didn't match values for row 2")
	}

	col := CellList(grid.Col(2))
	if !col.SameCol() {
		t.Log("The items in the col were not int he same col.")
		t.Fail()
	}

	if col.SameRow() {
		t.Log("For some reason we thought all the cells in a col were in the same row")
		t.Fail()
	}

	block := CellList(grid.Block(2))
	if !block.SameBlock() {
		t.Log("The items in the block were not int he same block.")
		t.Fail()
	}

	if block.SameRow() {
		t.Log("For some reason we thought all the cells in a col were in the same row")
		t.Fail()
	}

	if block.SameCol() {
		t.Log("For some reason we thought all the cells in a block were in the same col")
		t.Fail()
	}

	nums := row.CollectNums(func(cell *Cell) int {
		return cell.Row
	})

	if !nums.Same() {
		t.Log("Collecting rows gave us different numbers/.")
		t.Fail()
	}

	isZeroRow := func(cell *Cell) bool {
		return cell.Row == 0
	}

	cells := grid.Block(0).Filter(isZeroRow)

	if len(cells) != BLOCK_DIM {
		t.Log("We got back the wrong number of cells when filtering")
		t.Fail()
	}

	if !cells.SameRow() {
		t.Log("We got back cells not inthe same row.")
		t.Fail()
	}

	description := cells.Description()

	if description != "(0,0), (0,1), and (0,2)" {
		t.Log("Got wrong description of cellList: ", description)
		t.Fail()
	}

	unsortedList := CellList{grid.Cell(0, 1), grid.Cell(1, 2), grid.Cell(0, 0)}
	unsortedList.Sort()
	if unsortedList[0].Row != 0 || unsortedList[0].Col != 0 ||
		unsortedList[1].Row != 0 || unsortedList[1].Col != 1 ||
		unsortedList[2].Row != 1 || unsortedList[2].Col != 2 {
		t.Error("Cell List didn't get sorted: ", unsortedList)
	}

}

type chainTestConfiguration struct {
	name string
	one  []cellRef
	two  []cellRef
}

type chainTestResult struct {
	name          string
	originalIndex int
	value         float64
}

type chainTestResults []chainTestResult

func (self chainTestResults) Len() int {
	return len(self)
}

func (self chainTestResults) Less(i, j int) bool {
	return self[i].value < self[j].value
}

func (self chainTestResults) Swap(i, j int) {
	self[i], self[j] = self[j], self[i]
}

func TestChainDissimilarity(t *testing.T) {

	//The first bit is where we configure the tests.
	//We should add cases here in the order of similar to dissimilar. The test will then verify
	//they come out in that order.

	tests := []chainTestConfiguration{
		{
			"same row same block",
			[]cellRef{{0, 0}},
			[]cellRef{{0, 1}},
		},
		//These next two should be the same difficulty.
		//TODO: might need to generalize the test to allow me to say
		//that two can be equivalent.
		{
			"same block 2 in same row 2 in same col 2 total",
			[]cellRef{{0, 0}},
			[]cellRef{{0, 1}, {1, 0}},
		},
		{
			"same row different blocks",
			[]cellRef{{0, 0}, {0, 1}},
			[]cellRef{{0, 3}, {0, 4}},
		},
		{
			"same col different blocks",
			[]cellRef{{0, 0}, {1, 0}},
			[]cellRef{{3, 0}, {4, 0}},
		},
		{
			"single cell opposite corners",
			[]cellRef{{0, 0}},
			[]cellRef{{8, 8}},
		},
	}

	//Now run the tests

	grid := NewGrid()

	var results chainTestResults

	for i, test := range tests {
		var listOne CellList
		var listTwo CellList
		for _, ref := range test.one {
			listOne = append(listOne, ref.Cell(grid))
		}
		for _, ref := range test.two {
			listTwo = append(listTwo, ref.Cell(grid))
		}
		dissimilarity := listOne.ChainDissimilarity(listTwo)
		if dissimilarity < 0.0 {
			t.Fatal(test.name, "failed with a dissimilarity less than 0.0: ", dissimilarity)
		}
		if dissimilarity > 1.0 {
			t.Fatal(test.name, "failed with a dissimilarity great than 1.0:", dissimilarity)
		}
		result := chainTestResult{test.name, i, dissimilarity}
		results = append(results, result)
	}

	//sort them and see if their originalIndexes are now now in order.
	sort.Sort(results)

	for i, result := range results {
		if result.originalIndex != i {
			t.Error(result.name, "was in position", i, " but it was supposed to be in position", result.originalIndex, ". Value:", result.value)
		}
	}

}

func TestFilledNums(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	if !grid.LoadFromFile(puzzlePath("nakedpairblock1.sdk")) {
		t.Fatal("Couldn't load file")
	}

	filledNums := grid.Row(0).FilledNums()

	if !filledNums.SameContentAs(IntSlice{3, 7, 8, 9}) {
		t.Error("Filled nums had wrong nums", filledNums)
	}

}

func TestIntList(t *testing.T) {
	numArr := [...]int{1, 1, 1}
	if !IntSlice(numArr[:]).Same() {
		t.Log("We didn't think that a num list with all of the same ints was the same.")
		t.Fail()
	}
	differentNumArr := [...]int{1, 2, 1}
	if IntSlice(differentNumArr[:]).Same() {
		t.Log("We thought a list of different ints were the same")
		t.Fail()
	}
	description := IntSlice(numArr[:]).Description()
	if description != "1, 1, and 1" {
		t.Log("Didn't get right description: ", description)
		t.Fail()
	}

	unsortedList := IntSlice{3, 2, 1}
	unsortedList.Sort()
	if !unsortedList.SameAs(IntSlice{1, 2, 3}) {
		t.Error("IntSlice.Sort did not sort the list.")
	}

	oneList := IntSlice{1}

	description = oneList.Description()

	if description != "1" {
		t.Log("Didn't get the right description for a short intlist: ", description)
		t.Fail()
	}

	twoList := IntSlice{1, 1}

	description = twoList.Description()

	if description != "1 and 1" {
		t.Log("Did'get the the right description for a two-item intList: ", description)
		t.Fail()
	}
}

func TestInverseSubset(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	cells := grid.Row(0)

	indexes := IntSlice([]int{4, 6, 2})

	subset := cells.InverseSubset(indexes)

	if len(subset) != DIM-3 {
		t.Error("Inverse subset gave wrong number of results")
	}

	for _, cell := range subset {
		if cell.Col == 2 || cell.Col == 4 || cell.Col == 6 {
			t.Error("Inverse subset included cells it shouldn't have.")
		}
	}

}

func TestIntSliceIntersection(t *testing.T) {
	one := IntSlice([]int{1, 3, 2, 5})
	two := IntSlice([]int{2, 7, 6, 5})

	result := one.Intersection(two)

	if len(result) != 2 {
		t.Error("Intersection had wrong number of items")
	}

	sort.Ints(result)

	if result[0] != 2 || result[1] != 5 {
		t.Error("Intersection result was wrong.")
	}
}

func TestIntSliceDifference(t *testing.T) {
	one := IntSlice([]int{1, 2, 3, 4, 5, 6})
	two := IntSlice([]int{3, 4, 7})

	result := one.Difference(two)

	if !result.SameContentAs(IntSlice([]int{1, 2, 5, 6})) {
		t.Error("Int slice difference gave wrong result: ", result)
	}
}

func TestSameContentAs(t *testing.T) {
	one := IntSlice([]int{2, 3, 1})
	two := IntSlice([]int{2, 1, 3})

	if !one.SameContentAs(two) {
		t.Log("Didn't think two equivalent slices were the same.")
		t.Fail()
	}

	if !one.SameAs([]int{2, 3, 1}) {
		t.Log("We mutated one")
		t.Fail()
	}

	if !two.SameAs([]int{2, 1, 3}) {
		t.Log("We mutated two")
		t.Fail()
	}

	onePair := IntSlice([]int{3, 2})
	twoPair := IntSlice([]int{2, 3})

	if !onePair.SameContentAs(twoPair) {
		t.Log("Didn't think two equivalent pairs were the same.")
		t.Fail()
	}

}
<|endoftext|>


<commit_before>package goarken

import (
	"errors"
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"github.com/golang/glog"
	"time"
)

const (
	TIME_FORMAT = "2006-01-02 15:04:05"
)

// A Watcher loads and watch the etcd hierarchy for Domains and Services.
type Watcher struct {
	Client        *etcd.Client
	DomainPrefix  string
	ServicePrefix string
	Domains       map[string]*Domain
	Services      map[string]*ServiceCluster
	broadcaster   *Broadcaster
}

//Init Domains and Services.
func (w *Watcher) Init() {
	w.broadcaster = NewBroadcaster()
	setServicePrefix(w.ServicePrefix)
	setDomainPrefix(w.DomainPrefix)
	if w.Domains != nil {
		w.loadPrefix(w.DomainPrefix, w.registerDomain)
		go w.doWatch(w.DomainPrefix, w.registerDomain)
	}
	if w.Services != nil {
		w.loadPrefix(w.ServicePrefix, w.registerService)
		go w.doWatch(w.ServicePrefix, w.registerService)
	}

}

func (w *Watcher) Listen() chan interface{} {
	return w.broadcaster.Listen()
}

// Loads and watch an etcd directory to register objects like Domains, Services
// etc... The register function is passed the etcd Node that has been loaded.
func (w *Watcher) doWatch(etcdDir string, registerFunc func(*etcd.Node, string)) {
	stop := make(chan struct{})

	for {
		glog.Infof("Start watching %s", etcdDir)

		updateChannel := make(chan *etcd.Response, 10)
		go w.watch(updateChannel, stop, etcdDir, registerFunc)

		_, err := w.Client.Watch(etcdDir, (uint64)(0), true, updateChannel, nil)

		//If we are here, this means etcd watch ended in an error
		stop <- struct{}{}
		w.Client.CloseCURL()
		glog.Warningf("Error when watching %s : %v", etcdDir, err)
		glog.Warningf("Waiting 1 second and relaunch watch")
		time.Sleep(time.Second)

	}

}

func (w *Watcher) loadPrefix(etcDir string, registerFunc func(*etcd.Node, string)) {
	response, err := w.Client.Get(etcDir, true, true)
	if err == nil {
		for _, serviceNode := range response.Node.Nodes {
			registerFunc(serviceNode, response.Action)

		}
	}
}

func (w *Watcher) watch(updateChannel chan *etcd.Response, stop chan struct{}, key string, registerFunc func(*etcd.Node, string)) {
	for {
		select {
		case <-stop:
			glog.Warningf("Gracefully closing the etcd watch for %s", key)
			return
		case response := <-updateChannel:
			if response != nil {
				registerFunc(response.Node, response.Action)
			}
		default:
			// Don't slam the etcd server
			time.Sleep(time.Second)
		}
	}
}

func (w *Watcher) RemoveDomain(key string) {
	delete(w.Domains, key)

}

func (w *Watcher) RemoveEnv(serviceName string) {
	delete(w.Services, serviceName)
}

func GetServiceClusterFromPath(serviceClusterPath string, client *etcd.Client) (*ServiceCluster, error) {
	// Get service's root node instead of changed node.
	response, err := client.Get(serviceClusterPath, true, true)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Unable to get information for service %s from etcd", serviceClusterPath))
	}

	return GetServiceClusterFromNode(response.Node), nil
}

func GetServiceClusterFromNode(clusterNode *etcd.Node) *ServiceCluster {

	sc := NewServiceCluster("")
	for _, indexNode := range clusterNode.Nodes {
		service, err := NewService(indexNode)
		sc.Name = service.Name
		if err == nil {
			sc.Add(service)
		}
	}
	return sc

}

func (w *Watcher) registerDomain(node *etcd.Node, action string) {

	domainName := getDomainForNode(node)

	if action == "delete" || action == "expire" {
		w.RemoveDomain(domainName)
		return
	}

	domainKey := w.DomainPrefix + "/" + domainName
	response, err := w.Client.Get(domainKey, true, false)

	if err == nil {
		domain := NewDomain(response.Node)

		actualDomain := w.Domains[domainName]

		if domain.Typ != "" && domain.Value != "" && !domain.Equals(actualDomain) {
			w.Domains[domainName] = domain
			glog.Infof("Registered domain %s with (%s) %s", domainName, domain.Typ, domain.Value)

			//Broadcast the updated domain
			w.broadcaster.Write(domain)

		}

	}

}

func (w *Watcher) registerService(node *etcd.Node, action string) {

	serviceName := getEnvForNode(node)

	if action == "delete" && node.Key == w.ServicePrefix+"/"+serviceName {
		w.RemoveEnv(serviceName)
		return
	}

	// Get service's root node instead of changed node.
	response, err := w.Client.Get(w.ServicePrefix+"/"+serviceName, true, true)

	if err == nil {

		sc := GetServiceClusterFromNode(response.Node)

		if w.Services[sc.Name] == nil {
			w.Services[sc.Name] = sc
			w.broadcaster.Write(w.Services[serviceName])

		} else {
			for _, service := range sc.GetInstances() {
				actualEnv := w.Services[serviceName].Get(service.Index)
				if !actualEnv.Equals(service) {
					w.Services[serviceName].Add(service)
					if service.Location.Host != "" && service.Location.Port != 0 {
						glog.Infof("Registering service %s with location : http://%s:%d/", serviceName, service.Location.Host, service.Location.Port)
					} else {
						glog.Infof("Registering service %s without location", serviceName)
					}
					//Broadcast the updated object
					w.broadcaster.Write(w.Services[serviceName])
				}
			}
		}

	} else {
		glog.Errorf("Unable to get information for service %s from etcd", serviceName)
	}
}
<commit_msg>NXIO-345 prevent NPE<commit_after>package goarken

import (
	"errors"
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"github.com/golang/glog"
	"time"
)

const (
	TIME_FORMAT = "2006-01-02 15:04:05"
)

// A Watcher loads and watch the etcd hierarchy for Domains and Services.
type Watcher struct {
	Client        *etcd.Client
	DomainPrefix  string
	ServicePrefix string
	Domains       map[string]*Domain
	Services      map[string]*ServiceCluster
	broadcaster   *Broadcaster
}

//Init Domains and Services.
func (w *Watcher) Init() {
	w.broadcaster = NewBroadcaster()
	setServicePrefix(w.ServicePrefix)
	setDomainPrefix(w.DomainPrefix)
	if w.Domains != nil {
		w.loadPrefix(w.DomainPrefix, w.registerDomain)
		go w.doWatch(w.DomainPrefix, w.registerDomain)
	}
	if w.Services != nil {
		w.loadPrefix(w.ServicePrefix, w.registerService)
		go w.doWatch(w.ServicePrefix, w.registerService)
	}

}

func (w *Watcher) Listen() chan interface{} {
	return w.broadcaster.Listen()
}

// Loads and watch an etcd directory to register objects like Domains, Services
// etc... The register function is passed the etcd Node that has been loaded.
func (w *Watcher) doWatch(etcdDir string, registerFunc func(*etcd.Node, string)) {
	stop := make(chan struct{})

	for {
		glog.Infof("Start watching %s", etcdDir)

		updateChannel := make(chan *etcd.Response, 10)
		go w.watch(updateChannel, stop, etcdDir, registerFunc)

		_, err := w.Client.Watch(etcdDir, (uint64)(0), true, updateChannel, nil)

		//If we are here, this means etcd watch ended in an error
		stop <- struct{}{}
		w.Client.CloseCURL()
		glog.Warningf("Error when watching %s : %v", etcdDir, err)
		glog.Warningf("Waiting 1 second and relaunch watch")
		time.Sleep(time.Second)

	}

}

func (w *Watcher) loadPrefix(etcDir string, registerFunc func(*etcd.Node, string)) {
	response, err := w.Client.Get(etcDir, true, true)
	if err == nil {
		for _, serviceNode := range response.Node.Nodes {
			registerFunc(serviceNode, response.Action)

		}
	}
}

func (w *Watcher) watch(updateChannel chan *etcd.Response, stop chan struct{}, key string, registerFunc func(*etcd.Node, string)) {
	for {
		select {
		case <-stop:
			glog.Warningf("Gracefully closing the etcd watch for %s", key)
			return
		case response := <-updateChannel:
			if response != nil {
				registerFunc(response.Node, response.Action)
			}
		default:
			// Don't slam the etcd server
			time.Sleep(time.Second)
		}
	}
}

func (w *Watcher) RemoveDomain(key string) {
	delete(w.Domains, key)

}

func (w *Watcher) RemoveEnv(serviceName string) {
	delete(w.Services, serviceName)
}

func GetServiceClusterFromPath(serviceClusterPath string, client *etcd.Client) (*ServiceCluster, error) {
	// Get service's root node instead of changed node.
	response, err := client.Get(serviceClusterPath, true, true)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Unable to get information for service %s from etcd", serviceClusterPath))
	}

	return GetServiceClusterFromNode(response.Node), nil
}

func GetServiceClusterFromNode(clusterNode *etcd.Node) *ServiceCluster {

	sc := NewServiceCluster("")
	for _, indexNode := range clusterNode.Nodes {
		service, err := NewService(indexNode)
		if err == nil {
			sc.Add(service)
			sc.Name = service.Name
		}
	}
	return sc

}

func (w *Watcher) registerDomain(node *etcd.Node, action string) {

	domainName := getDomainForNode(node)

	if action == "delete" || action == "expire" {
		w.RemoveDomain(domainName)
		return
	}

	domainKey := w.DomainPrefix + "/" + domainName
	response, err := w.Client.Get(domainKey, true, false)

	if err == nil {
		domain := NewDomain(response.Node)

		actualDomain := w.Domains[domainName]

		if domain.Typ != "" && domain.Value != "" && !domain.Equals(actualDomain) {
			w.Domains[domainName] = domain
			glog.Infof("Registered domain %s with (%s) %s", domainName, domain.Typ, domain.Value)

			//Broadcast the updated domain
			w.broadcaster.Write(domain)

		}

	}

}

func (w *Watcher) registerService(node *etcd.Node, action string) {

	serviceName := getEnvForNode(node)

	if action == "delete" && node.Key == w.ServicePrefix+"/"+serviceName {
		w.RemoveEnv(serviceName)
		return
	}

	// Get service's root node instead of changed node.
	response, err := w.Client.Get(w.ServicePrefix+"/"+serviceName, true, true)

	if err == nil {

		sc := GetServiceClusterFromNode(response.Node)

		if w.Services[sc.Name] == nil {
			w.Services[sc.Name] = sc
			w.broadcaster.Write(w.Services[serviceName])

		} else {
			for _, service := range sc.GetInstances() {
				actualEnv := w.Services[serviceName].Get(service.Index)
				if !actualEnv.Equals(service) {
					w.Services[serviceName].Add(service)
					if service.Location.Host != "" && service.Location.Port != 0 {
						glog.Infof("Registering service %s with location : http://%s:%d/", serviceName, service.Location.Host, service.Location.Port)
					} else {
						glog.Infof("Registering service %s without location", serviceName)
					}
					//Broadcast the updated object
					w.broadcaster.Write(w.Services[serviceName])
				}
			}
		}

	} else {
		glog.Errorf("Unable to get information for service %s from etcd", serviceName)
	}
}
<|endoftext|>