<commit_before>package system

import (
  "io"
  "io/ioutil"
  "errors"
  "fmt"
  "os"
  "net/http"
  "crypto/tls"
  "crypto/md5"
  "runtime"

  "github.com/codegangsta/cli"
  "github.com/Sirupsen/logrus"
  "github.com/kardianos/osext"
)

const binaryURL = "https://artifactory1.eng.lancope.local/generic-dev-infrastructure/lc/lc-%s-%s-%s"

// CmdUpgrade will upgrade the current lc binary
func CmdUpgrade(c *cli.Context) error {
  version := c.String("version")
  if len(version) == 0 {
    return errors.New("upgrade command requires a version argument, none found")
  }

  platform := runtime.GOOS
  arch := runtime.GOARCH
  url := fmt.Sprintf(binaryURL, platform, arch, version)
  logrus.Debugf("using url: %s", url)

  // find location of lc currently running
  oldLc, err := getLcLocation()
  if err != nil {
    logrus.Errorf("could not find location of current lc")
    return err
  }

  // hash current binary for comparison with new binary
  oldMd5, err := computeMd5(oldLc)
  if err != nil {
    logrus.Debugf("could not compute md5 for old lc binary")
  }

  // rename current binary in preparation for replacing
  tmpDir, tmpLocation, err := mvLc(oldLc)
  if err != nil {
    return err
  }
  defer os.Remove(tmpDir)

  // do the upgrade
  if err := installNew(url, oldLc); err != nil {
    logrus.Errorf("failed upgrading lc, error: %q", err)
    if swapError := swap(tmpLocation, oldLc); swapError != nil {
      logrus.Errorf("failed replacing your lc, your old binary is located at %q, err: %q", tmpLocation, swapError)
    }
    return err
  }

  if newMd5, err := computeMd5(oldLc); err != nil {
    logrus.Debugf("could not compute md5 for new lc binary, not comparing them")
  } else {
    if oldMd5 != newMd5 {
      logrus.Infof("lc install finished, new lc binary installed")
    } else {
      logrus.Infof("lc install finished, lc binary was already the latest")
    }
  }
  return nil
}

// move src file into a temp location
// returns:
//  * temporary directory that should be deleted after the upgrade finishes
//  * filePath location of temporary location
func mvLc(src string) (string, string, error) {
  tmpDir, err := ioutil.TempDir("", "lcupgrade")
  if err != nil {
    logrus.Debugf("failed creating temp dir ", err)
    return "", "", err
  }
  tmpLocation := fmt.Sprintf("%s/%s", tmpDir, "lc.old")
  logrus.Debugf("moving binary '%s' to '%s'", src, tmpLocation)
  if err := swap(src, tmpLocation); err != nil {
      logrus.Debugf("failed moving binary ", err)
      return "", "", err
  }
  return tmpDir, tmpLocation, nil
}

func swap(src string, dst string) error {
  if err := os.Rename(src, dst); err != nil {
    logrus.Debugf("failed swapping '%s' to '%s'", src, dst, err)
    return err
  }
  return nil
}

func computeMd5(filePath string) (string, error) {
  file, err := os.Open(filePath)
  if err != nil {
    logrus.Debugf("could not open file at '%'", filePath, err)
    return "", err
  }
  defer file.Close()

  md5 := md5.New()
  if _, err := io.Copy(md5, file); err != nil {
    logrus.Debugf("could copy file at '%'", filePath, err)
    return "", err
  }

  var result []byte
  return string(md5.Sum(result)), nil
}

func getLcLocation() (string, error){
  // NOTE: if os.Args[0] is a symlink, this code will update the actual binary, not the link
  lcPath, err := osext.Executable()
  if err != nil {
    logrus.Debugf("lc not found", err)
    return "", err
  }
  return lcPath, nil
}

func installNew(url string, target string) error {
  tr := &http.Transport{
      TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
  }
  client := &http.Client{Transport: tr}

  resp, err := client.Get(url)
  if err != nil {
    logrus.Debugf("failed downloading binary", err)
  	return err
  }
  if resp.StatusCode != 200 {
    return fmt.Errorf("failed downloading binary, invalid http response: %d", resp.StatusCode)
  }
  defer resp.Body.Close()

  out, err := os.Create(target)
  if err != nil {
    logrus.Debugf("failed creating new lc file", err)
    return err
  }
  defer out.Close()

  if n, err := io.Copy(out, resp.Body); err != nil {
    logrus.Debugf("failed copying new lc file", err)
    return err
  } else {
    logrus.Debugf("successfully coppied %d bytes", n)
  }

  if err := os.Chmod(target, os.FileMode(0755)); err != nil {
    logrus.Debugf("failed making lc executable", err)
    return err
  }
  return nil
}
<commit_msg>download binary before replacing old binar during upgrade<commit_after>package system

import (
  "io"
  "io/ioutil"
  "errors"
  "fmt"
  "os"
  "net/http"
  "crypto/tls"
  "crypto/md5"
  "runtime"

  "github.com/codegangsta/cli"
  "github.com/Sirupsen/logrus"
  "github.com/kardianos/osext"
)

const tmpDirPrefix = "lcupgrade"
const binaryURL = "https://artifactory1.eng.lancope.local/generic-dev-infrastructure/lc/lc-%s-%s-%s"

// CmdUpgrade will upgrade the current lc binary
func CmdUpgrade(c *cli.Context) error {
  version := c.String("version")
  if len(version) == 0 {
    return errors.New("upgrade command requires a version argument, none found")
  }

  platform := runtime.GOOS
  arch := runtime.GOARCH
  url := fmt.Sprintf(binaryURL, platform, arch, version)
  logrus.Debugf("using url: %s", url)

  // find location of lc currently running
  lcPath, err := getLcLocation()
  if err != nil {
    logrus.Errorf("could not find location of current lc")
    return err
  }

  // hash current binary for comparison with new binary
  oldMd5, err := computeMd5(lcPath)
  if err != nil {
    logrus.Debugf("could not compute md5 for old lc binary")
  }

  //download new binary to staging location
  newTmpDir, newLcTmp, err := downloadNew(url)
  if err != nil {
    return err
  }
  defer os.Remove(newTmpDir)

  // rename current binary in preparation for replacing
  tmpDir, oldLcTmp, err := mvLc(lcPath)
  if err != nil {
    return err
  }
  defer os.Remove(tmpDir)

  //swap in new lc
  if err := swap(newLcTmp, lcPath); err != nil {
    logrus.Debugf("failed swaping new lc from %q to %q, err: %q", newLcTmp, lcPath, err)
    return fmt.Errorf("failed replacing your lc, your old binary is located at %q", oldLcTmp)
  }

  if newMd5, err := computeMd5(lcPath); err != nil {
    logrus.Debugf("could not compute md5 for new lc binary, not comparing them")
  } else {
    if oldMd5 != newMd5 {
      logrus.Infof("lc install finished, new lc binary installed")
    } else {
      logrus.Infof("lc install finished, lc binary was already the latest")
    }
  }
  return nil
}

// move src file into a temp location
// returns:
//  * temporary directory that should be deleted after the upgrade finishes
//  * filePath location of temporary location
func mvLc(src string) (string, string, error) {
  tmpDir, err := ioutil.TempDir("", tmpDirPrefix)
  if err != nil {
    logrus.Debugf("failed creating temp dir ", err)
    return "", "", err
  }
  tmpLocation := fmt.Sprintf("%s/%s", tmpDir, "lc.old")
  logrus.Debugf("moving binary '%s' to '%s'", src, tmpLocation)
  if err := swap(src, tmpLocation); err != nil {
      logrus.Debugf("failed moving binary ", err)
      return "", "", err
  }
  return tmpDir, tmpLocation, nil
}

// swap will rename the src file to the dst file
func swap(src string, dst string) error {
  if err := os.Rename(src, dst); err != nil {
    logrus.Debugf("failed swapping '%s' to '%s'", src, dst, err)
    return err
  }
  return nil
}

func computeMd5(filePath string) (string, error) {
  file, err := os.Open(filePath)
  if err != nil {
    logrus.Debugf("could not open file at '%'", filePath, err)
    return "", err
  }
  defer file.Close()

  md5 := md5.New()
  if _, err := io.Copy(md5, file); err != nil {
    logrus.Debugf("could copy file at '%'", filePath, err)
    return "", err
  }

  var result []byte
  return string(md5.Sum(result)), nil
}

func getLcLocation() (string, error){
  // NOTE: if os.Args[0] is a symlink, this code will update the actual binary, not the link
  lcPath, err := osext.Executable()
  if err != nil {
    logrus.Debugf("lc not found", err)
    return "", err
  }
  return lcPath, nil
}

// Will download the new binary from the given url to a temp location
// returns (dir holding binary, full path of binary)
func downloadNew(url string) (string, string, error) {
  tmpDir, err := ioutil.TempDir("", tmpDirPrefix)
  if err != nil {
    logrus.Debugf("failed creating temp dir ", err)
    return "", "", err
  }
  tmpLocation := fmt.Sprintf("%s/%s", tmpDir, "lc.new")
  logrus.Debugf("downloading new binary to '%s'", tmpLocation)
  if err := installNew(url, tmpLocation); err != nil {
    logrus.Debugf("failed downloading binary, err: %q", err)
    return "", "", err
  }
  return tmpDir, tmpLocation, nil
}

func installNew(url string, target string) error {
  tr := &http.Transport{
      TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
  }
  client := &http.Client{Transport: tr}

  resp, err := client.Get(url)
  if err != nil {
    logrus.Debugf("failed downloading binary", err)
  	return err
  }
  if resp.StatusCode != 200 {
    return fmt.Errorf("failed downloading binary, invalid http response: %d", resp.StatusCode)
  }
  defer resp.Body.Close()

  out, err := os.Create(target)
  if err != nil {
    logrus.Debugf("failed creating new lc file", err)
    return err
  }
  defer out.Close()

  if n, err := io.Copy(out, resp.Body); err != nil {
    logrus.Debugf("failed copying new lc file", err)
    return err
  } else {
    logrus.Debugf("successfully coppied %d bytes", n)
  }

  if err := os.Chmod(target, os.FileMode(0755)); err != nil {
    logrus.Debugf("failed making lc executable", err)
    return err
  }
  return nil
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"net"
	"os/exec"
	"sort"
	"strings"

	"github.com/hashicorp/nomad/api"
	"github.com/ryanuber/columnize"
)

type ServerMembersCommand struct {
	Meta
	Cmd *exec.Cmd
}

func (c *ServerMembersCommand) Help() string {
	helpText := `
Usage: maya omm-status [options]

  Display a list of the known servers and their status. Only Nomad servers are
  able to service this command.

General Options:

  ` + generalOptionsUsage() + `

Server Members Options:

  -detailed
    Show detailed information about each member. This dumps
    a raw set of tags which shows more information than the
    default output format.
`
	return strings.TrimSpace(helpText)
}

func (c *ServerMembersCommand) Synopsis() string {
	return "Display a list of known servers and their status"
}

func (c *ServerMembersCommand) Run(args []string) int {
	var detailed bool

	flags := c.Meta.FlagSet("server-members", FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.BoolVar(&detailed, "detailed", false, "Show detailed output")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	// Check for extra arguments
	args = flags.Args()
	if len(args) != 0 {
		c.Ui.Error(c.Help())
		return 1
	}

	// Get the HTTP client
	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	// Query the members
	srvMembers, err := client.Agent().Members()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error querying servers: %s", err))
		return 1
	}

	if srvMembers == nil {
		c.Ui.Error("Agent doesn't know about server members")
		return 0
	}

	// Sort the members
	sort.Sort(api.AgentMembersNameSort(srvMembers.Members))

	// Determine the leaders per region.
	leaders, err := regionLeaders(client, srvMembers.Members)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error determining leaders: %s", err))
		return 1
	}

	// Format the list
	var out []string
	if detailed {
		out = detailedOutput(srvMembers.Members)
	} else {
		out = standardOutput(srvMembers.Members, leaders)
	}

	// Dump the list
	c.Ui.Output(columnize.SimpleFormat(out))
	var runop int
	if runop = c.mserverstatus(); runop != 0 {
		return runop
	}
	return 0
}

func standardOutput(mem []*api.AgentMember, leaders map[string]string) []string {
	// Format the members list
	members := make([]string, len(mem)+1)
	members[0] = "Name|Address|Port|Status|Leader|Protocol|Build|Datacenter|Region"
	for i, member := range mem {
		reg := member.Tags["region"]
		regLeader, ok := leaders[reg]
		isLeader := false
		if ok {
			if regLeader == net.JoinHostPort(member.Addr, member.Tags["port"]) {

				isLeader = true
			}
		}

		members[i+1] = fmt.Sprintf("%s|%s|%d|%s|%t|%d|%s|%s|%s",
			member.Name,
			member.Addr,
			member.Port,
			member.Status,
			isLeader,
			member.ProtocolCur,
			member.Tags["build"],
			member.Tags["dc"],
			member.Tags["region"])
	}
	return members
}

func detailedOutput(mem []*api.AgentMember) []string {
	// Format the members list
	members := make([]string, len(mem)+1)
	members[0] = "Name|Address|Port|Tags"
	for i, member := range mem {
		// Format the tags
		tagPairs := make([]string, 0, len(member.Tags))
		for k, v := range member.Tags {
			tagPairs = append(tagPairs, fmt.Sprintf("%s=%s", k, v))
		}
		tags := strings.Join(tagPairs, ",")

		members[i+1] = fmt.Sprintf("%s|%s|%d|%s",
			member.Name,
			member.Addr,
			member.Port,
			tags)
	}
	return members
}

// regionLeaders returns a map of regions to the IP of the member that is the
// leader.
func regionLeaders(client *api.Client, mem []*api.AgentMember) (map[string]string, error) {
	// Determine the unique regions.
	leaders := make(map[string]string)
	regions := make(map[string]struct{})
	for _, m := range mem {
		regions[m.Tags["region"]] = struct{}{}
	}

	if len(regions) == 0 {
		return leaders, nil
	}

	status := client.Status()
	for reg := range regions {
		l, err := status.RegionLeader(reg)
		if err != nil {
			// This error means that region has no leader.
			if strings.Contains(err.Error(), "No cluster leader") {
				continue
			}
			return nil, err
		}

		leaders[reg] = l
	}

	return leaders, nil
}

// to get the status of mayaserver deamon,
// TODO proper CLI command once mayaserver have it's own
func (c *ServerMembersCommand) mserverstatus() int {
	var runop int = 0

	c.Cmd = exec.Command("systemctl", "status", "mayaserver")

	if runop := execute(c.Cmd, c.Ui); runop != 0 {
		c.Ui.Error("mayaserver not running")
	}

	return runop

}
<commit_msg>omm-status to show m-apiserver status<commit_after>package command

import (
	"fmt"
	"net"
	"os/exec"
	"sort"
	"strings"

	"github.com/hashicorp/nomad/api"
	"github.com/ryanuber/columnize"
)

type ServerMembersCommand struct {
	Meta
	Cmd *exec.Cmd
}

func (c *ServerMembersCommand) Help() string {
	helpText := `
Usage: maya omm-status [options]

  Display a list of the known servers and their status. Only Nomad servers are
  able to service this command.

General Options:

  ` + generalOptionsUsage() + `

Server Members Options:

  -detailed
    Show detailed information about each member. This dumps
    a raw set of tags which shows more information than the
    default output format.
`
	return strings.TrimSpace(helpText)
}

func (c *ServerMembersCommand) Synopsis() string {
	return "Display a list of known servers and their status"
}

func (c *ServerMembersCommand) Run(args []string) int {
	var detailed bool

	flags := c.Meta.FlagSet("server-members", FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.BoolVar(&detailed, "detailed", false, "Show detailed output")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	// Check for extra arguments
	args = flags.Args()
	if len(args) != 0 {
		c.Ui.Error(c.Help())
		return 1
	}

	// Get the HTTP client
	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	// Query the members
	srvMembers, err := client.Agent().Members()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error querying servers: %s", err))
		return 1
	}

	if srvMembers == nil {
		c.Ui.Error("Agent doesn't know about server members")
		return 0
	}

	// Sort the members
	sort.Sort(api.AgentMembersNameSort(srvMembers.Members))

	// Determine the leaders per region.
	leaders, err := regionLeaders(client, srvMembers.Members)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error determining leaders: %s", err))
		return 1
	}

	// Format the list
	var out []string
	if detailed {
		out = detailedOutput(srvMembers.Members)
	} else {
		out = standardOutput(srvMembers.Members, leaders)
	}

	// Dump the list
	c.Ui.Output(columnize.SimpleFormat(out))
	var runop int
	if runop = c.mserverstatus(); runop != 0 {
		return runop
	}
	return 0
}

func standardOutput(mem []*api.AgentMember, leaders map[string]string) []string {
	// Format the members list
	members := make([]string, len(mem)+1)
	members[0] = "Name|Address|Port|Status|Leader|Protocol|Build|Datacenter|Region"
	for i, member := range mem {
		reg := member.Tags["region"]
		regLeader, ok := leaders[reg]
		isLeader := false
		if ok {
			if regLeader == net.JoinHostPort(member.Addr, member.Tags["port"]) {

				isLeader = true
			}
		}

		members[i+1] = fmt.Sprintf("%s|%s|%d|%s|%t|%d|%s|%s|%s",
			member.Name,
			member.Addr,
			member.Port,
			member.Status,
			isLeader,
			member.ProtocolCur,
			member.Tags["build"],
			member.Tags["dc"],
			member.Tags["region"])
	}
	return members
}

func detailedOutput(mem []*api.AgentMember) []string {
	// Format the members list
	members := make([]string, len(mem)+1)
	members[0] = "Name|Address|Port|Tags"
	for i, member := range mem {
		// Format the tags
		tagPairs := make([]string, 0, len(member.Tags))
		for k, v := range member.Tags {
			tagPairs = append(tagPairs, fmt.Sprintf("%s=%s", k, v))
		}
		tags := strings.Join(tagPairs, ",")

		members[i+1] = fmt.Sprintf("%s|%s|%d|%s",
			member.Name,
			member.Addr,
			member.Port,
			tags)
	}
	return members
}

// regionLeaders returns a map of regions to the IP of the member that is the
// leader.
func regionLeaders(client *api.Client, mem []*api.AgentMember) (map[string]string, error) {
	// Determine the unique regions.
	leaders := make(map[string]string)
	regions := make(map[string]struct{})
	for _, m := range mem {
		regions[m.Tags["region"]] = struct{}{}
	}

	if len(regions) == 0 {
		return leaders, nil
	}

	status := client.Status()
	for reg := range regions {
		l, err := status.RegionLeader(reg)
		if err != nil {
			// This error means that region has no leader.
			if strings.Contains(err.Error(), "No cluster leader") {
				continue
			}
			return nil, err
		}

		leaders[reg] = l
	}

	return leaders, nil
}

// to get the status of mayaserver deamon,
// TODO proper CLI command once mayaserver have it's own
func (c *ServerMembersCommand) mserverstatus() int {
	var runop int = 0

	c.Cmd = exec.Command("systemctl", "status", "m-apiserver")

	if runop := execute(c.Cmd, c.Ui); runop != 0 {
		c.Ui.Error("m-apiserver not running")
	}

	return runop

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/url"
	"os"

	functions "github.com/funcy/functions_go"
	"github.com/urfave/cli"
)

// Version of Functions CLI
var Version = "0.3.28"

func version() cli.Command {
	r := versionCmd{VersionApi: functions.NewVersionApi()}
	return cli.Command{
		Name:   "version",
		Usage:  "displays fn and functions daemon versions",
		Action: r.version,
	}
}

type versionCmd struct {
	*functions.VersionApi
}

func (r *versionCmd) version(c *cli.Context) error {
	apiURL := os.Getenv("API_URL")
	if apiURL == "" {
		apiURL = "http://localhost:8080"
	}

	u, err := url.Parse(apiURL)
	if err != nil {
		return err
	}
	r.Configuration.BasePath = u.String()

	fmt.Println("Client version:", Version)
	v, _, err := r.VersionGet()
	if err != nil {
		return err
	}
	fmt.Println("Server version", v.Version)
	return nil
}
<commit_msg>fn tool: 0.3.29 release [skip ci]<commit_after>package main

import (
	"fmt"
	"net/url"
	"os"

	functions "github.com/funcy/functions_go"
	"github.com/urfave/cli"
)

// Version of Functions CLI
var Version = "0.3.29"

func version() cli.Command {
	r := versionCmd{VersionApi: functions.NewVersionApi()}
	return cli.Command{
		Name:   "version",
		Usage:  "displays fn and functions daemon versions",
		Action: r.version,
	}
}

type versionCmd struct {
	*functions.VersionApi
}

func (r *versionCmd) version(c *cli.Context) error {
	apiURL := os.Getenv("API_URL")
	if apiURL == "" {
		apiURL = "http://localhost:8080"
	}

	u, err := url.Parse(apiURL)
	if err != nil {
		return err
	}
	r.Configuration.BasePath = u.String()

	fmt.Println("Client version:", Version)
	v, _, err := r.VersionGet()
	if err != nil {
		return err
	}
	fmt.Println("Server version", v.Version)
	return nil
}
<|endoftext|>


<commit_before>package chain

// TODO
<commit_msg>Added tests (100% coverage).<commit_after>package chain_test

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/codemodus/chain"
	"golang.org/x/net/context"
)

var (
	bTxt0   = []byte("0")
	bTxt1   = []byte("1")
	bTxtA   = []byte("A")
	bTxtEnd = []byte("_END_")
)

func wrap0(n chain.Handler) chain.Handler {
	return chain.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		w.Write(bTxt0)
		n.ServeHTTPContext(ctx, w, r)
		w.Write(bTxt0)
	})
}

func wrap1(n chain.Handler) chain.Handler {
	return chain.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		w.Write(bTxt1)
		n.ServeHTTPContext(ctx, w, r)
		w.Write(bTxt1)
	})
}

func stdWrapA(n http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write(bTxtA)
		n.ServeHTTP(w, r)
		w.Write(bTxtA)
	})
}

func emptyWrap0(n chain.Handler) chain.Handler {
	return chain.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		n.ServeHTTPContext(ctx, w, r)
	})
}

func end0(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	w.Write(bTxtEnd)
	return
}

func TestChain(t *testing.T) {
	c0 := chain.New(context.Background(), wrap0)
	c1 := c0.Append(wrap1, chain.Meld(stdWrapA))
	m := http.NewServeMux()
	r0 := "/0"
	r1 := "/1"
	m.Handle(r0, c0.EndFn(end0))
	m.Handle(r1, c1.EndFn(end0))
	s := httptest.NewServer(m)

	re0, err := http.Get(s.URL + r0)
	if err != nil {
		t.Error(err)
	}
	defer re0.Body.Close()
	rb0, err := ioutil.ReadAll(re0.Body)
	if err != nil {
		t.Error(err)
	}

	re1, err := http.Get(s.URL + r1)
	if err != nil {
		t.Error(err)
	}
	defer re1.Body.Close()
	rb1, err := ioutil.ReadAll(re1.Body)
	if err != nil {
		t.Error(err)
	}

	bb := &bytes.Buffer{}
	bb.Write(bTxt0)
	bb.Write(bTxtEnd)
	bb.Write(bTxt0)
	bb.Write(bTxt0)
	bb.Write(bTxt1)
	bb.Write(bTxtA)
	bb.Write(bTxtEnd)
	bb.Write(bTxtA)
	bb.Write(bTxt1)
	bb.Write(bTxt0)

	want := string(bb.Bytes())
	got := string(rb0) + string(rb1)
	if got != want {
		t.Errorf("Body = %v, want %v", got, want)
	}
}

func TestNilEnd(t *testing.T) {
	c0 := chain.New(context.Background(), emptyWrap0)
	m := http.NewServeMux()
	r0 := "/0"
	r1 := "/1"
	m.Handle(r0, c0.End(nil))
	m.Handle(r1, c0.EndFn(nil))
	s := httptest.NewServer(m)

	re0, err := http.Get(s.URL + r0)
	if err != nil {
		t.Error(err)
	}
	defer re0.Body.Close()
	rs0 := re0.StatusCode

	want := http.StatusNoContent
	got := rs0
	if got != want {
		t.Errorf("Status Code = %v, want %v", got, want)
	}

	re1, err := http.Get(s.URL + r1)
	if err != nil {
		t.Error(err)
	}
	defer re1.Body.Close()
	rs1 := re1.StatusCode

	got = rs1
	if got != want {
		t.Errorf("Status Code = %v, want %v", got, want)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"

	contracts "github.com/estafette/estafette-ci-contracts"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"

	"github.com/rs/zerolog/log"
	"github.com/sethgrid/pester"
)

// EndOfLifeHelper has methods to shutdown the runner after a fatal or successful run
type EndOfLifeHelper interface {
	handleFatal(context.Context, contracts.BuildLog, error, string)
	sendBuildFinishedEvent(ctx context.Context, buildStatus string) error
	sendBuildCleanEvent(ctx context.Context, buildStatus string) error
	sendBuildJobLogEvent(ctx context.Context, buildLog contracts.BuildLog) error
}

type endOfLifeHelperImpl struct {
	runAsJob bool
	config   contracts.BuilderConfig
}

// NewEndOfLifeHelper returns a new EndOfLifeHelper
func NewEndOfLifeHelper(runAsJob bool, config contracts.BuilderConfig) EndOfLifeHelper {
	return &endOfLifeHelperImpl{
		runAsJob: runAsJob,
		config:   config,
	}
}

func (elh *endOfLifeHelperImpl) handleFatal(ctx context.Context, buildLog contracts.BuildLog, err error, message string) {

	// add error messages as step to show in logs
	fatalStep := contracts.BuildLogStep{
		Step:     "init",
		LogLines: []contracts.BuildLogLine{},
		ExitCode: -1,
		Status:   "FAILED",
	}
	lineNumber := 1

	if err != nil {
		fatalStep.LogLines = append(fatalStep.LogLines, contracts.BuildLogLine{
			LineNumber: lineNumber,
			Timestamp:  time.Now().UTC(),
			StreamType: "stderr",
			Text:       err.Error(),
		})
		lineNumber++
	}
	if message != "" {
		fatalStep.LogLines = append(fatalStep.LogLines, contracts.BuildLogLine{
			LineNumber: lineNumber,
			Timestamp:  time.Now().UTC(),
			StreamType: "stderr",
			Text:       message,
		})
		lineNumber++
	}

	buildLog.Steps = append(buildLog.Steps, fatalStep)

	_ = elh.sendBuildFinishedEvent(ctx, "failed")
	_ = elh.sendBuildJobLogEvent(ctx, buildLog)
	_ = elh.sendBuildCleanEvent(ctx, "failed")

	if elh.runAsJob {
		log.Error().Err(err).Msg(message)
		os.Exit(0)
	} else {
		log.Fatal().Err(err).Msg(message)
	}
}

func (elh *endOfLifeHelperImpl) sendBuildJobLogEvent(ctx context.Context, buildLog contracts.BuildLog) (err error) {

	err = elh.sendBuildJobLogEventCore(ctx, buildLog)

	if err == nil {
		return
	}

	// strip log lines from successful steps to reduce size of the logs and still keep the useful information
	slimBuildLog := buildLog
	slimBuildLog.Steps = []contracts.BuildLogStep{}
	for _, s := range buildLog.Steps {
		slimBuildLogStep := s
		if s.Status == "SUCCEEDED" {
			if len(s.LogLines) > 0 {
				slimBuildLogStep.LogLines = []contracts.BuildLogLine{
					contracts.BuildLogLine{
						LineNumber: s.LogLines[0].LineNumber,
						Timestamp:  s.LogLines[0].Timestamp,
						StreamType: "stdout",
						Text:       "Truncated logs for reducing total log size; to prevent this use less verbose logging",
					},
				}
			}
		}

		slimBuildLog.Steps = append(slimBuildLog.Steps, slimBuildLogStep)
	}

	return elh.sendBuildJobLogEventCore(ctx, slimBuildLog)
}

func (elh *endOfLifeHelperImpl) sendBuildJobLogEventCore(ctx context.Context, buildLog contracts.BuildLog) (err error) {

	span, ctx := opentracing.StartSpanFromContext(ctx, "SendLog")
	defer span.Finish()

	ciServerBuilderPostLogsURL := elh.config.CIServer.PostLogsURL
	ciAPIKey := elh.config.CIServer.APIKey
	jobName := *elh.config.JobName

	if ciServerBuilderPostLogsURL != "" && ciAPIKey != "" && jobName != "" {

		// convert BuildJobLogs to json
		var requestBody io.Reader

		var data []byte
		if *elh.config.Action == "release" {
			// copy buildLog to releaseLog and marshal that
			releaseLog := contracts.ReleaseLog{
				ID:         buildLog.ID,
				RepoSource: buildLog.RepoSource,
				RepoOwner:  buildLog.RepoOwner,
				RepoName:   buildLog.RepoName,
				ReleaseID:  strconv.Itoa(elh.config.ReleaseParams.ReleaseID),
				Steps:      buildLog.Steps,
				InsertedAt: buildLog.InsertedAt,
			}
			data, err = json.Marshal(releaseLog)
			if err != nil {
				log.Error().Err(err).Msgf("Failed marshalling ReleaseLog for job %v", jobName)
				return
			}
		} else {
			data, err = json.Marshal(buildLog)
			if err != nil {
				log.Error().Err(err).Msgf("Failed marshalling BuildLog for job %v", jobName)
				return
			}
		}

		requestBody = bytes.NewReader(data)

		// create client, in order to add headers
		client := pester.New()
		client.MaxRetries = 5
		client.Backoff = pester.ExponentialJitterBackoff
		client.KeepLog = true
		client.Timeout = time.Second * 120
		request, err := http.NewRequest("POST", ciServerBuilderPostLogsURL, requestBody)
		if err != nil {
			log.Error().Err(err).Msgf("Failed creating http client for job %v", jobName)
			return err
		}

		// add tracing context
		ext.SpanKindRPCClient.Set(span)
		ext.HTTPUrl.Set(span, request.URL.String())
		ext.HTTPMethod.Set(span, request.Method)
		span.Tracer().Inject(
			span.Context(),
			opentracing.HTTPHeaders,
			opentracing.HTTPHeadersCarrier(request.Header),
		)

		// add headers
		request.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ciAPIKey))
		request.Header.Add("Content-Type", "application/json")

		// perform actual request
		response, err := client.Do(request)
		if err != nil {
			log.Error().Err(err).Str("logs", client.LogString()).Msgf("Failed shipping logs to %v for job %v", ciServerBuilderPostLogsURL, jobName)
			return err
		}

		ext.HTTPStatusCode.Set(span, uint16(response.StatusCode))

		defer response.Body.Close()

		log.Debug().Str("logs", client.LogString()).Msgf("Successfully shipped logs to %v for job %v", ciServerBuilderPostLogsURL, jobName)
	}

	return nil
}

func (elh *endOfLifeHelperImpl) sendBuildFinishedEvent(ctx context.Context, buildStatus string) error {
	return elh.sendBuilderEvent(ctx, buildStatus, fmt.Sprintf("builder:%v", buildStatus))
}

func (elh *endOfLifeHelperImpl) sendBuildCleanEvent(ctx context.Context, buildStatus string) error {
	return elh.sendBuilderEvent(ctx, buildStatus, "builder:clean")
}

func (elh *endOfLifeHelperImpl) sendBuilderEvent(ctx context.Context, buildStatus, event string) (err error) {

	span, ctx := opentracing.StartSpanFromContext(ctx, "SendBuildStatus")
	defer span.Finish()
	span.SetTag("build-status", buildStatus)
	span.SetTag("event-type", event)

	ciServerBuilderEventsURL := elh.config.CIServer.BuilderEventsURL
	ciAPIKey := elh.config.CIServer.APIKey
	jobName := *elh.config.JobName

	if ciServerBuilderEventsURL != "" && ciAPIKey != "" && jobName != "" {
		// convert EstafetteCiBuilderEvent to json
		var requestBody io.Reader

		buildID := ""
		if elh.config.BuildParams != nil {
			buildID = strconv.Itoa(elh.config.BuildParams.BuildID)
		}

		releaseID := ""
		if elh.config.ReleaseParams != nil {
			releaseID = strconv.Itoa(elh.config.ReleaseParams.ReleaseID)
		}

		ciBuilderEvent := EstafetteCiBuilderEvent{
			JobName:      jobName,
			RepoSource:   elh.config.Git.RepoSource,
			RepoOwner:    elh.config.Git.RepoOwner,
			RepoName:     elh.config.Git.RepoName,
			RepoBranch:   elh.config.Git.RepoBranch,
			RepoRevision: elh.config.Git.RepoRevision,
			ReleaseID:    releaseID,
			BuildID:      buildID,
			BuildStatus:  buildStatus,
		}
		data, err := json.Marshal(ciBuilderEvent)
		if err != nil {
			log.Error().Err(err).Msgf("Failed marshalling EstafetteCiBuilderEvent for job %v", jobName)
			return err
		}
		requestBody = bytes.NewReader(data)

		// create client, in order to add headers
		client := pester.New()
		client.MaxRetries = 5
		client.Backoff = pester.ExponentialJitterBackoff
		client.KeepLog = true
		client.Timeout = time.Second * 10
		request, err := http.NewRequest("POST", ciServerBuilderEventsURL, requestBody)
		if err != nil {
			log.Error().Err(err).Msgf("Failed creating http client for job %v", jobName)
			return err
		}

		// add tracing context
		ext.SpanKindRPCClient.Set(span)
		ext.HTTPUrl.Set(span, request.URL.String())
		ext.HTTPMethod.Set(span, request.Method)
		span.Tracer().Inject(
			span.Context(),
			opentracing.HTTPHeaders,
			opentracing.HTTPHeadersCarrier(request.Header),
		)

		// add headers
		request.Header.Add("X-Estafette-Event", event)
		request.Header.Add("X-Estafette-Event-Job-Name", jobName)
		request.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ciAPIKey))

		// perform actual request
		response, err := client.Do(request)
		if err != nil {
			log.Error().Err(err).Str("pesterLogs", client.LogString()).Msgf("Failed performing http request to %v for job %v", ciServerBuilderEventsURL, jobName)
			return err
		}

		ext.HTTPStatusCode.Set(span, uint16(response.StatusCode))

		defer response.Body.Close()

		log.Debug().Str("pesterLogs", client.LogString()).Str("url", ciServerBuilderEventsURL).Msgf("Succesfully sent %v event to api", event)
	}

	return nil
}
<commit_msg>in case of error tag and add log to span<commit_after>package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"

	contracts "github.com/estafette/estafette-ci-contracts"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	tracingLog "github.com/opentracing/opentracing-go/log"

	"github.com/rs/zerolog/log"
	"github.com/sethgrid/pester"
)

// EndOfLifeHelper has methods to shutdown the runner after a fatal or successful run
type EndOfLifeHelper interface {
	handleFatal(context.Context, contracts.BuildLog, error, string)
	sendBuildFinishedEvent(ctx context.Context, buildStatus string) error
	sendBuildCleanEvent(ctx context.Context, buildStatus string) error
	sendBuildJobLogEvent(ctx context.Context, buildLog contracts.BuildLog) error
}

type endOfLifeHelperImpl struct {
	runAsJob bool
	config   contracts.BuilderConfig
}

// NewEndOfLifeHelper returns a new EndOfLifeHelper
func NewEndOfLifeHelper(runAsJob bool, config contracts.BuilderConfig) EndOfLifeHelper {
	return &endOfLifeHelperImpl{
		runAsJob: runAsJob,
		config:   config,
	}
}

func (elh *endOfLifeHelperImpl) handleFatal(ctx context.Context, buildLog contracts.BuildLog, err error, message string) {

	// add error messages as step to show in logs
	fatalStep := contracts.BuildLogStep{
		Step:     "init",
		LogLines: []contracts.BuildLogLine{},
		ExitCode: -1,
		Status:   "FAILED",
	}
	lineNumber := 1

	if err != nil {
		fatalStep.LogLines = append(fatalStep.LogLines, contracts.BuildLogLine{
			LineNumber: lineNumber,
			Timestamp:  time.Now().UTC(),
			StreamType: "stderr",
			Text:       err.Error(),
		})
		lineNumber++
	}
	if message != "" {
		fatalStep.LogLines = append(fatalStep.LogLines, contracts.BuildLogLine{
			LineNumber: lineNumber,
			Timestamp:  time.Now().UTC(),
			StreamType: "stderr",
			Text:       message,
		})
		lineNumber++
	}

	buildLog.Steps = append(buildLog.Steps, fatalStep)

	_ = elh.sendBuildFinishedEvent(ctx, "failed")
	_ = elh.sendBuildJobLogEvent(ctx, buildLog)
	_ = elh.sendBuildCleanEvent(ctx, "failed")

	if elh.runAsJob {
		log.Error().Err(err).Msg(message)
		os.Exit(0)
	} else {
		log.Fatal().Err(err).Msg(message)
	}
}

func (elh *endOfLifeHelperImpl) sendBuildJobLogEvent(ctx context.Context, buildLog contracts.BuildLog) (err error) {

	err = elh.sendBuildJobLogEventCore(ctx, buildLog)

	if err == nil {
		return
	}

	// strip log lines from successful steps to reduce size of the logs and still keep the useful information
	slimBuildLog := buildLog
	slimBuildLog.Steps = []contracts.BuildLogStep{}
	for _, s := range buildLog.Steps {
		slimBuildLogStep := s
		if s.Status == "SUCCEEDED" {
			if len(s.LogLines) > 0 {
				slimBuildLogStep.LogLines = []contracts.BuildLogLine{
					contracts.BuildLogLine{
						LineNumber: s.LogLines[0].LineNumber,
						Timestamp:  s.LogLines[0].Timestamp,
						StreamType: "stdout",
						Text:       "Truncated logs for reducing total log size; to prevent this use less verbose logging",
					},
				}
			}
		}

		slimBuildLog.Steps = append(slimBuildLog.Steps, slimBuildLogStep)
	}

	return elh.sendBuildJobLogEventCore(ctx, slimBuildLog)
}

func (elh *endOfLifeHelperImpl) sendBuildJobLogEventCore(ctx context.Context, buildLog contracts.BuildLog) (err error) {

	span, ctx := opentracing.StartSpanFromContext(ctx, "SendLog")
	defer span.Finish()

	ciServerBuilderPostLogsURL := elh.config.CIServer.PostLogsURL
	ciAPIKey := elh.config.CIServer.APIKey
	jobName := *elh.config.JobName

	if ciServerBuilderPostLogsURL != "" && ciAPIKey != "" && jobName != "" {

		// convert BuildJobLogs to json
		var requestBody io.Reader

		var data []byte
		if *elh.config.Action == "release" {
			// copy buildLog to releaseLog and marshal that
			releaseLog := contracts.ReleaseLog{
				ID:         buildLog.ID,
				RepoSource: buildLog.RepoSource,
				RepoOwner:  buildLog.RepoOwner,
				RepoName:   buildLog.RepoName,
				ReleaseID:  strconv.Itoa(elh.config.ReleaseParams.ReleaseID),
				Steps:      buildLog.Steps,
				InsertedAt: buildLog.InsertedAt,
			}
			data, err = json.Marshal(releaseLog)
			if err != nil {
				log.Error().Err(err).Msgf("Failed marshalling ReleaseLog for job %v", jobName)
				return
			}
		} else {
			data, err = json.Marshal(buildLog)
			if err != nil {
				log.Error().Err(err).Msgf("Failed marshalling BuildLog for job %v", jobName)
				return
			}
		}

		requestBody = bytes.NewReader(data)

		// create client, in order to add headers
		client := pester.New()
		client.MaxRetries = 5
		client.Backoff = pester.ExponentialJitterBackoff
		client.KeepLog = true
		client.Timeout = time.Second * 120
		request, err := http.NewRequest("POST", ciServerBuilderPostLogsURL, requestBody)
		if err != nil {
			log.Error().Err(err).Msgf("Failed creating http client for job %v", jobName)
			return err
		}

		// add tracing context
		ext.SpanKindRPCClient.Set(span)
		ext.HTTPUrl.Set(span, request.URL.String())
		ext.HTTPMethod.Set(span, request.Method)
		span.Tracer().Inject(
			span.Context(),
			opentracing.HTTPHeaders,
			opentracing.HTTPHeadersCarrier(request.Header),
		)

		// add headers
		request.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ciAPIKey))
		request.Header.Add("Content-Type", "application/json")

		// perform actual request
		response, err := client.Do(request)
		if err != nil {
			log.Error().Err(err).Str("logs", client.LogString()).Msgf("Failed shipping logs to %v for job %v", ciServerBuilderPostLogsURL, jobName)
			return err
		}

		ext.HTTPStatusCode.Set(span, uint16(response.StatusCode))

		defer response.Body.Close()

		log.Debug().Str("logs", client.LogString()).Msgf("Successfully shipped logs to %v for job %v", ciServerBuilderPostLogsURL, jobName)
	}

	return nil
}

func (elh *endOfLifeHelperImpl) sendBuildFinishedEvent(ctx context.Context, buildStatus string) error {
	return elh.sendBuilderEvent(ctx, buildStatus, fmt.Sprintf("builder:%v", buildStatus))
}

func (elh *endOfLifeHelperImpl) sendBuildCleanEvent(ctx context.Context, buildStatus string) error {
	return elh.sendBuilderEvent(ctx, buildStatus, "builder:clean")
}

func (elh *endOfLifeHelperImpl) sendBuilderEvent(ctx context.Context, buildStatus, event string) (err error) {

	span, ctx := opentracing.StartSpanFromContext(ctx, "SendBuildStatus")
	defer span.Finish()
	span.SetTag("build-status", buildStatus)
	span.SetTag("event-type", event)

	ciServerBuilderEventsURL := elh.config.CIServer.BuilderEventsURL
	ciAPIKey := elh.config.CIServer.APIKey
	jobName := *elh.config.JobName

	if ciServerBuilderEventsURL != "" && ciAPIKey != "" && jobName != "" {
		// convert EstafetteCiBuilderEvent to json
		var requestBody io.Reader

		buildID := ""
		if elh.config.BuildParams != nil {
			buildID = strconv.Itoa(elh.config.BuildParams.BuildID)
		}

		releaseID := ""
		if elh.config.ReleaseParams != nil {
			releaseID = strconv.Itoa(elh.config.ReleaseParams.ReleaseID)
		}

		ciBuilderEvent := EstafetteCiBuilderEvent{
			JobName:      jobName,
			RepoSource:   elh.config.Git.RepoSource,
			RepoOwner:    elh.config.Git.RepoOwner,
			RepoName:     elh.config.Git.RepoName,
			RepoBranch:   elh.config.Git.RepoBranch,
			RepoRevision: elh.config.Git.RepoRevision,
			ReleaseID:    releaseID,
			BuildID:      buildID,
			BuildStatus:  buildStatus,
		}
		data, err := json.Marshal(ciBuilderEvent)
		if err != nil {
			log.Error().Err(err).Msgf("Failed marshalling EstafetteCiBuilderEvent for job %v", jobName)
			return err
		}
		requestBody = bytes.NewReader(data)

		// create client, in order to add headers
		client := pester.New()
		client.MaxRetries = 5
		client.Backoff = pester.ExponentialJitterBackoff
		client.KeepLog = true
		client.Timeout = time.Second * 10
		request, err := http.NewRequest("POST", ciServerBuilderEventsURL, requestBody)
		if err != nil {
			log.Error().Err(err).Msgf("Failed creating http client for job %v", jobName)
			return err
		}

		// add tracing context
		ext.SpanKindRPCClient.Set(span)
		ext.HTTPUrl.Set(span, request.URL.String())
		ext.HTTPMethod.Set(span, request.Method)
		span.Tracer().Inject(
			span.Context(),
			opentracing.HTTPHeaders,
			opentracing.HTTPHeadersCarrier(request.Header),
		)

		// add headers
		request.Header.Add("X-Estafette-Event", event)
		request.Header.Add("X-Estafette-Event-Job-Name", jobName)
		request.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ciAPIKey))

		// perform actual request
		response, err := client.Do(request)
		if err != nil {
			span.SetTag("error", true)
			span.LogFields(
				tracingLog.String("error", err.Error()),
			)
			log.Error().Err(err).Str("pesterLogs", client.LogString()).Msgf("Failed performing http request to %v for job %v", ciServerBuilderEventsURL, jobName)
			return err
		}

		ext.HTTPStatusCode.Set(span, uint16(response.StatusCode))

		defer response.Body.Close()

		log.Debug().Str("pesterLogs", client.LogString()).Str("url", ciServerBuilderEventsURL).Msgf("Succesfully sent %v event to api", event)
	}

	return nil
}
<|endoftext|>


<commit_before>package tracing

import (
	"context"
	"fmt"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	zipkin "github.com/openzipkin/zipkin-go-opentracing"
	"github.com/sirupsen/logrus"
	"net/http"
)

// Tracer instance
var tracer opentracing.Tracer

// SetTracer can be used by unit tests to provide a NoopTracer instance. Real users should always
// use the InitTracing func.
func SetTracer(initializedTracer opentracing.Tracer) {
	tracer = initializedTracer
}

// InitTracing connects the calling service to Zipkin and initializes the tracer.
func InitTracing(zipkinURL string, serviceName string) {
	logrus.Infof("Connecting to zipkin server at %v", zipkinURL)
	collector, err := zipkin.NewHTTPCollector(
		fmt.Sprintf("%s/api/v1/spans", zipkinURL))
	if err != nil {
		logrus.Info("Error connecting to zipkin server at " +
			fmt.Sprintf("%s/api/v1/spans", zipkinURL) + ". Error: " + err.Error())
		logrus.Errorln("Error connecting to zipkin server at " +
			fmt.Sprintf("%s/api/v1/spans", zipkinURL) + ". Error: " + err.Error())
		panic("Error connecting to zipkin server at " +
			fmt.Sprintf("%s/api/v1/spans", zipkinURL) + ". Error: " + err.Error())
	}
	tracer, err = zipkin.NewTracer(
		zipkin.NewRecorder(collector, false, "127.0.0.1:0", serviceName))
	if err != nil {
		logrus.Errorln("Error starting new zipkin tracer. Error: " + err.Error())
		panic("Error starting new zipkin tracer. Error: " + err.Error())
	}
	logrus.Infof("Successfully started zipkin tracer for service '%v'", serviceName)
}

// StartHTTPTrace loads tracing information from an INCOMING HTTP request.
func StartHTTPTrace(r *http.Request, opName string) opentracing.Span {
	carrier := opentracing.HTTPHeadersCarrier(r.Header)
	clientContext, err := tracer.Extract(opentracing.HTTPHeaders, carrier)
	if err == nil {
		return tracer.StartSpan(
			opName, ext.RPCServerOption(clientContext))
	} else {
		return tracer.StartSpan(opName)
	}
}

// MapToCarrier converts a generic map to opentracing http headers carrier
func MapToCarrier(headers map[string]interface{}) opentracing.HTTPHeadersCarrier {
	carrier := make(opentracing.HTTPHeadersCarrier)
	for k, v := range headers {
		// delivery.Headers
		carrier.Set(k, v.(string))
	}
	return carrier
}

// CarrierToMap converts a TextMapCarrier to the amqp headers format
func CarrierToMap(values map[string]string) map[string]interface{} {
	headers := make(map[string]interface{})
	for k, v := range values {
		headers[k] = v
	}
	return headers
}

// StartTraceFromCarrier extracts tracing info from a generic map and starts a new span.
func StartTraceFromCarrier(carrier map[string]interface{}, spanName string) opentracing.Span {

	clientContext, err := tracer.Extract(opentracing.HTTPHeaders, MapToCarrier(carrier))
	var span opentracing.Span
	if err == nil {
		span = tracer.StartSpan(
			spanName, ext.RPCServerOption(clientContext))
	} else {
		span = tracer.StartSpan(spanName)
	}
	return span
}

// AddTracingToReq adds tracing information to an OUTGOING HTTP request
func AddTracingToReq(req *http.Request, span opentracing.Span) {
	carrier := opentracing.HTTPHeadersCarrier(req.Header)
	err := tracer.Inject(
		span.Context(),
		opentracing.HTTPHeaders,
		carrier)
	if err != nil {
		panic("Unable to inject tracing context: " + err.Error())
	}
}

// AddTracingToReqFromContext adds tracing information to an OUTGOING HTTP request
func AddTracingToReqFromContext(ctx context.Context, req *http.Request) {
	if ctx.Value("opentracing-span") == nil {
		return
	}
	carrier := opentracing.HTTPHeadersCarrier(req.Header)
	err := tracer.Inject(
		ctx.Value("opentracing-span").(opentracing.Span).Context(),
		opentracing.HTTPHeaders,
		carrier)
	if err != nil {
		panic("Unable to inject tracing context: " + err.Error())
	}
}

func AddTracingToTextMapCarrier(span opentracing.Span, val opentracing.TextMapCarrier) error {
	return tracer.Inject(span.Context(), opentracing.TextMap, val)
}

// StartSpanFromContext starts a span.
func StartSpanFromContext(ctx context.Context, opName string) opentracing.Span {
	span := ctx.Value("opentracing-span").(opentracing.Span)
	child := tracer.StartSpan(opName, ext.RPCServerOption(span.Context()))
	return child
}

// StartChildSpanFromContext starts a child span from span within the supplied context, if available.
func StartChildSpanFromContext(ctx context.Context, opName string) opentracing.Span {
	if ctx.Value("opentracing-span") == nil {
		return tracer.StartSpan(opName, ext.RPCServerOption(nil))
	}
	parent := ctx.Value("opentracing-span").(opentracing.Span)
	child := tracer.StartSpan(opName, opentracing.ChildOf(parent.Context()))
	return child
}

// StartSpanFromContextWithLogEvent starts span from context with logevent
func StartSpanFromContextWithLogEvent(ctx context.Context, opName string, logStatement string) opentracing.Span {
	span := ctx.Value("opentracing-span").(opentracing.Span)
	child := tracer.StartSpan(opName, ext.RPCServerOption(span.Context()))
	child.LogEvent(logStatement)
	return child
}

// CloseSpan logs event finishes span.
func CloseSpan(span opentracing.Span, event string) {
	span.LogEvent(event)
	span.Finish()
}

// LogEventToOngoingSpan extracts span from context and adds LogEvent
func LogEventToOngoingSpan(ctx context.Context, logMessage string) {
	if ctx.Value("opentracing-span") != nil {
		ctx.Value("opentracing-span").(opentracing.Span).LogEvent(logMessage)
	}
}

// UpdateContext updates the supplied context with the supplied span.
func UpdateContext(ctx context.Context, span opentracing.Span) context.Context {
	return context.WithValue(ctx, "opentracing-span", span)
}
<commit_msg>fixed missing warnf<commit_after>package tracing

import (
	"context"
	"fmt"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	zipkin "github.com/openzipkin/zipkin-go-opentracing"
	"github.com/sirupsen/logrus"
	"net/http"
)

// Tracer instance
var tracer opentracing.Tracer

// SetTracer can be used by unit tests to provide a NoopTracer instance. Real users should always
// use the InitTracing func.
func SetTracer(initializedTracer opentracing.Tracer) {
	tracer = initializedTracer
}

// InitTracing connects the calling service to Zipkin and initializes the tracer.
func InitTracing(zipkinURL string, serviceName string) {
	logrus.Infof("Connecting to zipkin server at %v", zipkinURL)
	collector, err := zipkin.NewHTTPCollector(
		fmt.Sprintf("%s/api/v1/spans", zipkinURL))
	if err != nil {
		logrus.Info("Error connecting to zipkin server at " +
			fmt.Sprintf("%s/api/v1/spans", zipkinURL) + ". Error: " + err.Error())
		logrus.Errorln("Error connecting to zipkin server at " +
			fmt.Sprintf("%s/api/v1/spans", zipkinURL) + ". Error: " + err.Error())
		panic("Error connecting to zipkin server at " +
			fmt.Sprintf("%s/api/v1/spans", zipkinURL) + ". Error: " + err.Error())
	}
	tracer, err = zipkin.NewTracer(
		zipkin.NewRecorder(collector, true, "127.0.0.1:0", serviceName))
	if err != nil {
		logrus.Errorln("Error starting new zipkin tracer. Error: " + err.Error())
		panic("Error starting new zipkin tracer. Error: " + err.Error())
	}
	logrus.Infof("Successfully started zipkin tracer for service '%v'", serviceName)
}

// StartHTTPTrace loads tracing information from an INCOMING HTTP request.
func StartHTTPTrace(r *http.Request, opName string) opentracing.Span {
	carrier := opentracing.HTTPHeadersCarrier(r.Header)
	clientContext, err := tracer.Extract(opentracing.HTTPHeaders, carrier)
	if err == nil {
		return tracer.StartSpan(
			opName, ext.RPCServerOption(clientContext))
	} else {
		return tracer.StartSpan(opName)
	}
}

// MapToCarrier converts a generic map to opentracing http headers carrier
func MapToCarrier(headers map[string]interface{}) opentracing.HTTPHeadersCarrier {
	carrier := make(opentracing.HTTPHeadersCarrier)
	for k, v := range headers {
		// delivery.Headers
		carrier.Set(k, v.(string))
	}
	return carrier
}

// CarrierToMap converts a TextMapCarrier to the amqp headers format
func CarrierToMap(values map[string]string) map[string]interface{} {
	headers := make(map[string]interface{})
	for k, v := range values {
		headers[k] = v
	}
	return headers
}

// StartTraceFromCarrier extracts tracing info from a generic map and starts a new span.
func StartTraceFromCarrier(carrier map[string]interface{}, spanName string) opentracing.Span {

	clientContext, err := tracer.Extract(opentracing.HTTPHeaders, MapToCarrier(carrier))
	var span opentracing.Span
	if err == nil {
		span = tracer.StartSpan(
			spanName, ext.RPCServerOption(clientContext))
	} else {
		span = tracer.StartSpan(spanName)
	}
	return span
}

// AddTracingToReq adds tracing information to an OUTGOING HTTP request
func AddTracingToReq(req *http.Request, span opentracing.Span) {
	carrier := opentracing.HTTPHeadersCarrier(req.Header)
	err := tracer.Inject(
		span.Context(),
		opentracing.HTTPHeaders,
		carrier)
	if err != nil {
		panic("Unable to inject tracing context: " + err.Error())
	}
}

// AddTracingToReqFromContext adds tracing information to an OUTGOING HTTP request
func AddTracingToReqFromContext(ctx context.Context, req *http.Request) {
	if ctx.Value("opentracing-span") == nil {
		logrus.Warnf("could not find 'opentracing-span' in outgoing context for request %v", req.RequestURI)
		return
	}
	carrier := opentracing.HTTPHeadersCarrier(req.Header)
	err := tracer.Inject(
		ctx.Value("opentracing-span").(opentracing.Span).Context(),
		opentracing.HTTPHeaders,
		carrier)
	if err != nil {
		panic("Unable to inject tracing context: " + err.Error())
	}
}

func AddTracingToTextMapCarrier(span opentracing.Span, val opentracing.TextMapCarrier) error {
	return tracer.Inject(span.Context(), opentracing.TextMap, val)
}

// StartSpanFromContext starts a span.
func StartSpanFromContext(ctx context.Context, opName string) opentracing.Span {
	span := ctx.Value("opentracing-span").(opentracing.Span)
	child := tracer.StartSpan(opName, ext.RPCServerOption(span.Context()))
	return child
}

// StartChildSpanFromContext starts a child span from span within the supplied context, if available.
func StartChildSpanFromContext(ctx context.Context, opName string) opentracing.Span {
	if ctx.Value("opentracing-span") == nil {
		return tracer.StartSpan(opName, ext.RPCServerOption(nil))
	}
	parent := ctx.Value("opentracing-span").(opentracing.Span)
	child := tracer.StartSpan(opName, opentracing.ChildOf(parent.Context()))
	return child
}

// StartSpanFromContextWithLogEvent starts span from context with logevent
func StartSpanFromContextWithLogEvent(ctx context.Context, opName string, logStatement string) opentracing.Span {
	span := ctx.Value("opentracing-span").(opentracing.Span)
	child := tracer.StartSpan(opName, ext.RPCServerOption(span.Context()))
	child.LogEvent(logStatement)
	return child
}

// CloseSpan logs event finishes span.
func CloseSpan(span opentracing.Span, event string) {
	span.LogEvent(event)
	span.Finish()
}

// LogEventToOngoingSpan extracts span from context and adds LogEvent
func LogEventToOngoingSpan(ctx context.Context, logMessage string) {
	if ctx.Value("opentracing-span") != nil {
		ctx.Value("opentracing-span").(opentracing.Span).LogEvent(logMessage)
	}
}

// UpdateContext updates the supplied context with the supplied span.
func UpdateContext(ctx context.Context, span opentracing.Span) context.Context {
	return context.WithValue(ctx, "opentracing-span", span)
}
<|endoftext|>


<commit_before>package buf

import (
	"os"
	"runtime"
	"strconv"
	"sync"
)

// Pool provides functionality to generate and recycle buffers on demand.
type Pool interface {
	// Allocate either returns a unused buffer from the pool, or generates a new one from system.
	Allocate() *Buffer
	// Free recycles the given buffer.
	Free(*Buffer)
}

// SyncPool is a buffer pool based on sync.Pool
type SyncPool struct {
	allocator *sync.Pool
}

// NewSyncPool creates a SyncPool with given buffer size.
func NewSyncPool(bufferSize uint32) *SyncPool {
	pool := &SyncPool{
		allocator: &sync.Pool{
			New: func() interface{} { return make([]byte, bufferSize) },
		},
	}
	return pool
}

// Allocate implements Pool.Allocate().
func (p *SyncPool) Allocate() *Buffer {
	return &Buffer{
		v:    p.allocator.Get().([]byte),
		pool: p,
	}
}

// Free implements Pool.Free().
func (p *SyncPool) Free(buffer *Buffer) {
	rawBuffer := buffer.v
	if rawBuffer == nil {
		return
	}
	p.allocator.Put(rawBuffer)
}

// BufferPool is a Pool that utilizes an internal cache.
type BufferPool struct {
	chain     chan []byte
	allocator *sync.Pool
}

// NewBufferPool creates a new BufferPool with given buffer size, and internal cache size.
func NewBufferPool(bufferSize, poolSize uint32) *BufferPool {
	pool := &BufferPool{
		chain: make(chan []byte, poolSize),
		allocator: &sync.Pool{
			New: func() interface{} { return make([]byte, bufferSize) },
		},
	}
	for i := uint32(0); i < poolSize; i++ {
		pool.chain <- make([]byte, bufferSize)
	}
	return pool
}

// Allocate implements Pool.Allocate().
func (p *BufferPool) Allocate() *Buffer {
	var b []byte
	select {
	case b = <-p.chain:
	default:
		b = p.allocator.Get().([]byte)
	}
	return &Buffer{
		v:    b,
		pool: p,
	}
}

// Free implements Pool.Free().
func (p *BufferPool) Free(buffer *Buffer) {
	rawBuffer := buffer.v
	if rawBuffer == nil {
		return
	}
	select {
	case p.chain <- rawBuffer:
	default:
		p.allocator.Put(rawBuffer)
	}
}

const (
	// Size of a regular buffer.
	Size = 2 * 1024

	poolSizeEnvKey = "v2ray.buffer.size"
)

var (
	mediumPool Pool
)

func getDefaultPoolSize() uint32 {
	switch runtime.GOARCH {
	case "amd64", "386":
		return 20
	default:
		return 5
	}
}

func init() {
	size := getDefaultPoolSize()
	sizeStr := os.Getenv(poolSizeEnvKey)
	if len(sizeStr) > 0 {
		customSize, err := strconv.ParseUint(sizeStr, 10, 32)
		if err == nil {
			size = uint32(customSize)
		}
	}
	if size > 0 {
		totalByteSize := size * 1024 * 1024
		mediumPool = NewBufferPool(Size, totalByteSize/Size)
	} else {
		mediumPool = NewSyncPool(Size)
	}
}
<commit_msg>simplify code<commit_after>package buf

import (
	"os"
	"runtime"
	"strconv"
	"sync"
)

// Pool provides functionality to generate and recycle buffers on demand.
type Pool interface {
	// Allocate either returns a unused buffer from the pool, or generates a new one from system.
	Allocate() *Buffer
	// Free recycles the given buffer.
	Free(*Buffer)
}

// SyncPool is a buffer pool based on sync.Pool
type SyncPool struct {
	allocator *sync.Pool
}

// NewSyncPool creates a SyncPool with given buffer size.
func NewSyncPool(bufferSize uint32) *SyncPool {
	pool := &SyncPool{
		allocator: &sync.Pool{
			New: func() interface{} { return make([]byte, bufferSize) },
		},
	}
	return pool
}

// Allocate implements Pool.Allocate().
func (p *SyncPool) Allocate() *Buffer {
	return &Buffer{
		v:    p.allocator.Get().([]byte),
		pool: p,
	}
}

// Free implements Pool.Free().
func (p *SyncPool) Free(buffer *Buffer) {
	if buffer.v != nil {
		p.allocator.Put(buffer.v)
	}
}

// BufferPool is a Pool that utilizes an internal cache.
type BufferPool struct {
	chain chan []byte
	sub   Pool
}

// NewBufferPool creates a new BufferPool with given buffer size, and internal cache size.
func NewBufferPool(bufferSize, poolSize uint32) *BufferPool {
	pool := &BufferPool{
		chain: make(chan []byte, poolSize),
		sub:   NewSyncPool(bufferSize),
	}
	for i := uint32(0); i < poolSize; i++ {
		pool.chain <- make([]byte, bufferSize)
	}
	return pool
}

// Allocate implements Pool.Allocate().
func (p *BufferPool) Allocate() *Buffer {
	select {
	case b := <-p.chain:
		return &Buffer{
			v:    b,
			pool: p,
		}
	default:
		return p.sub.Allocate()
	}
}

// Free implements Pool.Free().
func (p *BufferPool) Free(buffer *Buffer) {
	if buffer.v == nil {
		return
	}
	select {
	case p.chain <- buffer.v:
	default:
		p.sub.Free(buffer)
	}
}

const (
	// Size of a regular buffer.
	Size = 2 * 1024

	poolSizeEnvKey = "v2ray.buffer.size"
)

var (
	mediumPool Pool
)

func getDefaultPoolSize() uint32 {
	switch runtime.GOARCH {
	case "amd64", "386":
		return 20
	default:
		return 5
	}
}

func init() {
	size := getDefaultPoolSize()
	sizeStr := os.Getenv(poolSizeEnvKey)
	if len(sizeStr) > 0 {
		customSize, err := strconv.ParseUint(sizeStr, 10, 32)
		if err == nil {
			size = uint32(customSize)
		}
	}
	if size > 0 {
		totalByteSize := size * 1024 * 1024
		mediumPool = NewBufferPool(Size, totalByteSize/Size)
	} else {
		mediumPool = NewSyncPool(Size)
	}
}
<|endoftext|>


<commit_before>package pick

import (
	"io"
	"strings"
)

// input can be input/select/textarea
// TODO: now only get 1st value, need []string as result
// TODO: pick select option that has no value attr
// TODO: radio
func PickForm(reader io.Reader, attr *Attr) (input map[string][]string) {
	htmlArr := PickHtml(&Option{
		PageSource: reader,
		TagName:    "form",
		Attr:       attr,
	}, 1)
	if len(htmlArr) == 0 {
		return
	}

	html := htmlArr[0]
	input = make(map[string][]string)

	// pick input
	inputNameArr := PickAttr(&Option{
		PageSource: strings.NewReader(html),
		TagName:    "input",
		Attr:       nil,
	}, "name", 0)

	for i := 0; i < len(inputNameArr); i++ {
		values := PickAttr(&Option{
			PageSource: strings.NewReader(html),
			TagName:    "input",
			Attr: &Attr{
				Label: "name",
				Value: inputNameArr[i],
			},
		}, "value", 1)
		if len(values) == 0 {
			input[inputNameArr[i]] = []string{""}
			continue
		}

		input[inputNameArr[i]] = []string{values[0]}
	}

	// pick textarea
	textNameArr := PickAttr(&Option{
		PageSource: strings.NewReader(html),
		TagName:    "textarea",
		Attr:       nil,
	}, "name", 0)

	for i := 0; i < len(textNameArr); i++ {
		values := PickText(&Option{
			PageSource: strings.NewReader(html),
			TagName:    "textarea",
			Attr: &Attr{
				Label: "name",
				Value: textNameArr[i],
			},
		}, 1)
		if len(values) == 0 {
			input[textNameArr[i]] = []string{""}
			continue
		}

		input[textNameArr[i]] = []string{values[0]}
	}

	// pick select
	selectNameArr := PickAttr(&Option{
		PageSource: strings.NewReader(html),
		TagName:    "select",
		Attr:       nil,
	}, "name", 0)

	for i := 0; i < len(selectNameArr); i++ {
		selectHtmlArr := PickHtml(&Option{
			PageSource: strings.NewReader(html),
			TagName:    "select",
			Attr: &Attr{
				Label: "name",
				Value: selectNameArr[i],
			},
		}, 1)
		if len(selectHtmlArr) == 0 {
			continue
		}

		// get selected option
		values := PickAttr(&Option{
			PageSource: strings.NewReader(selectHtmlArr[0]),
			TagName:    "option",
			Attr: &Attr{
				Label: "selected",
				Value: "",
			},
		}, "value", 1)
		if len(values) != 0 {
			input[selectNameArr[i]] = []string{values[0]}
			continue
		}

		// if no select get the 1st option value
		values = PickAttr(&Option{
			PageSource: strings.NewReader(selectHtmlArr[0]),
			TagName:    "option",
			Attr:       nil,
		}, "value", 1)
		if len(values) == 0 {
			input[selectNameArr[i]] = []string{""}
			continue
		}

		input[selectNameArr[i]] = []string{values[0]}
	}

	return
}
<commit_msg>better #PickForm selected option<commit_after>package pick

import (
	"io"
	"strings"
)

// input can be input/select/textarea
// TODO: now only get 1st value, need []string as result
// TODO: pick select option that has no value attr
// TODO: radio
func PickForm(reader io.Reader, attr *Attr) (input map[string][]string) {
	htmlArr := PickHtml(&Option{
		PageSource: reader,
		TagName:    "form",
		Attr:       attr,
	}, 1)
	if len(htmlArr) == 0 {
		return
	}

	html := htmlArr[0]
	input = make(map[string][]string)

	// pick input
	inputNameArr := PickAttr(&Option{
		PageSource: strings.NewReader(html),
		TagName:    "input",
		Attr:       nil,
	}, "name", 0)

	for i := 0; i < len(inputNameArr); i++ {
		values := PickAttr(&Option{
			PageSource: strings.NewReader(html),
			TagName:    "input",
			Attr: &Attr{
				Label: "name",
				Value: inputNameArr[i],
			},
		}, "value", 1)
		if len(values) == 0 {
			input[inputNameArr[i]] = []string{""}
			continue
		}

		input[inputNameArr[i]] = []string{values[0]}
	}

	// pick textarea
	textNameArr := PickAttr(&Option{
		PageSource: strings.NewReader(html),
		TagName:    "textarea",
		Attr:       nil,
	}, "name", 0)

	for i := 0; i < len(textNameArr); i++ {
		values := PickText(&Option{
			PageSource: strings.NewReader(html),
			TagName:    "textarea",
			Attr: &Attr{
				Label: "name",
				Value: textNameArr[i],
			},
		}, 1)
		if len(values) == 0 {
			input[textNameArr[i]] = []string{""}
			continue
		}

		input[textNameArr[i]] = []string{values[0]}
	}

	// pick select
	selectNameArr := PickAttr(&Option{
		PageSource: strings.NewReader(html),
		TagName:    "select",
		Attr:       nil,
	}, "name", 0)

	for i := 0; i < len(selectNameArr); i++ {
		selectHtmlArr := PickHtml(&Option{
			PageSource: strings.NewReader(html),
			TagName:    "select",
			Attr: &Attr{
				Label: "name",
				Value: selectNameArr[i],
			},
		}, 1)
		if len(selectHtmlArr) == 0 {
			continue
		}

		// get selected option
		values := PickAttr(&Option{
			PageSource: strings.NewReader(selectHtmlArr[0]),
			TagName:    "option",
			Attr: &Attr{
				Label: "selected",
				Value: "",
			},
		}, "value", 1)
		if len(values) != 0 {
			input[selectNameArr[i]] = []string{values[0]}
			continue
		}

		values = PickAttr(&Option{
			PageSource: strings.NewReader(selectHtmlArr[0]),
			TagName:    "option",
			Attr: &Attr{
				Label: "selected",
				Value: "selected",
			},
		}, "value", 1)
		if len(values) != 0 {
			input[selectNameArr[i]] = []string{values[0]}
			continue
		}

		// if no select get the 1st option value
		values = PickAttr(&Option{
			PageSource: strings.NewReader(selectHtmlArr[0]),
			TagName:    "option",
			Attr:       nil,
		}, "value", 1)
		if len(values) == 0 {
			input[selectNameArr[i]] = []string{""}
			continue
		}

		input[selectNameArr[i]] = []string{values[0]}
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"context"

	"github.com/mafredri/cdp"
	"github.com/mafredri/cdp/protocol/dom"
	"github.com/mafredri/cdp/protocol/runtime"
)

func submitForm(ctx context.Context, c *cdp.Client, expression string) error {
	domContentEventFired, err := c.Page.DOMContentEventFired(ctx)
	if err != nil {
		return err
	}
	defer domContentEventFired.Close()

	_, err = c.Runtime.Evaluate(ctx, &runtime.EvaluateArgs{
		Expression: expression,
	})
	if err != nil {
		return err
	}

	_, err = domContentEventFired.Recv()
	return err
}

func setFormInputFiles(ctx context.Context, c *cdp.Client, selector string, files ...string) error {
	doc, err := c.DOM.GetDocument(ctx, dom.NewGetDocumentArgs())
	if err != nil {
		return err
	}

	fileInput, err := c.DOM.QuerySelector(ctx, &dom.QuerySelectorArgs{
		NodeID:   doc.Root.NodeID,
		Selector: selector,
	})
	if err != nil {
		return err
	}

	err = c.DOM.SetFileInputFiles(ctx, &dom.SetFileInputFilesArgs{
		NodeID: &fileInput.NodeID,
		Files:  files,
	})
	return err
}
<commit_msg>Handle eval exceptions in submitForm<commit_after>package main

import (
	"context"

	"github.com/mafredri/cdp"
	"github.com/mafredri/cdp/protocol/dom"
	"github.com/mafredri/cdp/protocol/runtime"
)

func submitForm(ctx context.Context, c *cdp.Client, expression string) error {
	domContentEventFired, err := c.Page.DOMContentEventFired(ctx)
	if err != nil {
		return err
	}
	defer domContentEventFired.Close()

	eval, err := c.Runtime.Evaluate(ctx, &runtime.EvaluateArgs{
		Expression: expression,
	})
	if err != nil {
		return err
	}
	if eval.ExceptionDetails != nil {
		return eval.ExceptionDetails
	}

	_, err = domContentEventFired.Recv()
	return err
}

func setFormInputFiles(ctx context.Context, c *cdp.Client, selector string, files ...string) error {
	doc, err := c.DOM.GetDocument(ctx, dom.NewGetDocumentArgs())
	if err != nil {
		return err
	}

	fileInput, err := c.DOM.QuerySelector(ctx, &dom.QuerySelectorArgs{
		NodeID:   doc.Root.NodeID,
		Selector: selector,
	})
	if err != nil {
		return err
	}

	err = c.DOM.SetFileInputFiles(ctx, &dom.SetFileInputFilesArgs{
		NodeID: &fileInput.NodeID,
		Files:  files,
	})
	return err
}
<|endoftext|>


<commit_before>package cobra

import (
	"bytes"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/spf13/pflag"
)

const (
	BashCompFilenameExt     = "cobra_annotation_bash_completion_filename_extentions"
	BashCompOneRequiredFlag = "cobra_annotation_bash_completion_one_required_flag"
	BashCompSubdirsInDir    = "cobra_annotation_bash_completion_subdirs_in_dir"
)

func preamble(out *bytes.Buffer) {
	fmt.Fprintf(out, `#!/bin/bash

__debug()
{
    if [[ -n ${BASH_COMP_DEBUG_FILE} ]]; then
        echo "$*" >> "${BASH_COMP_DEBUG_FILE}"
    fi
}

# Homebrew on Macs have version 1.3 of bash-completion which doesn't include
# _init_completion. This is a very minimal version of that function.
__my_init_completion()
{
    COMPREPLY=()
    _get_comp_words_by_ref cur prev words cword
}

__index_of_word()
{
    local w word=$1
    shift
    index=0
    for w in "$@"; do
        [[ $w = "$word" ]] && return
        index=$((index+1))
    done
    index=-1
}

__contains_word()
{
    local w word=$1; shift
    for w in "$@"; do
        [[ $w = "$word" ]] && return
    done
    return 1
}

__handle_reply()
{
    __debug "${FUNCNAME}"
    case $cur in
        -*)
            if [[ $(type -t compopt) = "builtin" ]]; then
                compopt -o nospace
            fi
            local allflags
            if [ ${#must_have_one_flag[@]} -ne 0 ]; then
                allflags=("${must_have_one_flag[@]}")
            else
                allflags=("${flags[*]} ${two_word_flags[*]}")
            fi
            COMPREPLY=( $(compgen -W "${allflags[*]}" -- "$cur") )
            if [[ $(type -t compopt) = "builtin" ]]; then
                [[ $COMPREPLY == *= ]] || compopt +o nospace
            fi
            return 0;
            ;;
    esac

    # check if we are handling a flag with special work handling
    local index
    __index_of_word "${prev}" "${flags_with_completion[@]}"
    if [[ ${index} -ge 0 ]]; then
        ${flags_completion[${index}]}
        return
    fi

    # we are parsing a flag and don't have a special handler, no completion
    if [[ ${cur} != "${words[cword]}" ]]; then
        return
    fi

    local completions
    if [[ ${#must_have_one_flag[@]} -ne 0 ]]; then
        completions=("${must_have_one_flag[@]}")
    elif [[ ${#must_have_one_noun[@]} -ne 0 ]]; then
        completions=("${must_have_one_noun[@]}")
    else
        completions=("${commands[@]}")
    fi
    COMPREPLY=( $(compgen -W "${completions[*]}" -- "$cur") )

    if [[ ${#COMPREPLY[@]} -eq 0 ]]; then
        declare -F __custom_func >/dev/null && __custom_func
    fi

    __ltrim_colon_completions "$cur"
}

# The arguments should be in the form "ext1|ext2|extn"
__handle_filename_extension_flag()
{
    local ext="$1"
    _filedir "@(${ext})"
}

__handle_subdirs_in_dir_flag()
{
    local dir="$1"
    pushd "${dir}" >/dev/null 2>&1 && _filedir -d && popd >/dev/null 2>&1
}

__handle_flag()
{
    __debug "${FUNCNAME}: c is $c words[c] is ${words[c]}"

    # if a command required a flag, and we found it, unset must_have_one_flag()
    local flagname=${words[c]}
    local flagvalue
    # if the word contained an =
    if [[ ${words[c]} == *"="* ]]; then
        flagvalue=${flagname#*=} # take in as flagvalue after the =
        flagname=${flagname%%=*} # strip everything after the =
        flagname="${flagname}=" # but put the = back
    fi
    __debug "${FUNCNAME}: looking for ${flagname}"
    if __contains_word "${flagname}" "${must_have_one_flag[@]}"; then
        must_have_one_flag=()
    fi

    # keep flag value with flagname as flaghash
    if [ ${flagvalue} ] ; then
        flaghash[${flagname}]=${flagvalue}
    elif [ ${words[ $((c+1)) ]} ] ; then
        flaghash[${flagname}]=${words[ $((c+1)) ]}
    else
        flaghash[${flagname}]="true" # pad "true" for bool flag
    fi

    # skip the argument to a two word flag
    if __contains_word "${words[c]}" "${two_word_flags[@]}"; then
        c=$((c+1))
        # if we are looking for a flags value, don't show commands
        if [[ $c -eq $cword ]]; then
            commands=()
        fi
    fi

    c=$((c+1))

}

__handle_noun()
{
    __debug "${FUNCNAME}: c is $c words[c] is ${words[c]}"

    if __contains_word "${words[c]}" "${must_have_one_noun[@]}"; then
        must_have_one_noun=()
    fi

    nouns+=("${words[c]}")
    c=$((c+1))
}

__handle_command()
{
    __debug "${FUNCNAME}: c is $c words[c] is ${words[c]}"

    local next_command
    if [[ -n ${last_command} ]]; then
        next_command="_${last_command}_${words[c]//:/__}"
    else
        next_command="_${words[c]//:/__}"
    fi
    c=$((c+1))
    __debug "${FUNCNAME}: looking for ${next_command}"
    declare -F $next_command >/dev/null && $next_command
}

__handle_word()
{
    if [[ $c -ge $cword ]]; then
        __handle_reply
        return
    fi
    __debug "${FUNCNAME}: c is $c words[c] is ${words[c]}"
    if [[ "${words[c]}" == -* ]]; then
        __handle_flag
    elif __contains_word "${words[c]}" "${commands[@]}"; then
        __handle_command
    else
        __handle_noun
    fi
    __handle_word
}

`)
}

func postscript(out *bytes.Buffer, name string) {
	name = strings.Replace(name, ":", "__", -1)
	fmt.Fprintf(out, "__start_%s()\n", name)
	fmt.Fprintf(out, `{
    local cur prev words cword
    declare -A flaghash
    if declare -F _init_completion >/dev/null 2>&1; then
        _init_completion -s || return
    else
        __my_init_completion || return
    fi

    local c=0
    local flags=()
    local two_word_flags=()
    local flags_with_completion=()
    local flags_completion=()
    local commands=("%s")
    local must_have_one_flag=()
    local must_have_one_noun=()
    local last_command
    local nouns=()

    __handle_word
}

`, name)
	fmt.Fprintf(out, `if [[ $(type -t compopt) = "builtin" ]]; then
    complete -o default -F __start_%s %s
else
    complete -o default -o nospace -F __start_%s %s
fi

`, name, name, name, name)
	fmt.Fprintf(out, "# ex: ts=4 sw=4 et filetype=sh\n")
}

func writeCommands(cmd *Command, out *bytes.Buffer) {
	fmt.Fprintf(out, "    commands=()\n")
	for _, c := range cmd.Commands() {
		if !c.IsAvailableCommand() || c == cmd.helpCommand {
			continue
		}
		fmt.Fprintf(out, "    commands+=(%q)\n", c.Name())
	}
	fmt.Fprintf(out, "\n")
}

func writeFlagHandler(name string, annotations map[string][]string, out *bytes.Buffer) {
	for key, value := range annotations {
		switch key {
		case BashCompFilenameExt:
			fmt.Fprintf(out, "    flags_with_completion+=(%q)\n", name)

			if len(value) > 0 {
				ext := "__handle_filename_extension_flag " + strings.Join(value, "|")
				fmt.Fprintf(out, "    flags_completion+=(%q)\n", ext)
			} else {
				ext := "_filedir"
				fmt.Fprintf(out, "    flags_completion+=(%q)\n", ext)
			}
		case BashCompSubdirsInDir:
			fmt.Fprintf(out, "    flags_with_completion+=(%q)\n", name)

			if len(value) == 1 {
				ext := "__handle_subdirs_in_dir_flag " + value[0]
				fmt.Fprintf(out, "    flags_completion+=(%q)\n", ext)
			} else {
				ext := "_filedir -d"
				fmt.Fprintf(out, "    flags_completion+=(%q)\n", ext)
			}
		}
	}
}

func writeShortFlag(flag *pflag.Flag, out *bytes.Buffer) {
	b := (flag.Value.Type() == "bool")
	name := flag.Shorthand
	format := "    "
	if !b {
		format += "two_word_"
	}
	format += "flags+=(\"-%s\")\n"
	fmt.Fprintf(out, format, name)
	writeFlagHandler("-"+name, flag.Annotations, out)
}

func writeFlag(flag *pflag.Flag, out *bytes.Buffer) {
	b := (flag.Value.Type() == "bool")
	name := flag.Name
	format := "    flags+=(\"--%s"
	if !b {
		format += "="
	}
	format += "\")\n"
	fmt.Fprintf(out, format, name)
	writeFlagHandler("--"+name, flag.Annotations, out)
}

func writeFlags(cmd *Command, out *bytes.Buffer) {
	fmt.Fprintf(out, `    flags=()
    two_word_flags=()
    flags_with_completion=()
    flags_completion=()

`)
	cmd.NonInheritedFlags().VisitAll(func(flag *pflag.Flag) {
		writeFlag(flag, out)
		if len(flag.Shorthand) > 0 {
			writeShortFlag(flag, out)
		}
	})
	cmd.InheritedFlags().VisitAll(func(flag *pflag.Flag) {
		writeFlag(flag, out)
		if len(flag.Shorthand) > 0 {
			writeShortFlag(flag, out)
		}
	})

	fmt.Fprintf(out, "\n")
}

func writeRequiredFlag(cmd *Command, out *bytes.Buffer) {
	fmt.Fprintf(out, "    must_have_one_flag=()\n")
	flags := cmd.NonInheritedFlags()
	flags.VisitAll(func(flag *pflag.Flag) {
		for key := range flag.Annotations {
			switch key {
			case BashCompOneRequiredFlag:
				format := "    must_have_one_flag+=(\"--%s"
				b := (flag.Value.Type() == "bool")
				if !b {
					format += "="
				}
				format += "\")\n"
				fmt.Fprintf(out, format, flag.Name)

				if len(flag.Shorthand) > 0 {
					fmt.Fprintf(out, "    must_have_one_flag+=(\"-%s\")\n", flag.Shorthand)
				}
			}
		}
	})
}

func writeRequiredNoun(cmd *Command, out *bytes.Buffer) {
	fmt.Fprintf(out, "    must_have_one_noun=()\n")
	sort.Sort(sort.StringSlice(cmd.ValidArgs))
	for _, value := range cmd.ValidArgs {
		fmt.Fprintf(out, "    must_have_one_noun+=(%q)\n", value)
	}
}

func gen(cmd *Command, out *bytes.Buffer) {
	for _, c := range cmd.Commands() {
		if !c.IsAvailableCommand() || c == cmd.helpCommand {
			continue
		}
		gen(c, out)
	}
	commandName := cmd.CommandPath()
	commandName = strings.Replace(commandName, " ", "_", -1)
	commandName = strings.Replace(commandName, ":", "__", -1)
	fmt.Fprintf(out, "_%s()\n{\n", commandName)
	fmt.Fprintf(out, "    last_command=%q\n", commandName)
	writeCommands(cmd, out)
	writeFlags(cmd, out)
	writeRequiredFlag(cmd, out)
	writeRequiredNoun(cmd, out)
	fmt.Fprintf(out, "}\n\n")
}

func (cmd *Command) GenBashCompletion(out *bytes.Buffer) {
	preamble(out)
	if len(cmd.BashCompletionFunction) > 0 {
		fmt.Fprintf(out, "%s\n", cmd.BashCompletionFunction)
	}
	gen(cmd, out)
	postscript(out, cmd.Name())
}

func (cmd *Command) GenBashCompletionFile(filename string) error {
	out := new(bytes.Buffer)

	cmd.GenBashCompletion(out)

	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	_, err = outFile.Write(out.Bytes())
	if err != nil {
		return err
	}
	return nil
}

// MarkFlagRequired adds the BashCompOneRequiredFlag annotation to the named flag, if it exists.
func (cmd *Command) MarkFlagRequired(name string) error {
	return MarkFlagRequired(cmd.Flags(), name)
}

// MarkPersistentFlagRequired adds the BashCompOneRequiredFlag annotation to the named persistent flag, if it exists.
func (cmd *Command) MarkPersistentFlagRequired(name string) error {
	return MarkFlagRequired(cmd.PersistentFlags(), name)
}

// MarkFlagRequired adds the BashCompOneRequiredFlag annotation to the named flag in the flag set, if it exists.
func MarkFlagRequired(flags *pflag.FlagSet, name string) error {
	return flags.SetAnnotation(name, BashCompOneRequiredFlag, []string{"true"})
}

// MarkFlagFilename adds the BashCompFilenameExt annotation to the named flag, if it exists.
// Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
func (cmd *Command) MarkFlagFilename(name string, extensions ...string) error {
	return MarkFlagFilename(cmd.Flags(), name, extensions...)
}

// MarkPersistentFlagFilename adds the BashCompFilenameExt annotation to the named persistent flag, if it exists.
// Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
func (cmd *Command) MarkPersistentFlagFilename(name string, extensions ...string) error {
	return MarkFlagFilename(cmd.PersistentFlags(), name, extensions...)
}

// MarkFlagFilename adds the BashCompFilenameExt annotation to the named flag in the flag set, if it exists.
// Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
func MarkFlagFilename(flags *pflag.FlagSet, name string, extensions ...string) error {
	return flags.SetAnnotation(name, BashCompFilenameExt, extensions)
}
<commit_msg>Replace *bytes.Buffer with io.Writer<commit_after>package cobra

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"

	"github.com/spf13/pflag"
)

const (
	BashCompFilenameExt     = "cobra_annotation_bash_completion_filename_extentions"
	BashCompOneRequiredFlag = "cobra_annotation_bash_completion_one_required_flag"
	BashCompSubdirsInDir    = "cobra_annotation_bash_completion_subdirs_in_dir"
)

func preamble(out io.Writer) {
	fmt.Fprintf(out, `#!/bin/bash

__debug()
{
    if [[ -n ${BASH_COMP_DEBUG_FILE} ]]; then
        echo "$*" >> "${BASH_COMP_DEBUG_FILE}"
    fi
}

# Homebrew on Macs have version 1.3 of bash-completion which doesn't include
# _init_completion. This is a very minimal version of that function.
__my_init_completion()
{
    COMPREPLY=()
    _get_comp_words_by_ref cur prev words cword
}

__index_of_word()
{
    local w word=$1
    shift
    index=0
    for w in "$@"; do
        [[ $w = "$word" ]] && return
        index=$((index+1))
    done
    index=-1
}

__contains_word()
{
    local w word=$1; shift
    for w in "$@"; do
        [[ $w = "$word" ]] && return
    done
    return 1
}

__handle_reply()
{
    __debug "${FUNCNAME}"
    case $cur in
        -*)
            if [[ $(type -t compopt) = "builtin" ]]; then
                compopt -o nospace
            fi
            local allflags
            if [ ${#must_have_one_flag[@]} -ne 0 ]; then
                allflags=("${must_have_one_flag[@]}")
            else
                allflags=("${flags[*]} ${two_word_flags[*]}")
            fi
            COMPREPLY=( $(compgen -W "${allflags[*]}" -- "$cur") )
            if [[ $(type -t compopt) = "builtin" ]]; then
                [[ $COMPREPLY == *= ]] || compopt +o nospace
            fi
            return 0;
            ;;
    esac

    # check if we are handling a flag with special work handling
    local index
    __index_of_word "${prev}" "${flags_with_completion[@]}"
    if [[ ${index} -ge 0 ]]; then
        ${flags_completion[${index}]}
        return
    fi

    # we are parsing a flag and don't have a special handler, no completion
    if [[ ${cur} != "${words[cword]}" ]]; then
        return
    fi

    local completions
    if [[ ${#must_have_one_flag[@]} -ne 0 ]]; then
        completions=("${must_have_one_flag[@]}")
    elif [[ ${#must_have_one_noun[@]} -ne 0 ]]; then
        completions=("${must_have_one_noun[@]}")
    else
        completions=("${commands[@]}")
    fi
    COMPREPLY=( $(compgen -W "${completions[*]}" -- "$cur") )

    if [[ ${#COMPREPLY[@]} -eq 0 ]]; then
        declare -F __custom_func >/dev/null && __custom_func
    fi

    __ltrim_colon_completions "$cur"
}

# The arguments should be in the form "ext1|ext2|extn"
__handle_filename_extension_flag()
{
    local ext="$1"
    _filedir "@(${ext})"
}

__handle_subdirs_in_dir_flag()
{
    local dir="$1"
    pushd "${dir}" >/dev/null 2>&1 && _filedir -d && popd >/dev/null 2>&1
}

__handle_flag()
{
    __debug "${FUNCNAME}: c is $c words[c] is ${words[c]}"

    # if a command required a flag, and we found it, unset must_have_one_flag()
    local flagname=${words[c]}
    local flagvalue
    # if the word contained an =
    if [[ ${words[c]} == *"="* ]]; then
        flagvalue=${flagname#*=} # take in as flagvalue after the =
        flagname=${flagname%%=*} # strip everything after the =
        flagname="${flagname}=" # but put the = back
    fi
    __debug "${FUNCNAME}: looking for ${flagname}"
    if __contains_word "${flagname}" "${must_have_one_flag[@]}"; then
        must_have_one_flag=()
    fi

    # keep flag value with flagname as flaghash
    if [ ${flagvalue} ] ; then
        flaghash[${flagname}]=${flagvalue}
    elif [ ${words[ $((c+1)) ]} ] ; then
        flaghash[${flagname}]=${words[ $((c+1)) ]}
    else
        flaghash[${flagname}]="true" # pad "true" for bool flag
    fi

    # skip the argument to a two word flag
    if __contains_word "${words[c]}" "${two_word_flags[@]}"; then
        c=$((c+1))
        # if we are looking for a flags value, don't show commands
        if [[ $c -eq $cword ]]; then
            commands=()
        fi
    fi

    c=$((c+1))

}

__handle_noun()
{
    __debug "${FUNCNAME}: c is $c words[c] is ${words[c]}"

    if __contains_word "${words[c]}" "${must_have_one_noun[@]}"; then
        must_have_one_noun=()
    fi

    nouns+=("${words[c]}")
    c=$((c+1))
}

__handle_command()
{
    __debug "${FUNCNAME}: c is $c words[c] is ${words[c]}"

    local next_command
    if [[ -n ${last_command} ]]; then
        next_command="_${last_command}_${words[c]//:/__}"
    else
        next_command="_${words[c]//:/__}"
    fi
    c=$((c+1))
    __debug "${FUNCNAME}: looking for ${next_command}"
    declare -F $next_command >/dev/null && $next_command
}

__handle_word()
{
    if [[ $c -ge $cword ]]; then
        __handle_reply
        return
    fi
    __debug "${FUNCNAME}: c is $c words[c] is ${words[c]}"
    if [[ "${words[c]}" == -* ]]; then
        __handle_flag
    elif __contains_word "${words[c]}" "${commands[@]}"; then
        __handle_command
    else
        __handle_noun
    fi
    __handle_word
}

`)
}

func postscript(w io.Writer, name string) {
	name = strings.Replace(name, ":", "__", -1)
	fmt.Fprintf(w, "__start_%s()\n", name)
	fmt.Fprintf(w, `{
    local cur prev words cword
    declare -A flaghash
    if declare -F _init_completion >/dev/null 2>&1; then
        _init_completion -s || return
    else
        __my_init_completion || return
    fi

    local c=0
    local flags=()
    local two_word_flags=()
    local flags_with_completion=()
    local flags_completion=()
    local commands=("%s")
    local must_have_one_flag=()
    local must_have_one_noun=()
    local last_command
    local nouns=()

    __handle_word
}

`, name)
	fmt.Fprintf(w, `if [[ $(type -t compopt) = "builtin" ]]; then
    complete -o default -F __start_%s %s
else
    complete -o default -o nospace -F __start_%s %s
fi

`, name, name, name, name)
	fmt.Fprintf(w, "# ex: ts=4 sw=4 et filetype=sh\n")
}

func writeCommands(cmd *Command, w io.Writer) {
	fmt.Fprintf(w, "    commands=()\n")
	for _, c := range cmd.Commands() {
		if !c.IsAvailableCommand() || c == cmd.helpCommand {
			continue
		}
		fmt.Fprintf(w, "    commands+=(%q)\n", c.Name())
	}
	fmt.Fprintf(w, "\n")
}

func writeFlagHandler(name string, annotations map[string][]string, w io.Writer) {
	for key, value := range annotations {
		switch key {
		case BashCompFilenameExt:
			fmt.Fprintf(w, "    flags_with_completion+=(%q)\n", name)

			if len(value) > 0 {
				ext := "__handle_filename_extension_flag " + strings.Join(value, "|")
				fmt.Fprintf(w, "    flags_completion+=(%q)\n", ext)
			} else {
				ext := "_filedir"
				fmt.Fprintf(w, "    flags_completion+=(%q)\n", ext)
			}
		case BashCompSubdirsInDir:
			fmt.Fprintf(w, "    flags_with_completion+=(%q)\n", name)

			if len(value) == 1 {
				ext := "__handle_subdirs_in_dir_flag " + value[0]
				fmt.Fprintf(w, "    flags_completion+=(%q)\n", ext)
			} else {
				ext := "_filedir -d"
				fmt.Fprintf(w, "    flags_completion+=(%q)\n", ext)
			}
		}
	}
}

func writeShortFlag(flag *pflag.Flag, w io.Writer) {
	b := (flag.Value.Type() == "bool")
	name := flag.Shorthand
	format := "    "
	if !b {
		format += "two_word_"
	}
	format += "flags+=(\"-%s\")\n"
	fmt.Fprintf(w, format, name)
	writeFlagHandler("-"+name, flag.Annotations, w)
}

func writeFlag(flag *pflag.Flag, w io.Writer) {
	b := (flag.Value.Type() == "bool")
	name := flag.Name
	format := "    flags+=(\"--%s"
	if !b {
		format += "="
	}
	format += "\")\n"
	fmt.Fprintf(w, format, name)
	writeFlagHandler("--"+name, flag.Annotations, w)
}

func writeFlags(cmd *Command, w io.Writer) {
	fmt.Fprintf(w, `    flags=()
    two_word_flags=()
    flags_with_completion=()
    flags_completion=()

`)
	cmd.NonInheritedFlags().VisitAll(func(flag *pflag.Flag) {
		writeFlag(flag, w)
		if len(flag.Shorthand) > 0 {
			writeShortFlag(flag, w)
		}
	})
	cmd.InheritedFlags().VisitAll(func(flag *pflag.Flag) {
		writeFlag(flag, w)
		if len(flag.Shorthand) > 0 {
			writeShortFlag(flag, w)
		}
	})

	fmt.Fprintf(w, "\n")
}

func writeRequiredFlag(cmd *Command, w io.Writer) {
	fmt.Fprintf(w, "    must_have_one_flag=()\n")
	flags := cmd.NonInheritedFlags()
	flags.VisitAll(func(flag *pflag.Flag) {
		for key := range flag.Annotations {
			switch key {
			case BashCompOneRequiredFlag:
				format := "    must_have_one_flag+=(\"--%s"
				b := (flag.Value.Type() == "bool")
				if !b {
					format += "="
				}
				format += "\")\n"
				fmt.Fprintf(w, format, flag.Name)

				if len(flag.Shorthand) > 0 {
					fmt.Fprintf(w, "    must_have_one_flag+=(\"-%s\")\n", flag.Shorthand)
				}
			}
		}
	})
}

func writeRequiredNoun(cmd *Command, w io.Writer) {
	fmt.Fprintf(w, "    must_have_one_noun=()\n")
	sort.Sort(sort.StringSlice(cmd.ValidArgs))
	for _, value := range cmd.ValidArgs {
		fmt.Fprintf(w, "    must_have_one_noun+=(%q)\n", value)
	}
}

func gen(cmd *Command, w io.Writer) {
	for _, c := range cmd.Commands() {
		if !c.IsAvailableCommand() || c == cmd.helpCommand {
			continue
		}
		gen(c, w)
	}
	commandName := cmd.CommandPath()
	commandName = strings.Replace(commandName, " ", "_", -1)
	commandName = strings.Replace(commandName, ":", "__", -1)
	fmt.Fprintf(w, "_%s()\n{\n", commandName)
	fmt.Fprintf(w, "    last_command=%q\n", commandName)
	writeCommands(cmd, w)
	writeFlags(cmd, w)
	writeRequiredFlag(cmd, w)
	writeRequiredNoun(cmd, w)
	fmt.Fprintf(w, "}\n\n")
}

func (cmd *Command) GenBashCompletion(w io.Writer) {
	preamble(w)
	if len(cmd.BashCompletionFunction) > 0 {
		fmt.Fprintf(w, "%s\n", cmd.BashCompletionFunction)
	}
	gen(cmd, w)
	postscript(w, cmd.Name())
}

func (cmd *Command) GenBashCompletionFile(filename string) error {
	out := new(bytes.Buffer)

	cmd.GenBashCompletion(out)

	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	_, err = outFile.Write(out.Bytes())
	if err != nil {
		return err
	}
	return nil
}

// MarkFlagRequired adds the BashCompOneRequiredFlag annotation to the named flag, if it exists.
func (cmd *Command) MarkFlagRequired(name string) error {
	return MarkFlagRequired(cmd.Flags(), name)
}

// MarkPersistentFlagRequired adds the BashCompOneRequiredFlag annotation to the named persistent flag, if it exists.
func (cmd *Command) MarkPersistentFlagRequired(name string) error {
	return MarkFlagRequired(cmd.PersistentFlags(), name)
}

// MarkFlagRequired adds the BashCompOneRequiredFlag annotation to the named flag in the flag set, if it exists.
func MarkFlagRequired(flags *pflag.FlagSet, name string) error {
	return flags.SetAnnotation(name, BashCompOneRequiredFlag, []string{"true"})
}

// MarkFlagFilename adds the BashCompFilenameExt annotation to the named flag, if it exists.
// Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
func (cmd *Command) MarkFlagFilename(name string, extensions ...string) error {
	return MarkFlagFilename(cmd.Flags(), name, extensions...)
}

// MarkPersistentFlagFilename adds the BashCompFilenameExt annotation to the named persistent flag, if it exists.
// Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
func (cmd *Command) MarkPersistentFlagFilename(name string, extensions ...string) error {
	return MarkFlagFilename(cmd.PersistentFlags(), name, extensions...)
}

// MarkFlagFilename adds the BashCompFilenameExt annotation to the named flag in the flag set, if it exists.
// Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
func MarkFlagFilename(flags *pflag.FlagSet, name string, extensions ...string) error {
	return flags.SetAnnotation(name, BashCompFilenameExt, extensions)
}
<|endoftext|>


<commit_before>package couchdb

import (
	"encoding/json"
	"testing"
)

const TESTDBNAME = "testdb"
const HOST = "127.0.0.1"
const PORT = 5984

func TestNewDb(t *testing.T) {
	client := NewClient(HOST, PORT)
	dbObj := client.DB(TESTDBNAME)
	status, err := dbObj.Exists()
	if err == nil {
		t.Log("Test ran without errors: ", TESTDBNAME, " exists is ", status)
	} else {
		t.Error("Error checking Exists for DB")
	}
}

// DBObject representation of the database under consideration.
var DBObject Database

func TestCreateDb(t *testing.T) {
	client := NewClient(HOST, PORT)
	dbObj := client.DB(TESTDBNAME)
	status, err := dbObj.Exists()
	if err == nil {
		if status == false {
			err = dbObj.Create()
			if err != nil {
				t.Error("Error creating DB ", err)
			}
		}
	} else {
		t.Error("Error running exists ", err)
	}
	DBObject = dbObj
}

type TestObj struct {
	Name string `json:"name"`
	Age  int    `json:"age"`
}

var Id string
var Rev string

func TestCreateDocument(t *testing.T) {

	testObj := &TestObj{
		Name: "Fred",
		Age:  18,
	}

	for i := 0; i < 10; i++ {
		strObj, err := json.Marshal(testObj)
		testObj.Age++
		t.Log("Saving Object:", strObj)
		if err != nil {
			t.Error("Error Marshalling testObj")
		}

		err, status := DBObject.CreateDocument(strObj)
		if err != nil {
			t.Error("Error creating Document ", err)
		} else {
			t.Log("Successfully created document " + TESTDBNAME)
			t.Log("Document Id " + status.Id)
			t.Log("Document Revision " + status.Rev)
			Id = status.Id
			Rev = status.Rev
		}
	}
}

func TestUpdateDocument(t *testing.T) {

	type UpdateObj struct {
		CouchWrapperUpdate
		TestObj
	}

	testObj := &UpdateObj{}

	testObj.Id = Id
	testObj.Rev = Rev

	testObj.Name = "Fred Updated"
	testObj.Age = 25

	strObj, err := json.Marshal(testObj)
	if err != nil {
		t.Error("Error Marshalling testObj")
	}

	err, status := DBObject.UpdateDocument(strObj)
	if err != nil {
		t.Error("Error Updating Document", err)
	} else {
		t.Log("Successfully updated document")
		t.Log("Document Id " + status.Id)
		t.Log("Revision " + status.Rev)
	}
}

func TestGetObject(t *testing.T) {

	err, jsonObj := DBObject.RetrieveDocument(Id)
	if err != nil {
		t.Error("Error ", err)
	}

	type UpdateObj struct {
		CouchWrapperUpdate
		TestObj
	}

	obj := &UpdateObj{}
	json.Unmarshal(jsonObj, obj)
	if obj.Id != Id {
		t.Error("Id should be the same as requested")
	}
}

func TestGetView(t *testing.T) {
	err, data := DBObject.GetView("_design/by_age", "by_age")
	if err != nil {
		t.Error("Error :", err)
	} else {
		t.Log(string(data))
	}
}

func TestDeleteDb(t *testing.T) {

	err := DBObject.Delete()
	if err == nil {
		t.Log("Deleted existing db " + TESTDBNAME + " Successful.")
	} else {
		t.Error("Error deleting "+TESTDBNAME, " ", err)
	}
}
<commit_msg>adding the createView stuff<commit_after>package couchdb

import (
	"encoding/json"
	"testing"
)

const TESTDBNAME = "testdb"
const HOST = "127.0.0.1"
const PORT = 5984

func TestNewDb(t *testing.T) {
	client := NewClient(HOST, PORT)
	dbObj := client.DB(TESTDBNAME)
	status, err := dbObj.Exists()
	if err == nil {
		t.Log("Test ran without errors: ", TESTDBNAME, " exists is ", status)
	} else {
		t.Error("Error checking Exists for DB")
	}
}

// DBObject representation of the database under consideration.
var DBObject Database

func TestCreateDb(t *testing.T) {
	client := NewClient(HOST, PORT)
	dbObj := client.DB(TESTDBNAME)
	status, err := dbObj.Exists()
	if err == nil {
		if status == false {
			err = dbObj.Create()
			if err != nil {
				t.Error("Error creating DB ", err)
			}
		}
	} else {
		t.Error("Error running exists ", err)
	}
	DBObject = dbObj
}

type TestObj struct {
	Name string `json:"name"`
	Age  int    `json:"age"`
}

var Id string
var Rev string

func TestCreateDocument(t *testing.T) {

	testObj := &TestObj{
		Name: "Fred",
		Age:  18,
	}

	for i := 0; i < 10; i++ {
		strObj, err := json.Marshal(testObj)
		testObj.Age++
		t.Log("Saving Object:", strObj)
		if err != nil {
			t.Error("Error Marshalling testObj")
		}

		err, status := DBObject.CreateDocument(strObj)
		if err != nil {
			t.Error("Error creating Document ", err)
		} else {
			t.Log("Successfully created document " + TESTDBNAME)
			t.Log("Document Id " + status.Id)
			t.Log("Document Revision " + status.Rev)
			Id = status.Id
			Rev = status.Rev
		}
	}
}

func TestUpdateDocument(t *testing.T) {

	type UpdateObj struct {
		CouchWrapperUpdate
		TestObj
	}

	testObj := &UpdateObj{}

	testObj.Id = Id
	testObj.Rev = Rev

	testObj.Name = "Fred Updated"
	testObj.Age = 25

	strObj, err := json.Marshal(testObj)
	if err != nil {
		t.Error("Error Marshalling testObj")
	}

	err, status := DBObject.UpdateDocument(strObj)
	if err != nil {
		t.Error("Error Updating Document", err)
	} else {
		t.Log("Successfully updated document")
		t.Log("Document Id " + status.Id)
		t.Log("Revision " + status.Rev)
	}
}

func TestGetObject(t *testing.T) {

	err, jsonObj := DBObject.RetrieveDocument(Id)
	if err != nil {
		t.Error("Error ", err)
	}

	type UpdateObj struct {
		CouchWrapperUpdate
		TestObj
	}

	obj := &UpdateObj{}
	json.Unmarshal(jsonObj, obj)
	if obj.Id != Id {
		t.Error("Id should be the same as requested")
	}
}

//TODO: Only one design can be added, extend to multiple as soon as possible.
func TestCreateView(t *testing.T) {
	err, _ := DBObject.CreateView("by_age", "doc", "doc.age < 22", "doc.age, doc.name")
	if err != nil {
		t.Error("Error creating view", err)
	}
}

func TestGetView(t *testing.T) {
	err, data := DBObject.GetView("_design/by_age", "by_age")
	if err != nil {
		t.Error("Error :", err)
	} else {
		t.Log(string(data))
	}
}

func TestDeleteDb(t *testing.T) {

	err := DBObject.Delete()
	if err == nil {
		t.Log("Deleted existing db " + TESTDBNAME + " Successful.")
	} else {
		t.Error("Error deleting "+TESTDBNAME, " ", err)
	}
}
<|endoftext|>