<commit_before>package deploy

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/git"
	hub "github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/github"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/google/go-github/github"
)

const (
	Name  = "deploy"
	Usage = "A command for creating GitHub deployments"
)

const (
	DefaultRef     = "master"
	DefaultTimeout = 20 * time.Second
)

var errTimeout = errors.New("Timed out waiting for build to start. Did you add a webhook to handle deployment events?")

func init() {
	cli.AppHelpTemplate = `USAGE:
   # Deploy the master branch of remind101/acme-inc to staging
   {{.Name}} --env=staging --ref=master remind101/acme-inc

   # Deploy HEAD of the current branch to staging
   {{.Name}} --env=staging remind101/acme-inc

   # Deploy the current GitHub repo to staging
   {{.Name}} --env=staging
{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`
}

var flags = []cli.Flag{
	cli.StringFlag{
		Name:  "ref, branch, commit, tag",
		Value: "",
		Usage: "The git ref to deploy. Can be a git commit, branch or tag.",
	},
	cli.StringFlag{
		Name:  "env, e",
		Value: "",
		Usage: "The environment to deploy to.",
	},
	cli.BoolFlag{
		Name:  "force, f",
		Usage: "Ignore commit status checks.",
	},
	cli.BoolFlag{
		Name:  "detached, d",
		Usage: "Don't wait for the deployment to complete.",
	},
	cli.BoolFlag{
		Name:  "quiet, q",
		Usage: "Silence any output to STDOUT.",
	},
}

// NewApp returns a new cli.App for the deploy command.
func NewApp() *cli.App {
	app := cli.NewApp()
	app.Version = Version
	app.Name = Name
	app.Usage = Usage
	app.Flags = flags
	app.Action = func(c *cli.Context) {
		if err := RunDeploy(c); err != nil {
			msg := err.Error()
			if err, ok := err.(*github.ErrorResponse); ok {
				if strings.HasPrefix(err.Message, "Conflict: Commit status checks failed for") {
					msg = "Commit status checks failed. You can bypass commit status checks with the --force flag."
				} else if strings.HasPrefix(err.Message, "No ref found for") {
					msg = fmt.Sprintf("%s. Did you push it to GitHub?", err.Message)
				} else {
					msg = err.Message
				}
			}

			fmt.Println(msg)
			os.Exit(-1)
		}
	}

	return app
}

// RunDeploy performs a deploy.
func RunDeploy(c *cli.Context) error {
	var w io.Writer
	if c.Bool("quiet") {
		w = ioutil.Discard
	} else {
		w = c.App.Writer
	}

	h, err := hub.CurrentConfig().PromptForHost("github.com")
	if err != nil {
		return err
	}

	client, err := newGitHubClient(h)
	if err != nil {
		return err
	}

	nwo, err := Repo(c.Args())
	if err != nil {
		return err
	}

	owner, repo, err := SplitRepo(nwo, os.Getenv("GITHUB_ORGANIZATION"))
	if err != nil {
		return fmt.Errorf("Invalid GitHub repo: %s", nwo)
	}

	r, err := newDeploymentRequest(c)
	if err != nil {
		return err
	}

	fmt.Fprintf(w, "Deploying %s/%s@%s to %s...\n", owner, repo, *r.Ref, *r.Environment)

	if c.Bool("detached") {
		return nil
	}

	d, _, err := client.Repositories.CreateDeployment(owner, repo, r)
	if err != nil {
		return err
	}

	started := make(chan *github.DeploymentStatus)
	completed := make(chan *github.DeploymentStatus)

	go func() {
		started <- waitState(pendingStates, owner, repo, *d.ID, client)
	}()

	go func() {
		completed <- waitState(completedStates, owner, repo, *d.ID, client)
	}()

	select {
	case <-time.After(DefaultTimeout):
		return errTimeout
	case status := <-started:
		var url string
		if status.TargetURL != nil {
			url = *status.TargetURL
		}
		fmt.Fprintf(w, "%s\n", url)
	}

	status := <-completed

	if isFailed(*status.State) {
		return errors.New("Failed to deploy")
	}

	return nil
}

func newDeploymentRequest(c *cli.Context) (*github.DeploymentRequest, error) {
	ref := Ref(c.String("ref"), git.Head)

	env := c.String("env")
	if env == "" {
		return nil, fmt.Errorf("--env flag is required")
	}

	var contexts *[]string
	if c.Bool("force") {
		s := []string{}
		contexts = &s
	}

	return &github.DeploymentRequest{
		Ref:              github.String(ref),
		Task:             github.String("deploy"),
		AutoMerge:        github.Bool(false),
		Environment:      github.String(env),
		RequiredContexts: contexts,
		Payload: map[string]interface{}{
			"force": c.Bool("force"),
		},
		// TODO Description:
	}, nil
}

var (
	pendingStates   = []string{"pending"}
	completedStates = []string{"success", "error", "failure"}
)

func isFailed(state string) bool {
	return state == "error" || state == "failure"
}

// waitState waits for a deployment status that matches the given states, then
// sends on the returned channel.
func waitState(states []string, owner, repo string, deploymentID int, c *github.Client) *github.DeploymentStatus {
	for {
		<-time.After(1 * time.Second)

		statuses, _, err := c.Repositories.ListDeploymentStatuses(owner, repo, deploymentID, nil)
		if err != nil {
			continue
		}

		status := firstStatus(states, statuses)
		if status != nil {
			return status
		}
	}
}

// firstStatus takes a slice of github.DeploymentStatus and returns the
// first status that matches the provided slice of states.
func firstStatus(states []string, statuses []github.DeploymentStatus) *github.DeploymentStatus {
	for _, ds := range statuses {
		for _, s := range states {
			if ds.State != nil && *ds.State == s {
				return &ds
			}
		}
	}

	return nil
}

// refRegex is a regular expression that matches a full git HEAD ref.
var refRegex = regexp.MustCompile(`^refs/heads/(.*)$`)

// Ref attempts to return the proper git ref to deploy. If a ref is provided,
// that will be returned. If not, it will fallback to calling headFunc. If an
// error is returned (not in a git repo), then it will fallback to DefaultRef.
func Ref(ref string, headFunc func() (string, error)) string {
	if ref != "" {
		return ref
	}

	ref, err := headFunc()
	if err != nil {
		// An error means that we're either not in a GitRepo or we're
		// not on a branch. In this case, we just fallback to the
		// DefaultRef.
		return DefaultRef
	}

	// Convert `refs/heads/test-deploy` => `test-deploy`
	return refRegex.ReplaceAllString(ref, "$1")
}

// Repo will determine the correct GitHub repo to deploy to, based on a set of
// arguments.
func Repo(arguments []string) (string, error) {
	if len(arguments) != 0 {
		return arguments[0], nil
	}

	remotes, err := hub.Remotes()
	if err != nil {
		return "", err
	}

	repo := GitHubRepo(remotes)
	if repo == "" {
		return repo, errors.New("no GitHub repo found in .git/config")
	}

	return repo, nil
}

// A regular expression that can convert a URL.Path into a GitHub repo name.
var remoteRegex = regexp.MustCompile(`^/(.*)\.git$`)

// GitHubRepo, given a list of git remotes, will determine what the GitHub repo
// is.
func GitHubRepo(remotes []hub.Remote) string {
	// We only want to look at the `origin` remote.
	remote := findRemote("origin", remotes)
	if remote == nil {
		return ""
	}

	// Remotes that are not pointed at a GitHub repo are not valid.
	if remote.URL.Host != "github.com" {
		return ""
	}

	// Convert `/remind101/acme-inc.git` => `remind101/acme-inc`.
	return remoteRegex.ReplaceAllString(remote.URL.Path, "$1")
}

func findRemote(name string, remotes []hub.Remote) *hub.Remote {
	for _, r := range remotes {
		if r.Name == name {
			return &r
		}
	}

	return nil
}

var errInvalidRepo = errors.New("invalid repo")

// SplitRepo splits a repo string in the form remind101/acme-inc into it's owner
// and repo components.
func SplitRepo(nwo, defaultOrg string) (owner string, repo string, err error) {
	parts := strings.Split(nwo, "/")

	// If we were only given a repo name, and a default organization is set,
	// we'll use the defaultOrg as the owner.
	if len(parts) == 1 && defaultOrg != "" && parts[0] != "" {
		owner = defaultOrg
		repo = parts[0]
		return
	}

	if len(parts) != 2 {
		err = errInvalidRepo
		return
	}

	owner = parts[0]
	repo = parts[1]

	return
}
<commit_msg>Ask for confirmation before deploying to prod. Fix #13.<commit_after>package deploy

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/git"
	hub "github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/github"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/google/go-github/github"
)

const (
	Name  = "deploy"
	Usage = "A command for creating GitHub deployments"
)

const (
	DefaultRef     = "master"
	DefaultTimeout = 20 * time.Second
)

var errTimeout = errors.New("Timed out waiting for build to start. Did you add a webhook to handle deployment events?")

func init() {
	cli.AppHelpTemplate = `USAGE:
   # Deploy the master branch of remind101/acme-inc to staging
   {{.Name}} --env=staging --ref=master remind101/acme-inc

   # Deploy HEAD of the current branch to staging
   {{.Name}} --env=staging remind101/acme-inc

   # Deploy the current GitHub repo to staging
   {{.Name}} --env=staging
{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`
}

var flags = []cli.Flag{
	cli.StringFlag{
		Name:  "ref, branch, commit, tag",
		Value: "",
		Usage: "The git ref to deploy. Can be a git commit, branch or tag.",
	},
	cli.StringFlag{
		Name:  "env, e",
		Value: "",
		Usage: "The environment to deploy to.",
	},
	cli.BoolFlag{
		Name:  "force, f",
		Usage: "Ignore commit status checks.",
	},
	cli.BoolFlag{
		Name:  "detached, d",
		Usage: "Don't wait for the deployment to complete.",
	},
	cli.BoolFlag{
		Name:  "quiet, q",
		Usage: "Silence any output to STDOUT.",
	},
}

// NewApp returns a new cli.App for the deploy command.
func NewApp() *cli.App {
	app := cli.NewApp()
	app.Version = Version
	app.Name = Name
	app.Usage = Usage
	app.Flags = flags
	app.Action = func(c *cli.Context) {
		if err := RunDeploy(c); err != nil {
			msg := err.Error()
			if err, ok := err.(*github.ErrorResponse); ok {
				if strings.HasPrefix(err.Message, "Conflict: Commit status checks failed for") {
					msg = "Commit status checks failed. You can bypass commit status checks with the --force flag."
				} else if strings.HasPrefix(err.Message, "No ref found for") {
					msg = fmt.Sprintf("%s. Did you push it to GitHub?", err.Message)
				} else {
					msg = err.Message
				}
			}

			fmt.Println(msg)
			os.Exit(-1)
		}
	}

	return app
}

// RunDeploy performs a deploy.
func RunDeploy(c *cli.Context) error {
	var w io.Writer
	if c.Bool("quiet") {
		w = ioutil.Discard
	} else {
		w = c.App.Writer
	}

	h, err := hub.CurrentConfig().PromptForHost("github.com")
	if err != nil {
		return err
	}

	client, err := newGitHubClient(h)
	if err != nil {
		return err
	}

	nwo, err := Repo(c.Args())
	if err != nil {
		return err
	}

	owner, repo, err := SplitRepo(nwo, os.Getenv("GITHUB_ORGANIZATION"))
	if err != nil {
		return fmt.Errorf("Invalid GitHub repo: %s", nwo)
	}

	r, err := newDeploymentRequest(c)
	if err != nil {
		return err
	}

	fmt.Fprintf(w, "Deploying %s/%s@%s to %s...\n", owner, repo, *r.Ref, *r.Environment)

	if c.Bool("detached") {
		return nil
	}

	d, _, err := client.Repositories.CreateDeployment(owner, repo, r)
	if err != nil {
		return err
	}

	started := make(chan *github.DeploymentStatus)
	completed := make(chan *github.DeploymentStatus)

	go func() {
		started <- waitState(pendingStates, owner, repo, *d.ID, client)
	}()

	go func() {
		completed <- waitState(completedStates, owner, repo, *d.ID, client)
	}()

	select {
	case <-time.After(DefaultTimeout):
		return errTimeout
	case status := <-started:
		var url string
		if status.TargetURL != nil {
			url = *status.TargetURL
		}
		fmt.Fprintf(w, "%s\n", url)
	}

	status := <-completed

	if isFailed(*status.State) {
		return errors.New("Failed to deploy")
	}

	return nil
}

func newDeploymentRequest(c *cli.Context) (*github.DeploymentRequest, error) {
	ref := Ref(c.String("ref"), git.Head)

	env := c.String("env")
	if env == "" {
		return nil, fmt.Errorf("--env flag is required")
	}

	if env == "production" {
		reader := bufio.NewReader(os.Stdin)
		fmt.Printf("Are you sure you want to push to production? (y/N)\n")
		a, _ := reader.ReadString('\n')
		if strings.ToUpper(a) == "N\n" {
			return nil, fmt.Errorf("Deployment aborted.")
		}
	}

	var contexts *[]string
	if c.Bool("force") {
		s := []string{}
		contexts = &s
	}

	return &github.DeploymentRequest{
		Ref:              github.String(ref),
		Task:             github.String("deploy"),
		AutoMerge:        github.Bool(false),
		Environment:      github.String(env),
		RequiredContexts: contexts,
		Payload: map[string]interface{}{
			"force": c.Bool("force"),
		},
		// TODO Description:
	}, nil
}

var (
	pendingStates   = []string{"pending"}
	completedStates = []string{"success", "error", "failure"}
)

func isFailed(state string) bool {
	return state == "error" || state == "failure"
}

// waitState waits for a deployment status that matches the given states, then
// sends on the returned channel.
func waitState(states []string, owner, repo string, deploymentID int, c *github.Client) *github.DeploymentStatus {
	for {
		<-time.After(1 * time.Second)

		statuses, _, err := c.Repositories.ListDeploymentStatuses(owner, repo, deploymentID, nil)
		if err != nil {
			continue
		}

		status := firstStatus(states, statuses)
		if status != nil {
			return status
		}
	}
}

// firstStatus takes a slice of github.DeploymentStatus and returns the
// first status that matches the provided slice of states.
func firstStatus(states []string, statuses []github.DeploymentStatus) *github.DeploymentStatus {
	for _, ds := range statuses {
		for _, s := range states {
			if ds.State != nil && *ds.State == s {
				return &ds
			}
		}
	}

	return nil
}

// refRegex is a regular expression that matches a full git HEAD ref.
var refRegex = regexp.MustCompile(`^refs/heads/(.*)$`)

// Ref attempts to return the proper git ref to deploy. If a ref is provided,
// that will be returned. If not, it will fallback to calling headFunc. If an
// error is returned (not in a git repo), then it will fallback to DefaultRef.
func Ref(ref string, headFunc func() (string, error)) string {
	if ref != "" {
		return ref
	}

	ref, err := headFunc()
	if err != nil {
		// An error means that we're either not in a GitRepo or we're
		// not on a branch. In this case, we just fallback to the
		// DefaultRef.
		return DefaultRef
	}

	// Convert `refs/heads/test-deploy` => `test-deploy`
	return refRegex.ReplaceAllString(ref, "$1")
}

// Repo will determine the correct GitHub repo to deploy to, based on a set of
// arguments.
func Repo(arguments []string) (string, error) {
	if len(arguments) != 0 {
		return arguments[0], nil
	}

	remotes, err := hub.Remotes()
	if err != nil {
		return "", err
	}

	repo := GitHubRepo(remotes)
	if repo == "" {
		return repo, errors.New("no GitHub repo found in .git/config")
	}

	return repo, nil
}

// A regular expression that can convert a URL.Path into a GitHub repo name.
var remoteRegex = regexp.MustCompile(`^/(.*)\.git$`)

// GitHubRepo, given a list of git remotes, will determine what the GitHub repo
// is.
func GitHubRepo(remotes []hub.Remote) string {
	// We only want to look at the `origin` remote.
	remote := findRemote("origin", remotes)
	if remote == nil {
		return ""
	}

	// Remotes that are not pointed at a GitHub repo are not valid.
	if remote.URL.Host != "github.com" {
		return ""
	}

	// Convert `/remind101/acme-inc.git` => `remind101/acme-inc`.
	return remoteRegex.ReplaceAllString(remote.URL.Path, "$1")
}

func findRemote(name string, remotes []hub.Remote) *hub.Remote {
	for _, r := range remotes {
		if r.Name == name {
			return &r
		}
	}

	return nil
}

var errInvalidRepo = errors.New("invalid repo")

// SplitRepo splits a repo string in the form remind101/acme-inc into it's owner
// and repo components.
func SplitRepo(nwo, defaultOrg string) (owner string, repo string, err error) {
	parts := strings.Split(nwo, "/")

	// If we were only given a repo name, and a default organization is set,
	// we'll use the defaultOrg as the owner.
	if len(parts) == 1 && defaultOrg != "" && parts[0] != "" {
		owner = defaultOrg
		repo = parts[0]
		return
	}

	if len(parts) != 2 {
		err = errInvalidRepo
		return
	}

	owner = parts[0]
	repo = parts[1]

	return
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"os"

	"github.com/remind101/emp/Godeps/_workspace/src/github.com/docker/docker/pkg/jsonmessage"
	"github.com/remind101/emp/Godeps/_workspace/src/github.com/docker/docker/pkg/term"
)

var cmdDeploy = &Command{
	Run:      runDeploy,
	Usage:    "deploy [<registry>]<image>:[<tag>]",
	Category: "deploy",
	Short:    "deploy a docker image",
	Long: `
Deploy is used to deploy a docker image to an app.
Examples:
    $ emp deploy remind101/acme-inc:latest
    Pulling repository remind101/acme-inc
    345c7524bc96: Download complete
    a1dd7097a8e8: Download complete
    23debee88b99: Download complete
    31862d352883: Download complete
    c7388ff7ab91: Download complete
    78fb106ed050: Download complete
    133fcef559c4: Download complete
    Status: Image is up to date for remind101/acme-inc:latest
    Status: Created new release v1 for acme-inc
    $ emp releases
    v1    Jan 1 12:55  Deploy remind101/acme-inc:latest
`,
}

type PostDeployForm struct {
	Image string `json:"image"`
}

func runDeploy(cmd *Command, args []string) {
	r, w := io.Pipe()

	if len(args) < 1 {
		printFatal("You must specify an image to deploy")
	}

	image := args[0]
	form := &PostDeployForm{Image: image}

	go func() {
		must(client.Post(w, "/deploys", form))
		must(w.Close())
	}()

	outFd, isTerminalOut := term.GetFdInfo(os.Stdout)
	must(jsonmessage.DisplayJSONMessagesStream(r, os.Stdout, outFd, isTerminalOut))
}
<commit_msg>Add `OptionalApp` support to deploy<commit_after>package main

import (
	"fmt"
	"io"
	"os"

	"github.com/remind101/emp/Godeps/_workspace/src/github.com/docker/docker/pkg/jsonmessage"
	"github.com/remind101/emp/Godeps/_workspace/src/github.com/docker/docker/pkg/term"
)

var cmdDeploy = &Command{
	Run:         runDeploy,
	Usage:       "deploy [<registry>]<image>:[<tag>]",
	OptionalApp: true,
	Category:    "deploy",
	Short:       "deploy a docker image",
	Long: `
Deploy is used to deploy a docker image to an app.
Examples:
    $ emp deploy remind101/acme-inc:latest
    Pulling repository remind101/acme-inc
    345c7524bc96: Download complete
    a1dd7097a8e8: Download complete
    23debee88b99: Download complete
    31862d352883: Download complete
    c7388ff7ab91: Download complete
    78fb106ed050: Download complete
    133fcef559c4: Download complete
    Status: Image is up to date for remind101/acme-inc:latest
    Status: Created new release v1 for acme-inc
    $ emp releases
    v1    Jan 1 12:55  Deploy remind101/acme-inc:latest
`,
}

type PostDeployForm struct {
	Image string `json:"image"`
}

func runDeploy(cmd *Command, args []string) {
	r, w := io.Pipe()

	if len(args) < 1 {
		printFatal("You must specify an image to deploy")
	}

	image := args[0]
	form := &PostDeployForm{Image: image}

	var endpoint string
	appName, _ := app()
	if appName != "" {
		endpoint = fmt.Sprintf("/apps/%s/deploys", appName)
	} else {
		endpoint = "/deploys"
	}

	go func() {
		must(client.Post(w, endpoint, form))
		must(w.Close())
	}()

	outFd, isTerminalOut := term.GetFdInfo(os.Stdout)
	must(jsonmessage.DisplayJSONMessagesStream(r, os.Stdout, outFd, isTerminalOut))
}
<|endoftext|>


<commit_before>package controller

import (
	. "eaciit/wfdemo-git/library/core"
	. "eaciit/wfdemo-git/library/models"
	"eaciit/wfdemo-git/web/helper"

	"github.com/eaciit/knot/knot.v1"
	tk "github.com/eaciit/toolkit"
	"github.com/eaciit/crowd"
	"sort"
	// "time"
)

func (m *AnalyticMeteorologyController) GetTurbulenceIntensity(k *knot.WebContext) interface{} {
	k.Config.OutputType = knot.OutputJson

	p := new(PayloadAnalytic)
	e := k.GetPayload(&p)
	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}
	tStart, tEnd, e := helper.GetStartEndDate(k, p.Period, p.DateStart, p.DateEnd)
	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}

	// tStart, _ = time.Parse("2006-01-02 15:04:05", "2016-08-21 00:00:00")
	// tEnd, _ = time.Parse("2006-01-02 15:04:05", "2016-08-23 00:00:00")

	turbine := p.Turbine
	var (
		query    []tk.M
		pipes    []tk.M
		pipesmet []tk.M
		results  tk.M
		datas    []tk.M
		dataSeries []tk.M
		sortTurbines []string
	)

	scadaHfds := make([]tk.M, 0)
	metTowers := make([]tk.M, 0)

	colors := []string{
		"#ED1C24", "#A3238E", "#00A65D", "#F58220", "#0066B3", "#5C2D91", "#FFF200", "#579835", "#CF3834", "#00B274", "#74489D",
		"#C06616", "#5565AF", "#CCBE00", "#390A5D", "#006D6F", "#65C295", "#F04E4D", "#407927", "#00599D", "#A09600", "#0D1F63",
		"#C38312", "#003D73", "#454FA1", "#BC312E",
	}

	query = append(query, tk.M{"_id": tk.M{"$ne": ""}})
	query = append(query, tk.M{"timestamp": tk.M{"$gte": tStart}})
	query = append(query, tk.M{"timestamp": tk.M{"$lte": tEnd}})

	pipes = append(pipes, tk.M{"$match": tk.M{"$and": query}})
	pipes = append(pipes, tk.M{"$group": tk.M{"_id": tk.M{"turbine": "$turbine", "windspeedbin": "$fast_windspeed_bin"},
		"avgws": tk.M{"$avg": "$fast_windspeed_ms"}, "avgwsstddev": tk.M{"$avg": "$fast_activepower_kw_stddev"}},
	})
	pipes = append(pipes, tk.M{"$sort": tk.M{"_id": 1}})

	pipesmet = append(pipesmet, tk.M{"$match": tk.M{"$and": query}})
	pipesmet = append(pipesmet, tk.M{"$group": tk.M{"_id": tk.M{"turbine": "Met Tower", "windspeedbin": "$windspeedbin"},
		"avgws": tk.M{"$avg": "$vhubws90mavg"}, "avgwsstddev": tk.M{"$avg": "$vhubws90mstddev"}},
	})
	pipesmet = append(pipesmet, tk.M{"$sort": tk.M{"_id": 1}})

	csr, e := DB().Connection.NewQuery().
		From(new(ScadaDataHFD).TableName()).
		Command("pipe", pipes).
		Cursor(nil)

	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}

	e = csr.Fetch(&scadaHfds, 0, false)

	csr.Close()

	csrt, e := DB().Connection.NewQuery().
		From(new(MetTower).TableName()).
		Command("pipe", pipesmet).
		Cursor(nil)

	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}

	e = csrt.Fetch(&metTowers, 0, false)

	csrt.Close()

	tk.Printf("metTowers : %s \n", len(metTowers))

	for _, m := range metTowers {
		iDs := m.Get("_id").(tk.M)
		turbine := iDs.GetString("turbine")
		windspeedbin := iDs.GetInt("windspeedbin")
		avgws := m.GetFloat64("avgws")
		avgwsstddev := m.GetFloat64("avgwsstddev")
		datas = append(datas, tk.M{
			"turbine":     turbine,
			"binws":       windspeedbin,
			"avgws":       avgws,
			"avgwsstddev": avgwsstddev,
			"tivalue":     tk.Div(avgwsstddev, avgws),
		})
	}

	for _, m := range scadaHfds {
		iDs := m.Get("_id").(tk.M)
		turbine := iDs.GetString("turbine")
		windspeedbin := iDs.GetInt("windspeedbin")
		avgws := m.GetFloat64("avgws")
		avgwsstddev := m.GetFloat64("avgwsstddev")
		datas = append(datas, tk.M{
			"turbine":     turbine,
			"binws":       windspeedbin,
			"avgws":       avgws,
			"avgwsstddev": avgwsstddev,
			"tivalue":     tk.Div(avgwsstddev, avgws),
		})
	}

	if len(p.Turbine) == 0 {
		for _, listVal := range datas {
			exist := false
			for _, val := range turbine {
				if listVal["turbine"] == val {
					exist = true
				}
			}
			if exist == false {
				turbine = append(turbine, listVal["turbine"])
			}
		}
	}

	for _, turX := range turbine {
		sortTurbines = append(sortTurbines, turX.(string))
	}
	sort.Strings(sortTurbines)

	selArr := 0
	for _, turbineX := range sortTurbines {

		exist := crowd.From(&datas).Where(func(x interface{}) interface{} {

			return x.(tk.M).GetString("turbine") == turbineX
		}).Exec().Result.Data().([]tk.M)

		var dts [][]float64
		turbineData := tk.M{}
		turbineData.Set("name", turbineX)
		turbineData.Set("type", "scatterLine")
		turbineData.Set("style", "smooth")
		turbineData.Set("dashType", "solid")
		turbineData.Set("markers", tk.M{"visible": false})
		turbineData.Set("width", 2)
		turbineData.Set("color", colors[selArr])
		turbineData.Set("idxseries", selArr)

		for _, val := range exist {
			dts = append(dts, []float64{val.GetFloat64("binws"), val.GetFloat64("tivalue")})
		}

		if len(dts) > 0 {
			turbineData.Set("data", dts)
		}

		dataSeries = append(dataSeries, turbineData)
		selArr++
	}

	if datas == nil {
		datas = make([]tk.M, 0)
	}

	results = tk.M{
		"Data":        datas,
		"ChartSeries": dataSeries,
	}

	return results
}
<commit_msg>Update column turbulance<commit_after>package controller

import (
	. "eaciit/wfdemo-git/library/core"
	. "eaciit/wfdemo-git/library/models"
	"eaciit/wfdemo-git/web/helper"

	"github.com/eaciit/knot/knot.v1"
	tk "github.com/eaciit/toolkit"
	"github.com/eaciit/crowd"
	"sort"
	// "time"
)

func (m *AnalyticMeteorologyController) GetTurbulenceIntensity(k *knot.WebContext) interface{} {
	k.Config.OutputType = knot.OutputJson

	p := new(PayloadAnalytic)
	e := k.GetPayload(&p)
	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}
	tStart, tEnd, e := helper.GetStartEndDate(k, p.Period, p.DateStart, p.DateEnd)
	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}

	// tStart, _ = time.Parse("2006-01-02 15:04:05", "2016-08-21 00:00:00")
	// tEnd, _ = time.Parse("2006-01-02 15:04:05", "2016-08-23 00:00:00")

	turbine := p.Turbine
	var (
		query    []tk.M
		pipes    []tk.M
		pipesmet []tk.M
		results  tk.M
		datas    []tk.M
		dataSeries []tk.M
		sortTurbines []string
	)

	scadaHfds := make([]tk.M, 0)
	metTowers := make([]tk.M, 0)

	colors := []string{
		"#ED1C24", "#A3238E", "#00A65D", "#F58220", "#0066B3", "#5C2D91", "#FFF200", "#579835", "#CF3834", "#00B274", "#74489D",
		"#C06616", "#5565AF", "#CCBE00", "#390A5D", "#006D6F", "#65C295", "#F04E4D", "#407927", "#00599D", "#A09600", "#0D1F63",
		"#C38312", "#003D73", "#454FA1", "#BC312E",
	}

	query = append(query, tk.M{"_id": tk.M{"$ne": ""}})
	query = append(query, tk.M{"timestamp": tk.M{"$gte": tStart}})
	query = append(query, tk.M{"timestamp": tk.M{"$lte": tEnd}})

	pipes = append(pipes, tk.M{"$match": tk.M{"$and": query}})
	pipes = append(pipes, tk.M{"$group": tk.M{"_id": tk.M{"turbine": "$turbine", "windspeedbin": "$fast_windspeed_bin"},
		"avgws": tk.M{"$avg": "$fast_windspeed_ms"}, "avgwsstddev": tk.M{"$avg": "$fast_windspeed_ms_stddev"}},
	})
	pipes = append(pipes, tk.M{"$sort": tk.M{"_id": 1}})

	pipesmet = append(pipesmet, tk.M{"$match": tk.M{"$and": query}})
	pipesmet = append(pipesmet, tk.M{"$group": tk.M{"_id": tk.M{"turbine": "Met Tower", "windspeedbin": "$windspeedbin"},
		"avgws": tk.M{"$avg": "$vhubws90mavg"}, "avgwsstddev": tk.M{"$avg": "$vhubws90mstddev"}},
	})
	pipesmet = append(pipesmet, tk.M{"$sort": tk.M{"_id": 1}})

	csr, e := DB().Connection.NewQuery().
		From(new(ScadaDataHFD).TableName()).
		Command("pipe", pipes).
		Cursor(nil)

	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}

	e = csr.Fetch(&scadaHfds, 0, false)

	csr.Close()

	csrt, e := DB().Connection.NewQuery().
		From(new(MetTower).TableName()).
		Command("pipe", pipesmet).
		Cursor(nil)

	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}

	e = csrt.Fetch(&metTowers, 0, false)

	csrt.Close()

	tk.Printf("metTowers : %s \n", len(metTowers))

	for _, m := range metTowers {
		iDs := m.Get("_id").(tk.M)
		turbine := iDs.GetString("turbine")
		windspeedbin := iDs.GetInt("windspeedbin")
		avgws := m.GetFloat64("avgws")
		avgwsstddev := m.GetFloat64("avgwsstddev")
		datas = append(datas, tk.M{
			"turbine":     turbine,
			"binws":       windspeedbin,
			"avgws":       avgws,
			"avgwsstddev": avgwsstddev,
			"tivalue":     tk.Div(avgwsstddev, avgws),
		})
	}

	for _, m := range scadaHfds {
		iDs := m.Get("_id").(tk.M)
		turbine := iDs.GetString("turbine")
		windspeedbin := iDs.GetInt("windspeedbin")
		avgws := m.GetFloat64("avgws")
		avgwsstddev := m.GetFloat64("avgwsstddev")
		datas = append(datas, tk.M{
			"turbine":     turbine,
			"binws":       windspeedbin,
			"avgws":       avgws,
			"avgwsstddev": avgwsstddev,
			"tivalue":     tk.Div(avgwsstddev, avgws),
		})
	}

	if len(p.Turbine) == 0 {
		for _, listVal := range datas {
			exist := false
			for _, val := range turbine {
				if listVal["turbine"] == val {
					exist = true
				}
			}
			if exist == false {
				turbine = append(turbine, listVal["turbine"])
			}
		}
	}

	for _, turX := range turbine {
		sortTurbines = append(sortTurbines, turX.(string))
	}
	sort.Strings(sortTurbines)

	selArr := 0
	for _, turbineX := range sortTurbines {

		exist := crowd.From(&datas).Where(func(x interface{}) interface{} {

			return x.(tk.M).GetString("turbine") == turbineX
		}).Exec().Result.Data().([]tk.M)

		var dts [][]float64
		turbineData := tk.M{}
		turbineData.Set("name", turbineX)
		turbineData.Set("type", "scatterLine")
		turbineData.Set("style", "smooth")
		turbineData.Set("dashType", "solid")
		turbineData.Set("markers", tk.M{"visible": false})
		turbineData.Set("width", 2)
		turbineData.Set("color", colors[selArr])
		turbineData.Set("idxseries", selArr)

		for _, val := range exist {
			dts = append(dts, []float64{val.GetFloat64("binws"), val.GetFloat64("tivalue")})
		}

		if len(dts) > 0 {
			turbineData.Set("data", dts)
		}

		dataSeries = append(dataSeries, turbineData)
		selArr++
	}

	if datas == nil {
		datas = make([]tk.M, 0)
	}

	results = tk.M{
		"Data":        datas,
		"ChartSeries": dataSeries,
	}

	return results
}
<|endoftext|>


<commit_before>package amazonec2

import (
	"errors"
)

type region struct {
	AmiId string
}

// Release 16.04 LTS 20160815
// See https://cloud-images.ubuntu.com/locator/ec2/
var regionDetails map[string]*region = map[string]*region{
	"ap-northeast-1":  {"ami-51f13330"},
	"ap-northeast-2":  {"ami-a3915acd"}, // 20160516.1
	"ap-southeast-1":  {"ami-fec51c9d"},
	"ap-southeast-2":  {"ami-a78ebac4"},
	"ap-south-1":      {"ami-7e94fe11"}, // 20160627
	"cn-north-1":      {"ami-2c3bf141"}, // 20160610
	"eu-central-1":    {"ami-004abc6f"},
	"eu-west-1":       {"ami-c06b1eb3"},
	"sa-east-1":       {"ami-a674e2ca"},
	"us-east-1":       {"ami-c60b90d1"},
	"us-west-1":       {"ami-1bf0b37b"},
	"us-west-2":       {"ami-f701cb97"},
	"us-gov-west-1":   {"ami-76f34a17"},
	"custom-endpoint": {""},
}

func awsRegionsList() []string {
	var list []string

	for k := range regionDetails {
		list = append(list, k)
	}

	return list
}

func validateAwsRegion(region string) (string, error) {
	for _, v := range awsRegionsList() {
		if v == region {
			return region, nil
		}
	}

	return "", errors.New("Invalid region specified")
}
<commit_msg>Add support for Ohio (us-east-2)<commit_after>package amazonec2

import (
	"errors"
)

type region struct {
	AmiId string
}

// Release 16.04 LTS 20161013
// See https://cloud-images.ubuntu.com/locator/ec2/
var regionDetails map[string]*region = map[string]*region{
	"ap-northeast-1":  {"ami-31892c50"},
	"ap-northeast-2":  {"ami-a3915acd"},
	"ap-southeast-1":  {"ami-18e7417b"},
	"ap-southeast-2":  {"ami-7be4d618"},
	"ap-south-1":      {"ami-7e94fe11"},
	"cn-north-1":      {"ami-d7c511ba"},
	"eu-central-1":    {"ami-597c8236"},
	"eu-west-1":       {"ami-c593deb6"},
	"sa-east-1":       {"ami-909b06fc"},
	"us-east-1":       {"ami-fd6e3bea"},
	"us-east-2":       {"ami-0a104a6f"},
	"us-west-1":       {"ami-73531b13"},
	"us-west-2":       {"ami-f1ca1091"},
	"us-gov-west-1":   {"ami-8df24aec"},
	"custom-endpoint": {""},
}

func awsRegionsList() []string {
	var list []string

	for k := range regionDetails {
		list = append(list, k)
	}

	return list
}

func validateAwsRegion(region string) (string, error) {
	for _, v := range awsRegionsList() {
		if v == region {
			return region, nil
		}
	}

	return "", errors.New("Invalid region specified")
}
<|endoftext|>


<commit_before>package mobi

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
)

type Metadata struct {
	Title       string
	Creator     string
	Publisher   string
	Description string
	Isbn        string
	Subject     []string
	Date        string
	Contributor string
	Rights      string
	SubjectCode string
	Type        string
	Source      string
}

type Mobi interface {
	Metadata() (Metadata, error)
	Close() error
}

type mobi struct {
	r readSeekCloser
	c io.Closer
}

type readSeekCloser interface {
	io.Reader
	io.ReaderAt
	io.Seeker
	io.Closer
}

func Open(file readSeekCloser) (Mobi, error) {
	return &mobi{file, file}, nil
}

func (m *mobi) Metadata() (Metadata, error) {
	const size = 4096
	metadata := Metadata{}
	buf := make([]byte, size)

	m.r.ReadAt(buf, 0)
	p := bytes.NewBuffer(buf)

	var header struct {
		_               [34]byte //  0
		Version         uint16   // 34
		CreationDate    uint32   // 36
		_               [12]byte // 40
		AppInfoId       uint32   // 52
		SortInfoId      uint32   // 56
		Type            [4]byte  // 60
		Creator         [4]byte  // 64
		_               [8]byte  // 68
		NumberOfRecords uint16   // 76
	}

	binary.Read(p, binary.BigEndian, &header)
	if string(header.Creator[:]) != "MOBI" || string(header.Type[:]) != "BOOK" {
		return Metadata{}, errors.New("not a book")
	}

	// 78: 8 * number of records
	p.Next(8 * int(header.NumberOfRecords))

	// xx: skip 2 zero bytes
	p.Next(2)

	var mobiHeader struct {
		_              [16]byte //   0  -- PalmDOC header
		Identifier     [4]byte  //  16
		Length         uint32   //  20  -- includes previous 4 bytes
		_              [60]byte //  24
		FullNameOffset uint32   //  84
		FullNameLength uint32   //  88
		_              [36]byte //  92
		ExthFlags      [4]byte  // 128
	}

	binary.Read(p, binary.BigEndian, &mobiHeader)
	if string(mobiHeader.Identifier[:]) != "MOBI" {
		return Metadata{}, errors.New("not a book")
	}

	p.Next(int(mobiHeader.Length) - 128 + 12) // +12 to skip PalmDOC header as well

	hasExthHeader := mobiHeader.ExthFlags[3]&0x40 != 0
	if hasExthHeader {
		var exthHeader struct {
			Identifier [4]byte // 0
			Length     uint32  // 4  -- includes previous 4 bytes
			Count      uint32  // 8
		}

		binary.Read(p, binary.BigEndian, &exthHeader)

		if string(exthHeader.Identifier[:]) != "EXTH" {
			return Metadata{}, errors.New("expected EXTH header")
		}

		type exthRecord struct {
			Type   uint32 // 0
			Length uint32 // 4
			Data   []byte // 8:Length-8
		}

		exthRecords := make([]exthRecord, exthHeader.Count)

		for i := 0; i < len(exthRecords); i++ {
			binary.Read(p, binary.BigEndian, &exthRecords[i].Type)
			binary.Read(p, binary.BigEndian, &exthRecords[i].Length)
			exthRecords[i].Data = make([]byte, exthRecords[i].Length-8)
			binary.Read(p, binary.BigEndian, &exthRecords[i].Data)
		}

		for _, rec := range exthRecords {
			val := string(rec.Data[:])
			switch rec.Type {
			case 100:
				metadata.Creator = val
			case 101:
				metadata.Publisher = val
			case 103:
				metadata.Description = val
			case 104:
				metadata.Isbn = val
			case 105:
				metadata.Subject = append(metadata.Subject, val)
			case 106:
				metadata.Date = val
			case 108:
				metadata.Contributor = val
			case 109:
				metadata.Rights = val
			case 110:
				metadata.SubjectCode = val
			case 111:
				metadata.Type = val
			case 112:
				metadata.Source = val
			}
		}

		// Skip padding
		for i := 0; i < 4; i++ {
			b, _ := p.ReadByte()
			if b != 0x00 {
				p.UnreadByte()
				break
			}
		}
	}

	name := make([]byte, mobiHeader.FullNameLength)
	binary.Read(p, binary.BigEndian, &name)
	metadata.Title = string(name[:])

	return metadata, nil
}

func (m *mobi) Close() error {
	return m.c.Close()
}
<commit_msg>Fix mobi metadata overflow problem<commit_after>package mobi

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
)

type Metadata struct {
	Title       string
	Creator     string
	Publisher   string
	Description string
	Isbn        string
	Subject     []string
	Date        string
	Contributor string
	Rights      string
	SubjectCode string
	Type        string
	Source      string
}

type Mobi interface {
	Metadata() (Metadata, error)
	Close() error
}

type mobi struct {
	r readSeekCloser
	c io.Closer
}

type readSeekCloser interface {
	io.Reader
	io.ReaderAt
	io.Seeker
	io.Closer
}

func Open(file readSeekCloser) (Mobi, error) {
	return &mobi{file, file}, nil
}

func (m *mobi) Metadata() (Metadata, error) {
	const size = 4096
	metadata := Metadata{}
	buf := make([]byte, size)

	m.r.ReadAt(buf, 0)
	p := bytes.NewBuffer(buf)

	var header struct {
		_               [34]byte //  0
		Version         uint16   // 34
		CreationDate    uint32   // 36
		_               [12]byte // 40
		AppInfoId       uint32   // 52
		SortInfoId      uint32   // 56
		Type            [4]byte  // 60
		Creator         [4]byte  // 64
		_               [8]byte  // 68
		NumberOfRecords uint16   // 76
	}

	binary.Read(p, binary.BigEndian, &header)
	if string(header.Creator[:]) != "MOBI" || string(header.Type[:]) != "BOOK" {
		return Metadata{}, errors.New("not a book")
	}

	// 78: 8 * number of records
	p.Next(8 * int(header.NumberOfRecords))

	// xx: skip 2 zero bytes
	p.Next(2)

	var mobiHeader struct {
		_              [16]byte //   0  -- PalmDOC header
		Identifier     [4]byte  //  16
		Length         uint32   //  20  -- includes previous 4 bytes
		_              [60]byte //  24
		FullNameOffset uint32   //  84
		FullNameLength uint32   //  88
		_              [36]byte //  92
		ExthFlags      [4]byte  // 128
	}

	binary.Read(p, binary.BigEndian, &mobiHeader)
	if string(mobiHeader.Identifier[:]) != "MOBI" {
		return Metadata{}, errors.New("not a book")
	}

	p.Next(int(mobiHeader.Length) - 128 + 12) // +12 to skip PalmDOC header as well

	hasExthHeader := mobiHeader.ExthFlags[3]&0x40 != 0
	if hasExthHeader {
		var exthHeader struct {
			Identifier [4]byte // 0
			Length     uint32  // 4  -- includes previous 4 bytes
			Count      uint32  // 8
		}

		binary.Read(p, binary.BigEndian, &exthHeader)

		if string(exthHeader.Identifier[:]) != "EXTH" {
			return Metadata{}, errors.New("expected EXTH header")
		}

		type exthRecord struct {
			Type   uint32 // 0
			Length uint32 // 4
			Data   []byte // 8:Length-8
		}

		exthRecords := make([]exthRecord, exthHeader.Count)

		for i := 0; i < len(exthRecords); i++ {
			binary.Read(p, binary.BigEndian, &exthRecords[i].Type)
			binary.Read(p, binary.BigEndian, &exthRecords[i].Length)
			recordLength := exthRecords[i].Length
			if exthRecords[i].Length > 8 {
				recordLength -= 8
			}
			exthRecords[i].Data = make([]byte, recordLength)
			binary.Read(p, binary.BigEndian, &exthRecords[i].Data)
		}

		for _, rec := range exthRecords {
			val := string(rec.Data[:])
			switch rec.Type {
			case 100:
				metadata.Creator = val
			case 101:
				metadata.Publisher = val
			case 103:
				metadata.Description = val
			case 104:
				metadata.Isbn = val
			case 105:
				metadata.Subject = append(metadata.Subject, val)
			case 106:
				metadata.Date = val
			case 108:
				metadata.Contributor = val
			case 109:
				metadata.Rights = val
			case 110:
				metadata.SubjectCode = val
			case 111:
				metadata.Type = val
			case 112:
				metadata.Source = val
			}
		}

		// Skip padding
		for i := 0; i < 4; i++ {
			b, _ := p.ReadByte()
			if b != 0x00 {
				p.UnreadByte()
				break
			}
		}
	}

	name := make([]byte, mobiHeader.FullNameLength)
	binary.Read(p, binary.BigEndian, &name)
	metadata.Title = string(name[:])

	return metadata, nil
}

func (m *mobi) Close() error {
	return m.c.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/Financial-Times/annotations-rw-neo4j/v4/annotations"

	logger "github.com/Financial-Times/go-logger/v2"
	"github.com/Financial-Times/kafka-client-go/kafka"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/suite"
)

const (
	knownUUID           = "12345"
	annotationLifecycle = "annotations-v1"
	platformVersion     = "v1"
)

type HttpHandlerTestSuite struct {
	suite.Suite
	body               []byte
	annotations        annotations.Annotations
	annotationsService *mockAnnotationsService
	forwarder          *mockForwarder
	message            kafka.FTMessage
	healthCheckHandler healthCheckHandler
	originMap          map[string]string
	lifecycleMap       map[string]string
	tid                string
	messageType        string
	log                *logger.UPPLogger
}

func (suite *HttpHandlerTestSuite) SetupTest() {
	suite.log = logger.NewUPPInfoLogger("annotations-rw")
	var err error
	suite.body, err = ioutil.ReadFile("annotations/examplePutBody.json")
	assert.NoError(suite.T(), err, "Unexpected error")

	suite.annotations, err = decode(bytes.NewReader(suite.body))
	assert.NoError(suite.T(), err, "Unexpected error")

	suite.annotationsService = new(mockAnnotationsService)
	suite.forwarder = new(mockForwarder)
	suite.tid = "tid_sample"

	suite.healthCheckHandler = healthCheckHandler{}
	suite.originMap, suite.lifecycleMap, suite.messageType, err = readConfigMap("annotation-config.json")

	assert.NoError(suite.T(), err, "Unexpected error")
}

func TestHttpHandlerTestSuite(t *testing.T) {
	suite.Run(t, new(HttpHandlerTestSuite))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_Success() {
	suite.annotationsService.On("Write", knownUUID, annotationLifecycle, platformVersion, suite.tid, suite.annotations).Return(nil)
	suite.forwarder.On("SendMessage", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", suite.body)
	request.Header.Add("X-Request-Id", "tid_sample")
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusCreated == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusCreated))
	assert.JSONEq(suite.T(), message("Annotations for content 12345 created"), rec.Body.String(), "Wrong body")
	suite.forwarder.AssertExpectations(suite.T())
}

func (suite *HttpHandlerTestSuite) TestPutHandler_ParseError() {
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", []byte(`{"id": "1234"}`))
	request.Header.Add("X-Request-Id", "tid_sample")
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusBadRequest == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusBadRequest))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_ValidationError() {
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", []byte(`"{"thing": {"prefLabel": "Apple"}`))
	request.Header.Add("X-Request-Id", "tid_sample")
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusBadRequest == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusBadRequest))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_NotJson() {
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "text/html", suite.body)
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusBadRequest == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusBadRequest))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_WriteFailed() {
	suite.annotationsService.On("Write", knownUUID, annotationLifecycle, platformVersion, suite.tid, suite.annotations).Return(errors.New("Write failed"))
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", suite.body)
	request.Header.Add("X-Request-Id", "tid_sample")
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusServiceUnavailable == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusServiceUnavailable))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_InvalidPredicate() {
	suite.annotationsService.On("Write", knownUUID, annotationLifecycle, platformVersion, suite.tid, suite.annotations).Return(annotations.UnsupportedPredicateErr)
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", suite.body)
	request.Header.Add("X-Request-Id", "tid_sample")
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusBadRequest == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusBadRequest))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_ForwardingFailed() {
	suite.annotationsService.On("Write", knownUUID, annotationLifecycle, platformVersion, suite.tid, suite.annotations).Return(nil)
	suite.forwarder.On("SendMessage", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("forwarding failed"))
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", suite.body)
	request.Header.Add("X-Request-Id", "tid_sample")
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusInternalServerError == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusInternalServerError))
	suite.forwarder.AssertExpectations(suite.T())
}

func (suite *HttpHandlerTestSuite) TestGetHandler_Success() {
	suite.annotationsService.On("Read", knownUUID, mock.Anything, annotationLifecycle).Return(suite.annotations, true, nil)
	request := newRequest("GET", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusOK == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusOK))
	expectedResponse, err := json.Marshal(suite.annotations)
	assert.NoError(suite.T(), err, "")
	assert.JSONEq(suite.T(), string(expectedResponse), rec.Body.String(), "Wrong body")
}

func (suite *HttpHandlerTestSuite) TestGetHandler_NotFound() {
	suite.annotationsService.On("Read", knownUUID, mock.Anything, annotationLifecycle).Return(nil, false, nil)
	request := newRequest("GET", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusNotFound == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusNotFound))
}

func (suite *HttpHandlerTestSuite) TestGetHandler_ReadError() {
	suite.annotationsService.On("Read", knownUUID, mock.Anything, annotationLifecycle).Return(nil, false, errors.New("Read error"))
	request := newRequest("GET", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusServiceUnavailable == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusServiceUnavailable))
}

func (suite *HttpHandlerTestSuite) TestDeleteHandler_Success() {
	suite.annotationsService.On("Delete", knownUUID, mock.Anything, annotationLifecycle).Return(true, nil)
	request := newRequest("DELETE", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusNoContent == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusNoContent))
}

func (suite *HttpHandlerTestSuite) TestDeleteHandler_NotFound() {
	suite.annotationsService.On("Delete", knownUUID, mock.Anything, annotationLifecycle).Return(false, nil)
	request := newRequest("DELETE", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusNotFound == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusNotFound))
}

func (suite *HttpHandlerTestSuite) TestDeleteHandler_DeleteError() {
	suite.annotationsService.On("Delete", knownUUID, mock.Anything, annotationLifecycle).Return(false, errors.New("Delete error"))
	request := newRequest("DELETE", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusServiceUnavailable == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusServiceUnavailable))
}

func (suite *HttpHandlerTestSuite) TestCount_Success() {
	suite.annotationsService.On("Count", annotationLifecycle, platformVersion).Return(10, nil)
	request := newRequest("GET", fmt.Sprintf("/content/annotations/%s/__count", annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusOK == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusOK))
}

func (suite *HttpHandlerTestSuite) TestCount_CountError() {
	suite.annotationsService.On("Count", annotationLifecycle, platformVersion).Return(0, errors.New("Count error"))
	request := newRequest("GET", fmt.Sprintf("/content/annotations/%s/__count", annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusServiceUnavailable == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusServiceUnavailable))
}

func newRequest(method, url, contentType string, body []byte) *http.Request {
	req, err := http.NewRequest(method, url, bytes.NewBuffer(body))
	if err != nil {
		panic(err)
	}
	req.Header.Add("Content-Type", contentType)
	return req
}

func message(errMsg string) string {
	return fmt.Sprintf("{\"message\": \"%s\"}\n", errMsg)
}
<commit_msg>Slightly improve tests<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/Financial-Times/annotations-rw-neo4j/v4/annotations"

	logger "github.com/Financial-Times/go-logger/v2"
	"github.com/Financial-Times/kafka-client-go/kafka"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/suite"
)

const (
	knownUUID           = "12345"
	annotationLifecycle = "annotations-v1"
	platformVersion     = "v1"
)

type HttpHandlerTestSuite struct {
	suite.Suite
	body               []byte
	annotations        annotations.Annotations
	annotationsService *mockAnnotationsService
	forwarder          *mockForwarder
	message            kafka.FTMessage
	healthCheckHandler healthCheckHandler
	originMap          map[string]string
	lifecycleMap       map[string]string
	tid                string
	messageType        string
	log                *logger.UPPLogger
}

func (suite *HttpHandlerTestSuite) SetupTest() {
	suite.log = logger.NewUPPInfoLogger("annotations-rw")
	var err error
	suite.body, err = ioutil.ReadFile("annotations/examplePutBody.json")
	assert.NoError(suite.T(), err, "Unexpected error")

	suite.annotations, err = decode(bytes.NewReader(suite.body))
	assert.NoError(suite.T(), err, "Unexpected error")

	suite.annotationsService = new(mockAnnotationsService)
	suite.forwarder = new(mockForwarder)
	suite.tid = "tid_sample"

	suite.healthCheckHandler = healthCheckHandler{}
	suite.originMap, suite.lifecycleMap, suite.messageType, err = readConfigMap("annotation-config.json")

	assert.NoError(suite.T(), err, "Unexpected error")
}

func TestHttpHandlerTestSuite(t *testing.T) {
	suite.Run(t, new(HttpHandlerTestSuite))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_Success() {
	suite.annotationsService.On("Write", knownUUID, annotationLifecycle, platformVersion, suite.tid, suite.annotations).Return(nil)
	suite.forwarder.On("SendMessage", suite.tid, "http://cmdb.ft.com/systems/methode-web-pub", knownUUID, suite.annotations).Return(nil).Once()
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", suite.body)
	request.Header.Add("X-Request-Id", suite.tid)
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusCreated == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusCreated))
	assert.JSONEq(suite.T(), message("Annotations for content 12345 created"), rec.Body.String(), "Wrong body")
	suite.forwarder.AssertExpectations(suite.T())
}

func (suite *HttpHandlerTestSuite) TestPutHandler_ParseError() {
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", []byte(`{"id": "1234"}`))
	request.Header.Add("X-Request-Id", suite.tid)
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusBadRequest == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusBadRequest))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_ValidationError() {
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", []byte(`"{"thing": {"prefLabel": "Apple"}`))
	request.Header.Add("X-Request-Id", suite.tid)
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusBadRequest == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusBadRequest))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_NotJson() {
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "text/html", suite.body)
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusBadRequest == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusBadRequest))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_WriteFailed() {
	suite.annotationsService.On("Write", knownUUID, annotationLifecycle, platformVersion, suite.tid, suite.annotations).Return(errors.New("Write failed"))
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", suite.body)
	request.Header.Add("X-Request-Id", suite.tid)
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusServiceUnavailable == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusServiceUnavailable))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_InvalidPredicate() {
	suite.annotationsService.On("Write", knownUUID, annotationLifecycle, platformVersion, suite.tid, suite.annotations).Return(annotations.UnsupportedPredicateErr)
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", suite.body)
	request.Header.Add("X-Request-Id", suite.tid)
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusBadRequest == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusBadRequest))
}

func (suite *HttpHandlerTestSuite) TestPutHandler_ForwardingFailed() {
	suite.annotationsService.On("Write", knownUUID, annotationLifecycle, platformVersion, suite.tid, suite.annotations).Return(nil)
	suite.forwarder.On("SendMessage", suite.tid, "http://cmdb.ft.com/systems/methode-web-pub", knownUUID, suite.annotations).Return(errors.New("forwarding failed"))
	request := newRequest("PUT", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", suite.body)
	request.Header.Add("X-Request-Id", suite.tid)
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	rec := httptest.NewRecorder()
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusInternalServerError == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusInternalServerError))
	suite.forwarder.AssertExpectations(suite.T())
}

func (suite *HttpHandlerTestSuite) TestGetHandler_Success() {
	suite.annotationsService.On("Read", knownUUID, mock.Anything, annotationLifecycle).Return(suite.annotations, true, nil)
	request := newRequest("GET", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusOK == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusOK))
	expectedResponse, err := json.Marshal(suite.annotations)
	assert.NoError(suite.T(), err, "")
	assert.JSONEq(suite.T(), string(expectedResponse), rec.Body.String(), "Wrong body")
}

func (suite *HttpHandlerTestSuite) TestGetHandler_NotFound() {
	suite.annotationsService.On("Read", knownUUID, mock.Anything, annotationLifecycle).Return(nil, false, nil)
	request := newRequest("GET", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusNotFound == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusNotFound))
}

func (suite *HttpHandlerTestSuite) TestGetHandler_ReadError() {
	suite.annotationsService.On("Read", knownUUID, mock.Anything, annotationLifecycle).Return(nil, false, errors.New("Read error"))
	request := newRequest("GET", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusServiceUnavailable == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusServiceUnavailable))
}

func (suite *HttpHandlerTestSuite) TestDeleteHandler_Success() {
	suite.annotationsService.On("Delete", knownUUID, mock.Anything, annotationLifecycle).Return(true, nil)
	request := newRequest("DELETE", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusNoContent == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusNoContent))
}

func (suite *HttpHandlerTestSuite) TestDeleteHandler_NotFound() {
	suite.annotationsService.On("Delete", knownUUID, mock.Anything, annotationLifecycle).Return(false, nil)
	request := newRequest("DELETE", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusNotFound == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusNotFound))
}

func (suite *HttpHandlerTestSuite) TestDeleteHandler_DeleteError() {
	suite.annotationsService.On("Delete", knownUUID, mock.Anything, annotationLifecycle).Return(false, errors.New("Delete error"))
	request := newRequest("DELETE", fmt.Sprintf("/content/%s/annotations/%s", knownUUID, annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusServiceUnavailable == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusServiceUnavailable))
}

func (suite *HttpHandlerTestSuite) TestCount_Success() {
	suite.annotationsService.On("Count", annotationLifecycle, platformVersion).Return(10, nil)
	request := newRequest("GET", fmt.Sprintf("/content/annotations/%s/__count", annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	router(&httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusOK == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusOK))
}

func (suite *HttpHandlerTestSuite) TestCount_CountError() {
	suite.annotationsService.On("Count", annotationLifecycle, platformVersion).Return(0, errors.New("Count error"))
	request := newRequest("GET", fmt.Sprintf("/content/annotations/%s/__count", annotationLifecycle), "application/json", nil)
	rec := httptest.NewRecorder()
	handler := httpHandler{suite.annotationsService, suite.forwarder, suite.originMap, suite.lifecycleMap, suite.messageType, suite.log}
	router(&handler, &suite.healthCheckHandler, suite.log).ServeHTTP(rec, request)
	assert.True(suite.T(), http.StatusServiceUnavailable == rec.Code, fmt.Sprintf("Wrong response code, was %d, should be %d", rec.Code, http.StatusServiceUnavailable))
}

func newRequest(method, url, contentType string, body []byte) *http.Request {
	req, err := http.NewRequest(method, url, bytes.NewBuffer(body))
	if err != nil {
		panic(err)
	}
	req.Header.Add("Content-Type", contentType)
	return req
}

func message(errMsg string) string {
	return fmt.Sprintf("{\"message\": \"%s\"}\n", errMsg)
}
<|endoftext|>


<commit_before>package scheduler

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"time"

	"github.com/chrislusf/glow/driver/cmd"
	"github.com/chrislusf/glow/resource"
	"github.com/chrislusf/glow/resource/service_discovery/client"
	"github.com/chrislusf/glow/util"
	"github.com/golang/protobuf/proto"
)

func NewStartRequest(path string, dir string, args []string, allocated resource.ComputeResource, envs []string, port int32) *cmd.ControlMessage {
	return &cmd.ControlMessage{
		Type: cmd.ControlMessage_StartRequest.Enum(),
		StartRequest: &cmd.StartRequest{
			Path: proto.String(path),
			Args: args,
			Dir:  proto.String(dir),
			Resource: &cmd.ComputeResource{
				CpuCount: proto.Int32(int32(allocated.CPUCount)),
				CpuLevel: proto.Int32(int32(allocated.CPULevel)),
				Memory:   proto.Int32(int32(allocated.MemoryMB)),
			},
			Envs: envs,
			Port: proto.Int32(port),
		},
	}
}

func NewDeleteDatasetShardRequest(name string) *cmd.ControlMessage {
	return &cmd.ControlMessage{
		Type: cmd.ControlMessage_DeleteDatasetShardRequest.Enum(),
		DeleteDatasetShardRequest: &cmd.DeleteDatasetShardRequest{
			Name: proto.String(name),
		},
	}
}

func RemoteDirectExecute(server string, command *cmd.ControlMessage) error {
	conn, err := getDirectCommandConnection(server)
	if err != nil {
		return err
	}
	defer conn.Close()

	return doExecute(conn, command)
}

func RemoteExecute(leader string, agentName string, command *cmd.ControlMessage) error {
	conn, err := getCommandConnection(leader, agentName)
	if err != nil {
		return err
	}
	defer conn.Close()

	return doExecute(conn, command)
}

func doExecute(conn net.Conn, command *cmd.ControlMessage) error {

	buf := make([]byte, 4)

	// serialize the commend
	data, err := proto.Marshal(command)
	if err != nil {
		log.Fatal("marshaling error: ", err)
	}

	remoteAddress := conn.RemoteAddr().String()

	// send the command
	if err = util.WriteData(conn, buf, []byte("CMD "), data); err != nil {
		println("failed to write to", remoteAddress, ":", err.Error())
		return err
	}

	// println("command sent")

	// read output and print it to stdout
	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		fmt.Printf("%s>%s\n", remoteAddress, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		return fmt.Errorf("Failed to scan output: %v", err)
	}

	return err
}

func getCommandConnection(leader string, agentName string) (net.Conn, error) {
	l := client.NewNameServiceProxy(leader)

	// looking for the agentName
	var target string
	for {
		locations := l.Find(agentName)
		if len(locations) > 0 {
			target = locations[0]
		}
		if target != "" {
			break
		} else {
			time.Sleep(time.Second)
			print("z")
		}
	}

	return getDirectCommandConnection(target)
}

func getDirectCommandConnection(target string) (net.Conn, error) {
	// connect to a TCP server
	network := "tcp"
	raddr, err := net.ResolveTCPAddr(network, target)
	if err != nil {
		return nil, fmt.Errorf("Fail to resolve %s:%v", target, err)
	}

	// println("dial tcp", raddr.String())
	conn, err := net.DialTCP(network, nil, raddr)
	if err != nil {
		return nil, fmt.Errorf("Fail to dial command %s:%v", raddr, err)
	}

	return conn, err
}
<commit_msg>add RemoteDirectCommand() to get response back<commit_after>package scheduler

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"time"

	"github.com/chrislusf/glow/driver/cmd"
	"github.com/chrislusf/glow/resource"
	"github.com/chrislusf/glow/resource/service_discovery/client"
	"github.com/chrislusf/glow/util"
	"github.com/golang/protobuf/proto"
)

func NewStartRequest(path string, dir string, args []string, allocated resource.ComputeResource, envs []string, port int32) *cmd.ControlMessage {
	request := &cmd.ControlMessage{
		Type: cmd.ControlMessage_StartRequest.Enum(),
		StartRequest: &cmd.StartRequest{
			Path: proto.String(path),
			Args: args,
			Dir:  proto.String(dir),
			Resource: &cmd.ComputeResource{
				CpuCount: proto.Int32(int32(allocated.CPUCount)),
				CpuLevel: proto.Int32(int32(allocated.CPULevel)),
				Memory:   proto.Int32(int32(allocated.MemoryMB)),
			},
			Envs: envs,
			Port: proto.Int32(port),
		},
	}

	// generate a unique hash code for the request
	data, err := proto.Marshal(request)
	if err != nil {
		log.Fatalf("marshaling start request error: %v", err)
		return nil
	}
	request.StartRequest.HashCode = proto.Int32(int32(util.Hash(data)))

	return request
}

func NewDeleteDatasetShardRequest(name string) *cmd.ControlMessage {
	return &cmd.ControlMessage{
		Type: cmd.ControlMessage_DeleteDatasetShardRequest.Enum(),
		DeleteDatasetShardRequest: &cmd.DeleteDatasetShardRequest{
			Name: proto.String(name),
		},
	}
}

func RemoteDirectExecute(server string, command *cmd.ControlMessage) error {
	conn, err := getDirectCommandConnection(server)
	if err != nil {
		return err
	}
	defer conn.Close()

	return doExecute(conn, command)
}

// doExecute() sends a request and expects the output from the connection
func doExecute(conn net.Conn, command *cmd.ControlMessage) error {

	buf := make([]byte, 4)

	// serialize the commend
	data, err := proto.Marshal(command)
	if err != nil {
		return fmt.Errorf("marshaling execute request error: %v", err)
	}

	remoteAddress := conn.RemoteAddr().String()

	// send the command
	if err = util.WriteData(conn, buf, []byte("CMD "), data); err != nil {
		return fmt.Errorf("failed to write to %s: %v", remoteAddress, err)
	}

	// println("command sent")

	// read output and print it to stdout
	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		fmt.Printf("%s>%s\n", remoteAddress, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		return fmt.Errorf("Failed to scan output: %v", err)
	}

	return err
}

func RemoteDirectCommand(server string, command *cmd.ControlMessage) (response *cmd.ControlMessage, err error) {
	conn, err := getDirectCommandConnection(server)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	return doCommand(conn, command)
}

// doCommand() sends a request and expects a response object
func doCommand(conn net.Conn, command *cmd.ControlMessage) (response *cmd.ControlMessage, err error) {

	buf := make([]byte, 4)

	// serialize the commend
	data, err := proto.Marshal(command)
	if err != nil {
		return nil, fmt.Errorf("marshaling command error: %v", err)
	}

	remoteAddress := conn.RemoteAddr().String()

	// send the command
	if err = util.WriteData(conn, buf, []byte("CMD "), data); err != nil {
		return nil, fmt.Errorf("failed to write command to %s: %v", remoteAddress, err)
	}

	// println("command sent")

	// read response
	replyBytes, err := ioutil.ReadAll(conn)
	if err != nil {
		return nil, fmt.Errorf("cmd response: %v", err)
	}

	// unmarshal the bytes
	response = &cmd.ControlMessage{}
	err = proto.Unmarshal(replyBytes, response)
	if err != nil {
		return nil, fmt.Errorf("unmarshaling error: %v", err)
	}

	return response, err
}

func getCommandConnection(leader string, agentName string) (net.Conn, error) {
	l := client.NewNameServiceProxy(leader)

	// looking for the agentName
	var target string
	for {
		locations := l.Find(agentName)
		if len(locations) > 0 {
			target = locations[0]
		}
		if target != "" {
			break
		} else {
			time.Sleep(time.Second)
			print("z")
		}
	}

	return getDirectCommandConnection(target)
}

func getDirectCommandConnection(target string) (net.Conn, error) {
	// connect to a TCP server
	network := "tcp"
	raddr, err := net.ResolveTCPAddr(network, target)
	if err != nil {
		return nil, fmt.Errorf("Fail to resolve %s:%v", target, err)
	}

	// println("dial tcp", raddr.String())
	conn, err := net.DialTCP(network, nil, raddr)
	if err != nil {
		return nil, fmt.Errorf("Fail to dial command %s:%v", raddr, err)
	}

	return conn, err
}
<|endoftext|>


<commit_before>// Copyright 2015 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build unit

package metrics

import (
	"errors"
	"reflect"
	"sync"
	"testing"
	"time"
)

type testClock struct {
	afterchan chan time.Time
}

func (c testClock) After(d time.Duration) <-chan time.Time {
	return c.afterchan
}

func TestCache(t *testing.T) {

	afterchan := make(chan time.Time)
	cache := MemoryUsageCache{
		Locks:  make(map[string]sync.Mutex),
		Usages: make(map[string][]MemoryUsageStats),
		TTL:    time.Minute,
		Clock:  testClock{afterchan},
	}

	memusage1 := &[]MemoryUsageStats{
		MemoryUsageStats{ServiceID: "memusage1"},
	}
	memusage2 := &[]MemoryUsageStats{
		MemoryUsageStats{ServiceID: "memusage2"},
	}

	getter1 := func() ([]MemoryUsageStats, error) {
		return memusage1, nil
	}

	getter2 := func() ([]MemoryUsageStats, error) {
		return memusage2, nil
	}

	errgetter := func() ([]MemoryUsageStats, error) {
		return nil, errors.New("")
	}

	// Cache is empty, test error propagation in getter
	x, err := cache.Get("first", errgetter)
	if err == nil {
		t.Errorf("Cache did not propagate error in getter properly")
	}

	// Cache still empty, test getter
	x, err = cache.Get("first", getter1)
	if !reflect.DeepEqual(x, memusage1) {
		t.Errorf("Empty cache did not return a new item when a key did not yet exist")
	}

	// Cache has value for key, try error getter again
	x, err = cache.Get("first", errgetter)
	if err != nil {
		t.Errorf("Cache returned an error when it should have returned a cached value")
	}

	// Cache has value for key, try a new key with a different getter
	x, err = cache.Get("second", getter2)
	if !reflect.DeepEqual(x, memusage2) {
		t.Errorf("Non-empty cache did not return a new item when a key did not yet exist")
	}

	// Cache has a value for key, try different getter
	x, err = cache.Get("first", getter2)
	if !reflect.DeepEqual(x, memusage1) {
		t.Errorf("Cache did not return the correct item")
	}

	// Force expiration
	afterchan <- time.Now()

	// Cache should no longer have a value for key, try with different getter
	x, err = cache.Get("first", getter2)
	if !reflect.DeepEqual(x, memusage2) {
		t.Errorf("Cache returned a value that should have expired")
	}
}
<commit_msg>CC-1286: Missed a pointer<commit_after>// Copyright 2015 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build unit

package metrics

import (
	"errors"
	"reflect"
	"sync"
	"testing"
	"time"
)

type testClock struct {
	afterchan chan time.Time
}

func (c testClock) After(d time.Duration) <-chan time.Time {
	return c.afterchan
}

func TestCache(t *testing.T) {

	afterchan := make(chan time.Time)
	cache := MemoryUsageCache{
		Locks:  make(map[string]sync.Mutex),
		Usages: make(map[string][]MemoryUsageStats),
		TTL:    time.Minute,
		Clock:  testClock{afterchan},
	}

	memusage1 := []MemoryUsageStats{
		MemoryUsageStats{ServiceID: "memusage1"},
	}
	memusage2 := []MemoryUsageStats{
		MemoryUsageStats{ServiceID: "memusage2"},
	}

	getter1 := func() ([]MemoryUsageStats, error) {
		return memusage1, nil
	}

	getter2 := func() ([]MemoryUsageStats, error) {
		return memusage2, nil
	}

	errgetter := func() ([]MemoryUsageStats, error) {
		return nil, errors.New("")
	}

	// Cache is empty, test error propagation in getter
	x, err := cache.Get("first", errgetter)
	if err == nil {
		t.Errorf("Cache did not propagate error in getter properly")
	}

	// Cache still empty, test getter
	x, err = cache.Get("first", getter1)
	if !reflect.DeepEqual(x, memusage1) {
		t.Errorf("Empty cache did not return a new item when a key did not yet exist")
	}

	// Cache has value for key, try error getter again
	x, err = cache.Get("first", errgetter)
	if err != nil {
		t.Errorf("Cache returned an error when it should have returned a cached value")
	}

	// Cache has value for key, try a new key with a different getter
	x, err = cache.Get("second", getter2)
	if !reflect.DeepEqual(x, memusage2) {
		t.Errorf("Non-empty cache did not return a new item when a key did not yet exist")
	}

	// Cache has a value for key, try different getter
	x, err = cache.Get("first", getter2)
	if !reflect.DeepEqual(x, memusage1) {
		t.Errorf("Cache did not return the correct item")
	}

	// Force expiration
	afterchan <- time.Now()

	// Cache should no longer have a value for key, try with different getter
	x, err = cache.Get("first", getter2)
	if !reflect.DeepEqual(x, memusage2) {
		t.Errorf("Cache returned a value that should have expired")
	}
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package poll

import (
	"fmt"
	"strings"
	"time"

	"github.com/ava-labs/avalanchego/utils/linkedhashmap"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/utils/metric"
)

type pollHolder interface {
	GetPoll() Poll
	StartTime() time.Time
}

type poll struct {
	Poll
	start time.Time
}

func (p poll) GetPoll() Poll {
	return p
}

func (p poll) StartTime() time.Time {
	return p.start
}

type set struct {
	log      logging.Logger
	numPolls prometheus.Gauge
	durPolls prometheus.Histogram
	factory  Factory
	// maps requestID -> poll
	polls linkedhashmap.LinkedHashmap
}

// NewSet returns a new empty set of polls
func NewSet(
	factory Factory,
	log logging.Logger,
	namespace string,
	registerer prometheus.Registerer,
) Set {
	numPolls := prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: namespace,
		Name:      "polls",
		Help:      "Number of pending network polls",
	})
	if err := registerer.Register(numPolls); err != nil {
		log.Error("failed to register polls statistics due to %s", err)
	}

	durPolls := prometheus.NewHistogram(prometheus.HistogramOpts{
		Namespace: namespace,
		Name:      "poll_duration",
		Help:      "Length of time the poll existed in milliseconds",
		Buckets:   metric.MillisecondsBuckets,
	})
	if err := registerer.Register(durPolls); err != nil {
		log.Error("failed to register poll_duration statistics due to %s", err)
	}

	return &set{
		log:      log,
		numPolls: numPolls,
		durPolls: durPolls,
		factory:  factory,
		polls:    linkedhashmap.New(),
	}
}

// Add to the current set of polls
// Returns true if the poll was registered correctly and the network sample
//         should be made.
func (s *set) Add(requestID uint32, vdrs ids.ShortBag) bool {
	if _, exists := s.polls.Get(requestID); exists {
		s.log.Debug("dropping poll due to duplicated requestID: %d", requestID)
		return false
	}

	s.log.Verbo("creating poll with requestID %d and validators %s",
		requestID,
		&vdrs)

	s.polls.Put(requestID, poll{
		Poll:  s.factory.New(vdrs), // create the new poll
		start: time.Now(),
	})
	s.numPolls.Inc() // increase the metrics
	return true
}

// Vote registers the connections response to a query for [id]. If there was no
// query, or the response has already be registered, nothing is performed.
func (s *set) Vote(requestID uint32, vdr ids.ShortID, vote ids.ID) ([]ids.Bag, bool) {
	pollHolderIntf, exists := s.polls.Get(requestID)
	if !exists {
		s.log.Verbo("dropping vote from %s to an unknown poll with requestID: %d",
			vdr,
			requestID)
		return []ids.Bag{}, false
	}

	holder := pollHolderIntf.(pollHolder)
	p := holder.GetPoll()

	s.log.Verbo("processing vote from %s in the poll with requestID: %d with the vote %s",
		vdr,
		requestID,
		vote)

	p.Vote(vdr, vote)
	if !p.Finished() {
		return []ids.Bag{}, false
	}

	s.log.Verbo("poll with requestID %d finished as %s", requestID, p)
	s.durPolls.Observe(float64(time.Since(holder.StartTime()).Milliseconds()))
	s.numPolls.Dec() // decrease the metrics

	var results []ids.Bag
	// check if previous poll(s) have finished
	if oldestRequestID, _, exists := s.polls.Oldest(); exists && oldestRequestID == requestID {
		// this is the oldest poll that has just finished
		// iterate from oldest to newest
		iter := s.polls.NewIterator()
		var keysToDelete []interface{}
		for iter.Next() {
			holder := iter.Value().(pollHolder)
			p = holder.GetPoll()
			if !p.Finished() {
				// since we're iterating from oldest to newest, if the next poll has not finished,
				// we can break and return what we have so far
				break
			}

			results = append(results, p.Result())
			keysToDelete = append(keysToDelete, iter.Key())
		}

		// delete the keys to be deleted
		for _, keyToDelete := range keysToDelete {
			s.polls.Delete(keyToDelete)
		}
	}

	// only gets here if the poll has finished
	// results will have values if this and other newer polls have finished
	return results, len(results) > 0
}

// Drop registers the connections response to a query for [id]. If there was no
// query, or the response has already be registered, nothing is performed.
func (s *set) Drop(requestID uint32, vdr ids.ShortID) ([]ids.Bag, bool) {
	pollHolderIntf, exists := s.polls.Get(requestID)
	if !exists {
		s.log.Verbo("dropping vote from %s to an unknown poll with requestID: %d",
			vdr,
			requestID)
		return []ids.Bag{}, false
	}

	s.log.Verbo("processing dropped vote from %s in the poll with requestID: %d",
		vdr,
		requestID)

	pollHolder := pollHolderIntf.(pollHolder)
	poll := pollHolder.GetPoll()

	poll.Drop(vdr)
	if !poll.Finished() {
		return []ids.Bag{}, false
	}

	s.log.Verbo("poll with requestID %d finished as %s", requestID, poll)

	s.polls.Delete(requestID) // remove the poll from the current set
	s.durPolls.Observe(float64(time.Since(pollHolder.StartTime()).Milliseconds()))
	s.numPolls.Dec() // decrease the metrics
	return []ids.Bag{poll.Result()}, true
}

// Len returns the number of outstanding polls
func (s *set) Len() int { return s.polls.Len() }

func (s *set) String() string {
	sb := strings.Builder{}
	sb.WriteString(fmt.Sprintf("current polls: (Size = %d)", s.polls.Len()))
	iter := s.polls.NewIterator()
	for iter.Next() {
		requestID := iter.Key()
		poll := iter.Value().(Poll)
		sb.WriteString(fmt.Sprintf("\n    %d: %s", requestID, poll.PrefixedString("    ")))
	}
	return sb.String()
}
<commit_msg>flip condition for clarity<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package poll

import (
	"fmt"
	"strings"
	"time"

	"github.com/ava-labs/avalanchego/utils/linkedhashmap"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/utils/metric"
)

type pollHolder interface {
	GetPoll() Poll
	StartTime() time.Time
}

type poll struct {
	Poll
	start time.Time
}

func (p poll) GetPoll() Poll {
	return p
}

func (p poll) StartTime() time.Time {
	return p.start
}

type set struct {
	log      logging.Logger
	numPolls prometheus.Gauge
	durPolls prometheus.Histogram
	factory  Factory
	// maps requestID -> poll
	polls linkedhashmap.LinkedHashmap
}

// NewSet returns a new empty set of polls
func NewSet(
	factory Factory,
	log logging.Logger,
	namespace string,
	registerer prometheus.Registerer,
) Set {
	numPolls := prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: namespace,
		Name:      "polls",
		Help:      "Number of pending network polls",
	})
	if err := registerer.Register(numPolls); err != nil {
		log.Error("failed to register polls statistics due to %s", err)
	}

	durPolls := prometheus.NewHistogram(prometheus.HistogramOpts{
		Namespace: namespace,
		Name:      "poll_duration",
		Help:      "Length of time the poll existed in milliseconds",
		Buckets:   metric.MillisecondsBuckets,
	})
	if err := registerer.Register(durPolls); err != nil {
		log.Error("failed to register poll_duration statistics due to %s", err)
	}

	return &set{
		log:      log,
		numPolls: numPolls,
		durPolls: durPolls,
		factory:  factory,
		polls:    linkedhashmap.New(),
	}
}

// Add to the current set of polls
// Returns true if the poll was registered correctly and the network sample
//         should be made.
func (s *set) Add(requestID uint32, vdrs ids.ShortBag) bool {
	if _, exists := s.polls.Get(requestID); exists {
		s.log.Debug("dropping poll due to duplicated requestID: %d", requestID)
		return false
	}

	s.log.Verbo("creating poll with requestID %d and validators %s",
		requestID,
		&vdrs)

	s.polls.Put(requestID, poll{
		Poll:  s.factory.New(vdrs), // create the new poll
		start: time.Now(),
	})
	s.numPolls.Inc() // increase the metrics
	return true
}

// Vote registers the connections response to a query for [id]. If there was no
// query, or the response has already be registered, nothing is performed.
func (s *set) Vote(requestID uint32, vdr ids.ShortID, vote ids.ID) ([]ids.Bag, bool) {
	pollHolderIntf, exists := s.polls.Get(requestID)
	if !exists {
		s.log.Verbo("dropping vote from %s to an unknown poll with requestID: %d",
			vdr,
			requestID)
		return []ids.Bag{}, false
	}

	holder := pollHolderIntf.(pollHolder)
	p := holder.GetPoll()

	s.log.Verbo("processing vote from %s in the poll with requestID: %d with the vote %s",
		vdr,
		requestID,
		vote)

	p.Vote(vdr, vote)
	if !p.Finished() {
		return []ids.Bag{}, false
	}

	s.log.Verbo("poll with requestID %d finished as %s", requestID, p)
	s.durPolls.Observe(float64(time.Since(holder.StartTime()).Milliseconds()))
	s.numPolls.Dec() // decrease the metrics

	var results []ids.Bag
	// If this is not the oldest poll, return as is.
	if oldestRequestID, _, _ := s.polls.Oldest(); oldestRequestID != requestID {
		return nil, false
	}

	// this is the oldest poll that has just finished
	// iterate from oldest to newest
	iter := s.polls.NewIterator()
	var keysToDelete []interface{}
	for iter.Next() {
		holder := iter.Value().(pollHolder)
		p = holder.GetPoll()
		if !p.Finished() {
			// since we're iterating from oldest to newest, if the next poll has not finished,
			// we can break and return what we have so far
			break
		}

		results = append(results, p.Result())
		keysToDelete = append(keysToDelete, iter.Key())
	}

	// delete the keys to be deleted
	for _, keyToDelete := range keysToDelete {
		s.polls.Delete(keyToDelete)
	}

	// only gets here if the poll has finished
	// results will have values if this and other newer polls have finished
	return results, len(results) > 0
}

// Drop registers the connections response to a query for [id]. If there was no
// query, or the response has already be registered, nothing is performed.
func (s *set) Drop(requestID uint32, vdr ids.ShortID) ([]ids.Bag, bool) {
	pollHolderIntf, exists := s.polls.Get(requestID)
	if !exists {
		s.log.Verbo("dropping vote from %s to an unknown poll with requestID: %d",
			vdr,
			requestID)
		return []ids.Bag{}, false
	}

	s.log.Verbo("processing dropped vote from %s in the poll with requestID: %d",
		vdr,
		requestID)

	pollHolder := pollHolderIntf.(pollHolder)
	poll := pollHolder.GetPoll()

	poll.Drop(vdr)
	if !poll.Finished() {
		return []ids.Bag{}, false
	}

	s.log.Verbo("poll with requestID %d finished as %s", requestID, poll)

	s.polls.Delete(requestID) // remove the poll from the current set
	s.durPolls.Observe(float64(time.Since(pollHolder.StartTime()).Milliseconds()))
	s.numPolls.Dec() // decrease the metrics
	return []ids.Bag{poll.Result()}, true
}

// Len returns the number of outstanding polls
func (s *set) Len() int { return s.polls.Len() }

func (s *set) String() string {
	sb := strings.Builder{}
	sb.WriteString(fmt.Sprintf("current polls: (Size = %d)", s.polls.Len()))
	iter := s.polls.NewIterator()
	for iter.Next() {
		requestID := iter.Key()
		poll := iter.Value().(Poll)
		sb.WriteString(fmt.Sprintf("\n    %d: %s", requestID, poll.PrefixedString("    ")))
	}
	return sb.String()
}
<|endoftext|>


<commit_before>// Copyright 2019, OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tests

import (
	"testing"

	"github.com/open-telemetry/opentelemetry-collector/testbed/testbed"
	scenarios "github.com/open-telemetry/opentelemetry-collector/testbed/tests"
)

func TestMetric10kDPS(t *testing.T) {
	tests := []struct {
		name         string
		sender       testbed.DataSender
		receiver     testbed.DataReceiver
		resourceSpec testbed.ResourceSpec
	}{
		{
			"OpenCensus",
			testbed.NewOCMetricDataSender(testbed.GetAvailablePort(t)),
			testbed.NewOCDataReceiver(testbed.GetAvailablePort(t)),
			testbed.ResourceSpec{
				ExpectedMaxCPU: 29,
				ExpectedMaxRAM: 59,
			},
		},
		{
			"Carbon",
			NewCarbonDataSender(testbed.GetAvailablePort(t)),
			NewCarbonDataReceiver(testbed.GetAvailablePort(t)),
			testbed.ResourceSpec{
				ExpectedMaxCPU: 115,
				ExpectedMaxRAM: 40,
			},
		},
		{
			"SignalFx",
			NewSFxMetricDataSender(testbed.GetAvailablePort(t)),
			NewSFxMetricsDataReceiver(testbed.GetAvailablePort(t)),
			testbed.ResourceSpec{
				ExpectedMaxCPU: 40,
				ExpectedMaxRAM: 58,
			},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			scenarios.Scenario10kItemsPerSecond(
				t,
				test.sender,
				test.receiver,
				testbed.LoadOptions{ItemsPerBatch: 100},
				test.resourceSpec,
			)
		})
	}

}
<commit_msg>Increase memory limit for Carbon E2E test (#121)<commit_after>// Copyright 2019, OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tests

import (
	"testing"

	"github.com/open-telemetry/opentelemetry-collector/testbed/testbed"
	scenarios "github.com/open-telemetry/opentelemetry-collector/testbed/tests"
)

func TestMetric10kDPS(t *testing.T) {
	tests := []struct {
		name         string
		sender       testbed.DataSender
		receiver     testbed.DataReceiver
		resourceSpec testbed.ResourceSpec
	}{
		{
			"OpenCensus",
			testbed.NewOCMetricDataSender(testbed.GetAvailablePort(t)),
			testbed.NewOCDataReceiver(testbed.GetAvailablePort(t)),
			testbed.ResourceSpec{
				ExpectedMaxCPU: 29,
				ExpectedMaxRAM: 59,
			},
		},
		{
			"Carbon",
			NewCarbonDataSender(testbed.GetAvailablePort(t)),
			NewCarbonDataReceiver(testbed.GetAvailablePort(t)),
			testbed.ResourceSpec{
				ExpectedMaxCPU: 115,
				ExpectedMaxRAM: 48,
			},
		},
		{
			"SignalFx",
			NewSFxMetricDataSender(testbed.GetAvailablePort(t)),
			NewSFxMetricsDataReceiver(testbed.GetAvailablePort(t)),
			testbed.ResourceSpec{
				ExpectedMaxCPU: 40,
				ExpectedMaxRAM: 58,
			},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			scenarios.Scenario10kItemsPerSecond(
				t,
				test.sender,
				test.receiver,
				testbed.LoadOptions{ItemsPerBatch: 100},
				test.resourceSpec,
			)
		})
	}

}
<|endoftext|>