<commit_before>package routes

import (
	"encoding/json"
	"github.com/danjac/photoshare/api/models"
	"github.com/danjac/photoshare/api/session"
	"github.com/gorilla/mux"
	"net/http"
)

type Result struct {
	Context *Context
	Code    int
	Payload interface{}
	Error   error
}

func (r *Result) Render() error {
	if r.Payload != nil {
		r.Context.Response.WriteHeader(r.Code)
		r.Context.Response.Header().Set("Content-type", "application/json")
		return json.NewEncoder(r.Context.Response).Encode(r.Payload)
	}
	if r.Error != nil {
		http.Error(r.Context.Response, "Sorry, an error has occurred", r.Code)
		return r.Error
	}
	return nil
}

type Context struct {
	*http.Request
	Response http.ResponseWriter
	Params   map[string]string
	User     *models.User
}

func (c *Context) Result(code int, payload interface{}, err error) *Result {
	return &Result{c, code, payload, err}
}

func (c *Context) Param(name string) string {
	return c.Params[name]
}

func (c *Context) GetCurrentUser() (*models.User, error) {
	var err error
	c.User, err = session.GetCurrentUser(c.Request)
	return c.User, err
}

func (c *Context) Login(user *models.User) error {
	c.User = user
	_, err := session.Login(c.Response, user)
	return err
}

func (c *Context) Logout() error {
	c.User = nil
	_, err := session.Logout(c.Response)
	return err
}

func (c *Context) OK(value interface{}) *Result {
	return c.Result(http.StatusOK, value, nil)
}

func (c *Context) Unauthorized(value interface{}) *Result {
	return c.Result(http.StatusUnauthorized, value, nil)
}

func (c *Context) Forbidden(value interface{}) *Result {
	return c.Result(http.StatusForbidden, value, nil)
}

func (c *Context) BadRequest(value interface{}) *Result {
	return c.Result(http.StatusBadRequest, value, nil)
}

func (c *Context) NotFound(value interface{}) *Result {
	return c.Result(http.StatusNotFound, value, nil)
}

func (c *Context) Error(err error) *Result {
	return c.Result(http.StatusInternalServerError, nil, err)
}

func (c *Context) ParseJSON(value interface{}) error {
	return json.NewDecoder(c.Request.Body).Decode(value)
}

func NewContext(w http.ResponseWriter, r *http.Request) *Context {
	return &Context{r, w, mux.Vars(r), nil}
}

func MakeAppHandler(fn AppHandlerFunc, loginRequired bool) http.HandlerFunc {

	// tbd defer ...
	return func(w http.ResponseWriter, r *http.Request) {
		c := NewContext(w, r)
		if loginRequired {
			if user, err := c.GetCurrentUser(); err != nil || user == nil {
				if err != nil {
					panic(err)
					return
				}
				c.Unauthorized("You must be logged in")
				return
			}
		}

		result := fn(c)
		if result != nil {
			if err := result.Render(); err != nil {
				panic(err)
			}
		}
	}

}

type AppHandlerFunc func(c *Context) *Result
<commit_msg>Recover from panic<commit_after>package routes

import (
	"encoding/json"
	"github.com/danjac/photoshare/api/models"
	"github.com/danjac/photoshare/api/session"
	"github.com/gorilla/mux"
	"net/http"
	"log"
)

type Result struct {
	Context *Context
	Code    int
	Payload interface{}
	Error   error
}

func (r *Result) Render() error {
	if r.Payload != nil {
		r.Context.Response.WriteHeader(r.Code)
		r.Context.Response.Header().Set("Content-type", "application/json")
		return json.NewEncoder(r.Context.Response).Encode(r.Payload)
	}
	if r.Error != nil {
		http.Error(r.Context.Response, "Sorry, an error has occurred", r.Code)
		return r.Error
	}
	return nil
}

type Context struct {
	*http.Request
	Response http.ResponseWriter
	Params   map[string]string
	User     *models.User
}

func (c *Context) Result(code int, payload interface{}, err error) *Result {
	return &Result{c, code, payload, err}
}

func (c *Context) Param(name string) string {
	return c.Params[name]
}

func (c *Context) GetCurrentUser() (*models.User, error) {
	var err error
	c.User, err = session.GetCurrentUser(c.Request)
	return c.User, err
}

func (c *Context) Login(user *models.User) error {
	c.User = user
	_, err := session.Login(c.Response, user)
	return err
}

func (c *Context) Logout() error {
	c.User = nil
	_, err := session.Logout(c.Response)
	return err
}

func (c *Context) OK(value interface{}) *Result {
	return c.Result(http.StatusOK, value, nil)
}

func (c *Context) Unauthorized(value interface{}) *Result {
	return c.Result(http.StatusUnauthorized, value, nil)
}

func (c *Context) Forbidden(value interface{}) *Result {
	return c.Result(http.StatusForbidden, value, nil)
}

func (c *Context) BadRequest(value interface{}) *Result {
	return c.Result(http.StatusBadRequest, value, nil)
}

func (c *Context) NotFound(value interface{}) *Result {
	return c.Result(http.StatusNotFound, value, nil)
}

func (c *Context) Error(err error) *Result {
	return c.Result(http.StatusInternalServerError, nil, err)
}

func (c *Context) ParseJSON(value interface{}) error {
	return json.NewDecoder(c.Request.Body).Decode(value)
}

func NewContext(w http.ResponseWriter, r *http.Request) *Context {
	return &Context{r, w, mux.Vars(r), nil}
}

type AppHandlerFunc func(c *Context) *Result

func MakeAppHandler(fn AppHandlerFunc, loginRequired bool) http.HandlerFunc {

	defer func () {
		if r := recover(); r != nil {
			log.Println(r)
		}
	}()
	return func(w http.ResponseWriter, r *http.Request) {
		c := NewContext(w, r)
		if loginRequired {
			if user, err := c.GetCurrentUser(); err != nil || user == nil {
				if err != nil {
					panic(err)
					return
				}
				c.Unauthorized("You must be logged in")
				return
			}
		}

		result := fn(c)
		if result != nil {
			if err := result.Render(); err != nil {
				panic(err)
			}
		}
	}

}

<|endoftext|>


<commit_before>package classic

import (
	"context"
	"fmt"

	"github.com/hashicorp/go-oracle-terraform/compute"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
)

type stepListImages struct{}

func (s *stepListImages) Run(_ context.Context, state multistep.StateBag) multistep.StepAction {
	// get variables from state
	ui := state.Get("ui").(packer.Ui)
	config := state.Get("config").(*Config)
	client := state.Get("client").(*compute.ComputeClient)
	ui.Say("Adding image to image list...")

	imageListClient := client.ImageList()
	getInput := compute.GetImageListInput{
		Name: config.DestImageList,
	}
	imList, err := imageListClient.GetImageList(&getInput)
	if err != nil {
		ui.Say(fmt.Sprintf(err.Error()))
		// If the list didn't exist, create it.
		ui.Say(fmt.Sprintf("Destination image list %s does not exist; Creating it...",
			config.DestImageList))

		ilInput := compute.CreateImageListInput{
			Name:        config.DestImageList,
			Description: "Packer-built image list",
		}

		// Load the packer-generated SSH key into the Oracle Compute cloud.
		imList, err = imageListClient.CreateImageList(&ilInput)
		if err != nil {
			err = fmt.Errorf("Problem creating image list: %s", err)
			ui.Error(err.Error())
			state.Put("error", err)
			return multistep.ActionHalt
		}
		ui.Message(fmt.Sprintf("Image list %s created!", imList.URI))
	}

	// Now create and image list entry for the image into that list.
	snap := state.Get("snapshot").(*compute.Snapshot)
	version := len(imList.Entries) + 1
	entriesClient := client.ImageListEntries()
	entriesInput := compute.CreateImageListEntryInput{
		Name:          config.DestImageList,
		MachineImages: []string{fmt.Sprintf("Compute-%s/%s/%s", config.IdentityDomain, config.Username, snap.MachineImage)},
		Version:       version,
	}
	entryInfo, err := entriesClient.CreateImageListEntry(&entriesInput)
	if err != nil {
		err = fmt.Errorf("Problem creating an image list entry: %s", err)
		ui.Error(err.Error())
		state.Put("error", err)
		return multistep.ActionHalt
	}
	state.Put("image_list_entry", entryInfo)
	ui.Message(fmt.Sprintf("created image list entry %s", entryInfo.Name))

	imList, err = imageListClient.GetImageList(&getInput)

	machineImagesClient := client.MachineImages()
	getImagesInput := compute.GetMachineImageInput{
		Name: config.ImageName,
	}

	// Grab info about the machine image to return with the artifact
	imInfo, err := machineImagesClient.GetMachineImage(&getImagesInput)
	if err != nil {
		err = fmt.Errorf("Problem getting machine image info: %s", err)
		ui.Error(err.Error())
		state.Put("error", err)
		return multistep.ActionHalt
	}
	state.Put("machine_image_file", imInfo.File)
	state.Put("machine_image_name", imInfo.Name)
	state.Put("image_list_version", version)

	return multistep.ActionContinue
}

func (s *stepListImages) Cleanup(state multistep.StateBag) {
	// Nothing to do
	return
}
<commit_msg>remove copypasta comment<commit_after>package classic

import (
	"context"
	"fmt"

	"github.com/hashicorp/go-oracle-terraform/compute"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
)

type stepListImages struct{}

func (s *stepListImages) Run(_ context.Context, state multistep.StateBag) multistep.StepAction {
	// get variables from state
	ui := state.Get("ui").(packer.Ui)
	config := state.Get("config").(*Config)
	client := state.Get("client").(*compute.ComputeClient)
	ui.Say("Adding image to image list...")

	imageListClient := client.ImageList()
	getInput := compute.GetImageListInput{
		Name: config.DestImageList,
	}
	imList, err := imageListClient.GetImageList(&getInput)
	if err != nil {
		// If the list didn't exist, create it.
		ui.Say(fmt.Sprintf(err.Error()))
		ui.Say(fmt.Sprintf("Destination image list %s does not exist; Creating it...",
			config.DestImageList))

		ilInput := compute.CreateImageListInput{
			Name:        config.DestImageList,
			Description: "Packer-built image list",
		}

		imList, err = imageListClient.CreateImageList(&ilInput)
		if err != nil {
			err = fmt.Errorf("Problem creating image list: %s", err)
			ui.Error(err.Error())
			state.Put("error", err)
			return multistep.ActionHalt
		}
		ui.Message(fmt.Sprintf("Image list %s created!", imList.URI))
	}

	// Now create and image list entry for the image into that list.
	snap := state.Get("snapshot").(*compute.Snapshot)
	version := len(imList.Entries) + 1
	entriesClient := client.ImageListEntries()
	entriesInput := compute.CreateImageListEntryInput{
		Name:          config.DestImageList,
		MachineImages: []string{fmt.Sprintf("Compute-%s/%s/%s", config.IdentityDomain, config.Username, snap.MachineImage)},
		Version:       version,
	}
	entryInfo, err := entriesClient.CreateImageListEntry(&entriesInput)
	if err != nil {
		err = fmt.Errorf("Problem creating an image list entry: %s", err)
		ui.Error(err.Error())
		state.Put("error", err)
		return multistep.ActionHalt
	}
	state.Put("image_list_entry", entryInfo)
	ui.Message(fmt.Sprintf("created image list entry %s", entryInfo.Name))

	imList, err = imageListClient.GetImageList(&getInput)

	machineImagesClient := client.MachineImages()
	getImagesInput := compute.GetMachineImageInput{
		Name: config.ImageName,
	}

	// Grab info about the machine image to return with the artifact
	imInfo, err := machineImagesClient.GetMachineImage(&getImagesInput)
	if err != nil {
		err = fmt.Errorf("Problem getting machine image info: %s", err)
		ui.Error(err.Error())
		state.Put("error", err)
		return multistep.ActionHalt
	}
	state.Put("machine_image_file", imInfo.File)
	state.Put("machine_image_name", imInfo.Name)
	state.Put("image_list_version", version)

	return multistep.ActionContinue
}

func (s *stepListImages) Cleanup(state multistep.StateBag) {
	// Nothing to do
	return
}
<|endoftext|>


<commit_before>package gnuplot

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	// "os/exec"
)

type Plotter struct {
	configures map[string]string
}

func (p *Plotter) Init() {
	p.configures = map[string]string{}
}

func (p *Plotter) Configure(key, val string) {
	p.configures[key] = val
}

func (p *Plotter) GetC(key string) string {
	return p.configures[key]
}

var DefaultFunction2dSplitNum int = 1000

type Function2d struct {
	plotter  Plotter
	splitNum int
	f        func(float64) float64
}

func (fun *Function2d) Init() {
	fun.splitNum = DefaultFunction2dSplitNum
	fun.plotter.configures = map[string]string{
		"_xMin": "-10.0",
		"_xMax": "10.0"}
}

func (fun *Function2d) UpdatePlotter(plotter *Plotter) {
	for key, val := range plotter.configures {
		fun.plotter.configures[key] = val
	}
}

func (fun *Function2d) GetData() [][2]float64 { // TODO: テスト書く
	xMin, _ := strconv.ParseFloat(fun.plotter.configures["_xMin"], 32)
	xMax, _ := strconv.ParseFloat(fun.plotter.configures["_xMax"], 32)
	var sep = float64(xMax-xMin) / float64(fun.splitNum-1)

	var a [][2]float64
	for j := 0; j < fun.splitNum; j++ {
		t := xMin + float64(j)*sep
		y := fun.f(t)
		a = append(a, [2]float64{t, y})
	}
	return a
}

func (fun *Function2d) getGnuData() string {
	var s string
	for _, xs := range fun.GetData() {
		s += fmt.Sprintf("%f %f\n", xs[0], xs[1])
	}
	return s
}

func (fun *Function2d) SetF(_f func(float64) float64) {
	fun.f = _f
}

func (fun Function2d) gnuplot(filename string) string {
	var s = fmt.Sprintf("plot \"%v\"", filename)
	for key, val := range fun.plotter.configures {
		if !strings.HasPrefix(key, "_") {
			s += fmt.Sprintf(" %v %v", key, val)
		}
	}
	return s + ";\n"
}

func (fun *Function2d) writeIntoGnufile(f os.File) {
	f.WriteString(fun.getGnuData())
}

var DefaultCurve2dSplitNum int = 100

type Curve2d struct {
	plotter  Plotter
	splitNum int
	c        func(float64) [2]float64
}

func (c *Curve2d) Init() {
	c.splitNum = DefaultCurve2dSplitNum
	c.plotter.configures = map[string]string{
		"_tMin": "-10.0",
		"_tMax": "10.0"}
}

func (c *Curve2d) UpdatePlotter(plotter *Plotter) {
	for key, val := range plotter.configures {
		c.plotter.Configure(key, val)
	}
}

func (c *Curve2d) GetData() [][2]float64 { // TODO: test
	tMin, _ := strconv.ParseFloat(c.plotter.configures["_tMin"], 32)
	tMax, _ := strconv.ParseFloat(c.plotter.configures["_tMax"], 32)
	var sep = float64(tMax-tMin) / float64(c.splitNum-1)

	var a [][2]float64
	for j := 0; j < c.splitNum; j++ {
		var t float64 = tMin + float64(j)*sep
		cs := c.c(tMin + t*float64(j))
		a = append(a, [2]float64{cs[0], cs[1]})
	}
	return a
}

func (c Curve2d) gnuplot(fileName string) string {
	var s = fmt.Sprintf("plot %v\n;", fileName)
	for key, val := range fun.plotter.configures {
		if !strings.HasPrefix(key, "_") {
			s += fmt.Sprintf(" %v %v", key, val)
		}
	}
	return s + ";\n"
}

// Graph
type Graph2d struct {
	plotter   Plotter
	functions []Function2d
	curves    []Curve2d
}

func (g *Graph2d) AppendFunc(f Function2d) {
	fmt.Println("before of AppendFunc")
	g.functions = append(g.functions, f)
	fmt.Println("after of AppendFunc")
}

func (g Graph2d) writeIntoFile(data string, f *os.File) {
	f.WriteString(data)
}

func (g Graph2d) exec_gnuplot() {
	// until
}

func (g Graph2d) gnuplot(funcFilenames []string, curveFilenames []string) string {
	var s string
	for j, _ := range g.functions {
		s += g.functions[j].gnuplot(funcFilenames[j])
	}
	s += "\n"
	for j, _ := range g.curves {
		s += g.curves[j].gnuplot(curveFilenames[j])
	}
	s += "pause -1;"
	s += "\n"
	return s
}

func (g *Graph2d) Run() {
	tmpDir := os.TempDir()

	// それぞれのfunctionのdataをtempファイルに書き込む
	// また, それらのファイルの名前を func_filenames []string に格納する
	var funcFilenames []string
	for _, fun := range g.functions {
		file, _ := ioutil.TempFile(tmpDir, "")
		defer func() {
			file.Close()
		}()
		g.writeIntoFile(fun.getGnuData(), file)
		funcFilenames = append(funcFilenames, file.Name())
	}

	// それぞれのcurveのdataをtempファイルに書き込む
	// また, それらのファイルの名前を curve_filenames []stringに格納する

	// 実行するgnuplotの実行ファイルをtempファイルに書き込む
	execFile, _ := os.OpenFile("exec_gnu.gnu", os.O_CREATE|os.O_WRONLY, 0666)
	defer func() {
		execFile.Close()
	}()
	fmt.Println(funcFilenames)
	execFile.WriteString(g.gnuplot(funcFilenames, []string{}))
}
<commit_msg>fix refactoring<commit_after>package gnuplot

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	// "os/exec"
)

type Plotter struct {
	configures map[string]string
}

func (p *Plotter) Init() {
	p.configures = map[string]string{}
}

func (p *Plotter) Configure(key, val string) {
	p.configures[key] = val
}

func (p *Plotter) GetC(key string) string {
	return p.configures[key]
}

var DefaultFunction2dSplitNum int = 1000

type Function2d struct {
	plotter  Plotter
	splitNum int
	f        func(float64) float64
}

func (fun *Function2d) Init() {
	fun.splitNum = DefaultFunction2dSplitNum
	fun.plotter.configures = map[string]string{
		"_xMin": "-10.0",
		"_xMax": "10.0"}
}

func (fun *Function2d) UpdatePlotter(plotter *Plotter) {
	for key, val := range plotter.configures {
		fun.plotter.configures[key] = val
	}
}

func (fun *Function2d) GetData() [][2]float64 { // TODO: テスト書く
	xMin, _ := strconv.ParseFloat(fun.plotter.configures["_xMin"], 32)
	xMax, _ := strconv.ParseFloat(fun.plotter.configures["_xMax"], 32)
	var sep = float64(xMax-xMin) / float64(fun.splitNum-1)

	var a [][2]float64
	for j := 0; j < fun.splitNum; j++ {
		t := xMin + float64(j)*sep
		y := fun.f(t)
		a = append(a, [2]float64{t, y})
	}
	return a
}

func (fun *Function2d) getGnuData() string {
	var s string
	for _, xs := range fun.GetData() {
		s += fmt.Sprintf("%f %f\n", xs[0], xs[1])
	}
	return s
}

func (fun *Function2d) SetF(_f func(float64) float64) {
	fun.f = _f
}

func (fun Function2d) gnuplot(filename string) string {
	var s = fmt.Sprintf("plot \"%v\"", filename)
	for key, val := range fun.plotter.configures {
		if !strings.HasPrefix(key, "_") {
			s += fmt.Sprintf(" %v %v", key, val)
		}
	}
	return s + ";\n"
}

func (fun *Function2d) writeIntoGnufile(f os.File) {
	f.WriteString(fun.getGnuData())
}

var DefaultCurve2dSplitNum int = 100

type Curve2d struct {
	plotter  Plotter
	splitNum int
	c        func(float64) [2]float64
}

func (c *Curve2d) Init() {
	c.splitNum = DefaultCurve2dSplitNum
	c.plotter.configures = map[string]string{
		"_tMin": "-10.0",
		"_tMax": "10.0"}
}

func (c *Curve2d) UpdatePlotter(plotter *Plotter) {
	for key, val := range plotter.configures {
		c.plotter.Configure(key, val)
	}
}

func (c *Curve2d) GetData() [][2]float64 { // TODO: test
	tMin, _ := strconv.ParseFloat(c.plotter.configures["_tMin"], 32)
	tMax, _ := strconv.ParseFloat(c.plotter.configures["_tMax"], 32)
	var sep = float64(tMax-tMin) / float64(c.splitNum-1)

	var a [][2]float64
	for j := 0; j < c.splitNum; j++ {
		var t float64 = tMin + float64(j)*sep
		cs := c.c(tMin + t*float64(j))
		a = append(a, [2]float64{cs[0], cs[1]})
	}
	return a
}

func (c Curve2d) gnuplot(fileName string) string {
	var s = fmt.Sprintf("plot %v\n;", fileName)
	for key, val := range c.plotter.configures {
		if !strings.HasPrefix(key, "_") {
			s += fmt.Sprintf(" %v %v", key, val)
		}
	}
	return s + ";\n"
}

// Graph
type Graph2d struct {
	plotter   Plotter
	functions []Function2d
	curves    []Curve2d
}

func (g *Graph2d) AppendFunc(f Function2d) {
	fmt.Println("before of AppendFunc")
	g.functions = append(g.functions, f)
	fmt.Println("after of AppendFunc")
}

func (g Graph2d) writeIntoFile(data string, f *os.File) {
	f.WriteString(data)
}

func (g Graph2d) exec_gnuplot() {
	// until
}

func (g Graph2d) gnuplot(funcFilenames []string, curveFilenames []string) string {
	var s string
	for j, _ := range g.functions {
		s += g.functions[j].gnuplot(funcFilenames[j])
	}
	s += "\n"
	for j, _ := range g.curves {
		s += g.curves[j].gnuplot(curveFilenames[j])
	}
	s += "pause -1;\n"
	return s
}

func (g *Graph2d) Run() {
	tmpDir := os.TempDir()

	// それぞれのfunctionのdataをtempファイルに書き込む
	// また, それらのファイルの名前を func_filenames []string に格納する
	var funcFilenames []string
	for _, fun := range g.functions {
		file, _ := ioutil.TempFile(tmpDir, "")
		defer func() {
			file.Close()
		}()
		g.writeIntoFile(fun.getGnuData(), file)
		funcFilenames = append(funcFilenames, file.Name())
	}

	// それぞれのcurveのdataをtempファイルに書き込む
	// また, それらのファイルの名前を curve_filenames []stringに格納する

	// 実行するgnuplotの実行ファイルをtempファイルに書き込む
	execFile, _ := os.OpenFile("exec_gnu.gnu", os.O_CREATE|os.O_WRONLY, 0666)
	defer func() {
		execFile.Close()
	}()
	fmt.Println(funcFilenames)
	execFile.WriteString(g.gnuplot(funcFilenames, []string{}))
}
<|endoftext|>


<commit_before>// Copyright 2016 The LUCI 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 buildbucket

import (
	"context"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"

	"google.golang.org/genproto/protobuf/field_mask"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"go.chromium.org/gae/service/datastore"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	gitpb "go.chromium.org/luci/common/proto/git"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/milo/api/config"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/milo/frontend/ui"
	"go.chromium.org/luci/server/auth"
	"go.chromium.org/luci/server/router"
)

var ErrNotFound = errors.Reason("Build not found").Tag(common.CodeNotFound).Err()

// BuildAddress constructs the build address of a buildbucketpb.Build.
// This is used as the key for the BuildSummary entity.
func BuildAddress(build *buildbucketpb.Build) string {
	if build == nil {
		return ""
	}
	num := strconv.FormatInt(build.Id, 10)
	if build.Number != 0 {
		num = strconv.FormatInt(int64(build.Number), 10)
	}
	b := build.Builder
	return fmt.Sprintf("luci.%s.%s/%s/%s", b.Project, b.Bucket, b.Builder, num)
}

// simplisticBlamelist returns a slice of ui.Blame for a build.
//
// HACK(iannucci) - Getting the frontend to render a proper blamelist will
// require some significant refactoring. To do this properly, we'll need:
//   * The frontend to get BuildSummary from the backend.
//   * BuildSummary to have a .PreviousBuild() API.
//   * The frontend to obtain the annotation streams itself (so it could see
//     the SourceManifest objects inside of them). Currently getRespBuild defers
//     to swarming's implementation of buildsource.ID.Get(), which only returns
//     the resp object.
func simplisticBlamelist(c context.Context, build *model.BuildSummary) (result []*ui.Commit) {
	var err error
	defer func() {
		if err != nil {
			msg := "Failed to fetch blame information\n"
			logging.WithError(err).Warningf(c, msg)
			if strings.Contains(err.Error(), "DEADLINE_EXCEEDED") || c.Err() != nil {
				// When the context deadline is reached, the error that is returned is
				// a bit racy.
				// Either the actual context error itself is returned, or the urlfetch variant.
				msg += "Gitiles was taking too long"
			} else {
				msg += err.Error()
			}
			result = append(result, &ui.Commit{Description: msg})
		}
	}()
	bs := build.GitilesCommit()
	if bs == nil {
		return
	}

	builds, commits, err := build.PreviousByGitilesCommit(c)
	switch {
	case err == nil || err == model.ErrUnknownPreviousBuild:
		// continue
	case status.Code(err) == codes.PermissionDenied:
		err = common.CodeUnauthorized.Tag().Apply(err)
		return
	default:
		return
	}

	repoURL := bs.RepoURL()
	result = make([]*ui.Commit, 0, len(commits)+1)
	for _, commit := range commits {
		result = append(result, uiCommit(commit, repoURL))
	}
	logging.Infof(c, "Fetched %d commit blamelist from Gitiles", len(result))

	// this means that there were more than 100 commits in-between.
	if len(builds) == 0 && len(commits) > 0 {
		result = append(result, &ui.Commit{
			Description: "<blame list capped at 100 commits>",
			Revision:    &ui.Link{},
			AuthorName:  "<blame list capped at 100 commits>",
		})
	}

	return
}

func uiCommit(commit *gitpb.Commit, repoURL string) *ui.Commit {
	res := &ui.Commit{
		AuthorName:  commit.Author.Name,
		AuthorEmail: commit.Author.Email,
		Repo:        repoURL,
		Description: commit.Message,

		// TODO(iannucci): this use of links is very sloppy; the frontend should
		// know how to render a Commit without having Links embedded in it.
		Revision: ui.NewLink(
			commit.Id,
			repoURL+"/+/"+commit.Id, fmt.Sprintf("commit by %s", commit.Author.Email)),
	}
	res.CommitTime, _ = ptypes.Timestamp(commit.Committer.Time)
	res.File = make([]string, 0, len(commit.TreeDiff))
	for _, td := range commit.TreeDiff {
		// If file is moved, there is both new and old path,
		// from which we take only new path.
		// If a file is deleted, its new path is /dev/null.
		// In that case, we're only interested in the old path.
		switch {
		case td.NewPath != "" && td.NewPath != "/dev/null":
			res.File = append(res.File, td.NewPath)
		case td.OldPath != "":
			res.File = append(res.File, td.OldPath)
		}
	}
	return res
}

// GetBuildSummary fetches a build summary where the Context URI matches the
// given address.
func GetBuildSummary(c context.Context, id int64) (*model.BuildSummary, error) {
	// The host is set to prod because buildbot is hardcoded to talk to prod.
	uri := fmt.Sprintf("buildbucket://cr-buildbucket.appspot.com/build/%d", id)
	bs := make([]*model.BuildSummary, 0, 1)
	q := datastore.NewQuery("BuildSummary").Eq("ContextURI", uri).Limit(1)
	switch err := datastore.GetAll(c, q, &bs); {
	case err != nil:
		return nil, common.ReplaceNSEWith(err.(errors.MultiError), ErrNotFound)
	case len(bs) == 0:
		return nil, ErrNotFound
	default:
		return bs[0], nil
	}
}

// getBlame fetches blame information from Gitiles.
// This requires the BuildSummary to be indexed in Milo.
func getBlame(c context.Context, host string, b *buildbucketpb.Build) []*ui.Commit {
	commit := b.GetInput().GetGitilesCommit()
	// No commit? No blamelist.
	if commit == nil {
		return nil
	}
	// TODO(hinoka): This converts a buildbucketpb.Commit into a string
	// and back into a buildbucketpb.Commit.  That's a bit silly.
	return simplisticBlamelist(c, &model.BuildSummary{
		BuildKey:  MakeBuildKey(c, host, BuildAddress(b)),
		BuildSet:  []string{commit.BuildSetString()},
		BuilderID: BuilderID{BuilderID: *b.Builder}.String(),
	})
}

func buildbucketClient(c context.Context, host string) (buildbucketpb.BuildsClient, error) {
	t, err := auth.GetRPCTransport(c, auth.AsUser)
	if err != nil {
		return nil, err
	}
	return buildbucketpb.NewBuildsPRPCClient(&prpc.Client{
		C:    &http.Client{Transport: t},
		Host: host,
	}), nil
}

// getBugLink attempts to formulate and return the build page bug link
// for the given build.
func getBugLink(c *router.Context, b *buildbucketpb.Build) (string, error) {
	project, err := common.GetProject(c.Context, b.Builder.GetProject())
	if err != nil || proto.Equal(&project.BuildBugTemplate, &config.BugTemplate{}) {
		return "", err
	}

	builderPath := fmt.Sprintf("/p/%s/builders/%s/%s", b.Builder.GetProject(), b.Builder.GetBucket(), b.Builder.GetBuilder())
	buildURL, err := c.Request.URL.Parse(builderPath + "/" + c.Params.ByName("numberOrId"))
	if err != nil {
		return "", errors.Annotate(err, "Unable to make build URL for build bug link.").Err()
	}
	builderURL, err := c.Request.URL.Parse(builderPath)
	if err != nil {
		return "", errors.Annotate(err, "Unable to make builder URL for build bug link.").Err()
	}

	return MakeBuildBugLink(&project.BuildBugTemplate, map[string]interface{}{
		"Build":          b,
		"MiloBuildUrl":   buildURL,
		"MiloBuilderUrl": builderURL,
	})
}

// searchBuildsRequest creates a searchBuildsRequest that looks for a buildset tag.
func searchBuildsRequest(buildset string) *buildbucketpb.SearchBuildsRequest {
	return &buildbucketpb.SearchBuildsRequest{
		Predicate: &buildbucketpb.BuildPredicate{
			Tags: []*buildbucketpb.StringPair{{Key: "buildset", Value: buildset}},
		},
		Fields:   summaryBuildMask,
		PageSize: 1000,
	}
}

// getRelatedBuilds fetches build summaries of builds with the same buildset as b.
func getRelatedBuilds(c context.Context, client buildbucketpb.BuildsClient, b *buildbucketpb.Build) ([]*ui.Build, error) {
	var bs []string
	for _, buildset := range b.Buildsets() {
		// HACK(hinoka): Remove the commit/git/ buildsets because we know they're redundant
		// with the commit/gitiles/ buildsets, and we don't need to ask Buildbucket twice.
		if strings.HasPrefix(buildset, "commit/git/") {
			continue
		}
		bs = append(bs, buildset)
	}
	if len(bs) == 0 {
		// No buildset? No builds.
		return nil, nil
	}

	// Do the search request.
	// Use multiple requests instead of a single batch request.
	// A single large request is CPU bound to a single GAE instance on the buildbucket side.
	// Multiple requests allows the use of multiple GAE instances, therefore more parallelism.
	resps := make([]*buildbucketpb.SearchBuildsResponse, len(bs))
	if err := parallel.FanOutIn(func(ch chan<- func() error) {
		for i, buildset := range bs {
			i := i
			buildset := buildset
			ch <- func() (err error) {
				logging.Debugf(c, "Searching for %s (%d)", buildset, i)
				resps[i], err = client.SearchBuilds(c, searchBuildsRequest(buildset))
				return
			}
		}
	}); err != nil {
		return nil, err
	}

	// Dedupe builds.
	// It's possible since we've made multiple requests that we got back the same builds
	// multiple times.
	seen := map[int64]bool{} // set of build IDs.
	result := []*ui.Build{}
	for _, resp := range resps {
		for _, rb := range resp.GetBuilds() {
			if seen[rb.Id] {
				continue
			}
			seen[rb.Id] = true
			result = append(result, &ui.Build{rb})
		}
	}

	// Sort builds by ID.
	sort.Slice(result, func(i, j int) bool { return result[i].Id < result[j].Id })

	return result, nil
}

var (
	fullBuildMask = &field_mask.FieldMask{
		// TODO(hinoka): Add statusReason here.
		Paths: []string{
			"id",
			"builder",
			"number",
			"created_by",
			"create_time",
			"start_time",
			"end_time",
			"update_time",
			"status",
			"input",
			"output",
			"steps",
			"infra",
			"tags",
		},
	}
	tagsAndGitilesMask = &field_mask.FieldMask{
		Paths: []string{
			"id",
			"number",
			"builder",
			"input.gitiles_commit",
			"tags",
		},
	}
	summaryBuildMask = &field_mask.FieldMask{
		Paths: []string{
			"builds.*.id",
			"builds.*.builder",
			"builds.*.number",
			"builds.*.create_time",
			"builds.*.start_time",
			"builds.*.end_time",
			"builds.*.update_time",
			"builds.*.status",
			"builds.*.summary_markdown",
		},
	}
)

// GetBuildPage fetches the full set of information for a Milo build page from Buildbucket.
// Including the blamelist and other auxiliary information.
func GetBuildPage(ctx *router.Context, br buildbucketpb.GetBuildRequest) (*ui.BuildPage, error) {
	c := ctx.Context
	host, err := getHost(c)
	if err != nil {
		return nil, err
	}
	client, err := buildbucketClient(c, host)
	if err != nil {
		return nil, err
	}

	var b *buildbucketpb.Build
	var relatedBuilds []*ui.Build
	var blame []*ui.Commit
	if err = parallel.FanOutIn(func(ch chan<- func() error) {
		ch <- func() (err error) {
			fullbr := br // Copy request
			fullbr.Fields = fullBuildMask
			b, err = client.GetBuild(c, &fullbr)
			return
		}

		// Fetch a small build with just a tiny bit of information.
		// We use this to get the Gitiles tag so that we can fetch
		// related builds and blamelist in parallel.
		smallbr := br // Copy request
		smallbr.Fields = tagsAndGitilesMask
		sb, err := client.GetBuild(c, &smallbr)
		if err != nil {
			return
		}
		ch <- func() (err error) {
			relatedBuilds, err = getRelatedBuilds(c, client, sb)
			return
		}
		ch <- func() error {
			// We wait at most 4s to grab a blamelist.
			nc, cancel := context.WithTimeout(c, 4*time.Second)
			defer cancel()
			blame = getBlame(nc, host, sb)
			return nil
		}
	}); err != nil {
		return nil, err
	}
	link, err := getBugLink(ctx, b)
	logging.Infof(c, "Got all the things")
	return &ui.BuildPage{
		Build:         ui.Build{b},
		Blame:         blame,
		RelatedBuilds: relatedBuilds,
		BuildBugLink:  link,
		Now:           clock.Now(c),
	}, err
}
<commit_msg>[milo] Raise blamelist timeout to 55s<commit_after>// Copyright 2016 The LUCI 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 buildbucket

import (
	"context"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"

	"google.golang.org/genproto/protobuf/field_mask"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"go.chromium.org/gae/service/datastore"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	gitpb "go.chromium.org/luci/common/proto/git"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/milo/api/config"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/milo/frontend/ui"
	"go.chromium.org/luci/server/auth"
	"go.chromium.org/luci/server/router"
)

var ErrNotFound = errors.Reason("Build not found").Tag(common.CodeNotFound).Err()

// BuildAddress constructs the build address of a buildbucketpb.Build.
// This is used as the key for the BuildSummary entity.
func BuildAddress(build *buildbucketpb.Build) string {
	if build == nil {
		return ""
	}
	num := strconv.FormatInt(build.Id, 10)
	if build.Number != 0 {
		num = strconv.FormatInt(int64(build.Number), 10)
	}
	b := build.Builder
	return fmt.Sprintf("luci.%s.%s/%s/%s", b.Project, b.Bucket, b.Builder, num)
}

// simplisticBlamelist returns a slice of ui.Blame for a build.
//
// HACK(iannucci) - Getting the frontend to render a proper blamelist will
// require some significant refactoring. To do this properly, we'll need:
//   * The frontend to get BuildSummary from the backend.
//   * BuildSummary to have a .PreviousBuild() API.
//   * The frontend to obtain the annotation streams itself (so it could see
//     the SourceManifest objects inside of them). Currently getRespBuild defers
//     to swarming's implementation of buildsource.ID.Get(), which only returns
//     the resp object.
func simplisticBlamelist(c context.Context, build *model.BuildSummary) (result []*ui.Commit) {
	var err error
	defer func() {
		if err != nil {
			msg := "Failed to fetch blame information\n"
			logging.WithError(err).Warningf(c, msg)
			if strings.Contains(err.Error(), "DEADLINE_EXCEEDED") || c.Err() != nil {
				// When the context deadline is reached, the error that is returned is
				// a bit racy.
				// Either the actual context error itself is returned, or the urlfetch variant.
				msg += "Gitiles was taking too long"
			} else {
				msg += err.Error()
			}
			result = append(result, &ui.Commit{Description: msg})
		}
	}()
	bs := build.GitilesCommit()
	if bs == nil {
		return
	}

	builds, commits, err := build.PreviousByGitilesCommit(c)
	switch {
	case err == nil || err == model.ErrUnknownPreviousBuild:
		// continue
	case status.Code(err) == codes.PermissionDenied:
		err = common.CodeUnauthorized.Tag().Apply(err)
		return
	default:
		return
	}

	repoURL := bs.RepoURL()
	result = make([]*ui.Commit, 0, len(commits)+1)
	for _, commit := range commits {
		result = append(result, uiCommit(commit, repoURL))
	}
	logging.Infof(c, "Fetched %d commit blamelist from Gitiles", len(result))

	// this means that there were more than 100 commits in-between.
	if len(builds) == 0 && len(commits) > 0 {
		result = append(result, &ui.Commit{
			Description: "<blame list capped at 100 commits>",
			Revision:    &ui.Link{},
			AuthorName:  "<blame list capped at 100 commits>",
		})
	}

	return
}

func uiCommit(commit *gitpb.Commit, repoURL string) *ui.Commit {
	res := &ui.Commit{
		AuthorName:  commit.Author.Name,
		AuthorEmail: commit.Author.Email,
		Repo:        repoURL,
		Description: commit.Message,

		// TODO(iannucci): this use of links is very sloppy; the frontend should
		// know how to render a Commit without having Links embedded in it.
		Revision: ui.NewLink(
			commit.Id,
			repoURL+"/+/"+commit.Id, fmt.Sprintf("commit by %s", commit.Author.Email)),
	}
	res.CommitTime, _ = ptypes.Timestamp(commit.Committer.Time)
	res.File = make([]string, 0, len(commit.TreeDiff))
	for _, td := range commit.TreeDiff {
		// If file is moved, there is both new and old path,
		// from which we take only new path.
		// If a file is deleted, its new path is /dev/null.
		// In that case, we're only interested in the old path.
		switch {
		case td.NewPath != "" && td.NewPath != "/dev/null":
			res.File = append(res.File, td.NewPath)
		case td.OldPath != "":
			res.File = append(res.File, td.OldPath)
		}
	}
	return res
}

// GetBuildSummary fetches a build summary where the Context URI matches the
// given address.
func GetBuildSummary(c context.Context, id int64) (*model.BuildSummary, error) {
	// The host is set to prod because buildbot is hardcoded to talk to prod.
	uri := fmt.Sprintf("buildbucket://cr-buildbucket.appspot.com/build/%d", id)
	bs := make([]*model.BuildSummary, 0, 1)
	q := datastore.NewQuery("BuildSummary").Eq("ContextURI", uri).Limit(1)
	switch err := datastore.GetAll(c, q, &bs); {
	case err != nil:
		return nil, common.ReplaceNSEWith(err.(errors.MultiError), ErrNotFound)
	case len(bs) == 0:
		return nil, ErrNotFound
	default:
		return bs[0], nil
	}
}

// getBlame fetches blame information from Gitiles.
// This requires the BuildSummary to be indexed in Milo.
func getBlame(c context.Context, host string, b *buildbucketpb.Build) []*ui.Commit {
	commit := b.GetInput().GetGitilesCommit()
	// No commit? No blamelist.
	if commit == nil {
		return nil
	}
	// TODO(hinoka): This converts a buildbucketpb.Commit into a string
	// and back into a buildbucketpb.Commit.  That's a bit silly.
	return simplisticBlamelist(c, &model.BuildSummary{
		BuildKey:  MakeBuildKey(c, host, BuildAddress(b)),
		BuildSet:  []string{commit.BuildSetString()},
		BuilderID: BuilderID{BuilderID: *b.Builder}.String(),
	})
}

func buildbucketClient(c context.Context, host string) (buildbucketpb.BuildsClient, error) {
	t, err := auth.GetRPCTransport(c, auth.AsUser)
	if err != nil {
		return nil, err
	}
	return buildbucketpb.NewBuildsPRPCClient(&prpc.Client{
		C:    &http.Client{Transport: t},
		Host: host,
	}), nil
}

// getBugLink attempts to formulate and return the build page bug link
// for the given build.
func getBugLink(c *router.Context, b *buildbucketpb.Build) (string, error) {
	project, err := common.GetProject(c.Context, b.Builder.GetProject())
	if err != nil || proto.Equal(&project.BuildBugTemplate, &config.BugTemplate{}) {
		return "", err
	}

	builderPath := fmt.Sprintf("/p/%s/builders/%s/%s", b.Builder.GetProject(), b.Builder.GetBucket(), b.Builder.GetBuilder())
	buildURL, err := c.Request.URL.Parse(builderPath + "/" + c.Params.ByName("numberOrId"))
	if err != nil {
		return "", errors.Annotate(err, "Unable to make build URL for build bug link.").Err()
	}
	builderURL, err := c.Request.URL.Parse(builderPath)
	if err != nil {
		return "", errors.Annotate(err, "Unable to make builder URL for build bug link.").Err()
	}

	return MakeBuildBugLink(&project.BuildBugTemplate, map[string]interface{}{
		"Build":          b,
		"MiloBuildUrl":   buildURL,
		"MiloBuilderUrl": builderURL,
	})
}

// searchBuildsRequest creates a searchBuildsRequest that looks for a buildset tag.
func searchBuildsRequest(buildset string) *buildbucketpb.SearchBuildsRequest {
	return &buildbucketpb.SearchBuildsRequest{
		Predicate: &buildbucketpb.BuildPredicate{
			Tags: []*buildbucketpb.StringPair{{Key: "buildset", Value: buildset}},
		},
		Fields:   summaryBuildMask,
		PageSize: 1000,
	}
}

// getRelatedBuilds fetches build summaries of builds with the same buildset as b.
func getRelatedBuilds(c context.Context, client buildbucketpb.BuildsClient, b *buildbucketpb.Build) ([]*ui.Build, error) {
	var bs []string
	for _, buildset := range b.Buildsets() {
		// HACK(hinoka): Remove the commit/git/ buildsets because we know they're redundant
		// with the commit/gitiles/ buildsets, and we don't need to ask Buildbucket twice.
		if strings.HasPrefix(buildset, "commit/git/") {
			continue
		}
		bs = append(bs, buildset)
	}
	if len(bs) == 0 {
		// No buildset? No builds.
		return nil, nil
	}

	// Do the search request.
	// Use multiple requests instead of a single batch request.
	// A single large request is CPU bound to a single GAE instance on the buildbucket side.
	// Multiple requests allows the use of multiple GAE instances, therefore more parallelism.
	resps := make([]*buildbucketpb.SearchBuildsResponse, len(bs))
	if err := parallel.FanOutIn(func(ch chan<- func() error) {
		for i, buildset := range bs {
			i := i
			buildset := buildset
			ch <- func() (err error) {
				logging.Debugf(c, "Searching for %s (%d)", buildset, i)
				resps[i], err = client.SearchBuilds(c, searchBuildsRequest(buildset))
				return
			}
		}
	}); err != nil {
		return nil, err
	}

	// Dedupe builds.
	// It's possible since we've made multiple requests that we got back the same builds
	// multiple times.
	seen := map[int64]bool{} // set of build IDs.
	result := []*ui.Build{}
	for _, resp := range resps {
		for _, rb := range resp.GetBuilds() {
			if seen[rb.Id] {
				continue
			}
			seen[rb.Id] = true
			result = append(result, &ui.Build{rb})
		}
	}

	// Sort builds by ID.
	sort.Slice(result, func(i, j int) bool { return result[i].Id < result[j].Id })

	return result, nil
}

var (
	fullBuildMask = &field_mask.FieldMask{
		// TODO(hinoka): Add statusReason here.
		Paths: []string{
			"id",
			"builder",
			"number",
			"created_by",
			"create_time",
			"start_time",
			"end_time",
			"update_time",
			"status",
			"input",
			"output",
			"steps",
			"infra",
			"tags",
		},
	}
	tagsAndGitilesMask = &field_mask.FieldMask{
		Paths: []string{
			"id",
			"number",
			"builder",
			"input.gitiles_commit",
			"tags",
		},
	}
	summaryBuildMask = &field_mask.FieldMask{
		Paths: []string{
			"builds.*.id",
			"builds.*.builder",
			"builds.*.number",
			"builds.*.create_time",
			"builds.*.start_time",
			"builds.*.end_time",
			"builds.*.update_time",
			"builds.*.status",
			"builds.*.summary_markdown",
		},
	}
)

// GetBuildPage fetches the full set of information for a Milo build page from Buildbucket.
// Including the blamelist and other auxiliary information.
func GetBuildPage(ctx *router.Context, br buildbucketpb.GetBuildRequest) (*ui.BuildPage, error) {
	c := ctx.Context
	host, err := getHost(c)
	if err != nil {
		return nil, err
	}
	client, err := buildbucketClient(c, host)
	if err != nil {
		return nil, err
	}

	var b *buildbucketpb.Build
	var relatedBuilds []*ui.Build
	var blame []*ui.Commit
	if err = parallel.FanOutIn(func(ch chan<- func() error) {
		ch <- func() (err error) {
			fullbr := br // Copy request
			fullbr.Fields = fullBuildMask
			b, err = client.GetBuild(c, &fullbr)
			return
		}

		// Fetch a small build with just a tiny bit of information.
		// We use this to get the Gitiles tag so that we can fetch
		// related builds and blamelist in parallel.
		smallbr := br // Copy request
		smallbr.Fields = tagsAndGitilesMask
		sb, err := client.GetBuild(c, &smallbr)
		if err != nil {
			return
		}
		ch <- func() (err error) {
			relatedBuilds, err = getRelatedBuilds(c, client, sb)
			return
		}
		ch <- func() error {
			nc, cancel := context.WithTimeout(c, 55*time.Second)
			defer cancel()
			blame = getBlame(nc, host, sb)
			return nil
		}
	}); err != nil {
		return nil, err
	}
	link, err := getBugLink(ctx, b)
	logging.Infof(c, "Got all the things")
	return &ui.BuildPage{
		Build:         ui.Build{b},
		Blame:         blame,
		RelatedBuilds: relatedBuilds,
		BuildBugLink:  link,
		Now:           clock.Now(c),
	}, err
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes 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 flag

import (
	goflag "flag"
	"strings"

	"github.com/golang/glog"
	"github.com/spf13/pflag"
)

// WordSepNormalizeFunc changes all flags that contain "_" separators
func WordSepNormalizeFunc(f *pflag.FlagSet, name string) pflag.NormalizedName {
	if strings.Contains(name, "_") {
		return pflag.NormalizedName(strings.Replace(name, "_", "-", -1))
	}
	return pflag.NormalizedName(name)
}

// WarnWordSepNormalizeFunc changes and warns for flags that contain "_" separators
func WarnWordSepNormalizeFunc(f *pflag.FlagSet, name string) pflag.NormalizedName {
	if strings.Contains(name, "_") {
		nname := strings.Replace(name, "_", "-", -1)
		glog.Warningf("%s is DEPRECATED and will be removed in a future version. Use %s instead.", name, nname)

		return pflag.NormalizedName(nname)
	}
	return pflag.NormalizedName(name)
}

// InitFlags normalizes, parses, then logs the command line flags
func InitFlags() {
	pflag.CommandLine.SetNormalizeFunc(WordSepNormalizeFunc)
	pflag.CommandLine.AddGoFlagSet(goflag.CommandLine)
	pflag.Parse()
	pflag.VisitAll(func(flag *pflag.Flag) {
		glog.V(4).Infof("FLAG: --%s=%q", flag.Name, flag.Value)
	})
}
<commit_msg>Ensure we log the flag apiserver starts with.<commit_after>/*
Copyright 2014 The Kubernetes 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 flag

import (
	goflag "flag"
	"strings"

	"github.com/golang/glog"
	"github.com/spf13/pflag"
)

// WordSepNormalizeFunc changes all flags that contain "_" separators
func WordSepNormalizeFunc(f *pflag.FlagSet, name string) pflag.NormalizedName {
	if strings.Contains(name, "_") {
		return pflag.NormalizedName(strings.Replace(name, "_", "-", -1))
	}
	return pflag.NormalizedName(name)
}

// WarnWordSepNormalizeFunc changes and warns for flags that contain "_" separators
func WarnWordSepNormalizeFunc(f *pflag.FlagSet, name string) pflag.NormalizedName {
	if strings.Contains(name, "_") {
		nname := strings.Replace(name, "_", "-", -1)
		glog.Warningf("%s is DEPRECATED and will be removed in a future version. Use %s instead.", name, nname)

		return pflag.NormalizedName(nname)
	}
	return pflag.NormalizedName(name)
}

// InitFlags normalizes, parses, then logs the command line flags
func InitFlags() {
	pflag.CommandLine.SetNormalizeFunc(WordSepNormalizeFunc)
	pflag.CommandLine.AddGoFlagSet(goflag.CommandLine)
	pflag.Parse()
	pflag.VisitAll(func(flag *pflag.Flag) {
		glog.V(2).Infof("FLAG: --%s=%q", flag.Name, flag.Value)
	})
}
<|endoftext|>


<commit_before>package event

import (
	"context"
	"fmt"
	"image"
	"sync"
	"testing"
	"time"

	"github.com/kylelemons/godebug/pretty"
	"github.com/mum4k/termdash/keyboard"
	"github.com/mum4k/termdash/terminalapi"
)

// receiver receives events from the distribution system.
type receiver struct {
	mu sync.Mutex

	// events are the received events.
	events []terminalapi.Event
}

// newReceiver returns a new event receiver.
func newReceiver() *receiver {
	return &receiver{}
}

// receive receives an event.
func (r *receiver) receive(ev terminalapi.Event) {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.events = append(r.events, ev)
}

// getEvents returns the received events.
func (r *receiver) getEvents() map[terminalapi.Event]bool {
	r.mu.Lock()
	defer r.mu.Unlock()

	res := map[terminalapi.Event]bool{}
	for _, ev := range r.events {
		res[ev] = true
	}
	return res
}

// waitFor waits until the receiver receives the specified number of events or
// the timeout.
// Returns the received events in an unspecified order.
func (r *receiver) waitFor(want int, timeout time.Duration) (map[terminalapi.Event]bool, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	tick := time.NewTimer(5 * time.Millisecond)
	defer tick.Stop()
	for {
		select {
		case <-tick.C:
			ev := r.getEvents()

			switch got := len(ev); {
			case got > want:
				return nil, fmt.Errorf("got %d events %v, want %d", got, ev, want)

			case got == want:
				return ev, nil
			}

		case <-ctx.Done():
			ev := r.getEvents()
			return nil, fmt.Errorf("while waiting for events, got %d so far: %v, err: %v", len(ev), ev, ctx.Err())
		}
	}
}

// subscriberCase holds test case specifics for one subscriber.
type subscriberCase struct {
	// filter is the subscribers filter.
	filter []terminalapi.Event

	// rec receives the events.
	rec *receiver

	// want are the expected events that should be delivered to this subscriber.
	want map[terminalapi.Event]bool

	// wantErr asserts whether we want an error from waitFor.
	wantErr bool
}

func TestDistributionSystem(t *testing.T) {
	tests := []struct {
		desc string
		// events will be sent down the distribution system.
		events []terminalapi.Event

		// subCase are the event subscribers and their expectations.
		subCase []*subscriberCase
	}{
		{
			desc: "no events and no subscribers",
		},
		{
			desc: "events and no subscribers",
			events: []terminalapi.Event{
				&terminalapi.Mouse{},
				&terminalapi.Keyboard{},
			},
		},
		{
			desc: "single subscriber, wants all events and gets them",
			events: []terminalapi.Event{
				&terminalapi.Keyboard{Key: keyboard.KeyEnter},
				&terminalapi.Mouse{Position: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{2, 2}},
				terminalapi.NewError("error"),
			},
			subCase: []*subscriberCase{
				{
					filter: nil,
					rec:    newReceiver(),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}:   true,
						&terminalapi.Mouse{Position: image.Point{1, 1}}: true,
						&terminalapi.Resize{Size: image.Point{2, 2}}:    true,
						terminalapi.NewError("error"):                   true,
					},
				},
			},
		},
		{
			desc: "single subscriber, filters events",
			events: []terminalapi.Event{
				&terminalapi.Keyboard{Key: keyboard.KeyEnter},
				&terminalapi.Mouse{Position: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{2, 2}},
			},
			subCase: []*subscriberCase{
				{
					filter: []terminalapi.Event{
						&terminalapi.Keyboard{},
						&terminalapi.Mouse{},
					},
					rec: newReceiver(),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}:   true,
						&terminalapi.Mouse{Position: image.Point{1, 1}}: true,
					},
				},
			},
		},
		{
			desc: "single subscriber, errors are always received",
			events: []terminalapi.Event{
				&terminalapi.Keyboard{Key: keyboard.KeyEnter},
				&terminalapi.Mouse{Position: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{2, 2}},
				terminalapi.NewError("error"),
			},
			subCase: []*subscriberCase{
				{
					filter: []terminalapi.Event{
						&terminalapi.Keyboard{},
						&terminalapi.Mouse{},
					},
					rec: newReceiver(),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}:   true,
						&terminalapi.Mouse{Position: image.Point{1, 1}}: true,
						terminalapi.NewError("error"):                   true,
					},
				},
			},
		},
		{
			desc: "multiple subscribers and events",
			events: []terminalapi.Event{
				&terminalapi.Keyboard{Key: keyboard.KeyEnter},
				&terminalapi.Keyboard{Key: keyboard.KeyEsc},
				&terminalapi.Mouse{Position: image.Point{0, 0}},
				&terminalapi.Mouse{Position: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{2, 2}},
				terminalapi.NewError("error1"),
				terminalapi.NewError("error2"),
			},
			subCase: []*subscriberCase{
				{
					filter: []terminalapi.Event{
						&terminalapi.Keyboard{},
					},
					rec: newReceiver(),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}: true,
						&terminalapi.Keyboard{Key: keyboard.KeyEsc}:   true,
						terminalapi.NewError("error1"):                true,
						terminalapi.NewError("error2"):                true,
					},
				},
				{
					filter: []terminalapi.Event{
						&terminalapi.Mouse{},
						&terminalapi.Resize{},
					},
					rec: newReceiver(),
					want: map[terminalapi.Event]bool{
						&terminalapi.Mouse{Position: image.Point{0, 0}}: true,
						&terminalapi.Mouse{Position: image.Point{1, 1}}: true,
						&terminalapi.Resize{Size: image.Point{1, 1}}:    true,
						&terminalapi.Resize{Size: image.Point{2, 2}}:    true,
						terminalapi.NewError("error1"):                  true,
						terminalapi.NewError("error2"):                  true,
					},
				},
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.desc, func(t *testing.T) {
			eds := NewDistributionSystem()
			for _, sc := range tc.subCase {
				stop := eds.Subscribe(sc.filter, sc.rec.receive)
				defer stop()
			}

			for _, ev := range tc.events {
				eds.Event(ev)
			}

			for i, sc := range tc.subCase {
				got, err := sc.rec.waitFor(len(sc.want), 5*time.Second)
				if err != nil {
					t.Fatalf("sc.rec.waitFor[%d] => unexpected error: %v", i, err)
				}

				if diff := pretty.Compare(sc.want, got); diff != "" {
					t.Errorf("sc.rec.waitFor[%d] => unexpected diff (-want, +got):\n%s", i, diff)
				}
			}
		})
	}
}
<commit_msg>Adding a test with misbehaving receiver.<commit_after>package event

import (
	"context"
	"fmt"
	"image"
	"sync"
	"testing"
	"time"

	"github.com/kylelemons/godebug/pretty"
	"github.com/mum4k/termdash/keyboard"
	"github.com/mum4k/termdash/terminalapi"
)

// receiverMode defines how the receiver behaves.
type receiverMode int

const (
	// receiverModeReceive tells the receiver to process the events
	receiverModeReceive receiverMode = iota

	// receiverModeBlock tells the receiver to block on the call to receive.
	receiverModeBlock
)

// receiver receives events from the distribution system.
type receiver struct {
	mu sync.Mutex

	// mode sets how the receiver behaves when receive(0 is called.
	mode receiverMode

	// events are the received events.
	events []terminalapi.Event
}

// newReceiver returns a new event receiver.
func newReceiver(mode receiverMode) *receiver {
	return &receiver{
		mode: mode,
	}
}

// receive receives an event.
func (r *receiver) receive(ev terminalapi.Event) {
	switch r.mode {
	case receiverModeBlock:
		for {
		}
	default:
		r.mu.Lock()
		defer r.mu.Unlock()

		r.events = append(r.events, ev)
	}
}

// getEvents returns the received events.
func (r *receiver) getEvents() map[terminalapi.Event]bool {
	r.mu.Lock()
	defer r.mu.Unlock()

	res := map[terminalapi.Event]bool{}
	for _, ev := range r.events {
		res[ev] = true
	}
	return res
}

// waitFor waits until the receiver receives the specified number of events or
// the timeout.
// Returns the received events in an unspecified order.
func (r *receiver) waitFor(want int, timeout time.Duration) (map[terminalapi.Event]bool, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	tick := time.NewTimer(5 * time.Millisecond)
	defer tick.Stop()
	for {
		select {
		case <-tick.C:
			ev := r.getEvents()

			switch got := len(ev); {
			case got > want:
				return nil, fmt.Errorf("got %d events %v, want %d", got, ev, want)

			case got == want:
				return ev, nil
			}

		case <-ctx.Done():
			ev := r.getEvents()
			return nil, fmt.Errorf("while waiting for events, got %d so far: %v, err: %v", len(ev), ev, ctx.Err())
		}
	}
}

// subscriberCase holds test case specifics for one subscriber.
type subscriberCase struct {
	// filter is the subscribers filter.
	filter []terminalapi.Event

	// rec receives the events.
	rec *receiver

	// want are the expected events that should be delivered to this subscriber.
	want map[terminalapi.Event]bool

	// wantErr asserts whether we want an error from waitFor.
	wantErr bool
}

func TestDistributionSystem(t *testing.T) {
	tests := []struct {
		desc string
		// events will be sent down the distribution system.
		events []terminalapi.Event

		// subCase are the event subscribers and their expectations.
		subCase []*subscriberCase
	}{
		{
			desc: "no events and no subscribers",
		},
		{
			desc: "events and no subscribers",
			events: []terminalapi.Event{
				&terminalapi.Mouse{},
				&terminalapi.Keyboard{},
			},
		},
		{
			desc: "single subscriber, wants all events and gets them",
			events: []terminalapi.Event{
				&terminalapi.Keyboard{Key: keyboard.KeyEnter},
				&terminalapi.Mouse{Position: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{2, 2}},
				terminalapi.NewError("error"),
			},
			subCase: []*subscriberCase{
				{
					filter: nil,
					rec:    newReceiver(receiverModeReceive),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}:   true,
						&terminalapi.Mouse{Position: image.Point{1, 1}}: true,
						&terminalapi.Resize{Size: image.Point{2, 2}}:    true,
						terminalapi.NewError("error"):                   true,
					},
				},
			},
		},
		{
			desc: "single subscriber, filters events",
			events: []terminalapi.Event{
				&terminalapi.Keyboard{Key: keyboard.KeyEnter},
				&terminalapi.Mouse{Position: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{2, 2}},
			},
			subCase: []*subscriberCase{
				{
					filter: []terminalapi.Event{
						&terminalapi.Keyboard{},
						&terminalapi.Mouse{},
					},
					rec: newReceiver(receiverModeReceive),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}:   true,
						&terminalapi.Mouse{Position: image.Point{1, 1}}: true,
					},
				},
			},
		},
		{
			desc: "single subscriber, errors are always received",
			events: []terminalapi.Event{
				&terminalapi.Keyboard{Key: keyboard.KeyEnter},
				&terminalapi.Mouse{Position: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{2, 2}},
				terminalapi.NewError("error"),
			},
			subCase: []*subscriberCase{
				{
					filter: []terminalapi.Event{
						&terminalapi.Keyboard{},
						&terminalapi.Mouse{},
					},
					rec: newReceiver(receiverModeReceive),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}:   true,
						&terminalapi.Mouse{Position: image.Point{1, 1}}: true,
						terminalapi.NewError("error"):                   true,
					},
				},
			},
		},
		{
			desc: "multiple subscribers and events",
			events: []terminalapi.Event{
				&terminalapi.Keyboard{Key: keyboard.KeyEnter},
				&terminalapi.Keyboard{Key: keyboard.KeyEsc},
				&terminalapi.Mouse{Position: image.Point{0, 0}},
				&terminalapi.Mouse{Position: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{1, 1}},
				&terminalapi.Resize{Size: image.Point{2, 2}},
				terminalapi.NewError("error1"),
				terminalapi.NewError("error2"),
			},
			subCase: []*subscriberCase{
				{
					filter: []terminalapi.Event{
						&terminalapi.Keyboard{},
					},
					rec: newReceiver(receiverModeReceive),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}: true,
						&terminalapi.Keyboard{Key: keyboard.KeyEsc}:   true,
						terminalapi.NewError("error1"):                true,
						terminalapi.NewError("error2"):                true,
					},
				},
				{
					filter: []terminalapi.Event{
						&terminalapi.Mouse{},
						&terminalapi.Resize{},
					},
					rec: newReceiver(receiverModeReceive),
					want: map[terminalapi.Event]bool{
						&terminalapi.Mouse{Position: image.Point{0, 0}}: true,
						&terminalapi.Mouse{Position: image.Point{1, 1}}: true,
						&terminalapi.Resize{Size: image.Point{1, 1}}:    true,
						&terminalapi.Resize{Size: image.Point{2, 2}}:    true,
						terminalapi.NewError("error1"):                  true,
						terminalapi.NewError("error2"):                  true,
					},
				},
			},
		},
		{
			desc: "a misbehaving receiver only blocks itself",
			events: []terminalapi.Event{
				&terminalapi.Keyboard{Key: keyboard.KeyEnter},
				&terminalapi.Keyboard{Key: keyboard.KeyEsc},
				terminalapi.NewError("error1"),
				terminalapi.NewError("error2"),
			},
			subCase: []*subscriberCase{
				{
					filter: []terminalapi.Event{
						&terminalapi.Keyboard{},
					},
					rec: newReceiver(receiverModeReceive),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}: true,
						&terminalapi.Keyboard{Key: keyboard.KeyEsc}:   true,
						terminalapi.NewError("error1"):                true,
						terminalapi.NewError("error2"):                true,
					},
				},
				{
					filter: []terminalapi.Event{
						&terminalapi.Keyboard{},
					},
					rec: newReceiver(receiverModeBlock),
					want: map[terminalapi.Event]bool{
						&terminalapi.Keyboard{Key: keyboard.KeyEnter}: true,
						&terminalapi.Keyboard{Key: keyboard.KeyEsc}:   true,
						terminalapi.NewError("error1"):                true,
						terminalapi.NewError("error2"):                true,
					},
					wantErr: true,
				},
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.desc, func(t *testing.T) {
			eds := NewDistributionSystem()
			for _, sc := range tc.subCase {
				stop := eds.Subscribe(sc.filter, sc.rec.receive)
				defer stop()
			}

			for _, ev := range tc.events {
				eds.Event(ev)
			}

			for i, sc := range tc.subCase {
				got, err := sc.rec.waitFor(len(sc.want), 5*time.Second)
				if (err != nil) != sc.wantErr {
					t.Errorf("sc.rec.waitFor[%d] => unexpected error: %v, wantErr:%v", i, err, sc.wantErr)
				}
				if err != nil {
					continue
				}

				if diff := pretty.Compare(sc.want, got); diff != "" {
					t.Errorf("sc.rec.waitFor[%d] => unexpected diff (-want, +got):\n%s", i, diff)
				}
			}
		})
	}
}
<|endoftext|>


<commit_before>// Command deploy-agent tracks single configuration against registry, updates
// local state as required and runs deploy script.
package main

import (
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/artyom/autoflags"
	"github.com/artyom/logger"
	"github.com/artyom/untar"
	"github.com/pkg/errors"
	"github.com/pkg/sftp"
)

func main() {
	args := &mainArgs{
		State:  "state.json",
		Key:    "id_ecdsa",
		Addr:   "localhost:2022",
		Dir:    ".",
		Script: "./deploy.sh",
	}
	autoflags.Define(args)
	flag.Parse()
	log := log.New(os.Stderr, "", log.LstdFlags)
	if err := run(args, log); err != nil {
		log.Fatal(err)
	}
}

type mainArgs struct {
	State    string `flag:"state,file to save state to"`
	Name     string `flag:"name,configuration to track"`
	Key      string `flag:"key,ssh private key to use"`
	Addr     string `flag:"addr,registry address (host:port)"`
	Fp       string `flag:"fp,registry server key fingerprint"`
	Dir      string `flag:"dir,directory to store downloaded and unpacked files"`
	Script   string `flag:"script,script to run on deploys"`
	CleanOld bool   `flag:"cleanold,try to remove previous state unpacked files after switching state"`
	Verbose  bool   `flag:"v,be more chatty about what's happening"`
}

func run(args *mainArgs, log logger.Interface) error {
	if args.Name == "" {
		return errors.New("empty configuration name")
	}
	if args.State == "" {
		return errors.New("no state file set")
	}
	cfg, err := newClientConfig(args.Key, args.Fp)
	if err != nil {
		return err
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go callOnSignal(func(s os.Signal) { log.Println(s); cancel() }, syscall.SIGINT, syscall.SIGTERM)
	for {
		if err := cycle(ctx, args, cfg, log); err != nil {
			log.Println(err)
		}
		select {
		case <-time.After(15 * time.Second):
		case <-ctx.Done():
			return nil
		}
	}
}

func cycle(ctx context.Context, args *mainArgs, cfg *ssh.ClientConfig, log logger.Interface) error {
	state, err := readLocalConfiguration(args.State)
	switch {
	case err == nil:
	case os.IsNotExist(err):
		state = &configuration{}
	default:
		return err
	}
	client, err := ssh.Dial("tcp", args.Addr, cfg)
	if err != nil {
		return err
	}
	defer client.Close()
	sconn, err := sftp.NewClient(client)
	if err != nil {
		return err
	}
	defer sconn.Close()
	newState, err := readRemoteConfiguration(path.Join("configs", args.Name+".json"), sconn)
	if err != nil {
		return err
	}
	if len(newState.Hash) != 64 {
		return errors.Errorf("invalid state hash value: %q, expecting sha256 sum", newState.Hash)
	}
	if newState.Hash == state.Hash {
		return nil
	}
	if args.Verbose {
		log.Println("configuration update available:", newState.Hash)
		for i, l := range newState.Layers {
			log.Printf("\tcomponent %d: %q, version: %q", i+1, l.Name, l.Version)
		}
	}
	if err := downloadMissing(ctx, sconn, filepath.Join(args.Dir, cacheDir), newState.Layers); err != nil {
		return err
	}
	sconn.Close()
	if err := unpackLayers(ctx, filepath.Join(args.Dir, unpackedDir, filepath.Base(newState.Hash)),
		filepath.Join(args.Dir, cacheDir), newState.Layers); err != nil {
		return err
	}
	if args.Script != "" {
		if err := runDeploy(ctx, args.Script, filepath.Join(args.Dir, unpackedDir), state, newState); err != nil {
			return err
		}
	}
	if err := saveState(args.State, 0644, newState); err != nil {
		return err
	}
	if args.Verbose {
		log.Println("switched to new configuration:", newState.Hash)
	}
	if err := cleanCache(filepath.Join(args.Dir, cacheDir), newState.Layers, time.Hour); err != nil {
		log.Println("error cleaning cache directory:", err)
	}
	if args.CleanOld && state.Hash != "" {
		oldDir := filepath.Join(args.Dir, unpackedDir, filepath.Base(state.Hash))
		if err := os.RemoveAll(oldDir); err != nil {
			log.Println("failed to remove old state directory:", err)
		}
		_ = os.Remove(oldDir + ".ok")
	}
	return nil
}

func cleanCache(dir string, layers []layer, age time.Duration) error {
	keep := make(map[string]struct{}, len(layers))
	for _, l := range layers {
		keep[l.Hash] = struct{}{}
	}
	fn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return errors.WithStack(err)
		}
		if _, ok := keep[filepath.Base(path)]; ok || !info.Mode().IsRegular() {
			return nil
		}
		if strings.HasPrefix(filepath.Base(path), tempDlPrefix) {
			return errors.WithStack(os.Remove(path))
		}
		if time.Now().Sub(info.ModTime()) < age {
			return nil
		}
		return errors.WithStack(os.Remove(path))
	}
	return filepath.Walk(dir, fn)
}

func runDeploy(ctx context.Context, script, dir string, oldState, newState *configuration) error {
	stateFile, err := saveStateTemp("", 0644, newState)
	if err != nil {
		return err
	}
	defer os.Remove(stateFile)
	cmd := exec.CommandContext(ctx, script)
	cmd.Env = []string{
		"OLDID=" + oldState.Hash,
		"OLDROOT=" + filepath.Join(dir, oldState.Hash),
		"NEWID=" + newState.Hash,
		"NEWROOT=" + filepath.Join(dir, newState.Hash),
		"STATEFILE=" + stateFile,
	}
	for _, name := range []string{"PATH", "HOME", "TMPDIR", "USER", "LOGNAME", "LANG", "LC_ALL", "SHELL"} {
		if val, ok := os.LookupEnv(name); ok {
			cmd.Env = append(cmd.Env, name+"="+val)
		}
	}
	for i, l := range newState.Layers {
		cmd.Env = append(cmd.Env,
			fmt.Sprintf("LAYER_%d_NAME=%s", i, l.Name),
			fmt.Sprintf("LAYER_%d_VERSION=%s", i, l.Version),
		)
	}
	return errors.Wrapf(cmd.Run(), "%q call failure", script)
}

func saveStateTemp(dir string, mode os.FileMode, state *configuration) (string, error) {
	f, err := ioutil.TempFile(dir, "json-state-")
	if err != nil {
		return "", err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	enc.SetIndent("", "\t")
	if err := enc.Encode(state); err != nil {
		os.Remove(f.Name())
		return "", err
	}
	if err := f.Chmod(mode); err != nil {
		os.Remove(f.Name())
		return "", err
	}
	return f.Name(), f.Close()
}

func saveState(dst string, mode os.FileMode, state *configuration) error {
	f, err := ioutil.TempFile(filepath.Dir(dst), "temp-state-")
	if err != nil {
		return err
	}
	defer f.Close()
	defer os.Remove(f.Name())
	enc := json.NewEncoder(f)
	enc.SetIndent("", "\t")
	if err := enc.Encode(state); err != nil {
		return err
	}
	if err := f.Chmod(mode); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return os.Rename(f.Name(), dst)
}

func unpackLayers(ctx context.Context, dst, filesDir string, layers []layer) error {
	touchfile := dst + ".ok"
	if _, err := os.Stat(touchfile); err == nil {
		if _, err := os.Stat(dst); err == nil {
			return err
		}
	}
	if err := os.MkdirAll(dst, 0750); err != nil {
		return err
	}
	for _, l := range layers {
		if err := ctx.Err(); err != nil {
			return err
		}
		if len(l.Hash) != 64 {
			return errors.Errorf("invalid length of layer hash: %v", l)
		}
		srcName := filepath.Join(filesDir, filepath.Base(l.Hash))
		if err := untarTo(srcName, dst); err != nil {
			return err
		}
	}
	_ = touch(touchfile)
	return nil
}

func untarTo(name, dir string) error {
	f, err := os.Open(name)
	if err != nil {
		return err
	}
	defer f.Close()
	gr, err := gzip.NewReader(f)
	if err != nil {
		return err
	}
	defer gr.Close()
	defer syscall.Umask(syscall.Umask(0))
	return untar.Untar(gr, dir)
}

func downloadMissing(ctx context.Context, sconn *sftp.Client, dir string, layers []layer) error {
	if err := os.MkdirAll(dir, 0750); err != nil {
		return err
	}
	for _, l := range layers {
		if err := ctx.Err(); err != nil {
			return err
		}
		if len(l.Hash) != 64 {
			return errors.Errorf("invalid length of layer hash: %v", l)
		}
		dstName := filepath.Join(dir, filepath.Base(l.Hash))
		if _, err := os.Stat(dstName); err == nil {
			continue
		}
		if err := downloadFile(sconn, dstName, l.File, l.Hash); err != nil {
			return err
		}
	}
	return nil
}

func downloadFile(sconn *sftp.Client, dst, src, contentHash string) error {
	dstFile, err := ioutil.TempFile(filepath.Dir(dst), tempDlPrefix)
	if err != nil {
		return err
	}
	defer dstFile.Close()
	defer os.Remove(dstFile.Name())
	srcFile, err := sconn.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	h := sha256.New()
	tr := io.TeeReader(srcFile, h)
	if _, err := io.Copy(dstFile, tr); err != nil {
		return err
	}
	if contentHash != fmt.Sprintf("%x", h.Sum(nil)) {
		return errors.Errorf("downloaded file hash mismatch: %v", src)
	}
	if err := dstFile.Close(); err != nil {
		return err
	}
	return os.Rename(dstFile.Name(), dst)
}

func readRemoteConfiguration(name string, sconn *sftp.Client) (*configuration, error) {
	f, err := sconn.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return decodeConfiguration(f)
}

func readLocalConfiguration(name string) (*configuration, error) {
	f, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return decodeConfiguration(f)
}

func decodeConfiguration(r io.Reader) (*configuration, error) {
	var cfg configuration
	if err := json.NewDecoder(r).Decode(&cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}

func newClientConfig(keyfile, fingerprint string) (*ssh.ClientConfig, error) {
	if !strings.HasPrefix(fingerprint, "SHA256:") {
		return nil, errors.New("invalid server key fingerprint, expecting SHA256 fingerprint")
	}
	key, err := readKey(keyfile)
	if err != nil {
		return nil, err
	}
	config := &ssh.ClientConfig{
		User: "deploy-agent",
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			if hostFp := ssh.FingerprintSHA256(key); hostFp != fingerprint {
				return errors.Errorf("host key fingerprint mismatch: %v", hostFp)
			}
			return nil
		},
		Auth:    []ssh.AuthMethod{ssh.PublicKeys(key)},
		Timeout: 30 * time.Second,
	}
	return config, nil
}

func readKey(name string) (ssh.Signer, error) {
	privateBytes, err := ioutil.ReadFile(name)
	if err != nil {
		return nil, err
	}
	return ssh.ParsePrivateKey(privateBytes)
}

func touch(name string) error {
	f, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE, 0640)
	if err != nil {
		return err
	}
	return f.Close()
}

func callOnSignal(fn func(os.Signal), sig ...os.Signal) {
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, sig...)
	defer signal.Stop(sigCh)
	fn(<-sigCh)
}

type configuration struct {
	Name   string
	Mtime  time.Time
	Hash   string
	Layers []layer
}

type layer struct {
	Name    string
	Version string
	File    string
	Hash    string
	Ctime   time.Time
}

const (
	cacheDir     = "cache"
	unpackedDir  = "unpacked"
	tempDlPrefix = "temp-"
)
<commit_msg>deploy-agent: explicitly return nil when skipping unpack step<commit_after>// Command deploy-agent tracks single configuration against registry, updates
// local state as required and runs deploy script.
package main

import (
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/artyom/autoflags"
	"github.com/artyom/logger"
	"github.com/artyom/untar"
	"github.com/pkg/errors"
	"github.com/pkg/sftp"
)

func main() {
	args := &mainArgs{
		State:  "state.json",
		Key:    "id_ecdsa",
		Addr:   "localhost:2022",
		Dir:    ".",
		Script: "./deploy.sh",
	}
	autoflags.Define(args)
	flag.Parse()
	log := log.New(os.Stderr, "", log.LstdFlags)
	if err := run(args, log); err != nil {
		log.Fatal(err)
	}
}

type mainArgs struct {
	State    string `flag:"state,file to save state to"`
	Name     string `flag:"name,configuration to track"`
	Key      string `flag:"key,ssh private key to use"`
	Addr     string `flag:"addr,registry address (host:port)"`
	Fp       string `flag:"fp,registry server key fingerprint"`
	Dir      string `flag:"dir,directory to store downloaded and unpacked files"`
	Script   string `flag:"script,script to run on deploys"`
	CleanOld bool   `flag:"cleanold,try to remove previous state unpacked files after switching state"`
	Verbose  bool   `flag:"v,be more chatty about what's happening"`
}

func run(args *mainArgs, log logger.Interface) error {
	if args.Name == "" {
		return errors.New("empty configuration name")
	}
	if args.State == "" {
		return errors.New("no state file set")
	}
	cfg, err := newClientConfig(args.Key, args.Fp)
	if err != nil {
		return err
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go callOnSignal(func(s os.Signal) { log.Println(s); cancel() }, syscall.SIGINT, syscall.SIGTERM)
	for {
		if err := cycle(ctx, args, cfg, log); err != nil {
			log.Println(err)
		}
		select {
		case <-time.After(15 * time.Second):
		case <-ctx.Done():
			return nil
		}
	}
}

func cycle(ctx context.Context, args *mainArgs, cfg *ssh.ClientConfig, log logger.Interface) error {
	state, err := readLocalConfiguration(args.State)
	switch {
	case err == nil:
	case os.IsNotExist(err):
		state = &configuration{}
	default:
		return err
	}
	client, err := ssh.Dial("tcp", args.Addr, cfg)
	if err != nil {
		return err
	}
	defer client.Close()
	sconn, err := sftp.NewClient(client)
	if err != nil {
		return err
	}
	defer sconn.Close()
	newState, err := readRemoteConfiguration(path.Join("configs", args.Name+".json"), sconn)
	if err != nil {
		return err
	}
	if len(newState.Hash) != 64 {
		return errors.Errorf("invalid state hash value: %q, expecting sha256 sum", newState.Hash)
	}
	if newState.Hash == state.Hash {
		return nil
	}
	if args.Verbose {
		log.Println("configuration update available:", newState.Hash)
		for i, l := range newState.Layers {
			log.Printf("\tcomponent %d: %q, version: %q", i+1, l.Name, l.Version)
		}
	}
	if err := downloadMissing(ctx, sconn, filepath.Join(args.Dir, cacheDir), newState.Layers); err != nil {
		return err
	}
	sconn.Close()
	if err := unpackLayers(ctx, filepath.Join(args.Dir, unpackedDir, filepath.Base(newState.Hash)),
		filepath.Join(args.Dir, cacheDir), newState.Layers); err != nil {
		return err
	}
	if args.Script != "" {
		if err := runDeploy(ctx, args.Script, filepath.Join(args.Dir, unpackedDir), state, newState); err != nil {
			return err
		}
	}
	if err := saveState(args.State, 0644, newState); err != nil {
		return err
	}
	if args.Verbose {
		log.Println("switched to new configuration:", newState.Hash)
	}
	if err := cleanCache(filepath.Join(args.Dir, cacheDir), newState.Layers, time.Hour); err != nil {
		log.Println("error cleaning cache directory:", err)
	}
	if args.CleanOld && state.Hash != "" {
		oldDir := filepath.Join(args.Dir, unpackedDir, filepath.Base(state.Hash))
		if err := os.RemoveAll(oldDir); err != nil {
			log.Println("failed to remove old state directory:", err)
		}
		_ = os.Remove(oldDir + ".ok")
	}
	return nil
}

func cleanCache(dir string, layers []layer, age time.Duration) error {
	keep := make(map[string]struct{}, len(layers))
	for _, l := range layers {
		keep[l.Hash] = struct{}{}
	}
	fn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return errors.WithStack(err)
		}
		if _, ok := keep[filepath.Base(path)]; ok || !info.Mode().IsRegular() {
			return nil
		}
		if strings.HasPrefix(filepath.Base(path), tempDlPrefix) {
			return errors.WithStack(os.Remove(path))
		}
		if time.Now().Sub(info.ModTime()) < age {
			return nil
		}
		return errors.WithStack(os.Remove(path))
	}
	return filepath.Walk(dir, fn)
}

func runDeploy(ctx context.Context, script, dir string, oldState, newState *configuration) error {
	stateFile, err := saveStateTemp("", 0644, newState)
	if err != nil {
		return err
	}
	defer os.Remove(stateFile)
	cmd := exec.CommandContext(ctx, script)
	cmd.Env = []string{
		"OLDID=" + oldState.Hash,
		"OLDROOT=" + filepath.Join(dir, oldState.Hash),
		"NEWID=" + newState.Hash,
		"NEWROOT=" + filepath.Join(dir, newState.Hash),
		"STATEFILE=" + stateFile,
	}
	for _, name := range []string{"PATH", "HOME", "TMPDIR", "USER", "LOGNAME", "LANG", "LC_ALL", "SHELL"} {
		if val, ok := os.LookupEnv(name); ok {
			cmd.Env = append(cmd.Env, name+"="+val)
		}
	}
	for i, l := range newState.Layers {
		cmd.Env = append(cmd.Env,
			fmt.Sprintf("LAYER_%d_NAME=%s", i, l.Name),
			fmt.Sprintf("LAYER_%d_VERSION=%s", i, l.Version),
		)
	}
	return errors.Wrapf(cmd.Run(), "%q call failure", script)
}

func saveStateTemp(dir string, mode os.FileMode, state *configuration) (string, error) {
	f, err := ioutil.TempFile(dir, "json-state-")
	if err != nil {
		return "", err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	enc.SetIndent("", "\t")
	if err := enc.Encode(state); err != nil {
		os.Remove(f.Name())
		return "", err
	}
	if err := f.Chmod(mode); err != nil {
		os.Remove(f.Name())
		return "", err
	}
	return f.Name(), f.Close()
}

func saveState(dst string, mode os.FileMode, state *configuration) error {
	f, err := ioutil.TempFile(filepath.Dir(dst), "temp-state-")
	if err != nil {
		return err
	}
	defer f.Close()
	defer os.Remove(f.Name())
	enc := json.NewEncoder(f)
	enc.SetIndent("", "\t")
	if err := enc.Encode(state); err != nil {
		return err
	}
	if err := f.Chmod(mode); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return os.Rename(f.Name(), dst)
}

func unpackLayers(ctx context.Context, dst, filesDir string, layers []layer) error {
	touchfile := dst + ".ok"
	if _, err := os.Stat(touchfile); err == nil {
		if _, err := os.Stat(dst); err == nil {
			return nil
		}
	}
	if err := os.MkdirAll(dst, 0750); err != nil {
		return err
	}
	for _, l := range layers {
		if err := ctx.Err(); err != nil {
			return err
		}
		if len(l.Hash) != 64 {
			return errors.Errorf("invalid length of layer hash: %v", l)
		}
		srcName := filepath.Join(filesDir, filepath.Base(l.Hash))
		if err := untarTo(srcName, dst); err != nil {
			return err
		}
	}
	_ = touch(touchfile)
	return nil
}

func untarTo(name, dir string) error {
	f, err := os.Open(name)
	if err != nil {
		return err
	}
	defer f.Close()
	gr, err := gzip.NewReader(f)
	if err != nil {
		return err
	}
	defer gr.Close()
	defer syscall.Umask(syscall.Umask(0))
	return untar.Untar(gr, dir)
}

func downloadMissing(ctx context.Context, sconn *sftp.Client, dir string, layers []layer) error {
	if err := os.MkdirAll(dir, 0750); err != nil {
		return err
	}
	for _, l := range layers {
		if err := ctx.Err(); err != nil {
			return err
		}
		if len(l.Hash) != 64 {
			return errors.Errorf("invalid length of layer hash: %v", l)
		}
		dstName := filepath.Join(dir, filepath.Base(l.Hash))
		if _, err := os.Stat(dstName); err == nil {
			continue
		}
		if err := downloadFile(sconn, dstName, l.File, l.Hash); err != nil {
			return err
		}
	}
	return nil
}

func downloadFile(sconn *sftp.Client, dst, src, contentHash string) error {
	dstFile, err := ioutil.TempFile(filepath.Dir(dst), tempDlPrefix)
	if err != nil {
		return err
	}
	defer dstFile.Close()
	defer os.Remove(dstFile.Name())
	srcFile, err := sconn.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	h := sha256.New()
	tr := io.TeeReader(srcFile, h)
	if _, err := io.Copy(dstFile, tr); err != nil {
		return err
	}
	if contentHash != fmt.Sprintf("%x", h.Sum(nil)) {
		return errors.Errorf("downloaded file hash mismatch: %v", src)
	}
	if err := dstFile.Close(); err != nil {
		return err
	}
	return os.Rename(dstFile.Name(), dst)
}

func readRemoteConfiguration(name string, sconn *sftp.Client) (*configuration, error) {
	f, err := sconn.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return decodeConfiguration(f)
}

func readLocalConfiguration(name string) (*configuration, error) {
	f, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return decodeConfiguration(f)
}

func decodeConfiguration(r io.Reader) (*configuration, error) {
	var cfg configuration
	if err := json.NewDecoder(r).Decode(&cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}

func newClientConfig(keyfile, fingerprint string) (*ssh.ClientConfig, error) {
	if !strings.HasPrefix(fingerprint, "SHA256:") {
		return nil, errors.New("invalid server key fingerprint, expecting SHA256 fingerprint")
	}
	key, err := readKey(keyfile)
	if err != nil {
		return nil, err
	}
	config := &ssh.ClientConfig{
		User: "deploy-agent",
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			if hostFp := ssh.FingerprintSHA256(key); hostFp != fingerprint {
				return errors.Errorf("host key fingerprint mismatch: %v", hostFp)
			}
			return nil
		},
		Auth:    []ssh.AuthMethod{ssh.PublicKeys(key)},
		Timeout: 30 * time.Second,
	}
	return config, nil
}

func readKey(name string) (ssh.Signer, error) {
	privateBytes, err := ioutil.ReadFile(name)
	if err != nil {
		return nil, err
	}
	return ssh.ParsePrivateKey(privateBytes)
}

func touch(name string) error {
	f, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE, 0640)
	if err != nil {
		return err
	}
	return f.Close()
}

func callOnSignal(fn func(os.Signal), sig ...os.Signal) {
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, sig...)
	defer signal.Stop(sigCh)
	fn(<-sigCh)
}

type configuration struct {
	Name   string
	Mtime  time.Time
	Hash   string
	Layers []layer
}

type layer struct {
	Name    string
	Version string
	File    string
	Hash    string
	Ctime   time.Time
}

const (
	cacheDir     = "cache"
	unpackedDir  = "unpacked"
	tempDlPrefix = "temp-"
)
<|endoftext|>


<commit_before>package main

import "testing"

func Test_commandConn_run(t *testing.T) {
	ctx := newContext()

	cmd, err := newCommandConn(nil, ctx)
	if err != nil {
		t.Errorf("err => %q, want => nil", err)
	}

	cmd.run()
}
<commit_msg>Update cmd/gmq-cli/command_conn_test.go<commit_after>package main

import "testing"

const testAddress = "iot.eclipse.org:1883"

func Test_commandConn_run_err(t *testing.T) {
	ctx := newContext()

	cmd, err := newCommandConn(nil, ctx)
	if err != nil {
		t.Errorf("err => %q, want => nil", err)
	}

	if err := cmd.run(); err == nil {
		t.Error("err => nil, want => not nil")
	}
}

func Test_commandConn_run(t *testing.T) {
	ctx := newContext()

	cmd, err := newCommandConn(nil, ctx)
	if err != nil {
		t.Errorf("err => %q, want => nil", err)
	}

	cmd.network = "tcp"
	cmd.address = testAddress

	if err := cmd.run(); err != nil {
		t.Error("err => %q, want => nil", err)
	}
}

func Test_commandConn_waitCONNACK_connack(t *testing.T) {
	ctx := newContext()

	cmd, err := newCommandConn(nil, ctx)
	if err != nil {
		t.Errorf("err => %q, want => nil", err)
	}

	cmd.ctx.wg.Add(1)
	go cmd.waitCONNACK()

	cmd.ctx.connack <- struct{}{}

	cmd.ctx.wg.Wait()
}

func Test_commandConn_waitCONNACK_timeout(t *testing.T) {
	ctx := newContext()

	cmd, err := newCommandConn(nil, ctx)
	if err != nil {
		t.Errorf("err => %q, want => nil", err)
	}

	cmd.connackTimeout = 1

	cmd.ctx.wg.Add(1)
	go cmd.waitCONNACK()

	cmd.ctx.wg.Wait()
}

func Test_commandConn_waitCONNACK_timeout_disconnDefault(t *testing.T) {
	ctx := newContext()

	cmd, err := newCommandConn(nil, ctx)
	if err != nil {
		t.Errorf("err => %q, want => nil", err)
	}

	cmd.connackTimeout = 1

	cmd.ctx.disconn <- struct{}{}

	cmd.ctx.wg.Add(1)
	go cmd.waitCONNACK()

	cmd.ctx.wg.Wait()
}

func Test_commandConn_waitCONNACK_connackEnd(t *testing.T) {
	ctx := newContext()

	cmd, err := newCommandConn(nil, ctx)
	if err != nil {
		t.Errorf("err => %q, want => nil", err)
	}

	cmd.ctx.wg.Add(1)
	go cmd.waitCONNACK()

	cmd.ctx.connackEnd <- struct{}{}

	cmd.ctx.wg.Wait()
}

func Test_commandCon_receive_ReceiveErr_disconnecting(t *testing.T) {
	ctx := newContext()

	cmd, err := newCommandConn(nil, ctx)
	if err != nil {
		t.Errorf("err => %q, want => nil", err)
	}

	cmd.ctx.disconnecting = true

	cmd.ctx.wg.Add(1)
	go cmd.receive()

	cmd.ctx.wg.Wait()
}

func Test_commandCon_receive_ReceiveErr(t *testing.T) {
	ctx := newContext()

	cmd, err := newCommandConn(nil, ctx)
	if err != nil {
		t.Errorf("err => %q, want => nil", err)
	}

	cmd.ctx.wg.Add(1)
	go cmd.receive()

	cmd.ctx.wg.Wait()
}
<|endoftext|>


<commit_before>package main

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

	"labix.org/v2/mgo"

	"github.com/stretchr/testify/assert"

	bsonScanner "github.com/Clever/oplog-replay/bson"
)

// simpleDocStruct is for inserting simple Mongo documents
type simpleDocStruct struct {
	key string
}

func TestComposingWithOplogReplay(t *testing.T) {
	// This test assumes that we're connecting to Mongo replica set. Otherwise, the oplog won't
	// be generated.

	unixTime := int(time.Now().Unix())

	time.Sleep(time.Duration(2) * time.Second)

	session, err := mgo.Dial("localhost")
	assert.Nil(t, err)
	db := session.DB("myTestDb")
	c := db.C("myCollection")
	assert.Nil(t, c.Insert(&simpleDocStruct{key: "key"}))

	time.Sleep(time.Duration(2) * time.Second)

	assert.Nil(t, c.Insert(&simpleDocStruct{key: "key2"}))
	// Dump at the time we started operations. Should get both operations
	dumpAtTime(t, unixTime, 2)
	// Three seconds later we should get only one.
	dumpAtTime(t, unixTime+3, 1)
}

func dumpAtTime(t *testing.T, unixTime, expectedResults int) {
	tempDir, err := ioutil.TempDir("/tmp", "oplogDumpTest")
	assert.Nil(t, err)
	defer os.RemoveAll(tempDir)
	assert.Nil(t, runDump(tempDir, "localhost", unixTime))

	file, err := os.Open(tempDir + "/local/oplog.rs.bson")
	assert.Nil(t, err)
	scanner := bsonScanner.New(file)
	count := 0
	for scanner.Scan() {
		count = count + 1
	}
	assert.Equal(t, expectedResults, count)
}

func TestCopyBsonFile(t *testing.T) {
	tempDir, err := ioutil.TempDir("/tmp", "oplogDumpTest")
	assert.Nil(t, err)
	defer os.RemoveAll(tempDir)

	// Create a directory structure that mirrors the oplog one the code expects
	err = os.Mkdir(tempDir+"/local", 0744)
	assert.Nil(t, err)
	file, err := os.Create(tempDir + "/local/oplog.rs.bson")
	assert.Nil(t, err)
	assert.Nil(t, ioutil.WriteFile(file.Name(), []byte("test-bson-file"), 0644))

	// Create a file to copy to
	toFile, err := ioutil.TempFile(tempDir, "bsonCopyTest")
	assert.Nil(t, err)
	assert.Nil(t, copyBsonFile(tempDir, toFile.Name()))

	// Check that the data matches
	fileData, err := ioutil.ReadFile(toFile.Name())
	assert.Nil(t, err)
	assert.Equal(t, "test-bson-file", string(fileData))
}
<commit_msg>Fix race condition in test<commit_after>package main

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

	"labix.org/v2/mgo"

	"github.com/stretchr/testify/assert"

	bsonScanner "github.com/Clever/oplog-replay/bson"
)

// simpleDocStruct is for inserting simple Mongo documents
type simpleDocStruct struct {
	key string
}

func TestComposingWithOplogReplay(t *testing.T) {
	// This test assumes that we're connecting to Mongo replica set. Otherwise, the oplog won't
	// be generated.

	unixTime := int(time.Now().Unix())

	time.Sleep(time.Duration(1) * time.Second)

	session, err := mgo.Dial("localhost")
	assert.Nil(t, err)
	db := session.DB("myTestDb")
	c := db.C("myCollection")
	assert.Nil(t, c.Insert(&simpleDocStruct{key: "key"}))

	time.Sleep(time.Duration(3) * time.Second)

	assert.Nil(t, c.Insert(&simpleDocStruct{key: "key2"}))
	// Dump at the time we started operations. Should get both operations
	dumpAtTime(t, unixTime, 2)
	// Three seconds later we should get only one.
	dumpAtTime(t, unixTime+3, 1)
}

func dumpAtTime(t *testing.T, unixTime, expectedResults int) {
	tempDir, err := ioutil.TempDir("/tmp", "oplogDumpTest")
	assert.Nil(t, err)
	defer os.RemoveAll(tempDir)
	assert.Nil(t, runDump(tempDir, "localhost", unixTime))

	file, err := os.Open(tempDir + "/local/oplog.rs.bson")
	assert.Nil(t, err)
	scanner := bsonScanner.New(file)
	count := 0
	for scanner.Scan() {
		count = count + 1
	}
	assert.Equal(t, expectedResults, count)
}

func TestCopyBsonFile(t *testing.T) {
	tempDir, err := ioutil.TempDir("/tmp", "oplogDumpTest")
	assert.Nil(t, err)
	defer os.RemoveAll(tempDir)

	// Create a directory structure that mirrors the oplog one the code expects
	err = os.Mkdir(tempDir+"/local", 0744)
	assert.Nil(t, err)
	file, err := os.Create(tempDir + "/local/oplog.rs.bson")
	assert.Nil(t, err)
	assert.Nil(t, ioutil.WriteFile(file.Name(), []byte("test-bson-file"), 0644))

	// Create a file to copy to
	toFile, err := ioutil.TempFile(tempDir, "bsonCopyTest")
	assert.Nil(t, err)
	assert.Nil(t, copyBsonFile(tempDir, toFile.Name()))

	// Check that the data matches
	fileData, err := ioutil.ReadFile(toFile.Name())
	assert.Nil(t, err)
	assert.Equal(t, "test-bson-file", string(fileData))
}
<|endoftext|>