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

package main

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

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

	"github.com/juju/juju/cmd/envcmd"
	cmdtesting "github.com/juju/juju/cmd/testing"
	"github.com/juju/juju/juju/osenv"
	_ "github.com/juju/juju/provider/dummy"
	"github.com/juju/juju/testing"
	"github.com/juju/juju/version"
)

type MainSuite struct {
	testing.FakeJujuHomeSuite
}

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

func deployHelpText() string {
	return cmdtesting.HelpText(envcmd.Wrap(&DeployCommand{}), "juju deploy")
}

func setHelpText() string {
	return cmdtesting.HelpText(envcmd.Wrap(&SetCommand{}), "juju set")
}

func syncToolsHelpText() string {
	return cmdtesting.HelpText(envcmd.Wrap(&SyncToolsCommand{}), "juju sync-tools")
}

func (s *MainSuite) TestRunMain(c *gc.C) {
	// The test array structure needs to be inline here as some of the
	// expected values below use deployHelpText().  This constructs the deploy
	// command and runs gets the help for it.  When the deploy command is
	// setting the flags (which is needed for the help text) it is accessing
	// osenv.JujuHome(), which panics if SetJujuHome has not been called.
	// The FakeHome from testing does this.
	for i, t := range []struct {
		summary string
		args    []string
		code    int
		out     string
	}{{
		summary: "no params shows help",
		args:    []string{},
		code:    0,
		out:     strings.TrimLeft(helpBasics, "\n"),
	}, {
		summary: "juju help is the same as juju",
		args:    []string{"help"},
		code:    0,
		out:     strings.TrimLeft(helpBasics, "\n"),
	}, {
		summary: "juju --help works too",
		args:    []string{"--help"},
		code:    0,
		out:     strings.TrimLeft(helpBasics, "\n"),
	}, {
		summary: "juju help basics is the same as juju",
		args:    []string{"help", "basics"},
		code:    0,
		out:     strings.TrimLeft(helpBasics, "\n"),
	}, {
		summary: "juju help foo doesn't exist",
		args:    []string{"help", "foo"},
		code:    1,
		out:     "ERROR unknown command or topic for foo\n",
	}, {
		summary: "juju help deploy shows the default help without global options",
		args:    []string{"help", "deploy"},
		code:    0,
		out:     deployHelpText(),
	}, {
		summary: "juju --help deploy shows the same help as 'help deploy'",
		args:    []string{"--help", "deploy"},
		code:    0,
		out:     deployHelpText(),
	}, {
		summary: "juju deploy --help shows the same help as 'help deploy'",
		args:    []string{"deploy", "--help"},
		code:    0,
		out:     deployHelpText(),
	}, {
		summary: "juju help set shows the default help without global options",
		args:    []string{"help", "set"},
		code:    0,
		out:     setHelpText(),
	}, {
		summary: "juju --help set shows the same help as 'help set'",
		args:    []string{"--help", "set"},
		code:    0,
		out:     setHelpText(),
	}, {
		summary: "juju set --help shows the same help as 'help set'",
		args:    []string{"set", "--help"},
		code:    0,
		out:     setHelpText(),
	}, {
		summary: "unknown command",
		args:    []string{"discombobulate"},
		code:    1,
		out:     "ERROR unrecognized command: juju discombobulate\n",
	}, {
		summary: "unknown option before command",
		args:    []string{"--cheese", "bootstrap"},
		code:    2,
		out:     "error: flag provided but not defined: --cheese\n",
	}, {
		summary: "unknown option after command",
		args:    []string{"bootstrap", "--cheese"},
		code:    2,
		out:     "error: flag provided but not defined: --cheese\n",
	}, {
		summary: "known option, but specified before command",
		args:    []string{"--environment", "blah", "bootstrap"},
		code:    2,
		out:     "error: flag provided but not defined: --environment\n",
	}, {
		summary: "juju sync-tools registered properly",
		args:    []string{"sync-tools", "--help"},
		code:    0,
		out:     syncToolsHelpText(),
	}, {
		summary: "check version command registered properly",
		args:    []string{"version"},
		code:    0,
		out:     version.Current.String() + "\n",
	},
	} {
		c.Logf("test %d: %s", i, t.summary)
		out := badrun(c, t.code, t.args...)
		c.Assert(out, gc.Equals, t.out)
	}
}

func (s *MainSuite) TestActualRunJujuArgOrder(c *gc.C) {
	logpath := filepath.Join(c.MkDir(), "log")
	tests := [][]string{
		{"--log-file", logpath, "--debug", "env"}, // global flags before
		{"env", "--log-file", logpath, "--debug"}, // after
		{"--log-file", logpath, "env", "--debug"}, // mixed
	}
	for i, test := range tests {
		c.Logf("test %d: %v", i, test)
		badrun(c, 0, test...)
		content, err := ioutil.ReadFile(logpath)
		c.Assert(err, gc.IsNil)
		c.Assert(string(content), gc.Matches, "(.|\n)*running juju(.|\n)*command finished(.|\n)*")
		err = os.Remove(logpath)
		c.Assert(err, gc.IsNil)
	}
}

var commandNames = []string{
	"add-machine",
	"add-relation",
	"add-unit",
	"api-endpoints",
	"authorised-keys", // alias for authorized-keys
	"authorized-keys",
	"backups",
	"bootstrap",
	"debug-hooks",
	"debug-log",
	"deploy",
	"destroy-environment",
	"destroy-machine",
	"destroy-relation",
	"destroy-service",
	"destroy-unit",
	"ensure-availability",
	"env", // alias for switch
	"expose",
	"generate-config", // alias for init
	"get",
	"get-constraints",
	"get-env", // alias for get-environment
	"get-environment",
	"help",
	"help-tool",
	"init",
	"publish",
	"remove-machine",  // alias for destroy-machine
	"remove-relation", // alias for destroy-relation
	"remove-service",  // alias for destroy-service
	"remove-unit",     // alias for destroy-unit
	"resolved",
	"retry-provisioning",
	"run",
	"scp",
	"set",
	"set-constraints",
	"set-env", // alias for set-environment
	"set-environment",
	"ssh",
	"stat", // alias for status
	"status",
	"switch",
	"sync-tools",
	"terminate-machine", // alias for destroy-machine
	"unexpose",
	"unset",
	"unset-env", // alias for unset-environment
	"unset-environment",
	"upgrade-charm",
	"upgrade-juju",
	"user",
	"version",
}

func (s *MainSuite) TestHelpCommands(c *gc.C) {
	// Check that we have correctly registered all the commands
	// by checking the help output.
	defer osenv.SetJujuHome(osenv.SetJujuHome(c.MkDir()))
	out := badrun(c, 0, "help", "commands")
	lines := strings.Split(out, "\n")
	var names []string
	for _, line := range lines {
		f := strings.Fields(line)
		if len(f) == 0 {
			continue
		}
		names = append(names, f[0])
	}
	// The names should be output in alphabetical order, so don't sort.
	c.Assert(names, jc.DeepEquals, commandNames)
}

var topicNames = []string{
	"azure-provider",
	"basics",
	"commands",
	"constraints",
	"ec2-provider",
	"global-options",
	"glossary",
	"hpcloud-provider",
	"local-provider",
	"logging",
	"openstack-provider",
	"plugins",
	"topics",
}

func (s *MainSuite) TestHelpTopics(c *gc.C) {
	// Check that we have correctly registered all the topics
	// by checking the help output.
	defer osenv.SetJujuHome(osenv.SetJujuHome(c.MkDir()))
	out := badrun(c, 0, "help", "topics")
	lines := strings.Split(out, "\n")
	var names []string
	for _, line := range lines {
		f := strings.Fields(line)
		if len(f) == 0 {
			continue
		}
		names = append(names, f[0])
	}
	// The names should be output in alphabetical order, so don't sort.
	c.Assert(names, gc.DeepEquals, topicNames)
}

var globalFlags = []string{
	"--debug .*",
	"--description .*",
	"-h, --help .*",
	"--log-file .*",
	"--logging-config .*",
	"-q, --quiet .*",
	"--show-log .*",
	"-v, --verbose .*",
}

func (s *MainSuite) TestHelpGlobalOptions(c *gc.C) {
	// Check that we have correctly registered all the topics
	// by checking the help output.
	defer osenv.SetJujuHome(osenv.SetJujuHome(c.MkDir()))
	out := badrun(c, 0, "help", "global-options")
	c.Assert(out, gc.Matches, `Global Options

These options may be used with any command, and may appear in front of any
command\.(.|\n)*`)
	lines := strings.Split(out, "\n")
	var flags []string
	for _, line := range lines {
		f := strings.Fields(line)
		if len(f) == 0 || line[0] != '-' {
			continue
		}
		flags = append(flags, line)
	}
	c.Assert(len(flags), gc.Equals, len(globalFlags))
	for i, line := range flags {
		c.Assert(line, gc.Matches, globalFlags[i])
	}
}

type commands []cmd.Command

func (r *commands) Register(c cmd.Command) {
	*r = append(*r, c)
}

func (s *MainSuite) TestEnvironCommands(c *gc.C) {
	var commands commands
	registerCommands(&commands, testing.Context(c))
	// There should not be any EnvironCommands registered.
	// EnvironCommands must be wrapped using envcmd.Wrap.
	for _, cmd := range commands {
		c.Logf("%v", cmd.Info().Name)
		c.Check(cmd, gc.Not(gc.FitsTypeOf), envcmd.EnvironCommand(&BootstrapCommand{}))
	}
}
<commit_msg>Forgot to add api-info to the help list test.<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

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

	"github.com/juju/juju/cmd/envcmd"
	cmdtesting "github.com/juju/juju/cmd/testing"
	"github.com/juju/juju/juju/osenv"
	_ "github.com/juju/juju/provider/dummy"
	"github.com/juju/juju/testing"
	"github.com/juju/juju/version"
)

type MainSuite struct {
	testing.FakeJujuHomeSuite
}

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

func deployHelpText() string {
	return cmdtesting.HelpText(envcmd.Wrap(&DeployCommand{}), "juju deploy")
}

func setHelpText() string {
	return cmdtesting.HelpText(envcmd.Wrap(&SetCommand{}), "juju set")
}

func syncToolsHelpText() string {
	return cmdtesting.HelpText(envcmd.Wrap(&SyncToolsCommand{}), "juju sync-tools")
}

func (s *MainSuite) TestRunMain(c *gc.C) {
	// The test array structure needs to be inline here as some of the
	// expected values below use deployHelpText().  This constructs the deploy
	// command and runs gets the help for it.  When the deploy command is
	// setting the flags (which is needed for the help text) it is accessing
	// osenv.JujuHome(), which panics if SetJujuHome has not been called.
	// The FakeHome from testing does this.
	for i, t := range []struct {
		summary string
		args    []string
		code    int
		out     string
	}{{
		summary: "no params shows help",
		args:    []string{},
		code:    0,
		out:     strings.TrimLeft(helpBasics, "\n"),
	}, {
		summary: "juju help is the same as juju",
		args:    []string{"help"},
		code:    0,
		out:     strings.TrimLeft(helpBasics, "\n"),
	}, {
		summary: "juju --help works too",
		args:    []string{"--help"},
		code:    0,
		out:     strings.TrimLeft(helpBasics, "\n"),
	}, {
		summary: "juju help basics is the same as juju",
		args:    []string{"help", "basics"},
		code:    0,
		out:     strings.TrimLeft(helpBasics, "\n"),
	}, {
		summary: "juju help foo doesn't exist",
		args:    []string{"help", "foo"},
		code:    1,
		out:     "ERROR unknown command or topic for foo\n",
	}, {
		summary: "juju help deploy shows the default help without global options",
		args:    []string{"help", "deploy"},
		code:    0,
		out:     deployHelpText(),
	}, {
		summary: "juju --help deploy shows the same help as 'help deploy'",
		args:    []string{"--help", "deploy"},
		code:    0,
		out:     deployHelpText(),
	}, {
		summary: "juju deploy --help shows the same help as 'help deploy'",
		args:    []string{"deploy", "--help"},
		code:    0,
		out:     deployHelpText(),
	}, {
		summary: "juju help set shows the default help without global options",
		args:    []string{"help", "set"},
		code:    0,
		out:     setHelpText(),
	}, {
		summary: "juju --help set shows the same help as 'help set'",
		args:    []string{"--help", "set"},
		code:    0,
		out:     setHelpText(),
	}, {
		summary: "juju set --help shows the same help as 'help set'",
		args:    []string{"set", "--help"},
		code:    0,
		out:     setHelpText(),
	}, {
		summary: "unknown command",
		args:    []string{"discombobulate"},
		code:    1,
		out:     "ERROR unrecognized command: juju discombobulate\n",
	}, {
		summary: "unknown option before command",
		args:    []string{"--cheese", "bootstrap"},
		code:    2,
		out:     "error: flag provided but not defined: --cheese\n",
	}, {
		summary: "unknown option after command",
		args:    []string{"bootstrap", "--cheese"},
		code:    2,
		out:     "error: flag provided but not defined: --cheese\n",
	}, {
		summary: "known option, but specified before command",
		args:    []string{"--environment", "blah", "bootstrap"},
		code:    2,
		out:     "error: flag provided but not defined: --environment\n",
	}, {
		summary: "juju sync-tools registered properly",
		args:    []string{"sync-tools", "--help"},
		code:    0,
		out:     syncToolsHelpText(),
	}, {
		summary: "check version command registered properly",
		args:    []string{"version"},
		code:    0,
		out:     version.Current.String() + "\n",
	},
	} {
		c.Logf("test %d: %s", i, t.summary)
		out := badrun(c, t.code, t.args...)
		c.Assert(out, gc.Equals, t.out)
	}
}

func (s *MainSuite) TestActualRunJujuArgOrder(c *gc.C) {
	logpath := filepath.Join(c.MkDir(), "log")
	tests := [][]string{
		{"--log-file", logpath, "--debug", "env"}, // global flags before
		{"env", "--log-file", logpath, "--debug"}, // after
		{"--log-file", logpath, "env", "--debug"}, // mixed
	}
	for i, test := range tests {
		c.Logf("test %d: %v", i, test)
		badrun(c, 0, test...)
		content, err := ioutil.ReadFile(logpath)
		c.Assert(err, gc.IsNil)
		c.Assert(string(content), gc.Matches, "(.|\n)*running juju(.|\n)*command finished(.|\n)*")
		err = os.Remove(logpath)
		c.Assert(err, gc.IsNil)
	}
}

var commandNames = []string{
	"add-machine",
	"add-relation",
	"add-unit",
	"api-endpoints",
	"api-info",
	"authorised-keys", // alias for authorized-keys
	"authorized-keys",
	"backups",
	"bootstrap",
	"debug-hooks",
	"debug-log",
	"deploy",
	"destroy-environment",
	"destroy-machine",
	"destroy-relation",
	"destroy-service",
	"destroy-unit",
	"ensure-availability",
	"env", // alias for switch
	"expose",
	"generate-config", // alias for init
	"get",
	"get-constraints",
	"get-env", // alias for get-environment
	"get-environment",
	"help",
	"help-tool",
	"init",
	"publish",
	"remove-machine",  // alias for destroy-machine
	"remove-relation", // alias for destroy-relation
	"remove-service",  // alias for destroy-service
	"remove-unit",     // alias for destroy-unit
	"resolved",
	"retry-provisioning",
	"run",
	"scp",
	"set",
	"set-constraints",
	"set-env", // alias for set-environment
	"set-environment",
	"ssh",
	"stat", // alias for status
	"status",
	"switch",
	"sync-tools",
	"terminate-machine", // alias for destroy-machine
	"unexpose",
	"unset",
	"unset-env", // alias for unset-environment
	"unset-environment",
	"upgrade-charm",
	"upgrade-juju",
	"user",
	"version",
}

func (s *MainSuite) TestHelpCommands(c *gc.C) {
	// Check that we have correctly registered all the commands
	// by checking the help output.
	defer osenv.SetJujuHome(osenv.SetJujuHome(c.MkDir()))
	out := badrun(c, 0, "help", "commands")
	lines := strings.Split(out, "\n")
	var names []string
	for _, line := range lines {
		f := strings.Fields(line)
		if len(f) == 0 {
			continue
		}
		names = append(names, f[0])
	}
	// The names should be output in alphabetical order, so don't sort.
	c.Assert(names, jc.DeepEquals, commandNames)
}

var topicNames = []string{
	"azure-provider",
	"basics",
	"commands",
	"constraints",
	"ec2-provider",
	"global-options",
	"glossary",
	"hpcloud-provider",
	"local-provider",
	"logging",
	"openstack-provider",
	"plugins",
	"topics",
}

func (s *MainSuite) TestHelpTopics(c *gc.C) {
	// Check that we have correctly registered all the topics
	// by checking the help output.
	defer osenv.SetJujuHome(osenv.SetJujuHome(c.MkDir()))
	out := badrun(c, 0, "help", "topics")
	lines := strings.Split(out, "\n")
	var names []string
	for _, line := range lines {
		f := strings.Fields(line)
		if len(f) == 0 {
			continue
		}
		names = append(names, f[0])
	}
	// The names should be output in alphabetical order, so don't sort.
	c.Assert(names, gc.DeepEquals, topicNames)
}

var globalFlags = []string{
	"--debug .*",
	"--description .*",
	"-h, --help .*",
	"--log-file .*",
	"--logging-config .*",
	"-q, --quiet .*",
	"--show-log .*",
	"-v, --verbose .*",
}

func (s *MainSuite) TestHelpGlobalOptions(c *gc.C) {
	// Check that we have correctly registered all the topics
	// by checking the help output.
	defer osenv.SetJujuHome(osenv.SetJujuHome(c.MkDir()))
	out := badrun(c, 0, "help", "global-options")
	c.Assert(out, gc.Matches, `Global Options

These options may be used with any command, and may appear in front of any
command\.(.|\n)*`)
	lines := strings.Split(out, "\n")
	var flags []string
	for _, line := range lines {
		f := strings.Fields(line)
		if len(f) == 0 || line[0] != '-' {
			continue
		}
		flags = append(flags, line)
	}
	c.Assert(len(flags), gc.Equals, len(globalFlags))
	for i, line := range flags {
		c.Assert(line, gc.Matches, globalFlags[i])
	}
}

type commands []cmd.Command

func (r *commands) Register(c cmd.Command) {
	*r = append(*r, c)
}

func (s *MainSuite) TestEnvironCommands(c *gc.C) {
	var commands commands
	registerCommands(&commands, testing.Context(c))
	// There should not be any EnvironCommands registered.
	// EnvironCommands must be wrapped using envcmd.Wrap.
	for _, cmd := range commands {
		c.Logf("%v", cmd.Info().Name)
		c.Check(cmd, gc.Not(gc.FitsTypeOf), envcmd.EnvironCommand(&BootstrapCommand{}))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/mattn/go-runewidth"
	"github.com/zyedidia/tcell"
)

func min(a, b int) int {
	if a <= b {
		return a
	}
	return b
}

func visualToCharPos(visualIndex int, lineN int, str string, buf *Buffer, tabsize int) (int, int, *tcell.Style) {
	charPos := 0
	var lineIdx int
	var lastWidth int
	var style *tcell.Style
	var width int
	var rw int
	for i, c := range str {
		// width := StringWidth(str[:i], tabsize)

		if group, ok := buf.Match(lineN)[charPos]; ok {
			s := GetColor(group.String())
			style = &s
		}

		if width >= visualIndex {
			return charPos, visualIndex - lastWidth, style
		}

		if i != 0 {
			charPos++
			lineIdx += rw
		}
		lastWidth = width
		rw = 0
		if c == '\t' {
			rw := tabsize - (lineIdx % tabsize)
			width += rw
		} else {
			rw = runewidth.RuneWidth(c)
			width += rw
		}
	}

	return -1, -1, style
}

type Char struct {
	visualLoc Loc
	realLoc   Loc
	char      rune
	// The actual character that is drawn
	// This is only different from char if it's for example hidden character
	drawChar rune
	style    tcell.Style
	width    int
}

type CellView struct {
	lines [][]*Char
}

func (c *CellView) Draw(buf *Buffer, top, height, left, width int) {
	tabsize := int(buf.Settings["tabsize"].(float64))
	softwrap := buf.Settings["softwrap"].(bool)
	indentchar := []rune(buf.Settings["indentchar"].(string))[0]

	start := buf.Cursor.Y
	if buf.Settings["syntax"].(bool) && buf.syntaxDef != nil {
		if start > 0 && buf.lines[start-1].rehighlight {
			buf.highlighter.ReHighlightLine(buf, start-1)
			buf.lines[start-1].rehighlight = false
		}

		buf.highlighter.ReHighlightStates(buf, start)

		buf.highlighter.HighlightMatches(buf, top, top+height)
	}

	c.lines = make([][]*Char, 0)

	viewLine := 0
	lineN := top

	curStyle := defStyle
	for viewLine < height {
		if lineN >= len(buf.lines) {
			break
		}

		lineStr := buf.Line(lineN)
		line := []rune(lineStr)

		colN, startOffset, startStyle := visualToCharPos(left, lineN, lineStr, buf, tabsize)
		if colN < 0 {
			colN = len(line)
		}
		viewCol := -startOffset
		if startStyle != nil {
			curStyle = *startStyle
		}

		// We'll either draw the length of the line, or the width of the screen
		// whichever is smaller
		lineLength := min(StringWidth(lineStr, tabsize), width)
		c.lines = append(c.lines, make([]*Char, lineLength))

		wrap := false
		// We only need to wrap if the length of the line is greater than the width of the terminal screen
		if softwrap && StringWidth(lineStr, tabsize) > width {
			wrap = true
			// We're going to draw the entire line now
			lineLength = StringWidth(lineStr, tabsize)
		}

		for viewCol < lineLength {
			if colN >= len(line) {
				break
			}
			if group, ok := buf.Match(lineN)[colN]; ok {
				curStyle = GetColor(group.String())
			}

			char := line[colN]

			if viewCol >= 0 {
				c.lines[viewLine][viewCol] = &Char{Loc{viewCol, viewLine}, Loc{colN, lineN}, char, char, curStyle, 1}
			}
			if char == '\t' {
				charWidth := tabsize - (viewCol+left)%tabsize
				if viewCol >= 0 {
					c.lines[viewLine][viewCol].drawChar = indentchar
					c.lines[viewLine][viewCol].width = charWidth

					indentStyle := curStyle
					if group, ok := colorscheme["indent-char"]; ok {
						indentStyle = group
					}

					c.lines[viewLine][viewCol].style = indentStyle
				}

				for i := 1; i < charWidth; i++ {
					viewCol++
					if viewCol >= 0 && viewCol < lineLength {
						c.lines[viewLine][viewCol] = &Char{Loc{viewCol, viewLine}, Loc{colN, lineN}, char, ' ', curStyle, 1}
					}
				}
				viewCol++
			} else if runewidth.RuneWidth(char) > 1 {
				charWidth := runewidth.RuneWidth(char)
				if viewCol >= 0 {
					c.lines[viewLine][viewCol].width = charWidth
				}
				for i := 1; i < charWidth; i++ {
					viewCol++
					if viewCol >= 0 && viewCol < lineLength {
						c.lines[viewLine][viewCol] = &Char{Loc{viewCol, viewLine}, Loc{colN, lineN}, char, ' ', curStyle, 1}
					}
				}
				viewCol++
			} else {
				viewCol++
			}
			colN++

			if wrap && viewCol >= width {
				viewLine++

				// If we go too far soft wrapping we have to cut off
				if viewLine >= height {
					break
				}

				nextLine := line[colN:]
				lineLength := min(StringWidth(string(nextLine), tabsize), width)
				c.lines = append(c.lines, make([]*Char, lineLength))

				viewCol = 0
			}

		}
		if group, ok := buf.Match(lineN)[len(line)]; ok {
			curStyle = GetColor(group.String())
		}

		// newline
		viewLine++
		lineN++
	}

	for i := top; i < top+height; i++ {
		if i >= buf.NumLines {
			break
		}
		buf.SetMatch(i, nil)
	}
}
<commit_msg>Fix rare problem with tabs and horizontal scrolling<commit_after>package main

import (
	"github.com/mattn/go-runewidth"
	"github.com/zyedidia/tcell"
)

func min(a, b int) int {
	if a <= b {
		return a
	}
	return b
}

func visualToCharPos(visualIndex int, lineN int, str string, buf *Buffer, tabsize int) (int, int, *tcell.Style) {
	charPos := 0
	var lineIdx int
	var lastWidth int
	var style *tcell.Style
	var width int
	var rw int
	for i, c := range str {
		// width := StringWidth(str[:i], tabsize)

		if group, ok := buf.Match(lineN)[charPos]; ok {
			s := GetColor(group.String())
			style = &s
		}

		if width >= visualIndex {
			return charPos, visualIndex - lastWidth, style
		}

		if i != 0 {
			charPos++
			lineIdx += rw
		}
		lastWidth = width
		rw = 0
		if c == '\t' {
			rw = tabsize - (lineIdx % tabsize)
			width += rw
		} else {
			rw = runewidth.RuneWidth(c)
			width += rw
		}
	}

	return -1, -1, style
}

type Char struct {
	visualLoc Loc
	realLoc   Loc
	char      rune
	// The actual character that is drawn
	// This is only different from char if it's for example hidden character
	drawChar rune
	style    tcell.Style
	width    int
}

type CellView struct {
	lines [][]*Char
}

func (c *CellView) Draw(buf *Buffer, top, height, left, width int) {
	tabsize := int(buf.Settings["tabsize"].(float64))
	softwrap := buf.Settings["softwrap"].(bool)
	indentchar := []rune(buf.Settings["indentchar"].(string))[0]

	start := buf.Cursor.Y
	if buf.Settings["syntax"].(bool) && buf.syntaxDef != nil {
		if start > 0 && buf.lines[start-1].rehighlight {
			buf.highlighter.ReHighlightLine(buf, start-1)
			buf.lines[start-1].rehighlight = false
		}

		buf.highlighter.ReHighlightStates(buf, start)

		buf.highlighter.HighlightMatches(buf, top, top+height)
	}

	c.lines = make([][]*Char, 0)

	viewLine := 0
	lineN := top

	curStyle := defStyle
	for viewLine < height {
		if lineN >= len(buf.lines) {
			break
		}

		lineStr := buf.Line(lineN)
		line := []rune(lineStr)

		colN, startOffset, startStyle := visualToCharPos(left, lineN, lineStr, buf, tabsize)
		if colN < 0 {
			colN = len(line)
		}
		viewCol := -startOffset
		if startStyle != nil {
			curStyle = *startStyle
		}

		// We'll either draw the length of the line, or the width of the screen
		// whichever is smaller
		lineLength := min(StringWidth(lineStr, tabsize), width)
		c.lines = append(c.lines, make([]*Char, lineLength))

		wrap := false
		// We only need to wrap if the length of the line is greater than the width of the terminal screen
		if softwrap && StringWidth(lineStr, tabsize) > width {
			wrap = true
			// We're going to draw the entire line now
			lineLength = StringWidth(lineStr, tabsize)
		}

		for viewCol < lineLength {
			if colN >= len(line) {
				break
			}
			if group, ok := buf.Match(lineN)[colN]; ok {
				curStyle = GetColor(group.String())
			}

			char := line[colN]

			if viewCol >= 0 {
				c.lines[viewLine][viewCol] = &Char{Loc{viewCol, viewLine}, Loc{colN, lineN}, char, char, curStyle, 1}
			}
			if char == '\t' {
				charWidth := tabsize - (viewCol+left)%tabsize
				if viewCol >= 0 {
					c.lines[viewLine][viewCol].drawChar = indentchar
					c.lines[viewLine][viewCol].width = charWidth

					indentStyle := curStyle
					if group, ok := colorscheme["indent-char"]; ok {
						indentStyle = group
					}

					c.lines[viewLine][viewCol].style = indentStyle
				}

				for i := 1; i < charWidth; i++ {
					viewCol++
					if viewCol >= 0 && viewCol < lineLength {
						c.lines[viewLine][viewCol] = &Char{Loc{viewCol, viewLine}, Loc{colN, lineN}, char, ' ', curStyle, 1}
					}
				}
				viewCol++
			} else if runewidth.RuneWidth(char) > 1 {
				charWidth := runewidth.RuneWidth(char)
				if viewCol >= 0 {
					c.lines[viewLine][viewCol].width = charWidth
				}
				for i := 1; i < charWidth; i++ {
					viewCol++
					if viewCol >= 0 && viewCol < lineLength {
						c.lines[viewLine][viewCol] = &Char{Loc{viewCol, viewLine}, Loc{colN, lineN}, char, ' ', curStyle, 1}
					}
				}
				viewCol++
			} else {
				viewCol++
			}
			colN++

			if wrap && viewCol >= width {
				viewLine++

				// If we go too far soft wrapping we have to cut off
				if viewLine >= height {
					break
				}

				nextLine := line[colN:]
				lineLength := min(StringWidth(string(nextLine), tabsize), width)
				c.lines = append(c.lines, make([]*Char, lineLength))

				viewCol = 0
			}

		}
		if group, ok := buf.Match(lineN)[len(line)]; ok {
			curStyle = GetColor(group.String())
		}

		// newline
		viewLine++
		lineN++
	}

	for i := top; i < top+height; i++ {
		if i >= buf.NumLines {
			break
		}
		buf.SetMatch(i, nil)
	}
}
<|endoftext|>


<commit_before>package main

import "log"

func main() {
	log.Println("Hello World")
}
<commit_msg>basic functionality<commit_after>package main

import (
	"encoding/xml"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"time"
)

type options struct {
	from   string
	link   string
	prefix string
	set    string
	until  string
}

type request struct {
	opts  options
	token string
	verb  string
}

type response struct {
	Date    string `xml:"responseDate"`
	Request struct {
		Verb  string `xml:"verb,attr"`
		Set   string `xml:"set,attr"`
		From  string `xml:"from,attr"`
		Until string `xml:"until,attr"`
		Link  string `xml:",chardata"`
	} `xml:"request"`
	Error struct {
		Code    string `xml:"code,attr"`
		Message string `xml:",chardata"`
	} `xml:"error"`
	ListRecords string `xml:",innerxml"`
}

func (r request) Link() string {
	v := url.Values{}
	v.Add("from", r.opts.from)
	v.Add("set", r.opts.set)
	v.Add("until", r.opts.until)
	v.Add("metadataPrefix", r.opts.prefix)
	v.Add("verb", r.verb)
	if r.token != "" {
		v.Add("resumptionToken", r.token)
	}
	return fmt.Sprintf("%s?%s", r.opts.link, v.Encode())
}

func main() {

	link := flag.String("link", "", "OAI provider URL")
	// output := flag.String("o", "", "output file")
	from := flag.String("f", "2000-01-01", "from parameter")
	until := flag.String("u", time.Now().Format("2006-01-02"), "until parameter")
	prefix := flag.String("p", "oai_dc", "metadata prefix")
	set := flag.String("s", "", "set name")
	verbose := flag.Bool("verbose", false, "be verbose")

	flag.Parse()

	opts := options{from: *from, until: *until, prefix: *prefix, set: *set, link: *link}
	oair := request{opts: opts, verb: "ListRecords"}

	client := http.Client{}

	if *verbose {
		log.Println(oair.Link())
	}

	req, err := http.NewRequest("GET", oair.Link(), nil)
	if err != nil {
		log.Fatal(err)
	}

	resp, err := client.Do(req)
	if err != nil {
		log.Fatal(err)
	}

	defer resp.Body.Close()
	decoder := xml.NewDecoder(resp.Body)

	for {
		t, err := decoder.Token()
		if t == nil {
			break
		}
		if err != nil {
			log.Fatal(err)
		}
		switch se := t.(type) {
		case xml.StartElement:
			if se.Name.Local == "OAI-PMH" {
				var resp response
				err := decoder.DecodeElement(&resp, &se)
				if err != nil {
					log.Fatal(err)
				}
				if resp.Error.Code != "" {
					log.Fatal(resp.Error.Message)
				}
				fmt.Printf("%+v", resp)
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015-2021 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/bcicen/jstream"
	"github.com/minio/minio-go/v7/pkg/set"
)

// startWithConds - map which indicates if a given condition supports starts-with policy operator
var startsWithConds = map[string]bool{
	"$acl":                     true,
	"$bucket":                  false,
	"$cache-control":           true,
	"$content-type":            true,
	"$content-disposition":     true,
	"$content-encoding":        true,
	"$expires":                 true,
	"$key":                     true,
	"$success_action_redirect": true,
	"$redirect":                true,
	"$success_action_status":   false,
	"$x-amz-algorithm":         false,
	"$x-amz-credential":        false,
	"$x-amz-date":              false,
}

// Add policy conditionals.
const (
	policyCondEqual         = "eq"
	policyCondStartsWith    = "starts-with"
	policyCondContentLength = "content-length-range"
)

// toString - Safely convert interface to string without causing panic.
func toString(val interface{}) string {
	switch v := val.(type) {
	case string:
		return v
	default:
		return ""
	}
}

// toLowerString - safely convert interface to lower string
func toLowerString(val interface{}) string {
	return strings.ToLower(toString(val))
}

// toInteger _ Safely convert interface to integer without causing panic.
func toInteger(val interface{}) (int64, error) {
	switch v := val.(type) {
	case float64:
		return int64(v), nil
	case int64:
		return v, nil
	case int:
		return int64(v), nil
	case string:
		i, err := strconv.Atoi(v)
		return int64(i), err
	default:
		return 0, errors.New("Invalid number format")
	}
}

// isString - Safely check if val is of type string without causing panic.
func isString(val interface{}) bool {
	_, ok := val.(string)
	return ok
}

// ContentLengthRange - policy content-length-range field.
type contentLengthRange struct {
	Min   int64
	Max   int64
	Valid bool // If content-length-range was part of policy
}

// PostPolicyForm provides strict static type conversion and validation for Amazon S3's POST policy JSON string.
type PostPolicyForm struct {
	Expiration time.Time // Expiration date and time of the POST policy.
	Conditions struct {  // Conditional policy structure.
		Policies []struct {
			Operator string
			Key      string
			Value    string
		}
		ContentLengthRange contentLengthRange
	}
}

// implemented to ensure that duplicate keys in JSON
// are merged together into a single JSON key, also
// to remove any extraneous JSON bodies.
//
// Go stdlib doesn't support parsing JSON with duplicate
// keys, so we need to use this technique to merge the
// keys.
func sanitizePolicy(r io.Reader) (io.Reader, error) {
	var buf bytes.Buffer
	e := json.NewEncoder(&buf)
	d := jstream.NewDecoder(r, 0).ObjectAsKVS()
	sset := set.NewStringSet()
	for mv := range d.Stream() {
		var kvs jstream.KVS
		if mv.ValueType == jstream.Object {
			// This is a JSON object type (that preserves key order)
			kvs = mv.Value.(jstream.KVS)
			for _, kv := range kvs {
				if sset.Contains(kv.Key) {
					// Reject duplicate conditions or expiration.
					return nil, fmt.Errorf("input policy has multiple %s, please fix your client code", kv.Key)
				}
				sset.Add(kv.Key)
			}
			e.Encode(kvs)
		}
	}
	return &buf, d.Err()
}

// parsePostPolicyForm - Parse JSON policy string into typed PostPolicyForm structure.
func parsePostPolicyForm(r io.Reader) (PostPolicyForm, error) {
	reader, err := sanitizePolicy(r)
	if err != nil {
		return PostPolicyForm{}, err
	}

	d := json.NewDecoder(reader)

	// Convert po into interfaces and
	// perform strict type conversion using reflection.
	var rawPolicy struct {
		Expiration string        `json:"expiration"`
		Conditions []interface{} `json:"conditions"`
	}

	d.DisallowUnknownFields()
	if err := d.Decode(&rawPolicy); err != nil {
		return PostPolicyForm{}, err
	}

	parsedPolicy := PostPolicyForm{}

	// Parse expiry time.
	parsedPolicy.Expiration, err = time.Parse(time.RFC3339Nano, rawPolicy.Expiration)
	if err != nil {
		return PostPolicyForm{}, err
	}

	// Parse conditions.
	for _, val := range rawPolicy.Conditions {
		switch condt := val.(type) {
		case map[string]interface{}: // Handle key:value map types.
			for k, v := range condt {
				if !isString(v) { // Pre-check value type.
					// All values must be of type string.
					return parsedPolicy, fmt.Errorf("Unknown type %s of conditional field value %s found in POST policy form", reflect.TypeOf(condt).String(), condt)
				}
				// {"acl": "public-read" } is an alternate way to indicate - [ "eq", "$acl", "public-read" ]
				// In this case we will just collapse this into "eq" for all use cases.
				parsedPolicy.Conditions.Policies = append(parsedPolicy.Conditions.Policies, struct {
					Operator string
					Key      string
					Value    string
				}{
					policyCondEqual, "$" + strings.ToLower(k), toString(v),
				})
			}
		case []interface{}: // Handle array types.
			if len(condt) != 3 { // Return error if we have insufficient elements.
				return parsedPolicy, fmt.Errorf("Malformed conditional fields %s of type %s found in POST policy form", condt, reflect.TypeOf(condt).String())
			}
			switch toLowerString(condt[0]) {
			case policyCondEqual, policyCondStartsWith:
				for _, v := range condt { // Pre-check all values for type.
					if !isString(v) {
						// All values must be of type string.
						return parsedPolicy, fmt.Errorf("Unknown type %s of conditional field value %s found in POST policy form", reflect.TypeOf(condt).String(), condt)
					}
				}
				operator, matchType, value := toLowerString(condt[0]), toLowerString(condt[1]), toString(condt[2])
				if !strings.HasPrefix(matchType, "$") {
					return parsedPolicy, fmt.Errorf("Invalid according to Policy: Policy Condition failed: [%s, %s, %s]", operator, matchType, value)
				}
				parsedPolicy.Conditions.Policies = append(parsedPolicy.Conditions.Policies, struct {
					Operator string
					Key      string
					Value    string
				}{
					operator, matchType, value,
				})
			case policyCondContentLength:
				min, err := toInteger(condt[1])
				if err != nil {
					return parsedPolicy, err
				}

				max, err := toInteger(condt[2])
				if err != nil {
					return parsedPolicy, err
				}

				parsedPolicy.Conditions.ContentLengthRange = contentLengthRange{
					Min:   min,
					Max:   max,
					Valid: true,
				}
			default:
				// Condition should be valid.
				return parsedPolicy, fmt.Errorf("Unknown type %s of conditional field value %s found in POST policy form",
					reflect.TypeOf(condt).String(), condt)
			}
		default:
			return parsedPolicy, fmt.Errorf("Unknown field %s of type %s found in POST policy form",
				condt, reflect.TypeOf(condt).String())
		}
	}
	return parsedPolicy, nil
}

// checkPolicyCond returns a boolean to indicate if a condition is satisified according
// to the passed operator
func checkPolicyCond(op string, input1, input2 string) bool {
	switch op {
	case policyCondEqual:
		return input1 == input2
	case policyCondStartsWith:
		return strings.HasPrefix(input1, input2)
	}
	return false
}

// checkPostPolicy - apply policy conditions and validate input values.
// (http://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-HTTPPOSTConstructPolicy.html)
func checkPostPolicy(formValues http.Header, postPolicyForm PostPolicyForm) error {
	// Check if policy document expiry date is still not reached
	if !postPolicyForm.Expiration.After(UTCNow()) {
		return fmt.Errorf("Invalid according to Policy: Policy expired")
	}
	// map to store the metadata
	metaMap := make(map[string]string)
	for _, policy := range postPolicyForm.Conditions.Policies {
		if strings.HasPrefix(policy.Key, "$x-amz-meta-") {
			formCanonicalName := http.CanonicalHeaderKey(strings.TrimPrefix(policy.Key, "$"))
			metaMap[formCanonicalName] = policy.Value
		}
	}
	// Check if any extra metadata field is passed as input
	for key := range formValues {
		if strings.HasPrefix(key, "X-Amz-Meta-") {
			if _, ok := metaMap[key]; !ok {
				return fmt.Errorf("Invalid according to Policy: Extra input fields: %s", key)
			}
		}
	}

	// Flag to indicate if all policies conditions are satisfied
	var condPassed bool

	// Iterate over policy conditions and check them against received form fields
	for _, policy := range postPolicyForm.Conditions.Policies {
		// Form fields names are in canonical format, convert conditions names
		// to canonical for simplification purpose, so `$key` will become `Key`
		formCanonicalName := http.CanonicalHeaderKey(strings.TrimPrefix(policy.Key, "$"))
		// Operator for the current policy condition
		op := policy.Operator
		// If the current policy condition is known
		if startsWithSupported, condFound := startsWithConds[policy.Key]; condFound {
			// Check if the current condition supports starts-with operator
			if op == policyCondStartsWith && !startsWithSupported {
				return fmt.Errorf("Invalid according to Policy: Policy Condition failed")
			}
			// Check if current policy condition is satisfied
			condPassed = checkPolicyCond(op, formValues.Get(formCanonicalName), policy.Value)
			if !condPassed {
				return fmt.Errorf("Invalid according to Policy: Policy Condition failed")
			}
		} else {
			// This covers all conditions X-Amz-Meta-* and X-Amz-*
			if strings.HasPrefix(policy.Key, "$x-amz-meta-") || strings.HasPrefix(policy.Key, "$x-amz-") {
				// Check if policy condition is satisfied
				condPassed = checkPolicyCond(op, formValues.Get(formCanonicalName), policy.Value)
				if !condPassed {
					return fmt.Errorf("Invalid according to Policy: Policy Condition failed: [%s, %s, %s]", op, policy.Key, policy.Value)
				}
			}
		}
	}

	return nil
}
<commit_msg>'starts-with' support for 'success_action_status' (#12698)<commit_after>// Copyright (c) 2015-2021 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/bcicen/jstream"
	"github.com/minio/minio-go/v7/pkg/set"
)

// startWithConds - map which indicates if a given condition supports starts-with policy operator
var startsWithConds = map[string]bool{
	"$acl":                     true,
	"$bucket":                  false,
	"$cache-control":           true,
	"$content-type":            true,
	"$content-disposition":     true,
	"$content-encoding":        true,
	"$expires":                 true,
	"$key":                     true,
	"$success_action_redirect": true,
	"$redirect":                true,
	"$success_action_status":   true,
	"$x-amz-algorithm":         false,
	"$x-amz-credential":        false,
	"$x-amz-date":              false,
}

// Add policy conditionals.
const (
	policyCondEqual         = "eq"
	policyCondStartsWith    = "starts-with"
	policyCondContentLength = "content-length-range"
)

// toString - Safely convert interface to string without causing panic.
func toString(val interface{}) string {
	switch v := val.(type) {
	case string:
		return v
	default:
		return ""
	}
}

// toLowerString - safely convert interface to lower string
func toLowerString(val interface{}) string {
	return strings.ToLower(toString(val))
}

// toInteger _ Safely convert interface to integer without causing panic.
func toInteger(val interface{}) (int64, error) {
	switch v := val.(type) {
	case float64:
		return int64(v), nil
	case int64:
		return v, nil
	case int:
		return int64(v), nil
	case string:
		i, err := strconv.Atoi(v)
		return int64(i), err
	default:
		return 0, errors.New("Invalid number format")
	}
}

// isString - Safely check if val is of type string without causing panic.
func isString(val interface{}) bool {
	_, ok := val.(string)
	return ok
}

// ContentLengthRange - policy content-length-range field.
type contentLengthRange struct {
	Min   int64
	Max   int64
	Valid bool // If content-length-range was part of policy
}

// PostPolicyForm provides strict static type conversion and validation for Amazon S3's POST policy JSON string.
type PostPolicyForm struct {
	Expiration time.Time // Expiration date and time of the POST policy.
	Conditions struct {  // Conditional policy structure.
		Policies []struct {
			Operator string
			Key      string
			Value    string
		}
		ContentLengthRange contentLengthRange
	}
}

// implemented to ensure that duplicate keys in JSON
// are merged together into a single JSON key, also
// to remove any extraneous JSON bodies.
//
// Go stdlib doesn't support parsing JSON with duplicate
// keys, so we need to use this technique to merge the
// keys.
func sanitizePolicy(r io.Reader) (io.Reader, error) {
	var buf bytes.Buffer
	e := json.NewEncoder(&buf)
	d := jstream.NewDecoder(r, 0).ObjectAsKVS()
	sset := set.NewStringSet()
	for mv := range d.Stream() {
		var kvs jstream.KVS
		if mv.ValueType == jstream.Object {
			// This is a JSON object type (that preserves key order)
			kvs = mv.Value.(jstream.KVS)
			for _, kv := range kvs {
				if sset.Contains(kv.Key) {
					// Reject duplicate conditions or expiration.
					return nil, fmt.Errorf("input policy has multiple %s, please fix your client code", kv.Key)
				}
				sset.Add(kv.Key)
			}
			e.Encode(kvs)
		}
	}
	return &buf, d.Err()
}

// parsePostPolicyForm - Parse JSON policy string into typed PostPolicyForm structure.
func parsePostPolicyForm(r io.Reader) (PostPolicyForm, error) {
	reader, err := sanitizePolicy(r)
	if err != nil {
		return PostPolicyForm{}, err
	}

	d := json.NewDecoder(reader)

	// Convert po into interfaces and
	// perform strict type conversion using reflection.
	var rawPolicy struct {
		Expiration string        `json:"expiration"`
		Conditions []interface{} `json:"conditions"`
	}

	d.DisallowUnknownFields()
	if err := d.Decode(&rawPolicy); err != nil {
		return PostPolicyForm{}, err
	}

	parsedPolicy := PostPolicyForm{}

	// Parse expiry time.
	parsedPolicy.Expiration, err = time.Parse(time.RFC3339Nano, rawPolicy.Expiration)
	if err != nil {
		return PostPolicyForm{}, err
	}

	// Parse conditions.
	for _, val := range rawPolicy.Conditions {
		switch condt := val.(type) {
		case map[string]interface{}: // Handle key:value map types.
			for k, v := range condt {
				if !isString(v) { // Pre-check value type.
					// All values must be of type string.
					return parsedPolicy, fmt.Errorf("Unknown type %s of conditional field value %s found in POST policy form", reflect.TypeOf(condt).String(), condt)
				}
				// {"acl": "public-read" } is an alternate way to indicate - [ "eq", "$acl", "public-read" ]
				// In this case we will just collapse this into "eq" for all use cases.
				parsedPolicy.Conditions.Policies = append(parsedPolicy.Conditions.Policies, struct {
					Operator string
					Key      string
					Value    string
				}{
					policyCondEqual, "$" + strings.ToLower(k), toString(v),
				})
			}
		case []interface{}: // Handle array types.
			if len(condt) != 3 { // Return error if we have insufficient elements.
				return parsedPolicy, fmt.Errorf("Malformed conditional fields %s of type %s found in POST policy form", condt, reflect.TypeOf(condt).String())
			}
			switch toLowerString(condt[0]) {
			case policyCondEqual, policyCondStartsWith:
				for _, v := range condt { // Pre-check all values for type.
					if !isString(v) {
						// All values must be of type string.
						return parsedPolicy, fmt.Errorf("Unknown type %s of conditional field value %s found in POST policy form", reflect.TypeOf(condt).String(), condt)
					}
				}
				operator, matchType, value := toLowerString(condt[0]), toLowerString(condt[1]), toString(condt[2])
				if !strings.HasPrefix(matchType, "$") {
					return parsedPolicy, fmt.Errorf("Invalid according to Policy: Policy Condition failed: [%s, %s, %s]", operator, matchType, value)
				}
				parsedPolicy.Conditions.Policies = append(parsedPolicy.Conditions.Policies, struct {
					Operator string
					Key      string
					Value    string
				}{
					operator, matchType, value,
				})
			case policyCondContentLength:
				min, err := toInteger(condt[1])
				if err != nil {
					return parsedPolicy, err
				}

				max, err := toInteger(condt[2])
				if err != nil {
					return parsedPolicy, err
				}

				parsedPolicy.Conditions.ContentLengthRange = contentLengthRange{
					Min:   min,
					Max:   max,
					Valid: true,
				}
			default:
				// Condition should be valid.
				return parsedPolicy, fmt.Errorf("Unknown type %s of conditional field value %s found in POST policy form",
					reflect.TypeOf(condt).String(), condt)
			}
		default:
			return parsedPolicy, fmt.Errorf("Unknown field %s of type %s found in POST policy form",
				condt, reflect.TypeOf(condt).String())
		}
	}
	return parsedPolicy, nil
}

// checkPolicyCond returns a boolean to indicate if a condition is satisified according
// to the passed operator
func checkPolicyCond(op string, input1, input2 string) bool {
	switch op {
	case policyCondEqual:
		return input1 == input2
	case policyCondStartsWith:
		return strings.HasPrefix(input1, input2)
	}
	return false
}

// checkPostPolicy - apply policy conditions and validate input values.
// (http://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-HTTPPOSTConstructPolicy.html)
func checkPostPolicy(formValues http.Header, postPolicyForm PostPolicyForm) error {
	// Check if policy document expiry date is still not reached
	if !postPolicyForm.Expiration.After(UTCNow()) {
		return fmt.Errorf("Invalid according to Policy: Policy expired")
	}
	// map to store the metadata
	metaMap := make(map[string]string)
	for _, policy := range postPolicyForm.Conditions.Policies {
		if strings.HasPrefix(policy.Key, "$x-amz-meta-") {
			formCanonicalName := http.CanonicalHeaderKey(strings.TrimPrefix(policy.Key, "$"))
			metaMap[formCanonicalName] = policy.Value
		}
	}
	// Check if any extra metadata field is passed as input
	for key := range formValues {
		if strings.HasPrefix(key, "X-Amz-Meta-") {
			if _, ok := metaMap[key]; !ok {
				return fmt.Errorf("Invalid according to Policy: Extra input fields: %s", key)
			}
		}
	}

	// Flag to indicate if all policies conditions are satisfied
	var condPassed bool

	// Iterate over policy conditions and check them against received form fields
	for _, policy := range postPolicyForm.Conditions.Policies {
		// Form fields names are in canonical format, convert conditions names
		// to canonical for simplification purpose, so `$key` will become `Key`
		formCanonicalName := http.CanonicalHeaderKey(strings.TrimPrefix(policy.Key, "$"))
		// Operator for the current policy condition
		op := policy.Operator
		// If the current policy condition is known
		if startsWithSupported, condFound := startsWithConds[policy.Key]; condFound {
			// Check if the current condition supports starts-with operator
			if op == policyCondStartsWith && !startsWithSupported {
				return fmt.Errorf("Invalid according to Policy: Policy Condition failed")
			}
			// Check if current policy condition is satisfied
			condPassed = checkPolicyCond(op, formValues.Get(formCanonicalName), policy.Value)
			if !condPassed {
				return fmt.Errorf("Invalid according to Policy: Policy Condition failed")
			}
		} else {
			// This covers all conditions X-Amz-Meta-* and X-Amz-*
			if strings.HasPrefix(policy.Key, "$x-amz-meta-") || strings.HasPrefix(policy.Key, "$x-amz-") {
				// Check if policy condition is satisfied
				condPassed = checkPolicyCond(op, formValues.Get(formCanonicalName), policy.Value)
				if !condPassed {
					return fmt.Errorf("Invalid according to Policy: Policy Condition failed: [%s, %s, %s]", op, policy.Key, policy.Value)
				}
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Home24 AG. All rights reserved.
// Proprietary license.
package particular

import (
	"encoding/json"
	"github.com/PuerkitoBio/goquery"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
)

const LG_URL = "http://www.lingvolive.com/api/Translation/Translate/"
const LG_PART_URL = "http://www.lingvolive.com/api/Translation/WordListPart/" //Fucking slow

var langsMap = map[string]string{
	"ru": "1049",
	"en": "1033",
	//	"da": "1030",?
	//	"zh": "1028",?
	//	"nl": "1035",?
	//	"fi": "1043",?
	"de": "32775",
	"fr": "1036",
}

type AbbyyLingvoLiveTranslator struct {
	client *http.Client
}

func (t *AbbyyLingvoLiveTranslator) SetClient(c *http.Client) {
	t.client = c
}

func (t *AbbyyLingvoLiveTranslator) TranslateOne(text string, language, to string) IParticularResponse {

	data := &LingvoLiveTranslatorResponseFull{}
	data.Lang = language
	_, ok1 := langsMap[language]
	_, ok2 := langsMap[to]
	if !ok1 || !ok2 {
		return data
	}
	if language != "ru" && to != "ru" {
		return data
	}
	reqUrl := LG_URL+"?"+t.getQueryStringFull(text, language, to)
	req, _ := http.NewRequest("GET", reqUrl, nil)
	data.Url = reqUrl
	resp, err := t.client.Do(req)
	defer resp.Body.Close()
	if err != nil {
		log.Println(err)
	}
	reader := ioutil.NopCloser(resp.Body)
	b, _ := ioutil.ReadAll(reader)
	str := string(b)

	// log.Println(str)
	if err := json.NewDecoder(strings.NewReader(str)).Decode(&data); err != nil {
		log.Println("error decode", err)
	}

	return data
}

func (t *AbbyyLingvoLiveTranslator) GetName() string {
	return "lingvo_live"
}

func (t *AbbyyLingvoLiveTranslator) getQueryStringFull(text, from, to string) string {
	form := url.Values{}
	form.Add("dstLang", langsMap[from])
	form.Add("srcLang", langsMap[to])
	form.Add("text", text)
	return form.Encode()
}

func (t *AbbyyLingvoLiveTranslator) getQueryStringPart(text, from, to string) string {
	form := url.Values{}
	form.Add("dstLang", from)
	form.Add("srcLang", to)
	form.Add("prefix", text)
	form.Add("pageSize", "10")
	form.Add("startIndex", "0")
	return form.Encode()
}

type LingvoLiveTranslatorResponseFull struct {
	Lang                  string
	Url                   string
	GlossaryUnits         interface{} `json:"glossaryUnits"`
	LanguagesReversed     bool        `json:"languagesReversed"`
	SeeAlsoWordForms      []string    `json:"seeAlsoWordForms"`
	Suggests              interface{} `json:"suggests"`
	WordByWordTranslation interface{} `json:"wordByWordTranslation"`
	Articles              []struct {
		Heading    string `json:"heading"`
		Dictionary string `json:"dictionary"`
		BodyHtml   string `json:"bodyHtml"`
	} `json:"lingvoArticles"`
}

type LingvoLiveTranslatorResponsePart struct {
	Lang string

	Items []struct {
		Heading      string `json:"heading"`
		Dictationary string `json:"lingvoDictionaryName"`
		Translations string `json:"lingvoTranslations"`
	} `json:"items"`
}

func (t *LingvoLiveTranslatorResponseFull) GetUrl() string {
	return t.Url
}

func (t *LingvoLiveTranslatorResponseFull) GetMeanings() []IParticularMeaning {
	meanings := []IParticularMeaning{}
	for _, v := range t.Articles {
		meaning := &Meaning{Dictationary: v.Dictionary}
		doc, _ := goquery.NewDocumentFromReader(strings.NewReader(v.BodyHtml))

		table := doc.Find(".article .article-body .article-body-items")
		table.Find(".article-body-items").Each(func(i int, s *goquery.Selection) {

			if s.Find(".paragraph-marker-top-level").Text() == "" {
				return
			}

			value := s.Find(".article-text-wrap").Text()
			value = strings.TrimSpace(value)
			if "" != value {
				meaning.All = append(meaning.All, value)
			}
			if len(meaning.All) > 0 && meaning.Text == "" {

				meaning.Text = meaning.All[0]
			}
		})
		meanings = append(meanings, meaning)
	}
	return meanings
}
<commit_msg>better parse<commit_after>// Copyright 2015 Home24 AG. All rights reserved.
// Proprietary license.
package particular

import (
	"encoding/json"
	"github.com/PuerkitoBio/goquery"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
)

const LG_URL = "http://www.lingvolive.com/api/Translation/Translate/"
const LG_PART_URL = "http://www.lingvolive.com/api/Translation/WordListPart/" //Fucking slow

var langsMap = map[string]string{
	"ru": "1049",
	"en": "1033",
	//	"da": "1030",?
	//	"zh": "1028",?
	//	"nl": "1035",?
	//	"fi": "1043",?
	"de": "32775",
	"fr": "1036",
}

type AbbyyLingvoLiveTranslator struct {
	client *http.Client
}

func (t *AbbyyLingvoLiveTranslator) SetClient(c *http.Client) {
	t.client = c
}

func (t *AbbyyLingvoLiveTranslator) TranslateOne(text string, language, to string) IParticularResponse {

	data := &LingvoLiveTranslatorResponseFull{}
	data.Lang = language
	_, ok1 := langsMap[language]
	_, ok2 := langsMap[to]
	if !ok1 || !ok2 {
		return data
	}
	if language != "ru" && to != "ru" {
		return data
	}
	reqUrl := LG_URL+"?"+t.getQueryStringFull(text, language, to)
	req, _ := http.NewRequest("GET", reqUrl, nil)
	data.Url = reqUrl
	resp, err := t.client.Do(req)
	defer resp.Body.Close()
	if err != nil {
		log.Println(err)
	}
	reader := ioutil.NopCloser(resp.Body)
	b, _ := ioutil.ReadAll(reader)
	str := string(b)

	// log.Println(str)
	if err := json.NewDecoder(strings.NewReader(str)).Decode(&data); err != nil {
		log.Println("error decode", err)
	}

	return data
}

func (t *AbbyyLingvoLiveTranslator) GetName() string {
	return "lingvo_live"
}

func (t *AbbyyLingvoLiveTranslator) getQueryStringFull(text, from, to string) string {
	form := url.Values{}
	form.Add("dstLang", langsMap[from])
	form.Add("srcLang", langsMap[to])
	form.Add("text", text)
	return form.Encode()
}

func (t *AbbyyLingvoLiveTranslator) getQueryStringPart(text, from, to string) string {
	form := url.Values{}
	form.Add("dstLang", from)
	form.Add("srcLang", to)
	form.Add("prefix", text)
	form.Add("pageSize", "10")
	form.Add("startIndex", "0")
	return form.Encode()
}

type LingvoLiveTranslatorResponseFull struct {
	Lang                  string
	Url                   string
	GlossaryUnits         interface{} `json:"glossaryUnits"`
	LanguagesReversed     bool        `json:"languagesReversed"`
	SeeAlsoWordForms      []string    `json:"seeAlsoWordForms"`
	Suggests              interface{} `json:"suggests"`
	WordByWordTranslation interface{} `json:"wordByWordTranslation"`
	Articles              []struct {
		Heading    string `json:"heading"`
		Dictionary string `json:"dictionary"`
		BodyHtml   string `json:"bodyHtml"`
	} `json:"lingvoArticles"`
}

type LingvoLiveTranslatorResponsePart struct {
	Lang string

	Items []struct {
		Heading      string `json:"heading"`
		Dictationary string `json:"lingvoDictionaryName"`
		Translations string `json:"lingvoTranslations"`
	} `json:"items"`
}

func (t *LingvoLiveTranslatorResponseFull) GetUrl() string {
	return t.Url
}

func (t *LingvoLiveTranslatorResponseFull) GetMeanings() []IParticularMeaning {
	meanings := []IParticularMeaning{}
	for _, v := range t.Articles {
		meaning := &Meaning{Dictationary: v.Dictionary}
		doc, _ := goquery.NewDocumentFromReader(strings.NewReader(v.BodyHtml))

		table := doc.Find(".article .article-body .article-body-items")
		table.Find(".article-body-items").Each(func(i int, s *goquery.Selection) {

			if s.Find(".paragraph-marker-top-level").Text() == "" {
				if s.Find(".parts-of-speech").Text() != "" && len(s.Find(".article-text").Nodes) == 0 {
					return
				}
			}

			value := s.Find(".article-text-wrap .article-text").Text()
			if value == "" {
				// maybe comment
				value = s.Find(".article-text-wrap .comment").Text()
			}
			value = strings.TrimLeft(value, "<-s, ->")
			value = strings.TrimSpace(value)
			if "" != value {
				meaning.All = append(meaning.All, value)
			}
			if len(meaning.All) > 0 && meaning.Text == "" {

				meaning.Text = meaning.All[0]
			}
		})
		meanings = append(meanings, meaning)
	}
	return meanings
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/loadimpact/speedboat"
	"github.com/loadimpact/speedboat/js"
	"github.com/loadimpact/speedboat/sampler"
	"github.com/loadimpact/speedboat/sampler/influxdb"
	"github.com/loadimpact/speedboat/simple"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	stdlog "log"
	"os"
	"strings"
	"sync"
	"time"
)

// Configure the global logger.
func configureLogging(c *cli.Context) {
	log.SetLevel(log.InfoLevel)
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}
}

// Configure the global sampler.
func configureSampler(c *cli.Context) {
	output := c.String("output")
	if output == "" {
		return
	}

	sampler.DefaultSampler.OnError = func(err error) {
		log.WithError(err).Error("[Sampler error]")
	}

	parts := strings.SplitN(output, "+", 2)
	switch parts[0] {
	case "influxdb":
		out, err := influxdb.NewFromURL(parts[1])
		if err != nil {
			log.WithError(err).Fatal("Couldn't create InfluxDB client")
		}
		sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
	}
}

func parse(cc *cli.Context) (conf Config, err error) {
	switch len(cc.Args()) {
	case 0:
		if !cc.IsSet("script") && !cc.IsSet("url") {
			return conf, errors.New("No config file, script or URL")
		}
	case 1:
		bytes, err := ioutil.ReadFile(cc.Args()[0])
		if err != nil {
			return conf, errors.New("Couldn't read config file")
		}
		if err := yaml.Unmarshal(bytes, &conf); err != nil {
			return conf, errors.New("Couldn't parse config file")
		}
	default:
		return conf, errors.New("Too many arguments!")
	}

	// Let commandline flags override config files
	if cc.IsSet("script") {
		conf.Script = cc.String("script")
	}
	if cc.IsSet("url") {
		conf.URL = cc.String("url")
	}
	if cc.IsSet("vus") {
		conf.VUs = cc.Int("vus")
	}
	if cc.IsSet("duration") {
		conf.Duration = cc.Duration("duration").String()
	}

	return conf, nil
}

func dumpTest(t *speedboat.Test) {
	log.WithFields(log.Fields{
		"script": t.Script,
		"url":    t.URL,
	}).Info("General")
	for i, stage := range t.Stages {
		log.WithFields(log.Fields{
			"#":        i,
			"duration": stage.Duration,
			"start":    stage.StartVUs,
			"end":      stage.EndVUs,
		}).Info("Stage")
	}
}

func headlessController(c context.Context, t *speedboat.Test) <-chan int {
	ch := make(chan int)

	go func() {
		defer close(ch)

		select {
		case ch <- t.VUsAt(0):
		case <-c.Done():
			return
		}

		startTime := time.Now()
		ticker := time.NewTicker(100 * time.Millisecond)
		for {
			select {
			case <-ticker.C:
				ch <- t.VUsAt(time.Since(startTime))
			case <-c.Done():
				return
			}
		}
	}()

	return ch
}

func action(cc *cli.Context) error {
	conf, err := parse(cc)
	if err != nil {
		log.WithError(err).Fatal("Invalid arguments; see --help")
	}

	t, err := conf.MakeTest()
	if err != nil {
		log.WithError(err).Fatal("Configuration error")
	}

	if cc.Bool("dump") {
		dumpTest(&t)
		return nil
	}

	// Inspect the test to find a suitable runner; additional ones can easily be added
	var runner speedboat.Runner
	switch {
	case t.URL != "":
		runner = simple.New()
	case strings.HasSuffix(t.Script, ".js"):
		src, err := ioutil.ReadFile(t.Script)
		if err != nil {
			log.WithError(err).Fatal("Couldn't read script")
		}
		runner = js.New(string(src))
	default:
		log.Fatal("No suitable runner found!")
	}

	// Global metrics
	mVUs := sampler.Gauge("vus")

	// Context that expires at the end of the test
	ctx, _ := context.WithTimeout(context.Background(), t.TotalDuration())

	// Output metrics appropriately; use a mutex to prevent garbled output
	logMetrics := cc.Bool("log")
	metricsLogger := stdlog.New(os.Stdout, "metrics: ", stdlog.Lmicroseconds)
	metricsMutex := sync.Mutex{}
	go func() {
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				if logMetrics {
					metricsMutex.Lock()
					printMetrics(metricsLogger)
					metricsMutex.Unlock()
				}
				commitMetrics()
			case <-ctx.Done():
				return
			}
		}
	}()

	// Use a "headless controller" to scale VUs by polling the test ramp
	vus := []context.CancelFunc{}
	for scale := range headlessController(ctx, &t) {
		for i := len(vus); i < scale; i++ {
			log.WithField("id", i).Debug("Spawning VU")
			vuCtx, cancel := context.WithCancel(ctx)
			vus = append(vus, cancel)
			go runner.RunVU(vuCtx, t, len(vus))
		}
		for i := len(vus); i > scale; i-- {
			log.WithField("id", i-1).Debug("Dropping VU")
			vus[i-1]()
			vus = vus[:i-1]
		}
		mVUs.Int(len(vus))
	}

	// Wait until the end of the test
	<-ctx.Done()

	// Print final metrics
	metricsMutex.Lock()
	printMetrics(metricsLogger)
	metricsMutex.Unlock()
	commitMetrics()

	return nil
}

func main() {
	// Free up -v and -h for our own flags
	cli.VersionFlag.Name = "version"
	cli.HelpFlag.Name = "help, ?"

	// Bootstrap using action-registered commandline flags
	app := cli.NewApp()
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "0.0.1a1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, v",
			Usage: "More verbose output",
		},
		cli.StringFlag{
			Name:  "script, s",
			Usage: "Script to run",
		},
		cli.StringFlag{
			Name:  "url",
			Usage: "URL to test",
		},
		cli.IntFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
			Value: 10,
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "Test duration",
			Value: time.Duration(10) * time.Second,
		},
		cli.StringFlag{
			Name:  "output, o",
			Usage: "Output metrics to a file or database",
		},
		cli.BoolFlag{
			Name:  "log, l",
			Usage: "Log metrics to stdout",
		},
		cli.BoolFlag{
			Name:  "dump",
			Usage: "Dump parsed test and exit",
		},
	}
	app.Before = func(c *cli.Context) error {
		configureLogging(c)
		configureSampler(c)
		return nil
	}
	app.Action = action
	app.Run(os.Args)
}
<commit_msg>[refactor] github.com/codegangsta/cli changed name to github.com/urfave/cli<commit_after>package main

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat"
	"github.com/loadimpact/speedboat/js"
	"github.com/loadimpact/speedboat/sampler"
	"github.com/loadimpact/speedboat/sampler/influxdb"
	"github.com/loadimpact/speedboat/simple"
	"github.com/urfave/cli"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	stdlog "log"
	"os"
	"strings"
	"sync"
	"time"
)

// Configure the global logger.
func configureLogging(c *cli.Context) {
	log.SetLevel(log.InfoLevel)
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}
}

// Configure the global sampler.
func configureSampler(c *cli.Context) {
	output := c.String("output")
	if output == "" {
		return
	}

	sampler.DefaultSampler.OnError = func(err error) {
		log.WithError(err).Error("[Sampler error]")
	}

	parts := strings.SplitN(output, "+", 2)
	switch parts[0] {
	case "influxdb":
		out, err := influxdb.NewFromURL(parts[1])
		if err != nil {
			log.WithError(err).Fatal("Couldn't create InfluxDB client")
		}
		sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
	}
}

func parse(cc *cli.Context) (conf Config, err error) {
	switch len(cc.Args()) {
	case 0:
		if !cc.IsSet("script") && !cc.IsSet("url") {
			return conf, errors.New("No config file, script or URL")
		}
	case 1:
		bytes, err := ioutil.ReadFile(cc.Args()[0])
		if err != nil {
			return conf, errors.New("Couldn't read config file")
		}
		if err := yaml.Unmarshal(bytes, &conf); err != nil {
			return conf, errors.New("Couldn't parse config file")
		}
	default:
		return conf, errors.New("Too many arguments!")
	}

	// Let commandline flags override config files
	if cc.IsSet("script") {
		conf.Script = cc.String("script")
	}
	if cc.IsSet("url") {
		conf.URL = cc.String("url")
	}
	if cc.IsSet("vus") {
		conf.VUs = cc.Int("vus")
	}
	if cc.IsSet("duration") {
		conf.Duration = cc.Duration("duration").String()
	}

	return conf, nil
}

func dumpTest(t *speedboat.Test) {
	log.WithFields(log.Fields{
		"script": t.Script,
		"url":    t.URL,
	}).Info("General")
	for i, stage := range t.Stages {
		log.WithFields(log.Fields{
			"#":        i,
			"duration": stage.Duration,
			"start":    stage.StartVUs,
			"end":      stage.EndVUs,
		}).Info("Stage")
	}
}

func headlessController(c context.Context, t *speedboat.Test) <-chan int {
	ch := make(chan int)

	go func() {
		defer close(ch)

		select {
		case ch <- t.VUsAt(0):
		case <-c.Done():
			return
		}

		startTime := time.Now()
		ticker := time.NewTicker(100 * time.Millisecond)
		for {
			select {
			case <-ticker.C:
				ch <- t.VUsAt(time.Since(startTime))
			case <-c.Done():
				return
			}
		}
	}()

	return ch
}

func action(cc *cli.Context) error {
	conf, err := parse(cc)
	if err != nil {
		log.WithError(err).Fatal("Invalid arguments; see --help")
	}

	t, err := conf.MakeTest()
	if err != nil {
		log.WithError(err).Fatal("Configuration error")
	}

	if cc.Bool("dump") {
		dumpTest(&t)
		return nil
	}

	// Inspect the test to find a suitable runner; additional ones can easily be added
	var runner speedboat.Runner
	switch {
	case t.URL != "":
		runner = simple.New()
	case strings.HasSuffix(t.Script, ".js"):
		src, err := ioutil.ReadFile(t.Script)
		if err != nil {
			log.WithError(err).Fatal("Couldn't read script")
		}
		runner = js.New(string(src))
	default:
		log.Fatal("No suitable runner found!")
	}

	// Global metrics
	mVUs := sampler.Gauge("vus")

	// Context that expires at the end of the test
	ctx, _ := context.WithTimeout(context.Background(), t.TotalDuration())

	// Output metrics appropriately; use a mutex to prevent garbled output
	logMetrics := cc.Bool("log")
	metricsLogger := stdlog.New(os.Stdout, "metrics: ", stdlog.Lmicroseconds)
	metricsMutex := sync.Mutex{}
	go func() {
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				if logMetrics {
					metricsMutex.Lock()
					printMetrics(metricsLogger)
					metricsMutex.Unlock()
				}
				commitMetrics()
			case <-ctx.Done():
				return
			}
		}
	}()

	// Use a "headless controller" to scale VUs by polling the test ramp
	vus := []context.CancelFunc{}
	for scale := range headlessController(ctx, &t) {
		for i := len(vus); i < scale; i++ {
			log.WithField("id", i).Debug("Spawning VU")
			vuCtx, cancel := context.WithCancel(ctx)
			vus = append(vus, cancel)
			go runner.RunVU(vuCtx, t, len(vus))
		}
		for i := len(vus); i > scale; i-- {
			log.WithField("id", i-1).Debug("Dropping VU")
			vus[i-1]()
			vus = vus[:i-1]
		}
		mVUs.Int(len(vus))
	}

	// Wait until the end of the test
	<-ctx.Done()

	// Print final metrics
	metricsMutex.Lock()
	printMetrics(metricsLogger)
	metricsMutex.Unlock()
	commitMetrics()

	return nil
}

func main() {
	// Free up -v and -h for our own flags
	cli.VersionFlag.Name = "version"
	cli.HelpFlag.Name = "help, ?"

	// Bootstrap using action-registered commandline flags
	app := cli.NewApp()
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "0.0.1a1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, v",
			Usage: "More verbose output",
		},
		cli.StringFlag{
			Name:  "script, s",
			Usage: "Script to run",
		},
		cli.StringFlag{
			Name:  "url",
			Usage: "URL to test",
		},
		cli.IntFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
			Value: 10,
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "Test duration",
			Value: time.Duration(10) * time.Second,
		},
		cli.StringFlag{
			Name:  "output, o",
			Usage: "Output metrics to a file or database",
		},
		cli.BoolFlag{
			Name:  "log, l",
			Usage: "Log metrics to stdout",
		},
		cli.BoolFlag{
			Name:  "dump",
			Usage: "Dump parsed test and exit",
		},
	}
	app.Before = func(c *cli.Context) error {
		configureLogging(c)
		configureSampler(c)
		return nil
	}
	app.Action = action
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>// Copyright 2015 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"log"
	"strconv"

	"github.com/tsuru/config"
	"github.com/tsuru/gnuflag"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/auth"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/db"
	"github.com/tsuru/tsuru/migration"
	"github.com/tsuru/tsuru/permission"
	"github.com/tsuru/tsuru/provision"
	"github.com/tsuru/tsuru/provision/docker"
	"gopkg.in/mgo.v2/bson"
)

func init() {
	err := migration.Register("migrate-docker-images", migrateImages)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
	err = migration.Register("migrate-pool", migratePool)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
	err = migration.Register("migrate-set-pool-to-app", setPoolToApps)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
	err = migration.Register("migrate-service-proxy-actions", migrateServiceProxyActions)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
	err = migration.RegisterOptional("migrate-roles", migrateRoles)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
}

func getProvisioner() (string, error) {
	provisioner, err := config.GetString("provisioner")
	if provisioner == "" {
		provisioner = "docker"
	}
	return provisioner, err
}

type migrationListCmd struct{}

func (*migrationListCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "migrate-list",
		Usage: "migrate-list",
		Desc:  "List available migration scripts from previous versions of tsurud",
	}
}

func (*migrationListCmd) Run(context *cmd.Context, client *cmd.Client) error {
	migrations, err := migration.List()
	if err != nil {
		return err
	}
	tbl := cmd.NewTable()
	tbl.Headers = cmd.Row{"Name", "Mandatory?", "Executed?"}
	for _, m := range migrations {
		tbl.AddRow(cmd.Row{m.Name, strconv.FormatBool(!m.Optional), strconv.FormatBool(m.Ran)})
	}
	fmt.Fprint(context.Stdout, tbl.String())
	return nil
}

type migrateCmd struct {
	fs    *gnuflag.FlagSet
	dry   bool
	force bool
	name  string
}

func (*migrateCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "migrate",
		Usage: "migrate [-n/--dry] [-f/--force] [--name name]",
		Desc:  "Runs migrations from previous versions of tsurud",
	}
}

func (c *migrateCmd) Run(context *cmd.Context, client *cmd.Client) error {
	return migration.Run(migration.RunArgs{
		Writer: context.Stdout,
		Dry:    c.dry,
		Name:   c.name,
		Force:  c.force,
	})
}

func (c *migrateCmd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("migrate", gnuflag.ExitOnError)
		dryMsg := "Do not run migrations, just print what would run"
		c.fs.BoolVar(&c.dry, "dry", false, dryMsg)
		c.fs.BoolVar(&c.dry, "n", false, dryMsg)
		forceMsg := "Force the execution of an already executed optional migration"
		c.fs.BoolVar(&c.force, "force", false, forceMsg)
		c.fs.BoolVar(&c.force, "f", false, forceMsg)
		c.fs.StringVar(&c.name, "name", "", "The name of an optional migration to run")
	}
	return c.fs
}

func migrateImages() error {
	provisioner, _ := getProvisioner()
	if provisioner == "docker" {
		p, err := provision.Get(provisioner)
		if err != nil {
			return err
		}
		err = p.(provision.InitializableProvisioner).Initialize()
		if err != nil {
			return err
		}
		return docker.MigrateImages()
	}
	return nil
}

func migratePool() error {
	db, err := db.Conn()
	if err != nil {
		return err
	}
	defer db.Close()
	poolColl := db.Collection("pool")
	if err != nil {
		return err
	}
	var pools []provision.Pool
	err = db.Collection("docker_scheduler").Find(nil).All(&pools)
	if err != nil {
		return err
	}
	for _, p := range pools {
		err = poolColl.Insert(p)
		if err != nil {
			return err
		}
	}
	return nil
}

func setPoolToApps() error {
	db, err := db.Conn()
	if err != nil {
		return err
	}
	defer db.Close()
	var apps []app.App
	var tooManyPoolsApps []app.App
	err = db.Apps().Find(nil).All(&apps)
	if err != nil {
		return err
	}
	for _, a := range apps {
		err = a.SetPool()
		if err != nil {
			tooManyPoolsApps = append(tooManyPoolsApps, a)
			continue
		}
		err = db.Apps().Update(bson.M{"name": a.Name}, bson.M{"$set": bson.M{"pool": a.Pool}})
		if err != nil {
			return err
		}
	}
	if len(tooManyPoolsApps) > 0 {
		fmt.Println("Apps bellow couldn't be migrated because they are in more than one pool.")
		fmt.Println("To fix this, please run `tsuru app-change-pool <pool_name> -a app` for each app.")
		fmt.Println("*****************************************")
		for _, a := range tooManyPoolsApps {
			fmt.Println(a.Name)
		}
	}
	return nil
}

func migrateServiceProxyActions() error {
	db, err := db.Conn()
	if err != nil {
		return err
	}
	defer db.Close()
	_, err = db.UserActions().UpdateAll(
		bson.M{"action": "service-proxy-status"},
		bson.M{"$set": bson.M{"action": "service-instance-proxy"}},
	)
	return err
}

func createRole(name, contextType string) (permission.Role, error) {
	role, err := permission.NewRole(name, contextType)
	if err == permission.ErrRoleAlreadyExists {
		role, err = permission.FindRole(name)
	}
	return role, err
}

func migrateRoles() error {
	adminTeam, err := config.GetString("admin-team")
	if err != nil {
		return err
	}
	adminRole, err := createRole("admin", "global")
	if err != nil {
		return err
	}
	err = adminRole.AddPermissions("*")
	if err != nil {
		return err
	}
	teamMember, err := createRole("team-member", "team")
	if err != nil {
		return err
	}
	err = teamMember.AddPermissions(permission.PermApp.FullName(),
		permission.PermTeam.FullName(),
		permission.PermServiceInstance.FullName())
	if err != nil {
		return err
	}
	err = teamMember.AddEvent(permission.RoleEventTeamCreate.String())
	if err != nil {
		return err
	}
	teamCreator, err := createRole("team-creator", "global")
	if err != nil {
		return err
	}
	err = teamCreator.AddPermissions(permission.PermTeamCreate.FullName())
	if err != nil {
		return err
	}
	err = teamCreator.AddEvent(permission.RoleEventUserCreate.String())
	if err != nil {
		return err
	}
	users, err := auth.ListUsers()
	if err != nil {
		return err
	}
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	for _, u := range users {
		var teams []auth.Team
		err := conn.Teams().Find(bson.M{"users": bson.M{"$in": []string{u.Email}}}).All(&teams)
		if err != nil {
			return err
		}
		for _, team := range teams {
			if team.Name == adminTeam {
				err := u.AddRole(adminRole.Name, "")
				if err != nil {
					fmt.Printf("%s\n", err.Error())
				}
				continue
			}
			err := u.AddRole(teamMember.Name, team.Name)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
			}
			err = u.AddRole(teamCreator.Name, "")
			if err != nil {
				fmt.Printf("%s\n", err.Error())
			}
		}
	}
	return nil
}
<commit_msg>cmd/tsurud: improve description of migrate command<commit_after>// Copyright 2015 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"log"
	"strconv"

	"github.com/tsuru/config"
	"github.com/tsuru/gnuflag"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/auth"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/db"
	"github.com/tsuru/tsuru/migration"
	"github.com/tsuru/tsuru/permission"
	"github.com/tsuru/tsuru/provision"
	"github.com/tsuru/tsuru/provision/docker"
	"gopkg.in/mgo.v2/bson"
)

func init() {
	err := migration.Register("migrate-docker-images", migrateImages)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
	err = migration.Register("migrate-pool", migratePool)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
	err = migration.Register("migrate-set-pool-to-app", setPoolToApps)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
	err = migration.Register("migrate-service-proxy-actions", migrateServiceProxyActions)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
	err = migration.RegisterOptional("migrate-roles", migrateRoles)
	if err != nil {
		log.Fatalf("unable to register migration: %s", err)
	}
}

func getProvisioner() (string, error) {
	provisioner, err := config.GetString("provisioner")
	if provisioner == "" {
		provisioner = "docker"
	}
	return provisioner, err
}

type migrationListCmd struct{}

func (*migrationListCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "migrate-list",
		Usage: "migrate-list",
		Desc:  "List available migration scripts from previous versions of tsurud",
	}
}

func (*migrationListCmd) Run(context *cmd.Context, client *cmd.Client) error {
	migrations, err := migration.List()
	if err != nil {
		return err
	}
	tbl := cmd.NewTable()
	tbl.Headers = cmd.Row{"Name", "Mandatory?", "Executed?"}
	for _, m := range migrations {
		tbl.AddRow(cmd.Row{m.Name, strconv.FormatBool(!m.Optional), strconv.FormatBool(m.Ran)})
	}
	fmt.Fprint(context.Stdout, tbl.String())
	return nil
}

type migrateCmd struct {
	fs    *gnuflag.FlagSet
	dry   bool
	force bool
	name  string
}

func (*migrateCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "migrate",
		Usage: "migrate [-n/--dry] [-f/--force] [--name name]",
		Desc: `Runs migrations from previous versions of tsurud. Only mandatory migrations
will be executed by default. To execute an optional migration the --name flag
must be informed.`,
	}
}

func (c *migrateCmd) Run(context *cmd.Context, client *cmd.Client) error {
	return migration.Run(migration.RunArgs{
		Writer: context.Stdout,
		Dry:    c.dry,
		Name:   c.name,
		Force:  c.force,
	})
}

func (c *migrateCmd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = gnuflag.NewFlagSet("migrate", gnuflag.ExitOnError)
		dryMsg := "Do not run migrations, just print what would run"
		c.fs.BoolVar(&c.dry, "dry", false, dryMsg)
		c.fs.BoolVar(&c.dry, "n", false, dryMsg)
		forceMsg := "Force the execution of an already executed optional migration"
		c.fs.BoolVar(&c.force, "force", false, forceMsg)
		c.fs.BoolVar(&c.force, "f", false, forceMsg)
		c.fs.StringVar(&c.name, "name", "", "The name of an optional migration to run")
	}
	return c.fs
}

func migrateImages() error {
	provisioner, _ := getProvisioner()
	if provisioner == "docker" {
		p, err := provision.Get(provisioner)
		if err != nil {
			return err
		}
		err = p.(provision.InitializableProvisioner).Initialize()
		if err != nil {
			return err
		}
		return docker.MigrateImages()
	}
	return nil
}

func migratePool() error {
	db, err := db.Conn()
	if err != nil {
		return err
	}
	defer db.Close()
	poolColl := db.Collection("pool")
	if err != nil {
		return err
	}
	var pools []provision.Pool
	err = db.Collection("docker_scheduler").Find(nil).All(&pools)
	if err != nil {
		return err
	}
	for _, p := range pools {
		err = poolColl.Insert(p)
		if err != nil {
			return err
		}
	}
	return nil
}

func setPoolToApps() error {
	db, err := db.Conn()
	if err != nil {
		return err
	}
	defer db.Close()
	var apps []app.App
	var tooManyPoolsApps []app.App
	err = db.Apps().Find(nil).All(&apps)
	if err != nil {
		return err
	}
	for _, a := range apps {
		err = a.SetPool()
		if err != nil {
			tooManyPoolsApps = append(tooManyPoolsApps, a)
			continue
		}
		err = db.Apps().Update(bson.M{"name": a.Name}, bson.M{"$set": bson.M{"pool": a.Pool}})
		if err != nil {
			return err
		}
	}
	if len(tooManyPoolsApps) > 0 {
		fmt.Println("Apps bellow couldn't be migrated because they are in more than one pool.")
		fmt.Println("To fix this, please run `tsuru app-change-pool <pool_name> -a app` for each app.")
		fmt.Println("*****************************************")
		for _, a := range tooManyPoolsApps {
			fmt.Println(a.Name)
		}
	}
	return nil
}

func migrateServiceProxyActions() error {
	db, err := db.Conn()
	if err != nil {
		return err
	}
	defer db.Close()
	_, err = db.UserActions().UpdateAll(
		bson.M{"action": "service-proxy-status"},
		bson.M{"$set": bson.M{"action": "service-instance-proxy"}},
	)
	return err
}

func createRole(name, contextType string) (permission.Role, error) {
	role, err := permission.NewRole(name, contextType)
	if err == permission.ErrRoleAlreadyExists {
		role, err = permission.FindRole(name)
	}
	return role, err
}

func migrateRoles() error {
	adminTeam, err := config.GetString("admin-team")
	if err != nil {
		return err
	}
	adminRole, err := createRole("admin", "global")
	if err != nil {
		return err
	}
	err = adminRole.AddPermissions("*")
	if err != nil {
		return err
	}
	teamMember, err := createRole("team-member", "team")
	if err != nil {
		return err
	}
	err = teamMember.AddPermissions(permission.PermApp.FullName(),
		permission.PermTeam.FullName(),
		permission.PermServiceInstance.FullName())
	if err != nil {
		return err
	}
	err = teamMember.AddEvent(permission.RoleEventTeamCreate.String())
	if err != nil {
		return err
	}
	teamCreator, err := createRole("team-creator", "global")
	if err != nil {
		return err
	}
	err = teamCreator.AddPermissions(permission.PermTeamCreate.FullName())
	if err != nil {
		return err
	}
	err = teamCreator.AddEvent(permission.RoleEventUserCreate.String())
	if err != nil {
		return err
	}
	users, err := auth.ListUsers()
	if err != nil {
		return err
	}
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	for _, u := range users {
		var teams []auth.Team
		err := conn.Teams().Find(bson.M{"users": bson.M{"$in": []string{u.Email}}}).All(&teams)
		if err != nil {
			return err
		}
		for _, team := range teams {
			if team.Name == adminTeam {
				err := u.AddRole(adminRole.Name, "")
				if err != nil {
					fmt.Printf("%s\n", err.Error())
				}
				continue
			}
			err := u.AddRole(teamMember.Name, team.Name)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
			}
			err = u.AddRole(teamCreator.Name, "")
			if err != nil {
				fmt.Printf("%s\n", err.Error())
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package api_test

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

	"github.com/lucapette/deloominator/pkg/api"
	"github.com/lucapette/deloominator/pkg/db"
	"github.com/lucapette/deloominator/pkg/testutil"
)

var update = flag.Bool("update", false, "update golden files")

var rows = db.QueryResult{
	Rows:    []db.Row{{db.Cell{Value: "42"}, db.Cell{Value: "Anna"}, db.Cell{Value: "Torv"}}},
	Columns: []db.Column{{Name: "actor_id"}, {Name: "first_name"}, {Name: "last_name"}},
}

type test struct {
	query   string
	code    int
	fixture string
}

func graphqlPayload(t *testing.T, query string) string {
	payload := struct {
		Query string `json:"query"`
	}{Query: query}

	json, err := json.Marshal(payload)
	if err != nil {
		t.Fatalf(err.Error())
	}

	return string(json)
}

func TestGraphQLDataSources(t *testing.T) {
	dsn, cleanup := testutil.SetupPG(t)
	cfg := testutil.InitConfig(t, map[string]string{
		"DATA_SOURCES": dsn,
	})
	dataSources, err := db.NewDataSources(cfg.Sources)
	if err != nil {
		t.Fatal(err)
	}

	defer func() {
		dataSources.Shutdown()
		cleanup()
	}()

	tests := []test{
		{query: "{ notAQuery }", code: 400, fixture: "wrong_query.json"},
		{query: "{ dataSources {name} }", code: 200, fixture: "data_sources.json"},
		{query: "{ dataSources {name tables {name}}}", code: 200, fixture: "data_sources_with_tables.json"},
	}

	for _, dataSource := range dataSources {
		testutil.LoadData(t, dataSource, "actor", rows)

		for _, test := range tests {
			t.Run(test.fixture, func(t *testing.T) {
				req := httptest.NewRequest("POST", "http://example.com/graphql", strings.NewReader(graphqlPayload(t, test.query)))
				w := httptest.NewRecorder()

				api.GraphQLHandler(dataSources)(w, req)

				resp, err := ioutil.ReadAll(w.Body)
				if err != nil {
					t.Fatal(err)
				}
				actual := string(resp)

				if w.Code != test.code {
					t.Fatalf("expected code %d, got: %d. Resp: %s", test.code, w.Code, actual)
				}

				var expected bytes.Buffer
				testutil.ParseFixture(t, &expected, test.fixture, testutil.DBTemplate{Name: dataSource.Name()})
				if *update {
					testutil.WriteFixture(t, test.fixture, actual)
				}

				if !reflect.DeepEqual(strings.TrimSuffix(expected.String(), "\n"), actual) {
					t.Fatalf("Unexpected result, diff: %v", testutil.Diff(expected.String(), actual))
				}
			})
		}
	}
}

var graphQLQuery = `
{ query(source: "%s", input: "%s") {
	... on results {
		chartName
	    columns { name type }
	    rows { cells { value } }
    }
    ... on queryError { message }
  }
}
`

var queryTests = []test{
	{
		query:   `select * from table_that_does_not_exist`,
		fixture: "query_error.json",
	},
	{
		query:   `select actor_id, first_name, last_name from actor`,
		fixture: "query_raw_results.json",
	},
	{
		query:   `select substr(first_name, 1, 1) initial, count(*)  from actor group by 1`,
		fixture: "query_simple_bar_detected.json",
	},
}

func TestGraphQLQuery(t *testing.T) {
	dsn, cleanup := testutil.SetupPG(t)
	cfg := testutil.InitConfig(t, map[string]string{
		"DATA_SOURCES": dsn,
	})
	dataSources, err := db.NewDataSources(cfg.Sources)
	if err != nil {
		t.Fatal(err)
	}

	defer func() {
		dataSources.Shutdown()
		cleanup()
	}()

	for _, dataSource := range dataSources {
		testutil.LoadData(t, dataSource, "actor", rows)

		for _, test := range queryTests {
			t.Run(test.fixture, func(t *testing.T) {
				query := graphqlPayload(t, fmt.Sprintf(graphQLQuery, dataSource.Name(), test.query))
				req := httptest.NewRequest("POST", "http://example.com/graphql", strings.NewReader(query))
				w := httptest.NewRecorder()

				api.GraphQLHandler(dataSources)(w, req)

				resp, err := ioutil.ReadAll(w.Body)
				if err != nil {
					t.Fatal(err)
				}
				actual := string(resp)

				if w.Code != 200 {
					t.Fatalf("expected code %d, got: %d. Resp: %s", 200, w.Code, actual)
				}

				var expected bytes.Buffer
				testutil.ParseFixture(t, &expected, test.fixture, testutil.DBTemplate{Name: dataSource.Name()})
				if *update {
					testutil.WriteFixture(t, test.fixture, actual)
				}

				if !reflect.DeepEqual(strings.TrimSuffix(expected.String(), "\n"), actual) {
					t.Fatalf("Unexpected result, diff: %v", testutil.Diff(expected.String(), actual))
				}
			})
		}
	}
}
<commit_msg>Simpler code, same result<commit_after>package api_test

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

	"github.com/lucapette/deloominator/pkg/api"
	"github.com/lucapette/deloominator/pkg/db"
	"github.com/lucapette/deloominator/pkg/testutil"
)

var update = flag.Bool("update", false, "update golden files")

var rows = db.QueryResult{
	Rows:    []db.Row{{db.Cell{Value: "42"}, db.Cell{Value: "Anna"}, db.Cell{Value: "Torv"}}},
	Columns: []db.Column{{Name: "actor_id"}, {Name: "first_name"}, {Name: "last_name"}},
}

type test struct {
	query   string
	code    int
	fixture string
}

func graphqlPayload(t *testing.T, query string) string {
	payload := struct {
		Query string `json:"query"`
	}{Query: query}

	json, err := json.Marshal(payload)
	if err != nil {
		t.Fatalf(err.Error())
	}

	return string(json)
}

func TestGraphQLDataSources(t *testing.T) {
	dsnPG, cleanupPG := testutil.SetupPG(t)
	cfg := testutil.InitConfig(t, map[string]string{
		"DATA_SOURCES": dsnPG,
	})
	dataSources, err := db.NewDataSources(cfg.Sources)
	if err != nil {
		t.Fatal(err)
	}

	defer func() {
		dataSources.Shutdown()
		cleanupPG()
	}()

	tests := []test{
		{query: "{ notAQuery }", code: 400, fixture: "wrong_query.json"},
		{query: "{ dataSources {name} }", code: 200, fixture: "data_sources.json"},
		{query: "{ dataSources {name tables {name}}}", code: 200, fixture: "data_sources_with_tables.json"},
	}

	for _, dataSource := range dataSources {
		testutil.LoadData(t, dataSource, "actor", rows)

		for _, test := range tests {
			t.Run(test.fixture, func(t *testing.T) {
				req := httptest.NewRequest("POST", "http://example.com/graphql", strings.NewReader(graphqlPayload(t, test.query)))
				w := httptest.NewRecorder()

				api.GraphQLHandler(dataSources)(w, req)

				resp, err := ioutil.ReadAll(w.Body)
				if err != nil {
					t.Fatal(err)
				}
				actual := string(resp)

				if w.Code != test.code {
					t.Fatalf("expected code %d, got: %d. Resp: %s", test.code, w.Code, actual)
				}

				var expected bytes.Buffer
				testutil.ParseFixture(t, &expected, test.fixture, testutil.DBTemplate{Name: dataSource.Name()})
				if *update {
					testutil.WriteFixture(t, test.fixture, actual)
				}

				if !reflect.DeepEqual(strings.TrimSuffix(expected.String(), "\n"), actual) {
					t.Fatalf("Unexpected result, diff: %v", testutil.Diff(expected.String(), actual))
				}
			})
		}
	}
}

var graphQLQuery = `
{ query(source: "%s", input: "%s") {
	... on results {
		chartName
	    columns { name type }
	    rows { cells { value } }
    }
    ... on queryError { message }
  }
}
`

var queryTests = []test{
	{
		query:   `select * from table_that_does_not_exist`,
		fixture: "query_error.json",
	},
	{
		query:   `select actor_id, first_name, last_name from actor`,
		fixture: "query_raw_results.json",
	},
	{
		query:   `select substr(first_name, 1, 1) initial, count(*)  from actor group by 1`,
		fixture: "query_simple_bar_detected.json",
	},
}

func TestGraphQLQuery(t *testing.T) {
	dsn, cleanup := testutil.SetupPG(t)
	cfg := testutil.InitConfig(t, map[string]string{
		"DATA_SOURCES": dsn,
	})
	dataSources, err := db.NewDataSources(cfg.Sources)
	if err != nil {
		t.Fatal(err)
	}

	defer func() {
		dataSources.Shutdown()
		cleanup()
	}()

	for _, dataSource := range dataSources {
		testutil.LoadData(t, dataSource, "actor", rows)

		for _, test := range queryTests {
			t.Run(test.fixture, func(t *testing.T) {
				query := graphqlPayload(t, fmt.Sprintf(graphQLQuery, dataSource.Name(), test.query))
				req := httptest.NewRequest("POST", "http://example.com/graphql", strings.NewReader(query))
				w := httptest.NewRecorder()

				api.GraphQLHandler(dataSources)(w, req)

				resp, err := ioutil.ReadAll(w.Body)
				if err != nil {
					t.Fatal(err)
				}
				actual := string(resp)

				if w.Code != 200 {
					t.Fatalf("expected code %d, got: %d. Resp: %s", 200, w.Code, actual)
				}

				expected := testutil.LoadFixture(t, test.fixture)
				if *update {
					testutil.WriteFixture(t, test.fixture, actual)
				}

				if !reflect.DeepEqual(strings.TrimSuffix(expected, "\n"), actual) {
					t.Fatalf("Unexpected result, diff: %v", testutil.Diff(expected, actual))
				}
			})
		}
	}
}
<|endoftext|>