<commit_before>package arg

import (
	"os"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func parse(cmdline string, dest interface{}) error {
	p, err := NewParser(dest)
	if err != nil {
		return err
	}
	return p.Parse(strings.Split(cmdline, " "))
}

func TestString(t *testing.T) {
	var args struct {
		Foo string
	}
	err := parse("--foo bar", &args)
	require.NoError(t, err)
	assert.Equal(t, "bar", args.Foo)
}

func TestInt(t *testing.T) {
	var args struct {
		Foo int
	}
	err := parse("--foo 7", &args)
	require.NoError(t, err)
	assert.EqualValues(t, 7, args.Foo)
}

func TestUint(t *testing.T) {
	var args struct {
		Foo uint
	}
	err := parse("--foo 7", &args)
	require.NoError(t, err)
	assert.EqualValues(t, 7, args.Foo)
}

func TestFloat(t *testing.T) {
	var args struct {
		Foo float32
	}
	err := parse("--foo 3.4", &args)
	require.NoError(t, err)
	assert.EqualValues(t, 3.4, args.Foo)
}

func TestDuration(t *testing.T) {
	var args struct {
		Foo time.Duration
	}
	err := parse("--foo 3ms", &args)
	require.NoError(t, err)
	assert.Equal(t, 3*time.Millisecond, args.Foo)
}

func TestMixed(t *testing.T) {
	var args struct {
		Foo  string `arg:"-f"`
		Bar  int
		Baz  uint `arg:"positional"`
		Ham  bool
		Spam float32
	}
	args.Bar = 3
	err := parse("123 -spam=1.2 -ham -f xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)
	assert.Equal(t, 3, args.Bar)
	assert.Equal(t, uint(123), args.Baz)
	assert.Equal(t, true, args.Ham)
	assert.EqualValues(t, 1.2, args.Spam)
}

func TestRequired(t *testing.T) {
	var args struct {
		Foo string `arg:"required"`
	}
	err := parse("", &args)
	require.Error(t, err, "--foo is required")
}

func TestShortFlag(t *testing.T) {
	var args struct {
		Foo string `arg:"-f"`
	}

	err := parse("-f xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)

	err = parse("-foo xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)

	err = parse("--foo xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)
}

func TestInvalidShortFlag(t *testing.T) {
	var args struct {
		Foo string `arg:"-foo"`
	}
	err := parse("", &args)
	assert.Error(t, err)
}

func TestLongFlag(t *testing.T) {
	var args struct {
		Foo string `arg:"--abc"`
	}

	err := parse("-abc xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)

	err = parse("--abc xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)
}

func TestCaseSensitive(t *testing.T) {
	var args struct {
		Lower bool `arg:"-v"`
		Upper bool `arg:"-V"`
	}

	err := parse("-v", &args)
	require.NoError(t, err)
	assert.True(t, args.Lower)
	assert.False(t, args.Upper)
}

func TestCaseSensitive2(t *testing.T) {
	var args struct {
		Lower bool `arg:"-v"`
		Upper bool `arg:"-V"`
	}

	err := parse("-V", &args)
	require.NoError(t, err)
	assert.False(t, args.Lower)
	assert.True(t, args.Upper)
}

func TestPositional(t *testing.T) {
	var args struct {
		Input  string `arg:"positional"`
		Output string `arg:"positional"`
	}
	err := parse("foo", &args)
	require.NoError(t, err)
	assert.Equal(t, "foo", args.Input)
	assert.Equal(t, "", args.Output)
}

func TestPositionalPointer(t *testing.T) {
	var args struct {
		Input  string    `arg:"positional"`
		Output []*string `arg:"positional"`
	}
	err := parse("foo bar baz", &args)
	require.NoError(t, err)
	assert.Equal(t, "foo", args.Input)
	bar := "bar"
	baz := "baz"
	assert.Equal(t, []*string{&bar, &baz}, args.Output)
}

func TestRequiredPositional(t *testing.T) {
	var args struct {
		Input  string `arg:"positional"`
		Output string `arg:"positional,required"`
	}
	err := parse("foo", &args)
	assert.Error(t, err)
}

func TestTooManyPositional(t *testing.T) {
	var args struct {
		Input  string `arg:"positional"`
		Output string `arg:"positional"`
	}
	err := parse("foo bar baz", &args)
	assert.Error(t, err)
}

func TestMultiple(t *testing.T) {
	var args struct {
		Foo []int
		Bar []string
	}
	err := parse("--foo 1 2 3 --bar x y z", &args)
	require.NoError(t, err)
	assert.Equal(t, []int{1, 2, 3}, args.Foo)
	assert.Equal(t, []string{"x", "y", "z"}, args.Bar)
}

func TestMultipleWithEq(t *testing.T) {
	var args struct {
		Foo []int
		Bar []string
	}
	err := parse("--foo 1 2 3 --bar=x", &args)
	require.NoError(t, err)
	assert.Equal(t, []int{1, 2, 3}, args.Foo)
	assert.Equal(t, []string{"x"}, args.Bar)
}

func TestExemptField(t *testing.T) {
	var args struct {
		Foo string
		Bar interface{} `arg:"-"`
	}
	err := parse("--foo xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)
}

func TestUnknownField(t *testing.T) {
	var args struct {
		Foo string
	}
	err := parse("--bar xyz", &args)
	assert.Error(t, err)
}

func TestMissingRequired(t *testing.T) {
	var args struct {
		Foo string   `arg:"required"`
		X   []string `arg:"positional"`
	}
	err := parse("x", &args)
	assert.Error(t, err)
}

func TestMissingValue(t *testing.T) {
	var args struct {
		Foo string
	}
	err := parse("--foo", &args)
	assert.Error(t, err)
}

func TestInvalidInt(t *testing.T) {
	var args struct {
		Foo int
	}
	err := parse("--foo=xyz", &args)
	assert.Error(t, err)
}

func TestInvalidUint(t *testing.T) {
	var args struct {
		Foo uint
	}
	err := parse("--foo=xyz", &args)
	assert.Error(t, err)
}

func TestInvalidFloat(t *testing.T) {
	var args struct {
		Foo float64
	}
	err := parse("--foo xyz", &args)
	require.Error(t, err)
}

func TestInvalidBool(t *testing.T) {
	var args struct {
		Foo bool
	}
	err := parse("--foo=xyz", &args)
	require.Error(t, err)
}

func TestInvalidIntSlice(t *testing.T) {
	var args struct {
		Foo []int
	}
	err := parse("--foo 1 2 xyz", &args)
	require.Error(t, err)
}

func TestInvalidPositional(t *testing.T) {
	var args struct {
		Foo int `arg:"positional"`
	}
	err := parse("xyz", &args)
	require.Error(t, err)
}

func TestInvalidPositionalSlice(t *testing.T) {
	var args struct {
		Foo []int `arg:"positional"`
	}
	err := parse("1 2 xyz", &args)
	require.Error(t, err)
}

func TestNoMoreOptions(t *testing.T) {
	var args struct {
		Foo string
		Bar []string `arg:"positional"`
	}
	err := parse("abc -- --foo xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "", args.Foo)
	assert.Equal(t, []string{"abc", "--foo", "xyz"}, args.Bar)
}

func TestHelpFlag(t *testing.T) {
	var args struct {
		Foo string
		Bar interface{} `arg:"-"`
	}
	err := parse("--help", &args)
	assert.Equal(t, ErrHelp, err)
}

func TestPanicOnNonPointer(t *testing.T) {
	var args struct{}
	assert.Panics(t, func() {
		parse("", args)
	})
}

func TestPanicOnNonStruct(t *testing.T) {
	var args string
	assert.Panics(t, func() {
		parse("", &args)
	})
}

func TestUnsupportedType(t *testing.T) {
	var args struct {
		Foo interface{}
	}
	err := parse("--foo", &args)
	assert.Error(t, err)
}

func TestUnsupportedSliceElement(t *testing.T) {
	var args struct {
		Foo []interface{}
	}
	err := parse("--foo", &args)
	assert.Error(t, err)
}

func TestUnknownTag(t *testing.T) {
	var args struct {
		Foo string `arg:"this_is_not_valid"`
	}
	err := parse("--foo xyz", &args)
	assert.Error(t, err)
}

func TestParse(t *testing.T) {
	var args struct {
		Foo string
	}
	os.Args = []string{"example", "--foo", "bar"}
	err := Parse(&args)
	require.NoError(t, err)
	assert.Equal(t, "bar", args.Foo)
}

func TestParseError(t *testing.T) {
	var args struct {
		Foo string `arg:"this_is_not_valid"`
	}
	os.Args = []string{"example", "--bar"}
	err := Parse(&args)
	assert.Error(t, err)
}

func TestMustParse(t *testing.T) {
	var args struct {
		Foo string
	}
	os.Args = []string{"example", "--foo", "bar"}
	parser := MustParse(&args)
	assert.Equal(t, "bar", args.Foo)
	assert.NotNil(t, parser)
}

func TestEnvironmentVariable(t *testing.T) {
	var args struct {
		Foo string `arg:"env"`
	}
	os.Setenv("FOO", "bar")
	os.Args = []string{"example"}
	MustParse(&args)
	assert.Equal(t, "bar", args.Foo)
}

func TestEnvironmentVariableOverrideName(t *testing.T) {
	var args struct {
		Foo string `arg:"env:BAZ"`
	}
	os.Setenv("BAZ", "bar")
	os.Args = []string{"example"}
	MustParse(&args)
	assert.Equal(t, "bar", args.Foo)
}

func TestEnvironmentVariableOverrideArgument(t *testing.T) {
	var args struct {
		Foo string `arg:"env"`
	}
	os.Setenv("FOO", "bar")
	os.Args = []string{"example", "--foo", "baz"}
	MustParse(&args)
	assert.Equal(t, "baz", args.Foo)
}

func TestEnvironmentVariableError(t *testing.T) {
	var args struct {
		Foo int `arg:"env"`
	}
	os.Setenv("FOO", "bar")
	os.Args = []string{"example"}
	err := Parse(&args)
	assert.Error(t, err)
}

func TestEnvironmentVariableRequired(t *testing.T) {
	var args struct {
		Foo string `arg:"env,required"`
	}
	os.Setenv("FOO", "bar")
	os.Args = []string{"example"}
	MustParse(&args)
	assert.Equal(t, "bar", args.Foo)
}
<commit_msg>add test for invalid durations<commit_after>package arg

import (
	"os"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func parse(cmdline string, dest interface{}) error {
	p, err := NewParser(dest)
	if err != nil {
		return err
	}
	return p.Parse(strings.Split(cmdline, " "))
}

func TestString(t *testing.T) {
	var args struct {
		Foo string
	}
	err := parse("--foo bar", &args)
	require.NoError(t, err)
	assert.Equal(t, "bar", args.Foo)
}

func TestInt(t *testing.T) {
	var args struct {
		Foo int
	}
	err := parse("--foo 7", &args)
	require.NoError(t, err)
	assert.EqualValues(t, 7, args.Foo)
}

func TestUint(t *testing.T) {
	var args struct {
		Foo uint
	}
	err := parse("--foo 7", &args)
	require.NoError(t, err)
	assert.EqualValues(t, 7, args.Foo)
}

func TestFloat(t *testing.T) {
	var args struct {
		Foo float32
	}
	err := parse("--foo 3.4", &args)
	require.NoError(t, err)
	assert.EqualValues(t, 3.4, args.Foo)
}

func TestDuration(t *testing.T) {
	var args struct {
		Foo time.Duration
	}
	err := parse("--foo 3ms", &args)
	require.NoError(t, err)
	assert.Equal(t, 3*time.Millisecond, args.Foo)
}

func TestInvalidDuration(t *testing.T) {
	var args struct {
		Foo time.Duration
	}
	err := parse("--foo xxx", &args)
	require.Error(t, err)
}

func TestMixed(t *testing.T) {
	var args struct {
		Foo  string `arg:"-f"`
		Bar  int
		Baz  uint `arg:"positional"`
		Ham  bool
		Spam float32
	}
	args.Bar = 3
	err := parse("123 -spam=1.2 -ham -f xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)
	assert.Equal(t, 3, args.Bar)
	assert.Equal(t, uint(123), args.Baz)
	assert.Equal(t, true, args.Ham)
	assert.EqualValues(t, 1.2, args.Spam)
}

func TestRequired(t *testing.T) {
	var args struct {
		Foo string `arg:"required"`
	}
	err := parse("", &args)
	require.Error(t, err, "--foo is required")
}

func TestShortFlag(t *testing.T) {
	var args struct {
		Foo string `arg:"-f"`
	}

	err := parse("-f xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)

	err = parse("-foo xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)

	err = parse("--foo xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)
}

func TestInvalidShortFlag(t *testing.T) {
	var args struct {
		Foo string `arg:"-foo"`
	}
	err := parse("", &args)
	assert.Error(t, err)
}

func TestLongFlag(t *testing.T) {
	var args struct {
		Foo string `arg:"--abc"`
	}

	err := parse("-abc xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)

	err = parse("--abc xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)
}

func TestCaseSensitive(t *testing.T) {
	var args struct {
		Lower bool `arg:"-v"`
		Upper bool `arg:"-V"`
	}

	err := parse("-v", &args)
	require.NoError(t, err)
	assert.True(t, args.Lower)
	assert.False(t, args.Upper)
}

func TestCaseSensitive2(t *testing.T) {
	var args struct {
		Lower bool `arg:"-v"`
		Upper bool `arg:"-V"`
	}

	err := parse("-V", &args)
	require.NoError(t, err)
	assert.False(t, args.Lower)
	assert.True(t, args.Upper)
}

func TestPositional(t *testing.T) {
	var args struct {
		Input  string `arg:"positional"`
		Output string `arg:"positional"`
	}
	err := parse("foo", &args)
	require.NoError(t, err)
	assert.Equal(t, "foo", args.Input)
	assert.Equal(t, "", args.Output)
}

func TestPositionalPointer(t *testing.T) {
	var args struct {
		Input  string    `arg:"positional"`
		Output []*string `arg:"positional"`
	}
	err := parse("foo bar baz", &args)
	require.NoError(t, err)
	assert.Equal(t, "foo", args.Input)
	bar := "bar"
	baz := "baz"
	assert.Equal(t, []*string{&bar, &baz}, args.Output)
}

func TestRequiredPositional(t *testing.T) {
	var args struct {
		Input  string `arg:"positional"`
		Output string `arg:"positional,required"`
	}
	err := parse("foo", &args)
	assert.Error(t, err)
}

func TestTooManyPositional(t *testing.T) {
	var args struct {
		Input  string `arg:"positional"`
		Output string `arg:"positional"`
	}
	err := parse("foo bar baz", &args)
	assert.Error(t, err)
}

func TestMultiple(t *testing.T) {
	var args struct {
		Foo []int
		Bar []string
	}
	err := parse("--foo 1 2 3 --bar x y z", &args)
	require.NoError(t, err)
	assert.Equal(t, []int{1, 2, 3}, args.Foo)
	assert.Equal(t, []string{"x", "y", "z"}, args.Bar)
}

func TestMultipleWithEq(t *testing.T) {
	var args struct {
		Foo []int
		Bar []string
	}
	err := parse("--foo 1 2 3 --bar=x", &args)
	require.NoError(t, err)
	assert.Equal(t, []int{1, 2, 3}, args.Foo)
	assert.Equal(t, []string{"x"}, args.Bar)
}

func TestExemptField(t *testing.T) {
	var args struct {
		Foo string
		Bar interface{} `arg:"-"`
	}
	err := parse("--foo xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "xyz", args.Foo)
}

func TestUnknownField(t *testing.T) {
	var args struct {
		Foo string
	}
	err := parse("--bar xyz", &args)
	assert.Error(t, err)
}

func TestMissingRequired(t *testing.T) {
	var args struct {
		Foo string   `arg:"required"`
		X   []string `arg:"positional"`
	}
	err := parse("x", &args)
	assert.Error(t, err)
}

func TestMissingValue(t *testing.T) {
	var args struct {
		Foo string
	}
	err := parse("--foo", &args)
	assert.Error(t, err)
}

func TestInvalidInt(t *testing.T) {
	var args struct {
		Foo int
	}
	err := parse("--foo=xyz", &args)
	assert.Error(t, err)
}

func TestInvalidUint(t *testing.T) {
	var args struct {
		Foo uint
	}
	err := parse("--foo=xyz", &args)
	assert.Error(t, err)
}

func TestInvalidFloat(t *testing.T) {
	var args struct {
		Foo float64
	}
	err := parse("--foo xyz", &args)
	require.Error(t, err)
}

func TestInvalidBool(t *testing.T) {
	var args struct {
		Foo bool
	}
	err := parse("--foo=xyz", &args)
	require.Error(t, err)
}

func TestInvalidIntSlice(t *testing.T) {
	var args struct {
		Foo []int
	}
	err := parse("--foo 1 2 xyz", &args)
	require.Error(t, err)
}

func TestInvalidPositional(t *testing.T) {
	var args struct {
		Foo int `arg:"positional"`
	}
	err := parse("xyz", &args)
	require.Error(t, err)
}

func TestInvalidPositionalSlice(t *testing.T) {
	var args struct {
		Foo []int `arg:"positional"`
	}
	err := parse("1 2 xyz", &args)
	require.Error(t, err)
}

func TestNoMoreOptions(t *testing.T) {
	var args struct {
		Foo string
		Bar []string `arg:"positional"`
	}
	err := parse("abc -- --foo xyz", &args)
	require.NoError(t, err)
	assert.Equal(t, "", args.Foo)
	assert.Equal(t, []string{"abc", "--foo", "xyz"}, args.Bar)
}

func TestHelpFlag(t *testing.T) {
	var args struct {
		Foo string
		Bar interface{} `arg:"-"`
	}
	err := parse("--help", &args)
	assert.Equal(t, ErrHelp, err)
}

func TestPanicOnNonPointer(t *testing.T) {
	var args struct{}
	assert.Panics(t, func() {
		parse("", args)
	})
}

func TestPanicOnNonStruct(t *testing.T) {
	var args string
	assert.Panics(t, func() {
		parse("", &args)
	})
}

func TestUnsupportedType(t *testing.T) {
	var args struct {
		Foo interface{}
	}
	err := parse("--foo", &args)
	assert.Error(t, err)
}

func TestUnsupportedSliceElement(t *testing.T) {
	var args struct {
		Foo []interface{}
	}
	err := parse("--foo", &args)
	assert.Error(t, err)
}

func TestUnknownTag(t *testing.T) {
	var args struct {
		Foo string `arg:"this_is_not_valid"`
	}
	err := parse("--foo xyz", &args)
	assert.Error(t, err)
}

func TestParse(t *testing.T) {
	var args struct {
		Foo string
	}
	os.Args = []string{"example", "--foo", "bar"}
	err := Parse(&args)
	require.NoError(t, err)
	assert.Equal(t, "bar", args.Foo)
}

func TestParseError(t *testing.T) {
	var args struct {
		Foo string `arg:"this_is_not_valid"`
	}
	os.Args = []string{"example", "--bar"}
	err := Parse(&args)
	assert.Error(t, err)
}

func TestMustParse(t *testing.T) {
	var args struct {
		Foo string
	}
	os.Args = []string{"example", "--foo", "bar"}
	parser := MustParse(&args)
	assert.Equal(t, "bar", args.Foo)
	assert.NotNil(t, parser)
}

func TestEnvironmentVariable(t *testing.T) {
	var args struct {
		Foo string `arg:"env"`
	}
	os.Setenv("FOO", "bar")
	os.Args = []string{"example"}
	MustParse(&args)
	assert.Equal(t, "bar", args.Foo)
}

func TestEnvironmentVariableOverrideName(t *testing.T) {
	var args struct {
		Foo string `arg:"env:BAZ"`
	}
	os.Setenv("BAZ", "bar")
	os.Args = []string{"example"}
	MustParse(&args)
	assert.Equal(t, "bar", args.Foo)
}

func TestEnvironmentVariableOverrideArgument(t *testing.T) {
	var args struct {
		Foo string `arg:"env"`
	}
	os.Setenv("FOO", "bar")
	os.Args = []string{"example", "--foo", "baz"}
	MustParse(&args)
	assert.Equal(t, "baz", args.Foo)
}

func TestEnvironmentVariableError(t *testing.T) {
	var args struct {
		Foo int `arg:"env"`
	}
	os.Setenv("FOO", "bar")
	os.Args = []string{"example"}
	err := Parse(&args)
	assert.Error(t, err)
}

func TestEnvironmentVariableRequired(t *testing.T) {
	var args struct {
		Foo string `arg:"env,required"`
	}
	os.Setenv("FOO", "bar")
	os.Args = []string{"example"}
	MustParse(&args)
	assert.Equal(t, "bar", args.Foo)
}
<|endoftext|>


<commit_before>package jsonutil

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/ioutil"
)

var (
	MarshalPrefix = ""
	MarshalIndent = "    "
)

type mustMarhshalError struct {
	MustMarhshalError string `json:"must_marshal_error"`
}

func MustMarshal(i interface{}, embedError bool) []byte {
	bytes, err := json.Marshal(i)
	if err != nil {
		if embedError {
			e := mustMarhshalError{
				MustMarhshalError: err.Error(),
			}
			bytes, err := json.Marshal(e)
			if err != nil {
				panic(err)
			}
			return bytes
		}
		panic(err)
	}
	return bytes
}

func MustMarshalString(i interface{}, embedError bool) string {
	return string(MustMarshal(i, embedError))
}

func MustMarshalIndent(i interface{}, prefix, indent string, embedError bool) []byte {
	bytes, err := json.MarshalIndent(i, prefix, indent)
	if err != nil {
		panic(err)
	}
	return bytes
}

func PrettyPrint(b []byte) ([]byte, error) {
	var out bytes.Buffer
	err := json.Indent(&out, b, MarshalPrefix, MarshalIndent)
	return out.Bytes(), err
}

func MarshalBase64(i interface{}) (string, error) {
	data, err := json.Marshal(i)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}

func UnmarshalIoReader(r io.Reader, i interface{}) ([]byte, error) {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return b, err
	}
	err = json.Unmarshal(b, i)
	return b, err
}
<commit_msg>streamline code<commit_after>package jsonutil

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/ioutil"
)

var (
	MarshalPrefix = ""
	MarshalIndent = "    "
)

type mustMarhshalError struct {
	MustMarhshalError string `json:"must_marshal_error"`
}

func MustMarshal(i interface{}, embedError bool) []byte {
	bytes, err := json.Marshal(i)
	if err != nil {
		if embedError {
			e := mustMarhshalError{
				MustMarhshalError: err.Error(),
			}
			bytes, err := json.Marshal(e)
			if err != nil {
				panic(err)
			}
			return bytes
		}
		panic(err)
	}
	return bytes
}

func MustMarshalString(i interface{}, embedError bool) string {
	return string(MustMarshal(i, embedError))
}

func MustMarshalIndent(i interface{}, prefix, indent string, embedError bool) []byte {
	bytes, err := json.MarshalIndent(i, prefix, indent)
	if err != nil {
		panic(err)
	}
	return bytes
}

func PrettyPrint(b []byte) ([]byte, error) {
	var out bytes.Buffer
	err := json.Indent(&out, b, MarshalPrefix, MarshalIndent)
	return out.Bytes(), err
}

func MarshalBase64(i interface{}) (string, error) {
	data, err := json.Marshal(i)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}

func UnmarshalIoReader(r io.Reader, i interface{}) ([]byte, error) {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return b, err
	}
	return b, json.Unmarshal(b, i)
}
<|endoftext|>


<commit_before>package main_test

import (
	"crypto/md5"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	. "github.com/cloudfoundry-incubator/buildpack_app_lifecycle/Godeps/_workspace/src/github.com/onsi/ginkgo"
	. "github.com/cloudfoundry-incubator/buildpack_app_lifecycle/Godeps/_workspace/src/github.com/onsi/gomega"
	"github.com/cloudfoundry-incubator/buildpack_app_lifecycle/Godeps/_workspace/src/github.com/onsi/gomega/gbytes"
	"github.com/cloudfoundry-incubator/buildpack_app_lifecycle/Godeps/_workspace/src/github.com/onsi/gomega/gexec"
)

var _ = Describe("Building", func() {
	buildpackFixtures := "fixtures/buildpacks"
	appFixtures := "fixtures/apps"

	var (
		builderCmd *exec.Cmd

		tmpDir                    string
		buildDir                  string
		buildpacksDir             string
		outputDroplet             string
		buildpackOrder            string
		buildArtifactsCacheDir    string
		outputMetadata            string
		outputBuildArtifactsCache string
	)

	builder := func() *gexec.Session {
		session, err := gexec.Start(
			builderCmd,
			GinkgoWriter,
			GinkgoWriter,
		)
		Ω(err).ShouldNot(HaveOccurred())

		return session
	}

	cpBuildpack := func(buildpack string) {
		hash := fmt.Sprintf("%x", md5.Sum([]byte(buildpack)))
		cp(path.Join(buildpackFixtures, buildpack), path.Join(buildpacksDir, hash))
	}

	BeforeEach(func() {
		var err error

		tmpDir, err = ioutil.TempDir("", "building-tmp")
		buildDir, err = ioutil.TempDir(tmpDir, "building-app")
		Ω(err).ShouldNot(HaveOccurred())

		buildpacksDir, err = ioutil.TempDir(tmpDir, "building-buildpacks")
		Ω(err).ShouldNot(HaveOccurred())

		outputDropletFile, err := ioutil.TempFile(tmpDir, "building-droplet")
		Ω(err).ShouldNot(HaveOccurred())
		outputDroplet = outputDropletFile.Name()

		outputBuildArtifactsCacheDir, err := ioutil.TempDir(tmpDir, "building-cache-output")
		Ω(err).ShouldNot(HaveOccurred())
		outputBuildArtifactsCache = filepath.Join(outputBuildArtifactsCacheDir, "cache.tgz")

		buildArtifactsCacheDir, err = ioutil.TempDir(tmpDir, "building-cache")
		Ω(err).ShouldNot(HaveOccurred())

		outputMetadataFile, err := ioutil.TempFile(tmpDir, "building-result")
		Ω(err).ShouldNot(HaveOccurred())
		outputMetadata = outputMetadataFile.Name()

		buildpackOrder = ""
	})

	AfterEach(func() {
		os.RemoveAll(tmpDir)
	})

	JustBeforeEach(func() {
		builderCmd = exec.Command(builderPath,
			"-buildDir", buildDir,
			"-buildpacksDir", buildpacksDir,
			"-outputDroplet", outputDroplet,
			"-outputBuildArtifactsCache", outputBuildArtifactsCache,
			"-buildArtifactsCacheDir", buildArtifactsCacheDir,
			"-buildpackOrder", buildpackOrder,
			"-outputMetadata", outputMetadata,
		)

		env := os.Environ()
		builderCmd.Env = append(env, "TMPDIR="+tmpDir)
	})

	resultJSON := func() []byte {
		resultInfo, err := ioutil.ReadFile(outputMetadata)
		Ω(err).ShouldNot(HaveOccurred())

		return resultInfo
	}

	Context("with a normal buildpack", func() {
		BeforeEach(func() {
			buildpackOrder = "always-detects,also-always-detects"

			cpBuildpack("always-detects")
			cpBuildpack("also-always-detects")
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		JustBeforeEach(func() {
			Eventually(builder()).Should(gexec.Exit(0))
		})

		Describe("the contents of the output tgz", func() {
			var files []string

			JustBeforeEach(func() {
				result, err := exec.Command("tar", "-tzf", outputDroplet).Output()
				Ω(err).ShouldNot(HaveOccurred())

				files = strings.Split(string(result), "\n")
			})

			It("should contain an /app dir with the contents of the compilation", func() {
				Ω(files).Should(ContainElement("./app/"))
				Ω(files).Should(ContainElement("./app/app.sh"))
				Ω(files).Should(ContainElement("./app/compiled"))
			})

			It("should contain an empty /tmp directory", func() {
				Ω(files).Should(ContainElement("./tmp/"))
				Ω(files).ShouldNot(ContainElement(MatchRegexp("\\./tmp/.+")))
			})

			It("should contain an empty /logs directory", func() {
				Ω(files).Should(ContainElement("./logs/"))
				Ω(files).ShouldNot(ContainElement(MatchRegexp("\\./logs/.+")))
			})

			It("should contain a staging_info.yml with the detected buildpack", func() {
				stagingInfo, err := exec.Command("tar", "-xzf", outputDroplet, "-O", "./staging_info.yml").Output()
				Ω(err).ShouldNot(HaveOccurred())

				expectedYAML := `detected_buildpack: Always Matching
start_command: the start command
`
				Ω(string(stagingInfo)).Should(Equal(expectedYAML))
			})
		})

		Describe("the build artifacts cache output tgz", func() {
			BeforeEach(func() {
				buildpackOrder = "always-detects-creates-build-artifacts"

				cpBuildpack("always-detects-creates-build-artifacts")
			})

			It("gets created", func() {
				result, err := exec.Command("tar", "-tzf", outputBuildArtifactsCache).Output()
				Ω(err).ShouldNot(HaveOccurred())

				Ω(strings.Split(string(result), "\n")).Should(ContainElement("./build-artifact"))
			})
		})

		Describe("the result.json, which is used to communicate back to the stager", func() {
			It("exists, and contains the detected buildpack", func() {
				Ω(resultJSON()).Should(MatchJSON(`{
					"detected_buildpack": "Always Matching",
					"execution_metadata": "{\"start_command\":\"the start command\"}",
					"buildpack_key": "always-detects",
					"detected_start_command":{"web":"the start command"}
				}`))
			})

			Context("when the app has a Procfile", func() {
				Context("with web defined", func() {
					BeforeEach(func() {
						cp(path.Join(appFixtures, "with-procfile-with-web", "Procfile"), buildDir)
					})

					It("chooses the Procfile-provided command", func() {
						Ω(resultJSON()).Should(MatchJSON(`{
					"detected_buildpack": "Always Matching",
					"execution_metadata": "{\"start_command\":\"procfile-provided start-command\"}",
					"buildpack_key": "always-detects",
					"detected_start_command":{"web":"procfile-provided start-command"}
				}`))
					})
				})

				Context("without web", func() {
					BeforeEach(func() {
						cp(path.Join(appFixtures, "with-procfile", "Procfile"), buildDir)
					})

					It("chooses the buildpack-provided command", func() {
						Ω(resultJSON()).Should(MatchJSON(`{
					"detected_buildpack": "Always Matching",
					"execution_metadata": "{\"start_command\":\"the start command\"}",
					"buildpack_key": "always-detects",
					"detected_start_command":{"web":"the start command"}
				}`))
					})
				})
			})
		})
	})

	Context("with a buildpack that does not determine a start command", func() {
		BeforeEach(func() {
			buildpackOrder = "release-without-command"
			cpBuildpack("release-without-command")
		})

		Context("when the app has a Procfile", func() {
			Context("with web defined", func() {
				JustBeforeEach(func() {
					Eventually(builder()).Should(gexec.Exit(0))
				})

				BeforeEach(func() {
					cp(path.Join(appFixtures, "with-procfile-with-web", "Procfile"), buildDir)
				})

				It("uses the command defined by web in the Procfile", func() {
					Ω(resultJSON()).Should(MatchJSON(`{
						"detected_buildpack": "Release Without Command",
						"execution_metadata": "{\"start_command\":\"procfile-provided start-command\"}",
						"buildpack_key": "release-without-command",
						"detected_start_command":{"web":"procfile-provided start-command"}
					}`))
				})
			})

			Context("without web", func() {
				BeforeEach(func() {
					cp(path.Join(appFixtures, "with-procfile", "Procfile"), buildDir)
				})

				It("fails", func() {
					session := builder()
					Eventually(session.Err).Should(gbytes.Say("No start command detected"))
					Eventually(session).Should(gexec.Exit(0))
				})
			})
		})

		Context("and the app has no Procfile", func() {
			BeforeEach(func() {
				cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
			})

			It("fails", func() {
				session := builder()
				Eventually(session.Err).Should(gbytes.Say("No start command detected"))
				Eventually(session).Should(gexec.Exit(0))
			})
		})
	})

	Context("with an app with an invalid Procfile", func() {
		BeforeEach(func() {
			buildpackOrder = "always-detects,also-always-detects"

			cpBuildpack("always-detects")
			cpBuildpack("also-always-detects")

			cp(path.Join(appFixtures, "bogus-procfile", "Procfile"), buildDir)
		})

		It("fails", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("Failed to read command from Procfile: invalid YAML"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("when no buildpacks match", func() {
		BeforeEach(func() {
			buildpackOrder = "always-fails"

			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
			cpBuildpack("always-fails")
		})

		It("should exit with an error", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("None of the buildpacks detected a compatible application"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("when the buildpack fails in compile", func() {
		BeforeEach(func() {
			buildpackOrder = "fails-to-compile"

			cpBuildpack("fails-to-compile")
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		It("should exit with an error", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("Failed to compile droplet: exit status 1"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("when the buildpack release generates invalid yaml", func() {
		BeforeEach(func() {
			buildpackOrder = "release-generates-bad-yaml"

			cpBuildpack("release-generates-bad-yaml")
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		It("should exit with an error", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("buildpack's release output invalid"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("when the buildpack fails to release", func() {
		BeforeEach(func() {
			buildpackOrder = "fails-to-release"

			cpBuildpack("fails-to-release")
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		It("should exit with an error", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("Failed to build droplet release: exit status 1"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("with a nested buildpack", func() {
		BeforeEach(func() {
			nestedBuildpack := "nested-buildpack"
			buildpackOrder = nestedBuildpack

			nestedBuildpackHash := "70d137ae4ee01fbe39058ccdebf48460"

			nestedBuildpackDir := path.Join(buildpacksDir, nestedBuildpackHash)
			err := os.MkdirAll(nestedBuildpackDir, 0777)
			Ω(err).ShouldNot(HaveOccurred())

			cp(path.Join(buildpackFixtures, "always-detects"), nestedBuildpackDir)
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		It("should detect the nested buildpack", func() {
			Eventually(builder()).Should(gexec.Exit(0))
		})
	})
})

func cp(src string, dst string) {
	session, err := gexec.Start(
		exec.Command("cp", "-a", src, dst),
		GinkgoWriter,
		GinkgoWriter,
	)
	Ω(err).ShouldNot(HaveOccurred())
	Eventually(session).Should(gexec.Exit(0))
}
<commit_msg>wait a bit longer for builder to execute & finish<commit_after>package main_test

import (
	"crypto/md5"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"

	. "github.com/cloudfoundry-incubator/buildpack_app_lifecycle/Godeps/_workspace/src/github.com/onsi/ginkgo"
	. "github.com/cloudfoundry-incubator/buildpack_app_lifecycle/Godeps/_workspace/src/github.com/onsi/gomega"
	"github.com/cloudfoundry-incubator/buildpack_app_lifecycle/Godeps/_workspace/src/github.com/onsi/gomega/gbytes"
	"github.com/cloudfoundry-incubator/buildpack_app_lifecycle/Godeps/_workspace/src/github.com/onsi/gomega/gexec"
)

var _ = Describe("Building", func() {
	buildpackFixtures := "fixtures/buildpacks"
	appFixtures := "fixtures/apps"

	var (
		builderCmd *exec.Cmd

		tmpDir                    string
		buildDir                  string
		buildpacksDir             string
		outputDroplet             string
		buildpackOrder            string
		buildArtifactsCacheDir    string
		outputMetadata            string
		outputBuildArtifactsCache string
	)

	builder := func() *gexec.Session {
		session, err := gexec.Start(
			builderCmd,
			GinkgoWriter,
			GinkgoWriter,
		)
		Ω(err).ShouldNot(HaveOccurred())

		return session
	}

	cpBuildpack := func(buildpack string) {
		hash := fmt.Sprintf("%x", md5.Sum([]byte(buildpack)))
		cp(path.Join(buildpackFixtures, buildpack), path.Join(buildpacksDir, hash))
	}

	BeforeEach(func() {
		var err error

		tmpDir, err = ioutil.TempDir("", "building-tmp")
		buildDir, err = ioutil.TempDir(tmpDir, "building-app")
		Ω(err).ShouldNot(HaveOccurred())

		buildpacksDir, err = ioutil.TempDir(tmpDir, "building-buildpacks")
		Ω(err).ShouldNot(HaveOccurred())

		outputDropletFile, err := ioutil.TempFile(tmpDir, "building-droplet")
		Ω(err).ShouldNot(HaveOccurred())
		outputDroplet = outputDropletFile.Name()

		outputBuildArtifactsCacheDir, err := ioutil.TempDir(tmpDir, "building-cache-output")
		Ω(err).ShouldNot(HaveOccurred())
		outputBuildArtifactsCache = filepath.Join(outputBuildArtifactsCacheDir, "cache.tgz")

		buildArtifactsCacheDir, err = ioutil.TempDir(tmpDir, "building-cache")
		Ω(err).ShouldNot(HaveOccurred())

		outputMetadataFile, err := ioutil.TempFile(tmpDir, "building-result")
		Ω(err).ShouldNot(HaveOccurred())
		outputMetadata = outputMetadataFile.Name()

		buildpackOrder = ""
	})

	AfterEach(func() {
		os.RemoveAll(tmpDir)
	})

	JustBeforeEach(func() {
		builderCmd = exec.Command(builderPath,
			"-buildDir", buildDir,
			"-buildpacksDir", buildpacksDir,
			"-outputDroplet", outputDroplet,
			"-outputBuildArtifactsCache", outputBuildArtifactsCache,
			"-buildArtifactsCacheDir", buildArtifactsCacheDir,
			"-buildpackOrder", buildpackOrder,
			"-outputMetadata", outputMetadata,
		)

		env := os.Environ()
		builderCmd.Env = append(env, "TMPDIR="+tmpDir)
	})

	resultJSON := func() []byte {
		resultInfo, err := ioutil.ReadFile(outputMetadata)
		Ω(err).ShouldNot(HaveOccurred())

		return resultInfo
	}

	Context("with a normal buildpack", func() {
		BeforeEach(func() {
			buildpackOrder = "always-detects,also-always-detects"

			cpBuildpack("always-detects")
			cpBuildpack("also-always-detects")
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		JustBeforeEach(func() {
			Eventually(builder(), 5*time.Second).Should(gexec.Exit(0))
		})

		Describe("the contents of the output tgz", func() {
			var files []string

			JustBeforeEach(func() {
				result, err := exec.Command("tar", "-tzf", outputDroplet).Output()
				Ω(err).ShouldNot(HaveOccurred())

				files = strings.Split(string(result), "\n")
			})

			It("should contain an /app dir with the contents of the compilation", func() {
				Ω(files).Should(ContainElement("./app/"))
				Ω(files).Should(ContainElement("./app/app.sh"))
				Ω(files).Should(ContainElement("./app/compiled"))
			})

			It("should contain an empty /tmp directory", func() {
				Ω(files).Should(ContainElement("./tmp/"))
				Ω(files).ShouldNot(ContainElement(MatchRegexp("\\./tmp/.+")))
			})

			It("should contain an empty /logs directory", func() {
				Ω(files).Should(ContainElement("./logs/"))
				Ω(files).ShouldNot(ContainElement(MatchRegexp("\\./logs/.+")))
			})

			It("should contain a staging_info.yml with the detected buildpack", func() {
				stagingInfo, err := exec.Command("tar", "-xzf", outputDroplet, "-O", "./staging_info.yml").Output()
				Ω(err).ShouldNot(HaveOccurred())

				expectedYAML := `detected_buildpack: Always Matching
start_command: the start command
`
				Ω(string(stagingInfo)).Should(Equal(expectedYAML))
			})
		})

		Describe("the build artifacts cache output tgz", func() {
			BeforeEach(func() {
				buildpackOrder = "always-detects-creates-build-artifacts"

				cpBuildpack("always-detects-creates-build-artifacts")
			})

			It("gets created", func() {
				result, err := exec.Command("tar", "-tzf", outputBuildArtifactsCache).Output()
				Ω(err).ShouldNot(HaveOccurred())

				Ω(strings.Split(string(result), "\n")).Should(ContainElement("./build-artifact"))
			})
		})

		Describe("the result.json, which is used to communicate back to the stager", func() {
			It("exists, and contains the detected buildpack", func() {
				Ω(resultJSON()).Should(MatchJSON(`{
					"detected_buildpack": "Always Matching",
					"execution_metadata": "{\"start_command\":\"the start command\"}",
					"buildpack_key": "always-detects",
					"detected_start_command":{"web":"the start command"}
				}`))
			})

			Context("when the app has a Procfile", func() {
				Context("with web defined", func() {
					BeforeEach(func() {
						cp(path.Join(appFixtures, "with-procfile-with-web", "Procfile"), buildDir)
					})

					It("chooses the Procfile-provided command", func() {
						Ω(resultJSON()).Should(MatchJSON(`{
					"detected_buildpack": "Always Matching",
					"execution_metadata": "{\"start_command\":\"procfile-provided start-command\"}",
					"buildpack_key": "always-detects",
					"detected_start_command":{"web":"procfile-provided start-command"}
				}`))
					})
				})

				Context("without web", func() {
					BeforeEach(func() {
						cp(path.Join(appFixtures, "with-procfile", "Procfile"), buildDir)
					})

					It("chooses the buildpack-provided command", func() {
						Ω(resultJSON()).Should(MatchJSON(`{
					"detected_buildpack": "Always Matching",
					"execution_metadata": "{\"start_command\":\"the start command\"}",
					"buildpack_key": "always-detects",
					"detected_start_command":{"web":"the start command"}
				}`))
					})
				})
			})
		})
	})

	Context("with a buildpack that does not determine a start command", func() {
		BeforeEach(func() {
			buildpackOrder = "release-without-command"
			cpBuildpack("release-without-command")
		})

		Context("when the app has a Procfile", func() {
			Context("with web defined", func() {
				JustBeforeEach(func() {
					Eventually(builder() * time.Second).Should(gexec.Exit(0))
				})

				BeforeEach(func() {
					cp(path.Join(appFixtures, "with-procfile-with-web", "Procfile"), buildDir)
				})

				It("uses the command defined by web in the Procfile", func() {
					Ω(resultJSON()).Should(MatchJSON(`{
						"detected_buildpack": "Release Without Command",
						"execution_metadata": "{\"start_command\":\"procfile-provided start-command\"}",
						"buildpack_key": "release-without-command",
						"detected_start_command":{"web":"procfile-provided start-command"}
					}`))
				})
			})

			Context("without web", func() {
				BeforeEach(func() {
					cp(path.Join(appFixtures, "with-procfile", "Procfile"), buildDir)
				})

				It("fails", func() {
					session := builder()
					Eventually(session.Err).Should(gbytes.Say("No start command detected"))
					Eventually(session).Should(gexec.Exit(0))
				})
			})
		})

		Context("and the app has no Procfile", func() {
			BeforeEach(func() {
				cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
			})

			It("fails", func() {
				session := builder()
				Eventually(session.Err).Should(gbytes.Say("No start command detected"))
				Eventually(session).Should(gexec.Exit(0))
			})
		})
	})

	Context("with an app with an invalid Procfile", func() {
		BeforeEach(func() {
			buildpackOrder = "always-detects,also-always-detects"

			cpBuildpack("always-detects")
			cpBuildpack("also-always-detects")

			cp(path.Join(appFixtures, "bogus-procfile", "Procfile"), buildDir)
		})

		It("fails", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("Failed to read command from Procfile: invalid YAML"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("when no buildpacks match", func() {
		BeforeEach(func() {
			buildpackOrder = "always-fails"

			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
			cpBuildpack("always-fails")
		})

		It("should exit with an error", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("None of the buildpacks detected a compatible application"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("when the buildpack fails in compile", func() {
		BeforeEach(func() {
			buildpackOrder = "fails-to-compile"

			cpBuildpack("fails-to-compile")
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		It("should exit with an error", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("Failed to compile droplet: exit status 1"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("when the buildpack release generates invalid yaml", func() {
		BeforeEach(func() {
			buildpackOrder = "release-generates-bad-yaml"

			cpBuildpack("release-generates-bad-yaml")
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		It("should exit with an error", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("buildpack's release output invalid"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("when the buildpack fails to release", func() {
		BeforeEach(func() {
			buildpackOrder = "fails-to-release"

			cpBuildpack("fails-to-release")
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		It("should exit with an error", func() {
			session := builder()
			Eventually(session.Err).Should(gbytes.Say("Failed to build droplet release: exit status 1"))
			Eventually(session).Should(gexec.Exit(1))
		})
	})

	Context("with a nested buildpack", func() {
		BeforeEach(func() {
			nestedBuildpack := "nested-buildpack"
			buildpackOrder = nestedBuildpack

			nestedBuildpackHash := "70d137ae4ee01fbe39058ccdebf48460"

			nestedBuildpackDir := path.Join(buildpacksDir, nestedBuildpackHash)
			err := os.MkdirAll(nestedBuildpackDir, 0777)
			Ω(err).ShouldNot(HaveOccurred())

			cp(path.Join(buildpackFixtures, "always-detects"), nestedBuildpackDir)
			cp(path.Join(appFixtures, "bash-app", "app.sh"), buildDir)
		})

		It("should detect the nested buildpack", func() {
			Eventually(builder()).Should(gexec.Exit(0))
		})
	})
})

func cp(src string, dst string) {
	session, err := gexec.Start(
		exec.Command("cp", "-a", src, dst),
		GinkgoWriter,
		GinkgoWriter,
	)
	Ω(err).ShouldNot(HaveOccurred())
	Eventually(session).Should(gexec.Exit(0))
}
<|endoftext|>


<commit_before>package pgbroadcast

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/lib/pq"
)

type pgnotification struct {
	Table  string                 `json:"table"`
	Action string                 `json:"action"`
	Data   map[string]interface{} `json:"data"`
}

func Run(pgconninfo string, pglistenchannel string) error {
	go h.run()
	err := startPgListener(pgconninfo, pglistenchannel, h.broadcast)
	return err
}

func startPgListener(pgconninfo string, pglistenchannel string, broadcastingchannel chan pgnotification) error {
	// callback func
	pgEventCallback := func(ev pq.ListenerEventType, err error) {
		if err != nil {
			fmt.Println("pgbroadcast: ", err.Error())
		}
	}

	// create listener and start listening
	l := pq.NewListener(pgconninfo, 10*time.Second, time.Minute, pgEventCallback)

	err := l.Listen(pglistenchannel)
	if err != nil {
		return err
	}
	// Wait for notifications in goroutine, pass them to the broadcastingchannel.
	go waitForNotification(l)
	fmt.Println("pgbroadcast: listening for notifications")
	return nil
}

func waitForNotification(l *pq.Listener) {
	for {
		select {
		case n := <-l.Notify:
			// Unmarshal JSON in pgnotification struct
			var pgn pgnotification
			err := json.Unmarshal([]byte(n.Extra), &pgn)
			if err != nil {
				fmt.Println("pgbroadcast: error processing JSON: ", err)
			} else {
				h.broadcast <- pgn
			}
		case <-time.After(60 * time.Second):
			// received no events for 60 seconds, ping connection")
			go func() {
				l.Ping()
			}()
		}
	}
}
<commit_msg>Check for nil notifications causing panic.<commit_after>package pgbroadcast

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/lib/pq"
)

type pgnotification struct {
	Table  string                 `json:"table"`
	Action string                 `json:"action"`
	Data   map[string]interface{} `json:"data"`
}

func Run(pgconninfo string, pglistenchannel string) error {
	go h.run()
	err := startPgListener(pgconninfo, pglistenchannel, h.broadcast)
	return err
}

func startPgListener(pgconninfo string, pglistenchannel string, broadcastingchannel chan pgnotification) error {
	// callback func
	pgEventCallback := func(ev pq.ListenerEventType, err error) {
		if err != nil {
			fmt.Println("pgbroadcast: ", err.Error())
		}
	}

	// create listener and start listening
	l := pq.NewListener(pgconninfo, 10*time.Second, time.Minute, pgEventCallback)

	err := l.Listen(pglistenchannel)
	if err != nil {
		return err
	}
	// Wait for notifications in goroutine, pass them to the broadcastingchannel.
	go waitForNotification(l)
	fmt.Println("pgbroadcast: listening for notifications")
	return nil
}

func waitForNotification(l *pq.Listener) {
	for {
		select {
		case n := <-l.Notify:
			//
			if n == nil {
				continue
			}
			// Unmarshal JSON in pgnotification struct
			var pgn pgnotification
			err := json.Unmarshal([]byte(n.Extra), &pgn)
			if err != nil {
				fmt.Println("pgbroadcast: error processing JSON: ", err)
			} else {
				h.broadcast <- pgn
			}
		case <-time.After(60 * time.Second):
			// received no events for 60 seconds, ping connection")
			go func() {
				l.Ping()
			}()
		}
	}
}
<|endoftext|>


<commit_before>package api

import (
	"archive/tar"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/gorilla/mux"
	"gopkg.in/yaml.v2"

	"github.com/ovh/cds/engine/api/application"
	"github.com/ovh/cds/engine/api/environment"
	"github.com/ovh/cds/engine/api/pipeline"
	"github.com/ovh/cds/engine/api/project"
	"github.com/ovh/cds/engine/api/workflow"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/exportentities"
	"github.com/ovh/cds/sdk/log"
)

func (api *API) postWorkflowImportHandler() Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		// Get project name in URL
		vars := mux.Vars(r)
		key := vars["permProjectKey"]
		force := FormBool(r, "force")

		//Load project
		proj, errp := project.Load(api.mustDB(), api.Cache, key, getUser(ctx),
			project.LoadOptions.WithGroups,
			project.LoadOptions.WithApplications,
			project.LoadOptions.WithEnvironments,
			project.LoadOptions.WithPipelines,
		)
		if errp != nil {
			return sdk.WrapError(errp, "postWorkflowImportHandler>> Unable load project")
		}

		body, errr := ioutil.ReadAll(r.Body)
		if errr != nil {
			return sdk.NewError(sdk.ErrWrongRequest, errr)
		}
		defer r.Body.Close()

		contentType := r.Header.Get("Content-Type")
		if contentType == "" {
			contentType = http.DetectContentType(body)
		}

		var ew = new(exportentities.Workflow)
		var errw error
		switch contentType {
		case "application/json":
			errw = json.Unmarshal(body, ew)
		case "application/x-yaml", "text/x-yam":
			errw = yaml.Unmarshal(body, ew)
		default:
			return sdk.NewError(sdk.ErrWrongRequest, fmt.Errorf("unsupported content-type: %s", contentType))
		}

		if errw != nil {
			return sdk.NewError(sdk.ErrWrongRequest, errw)
		}

		tx, errtx := api.mustDB().Begin()
		if errtx != nil {
			return sdk.WrapError(errtx, "postWorkflowImportHandler> Unable to start tx")
		}
		defer tx.Rollback()

		msgList, globalError := workflow.ParseAndImport(tx, api.Cache, proj, ew, force, getUser(ctx))
		msgListString := translate(r, msgList)

		if globalError != nil {
			myError, ok := globalError.(sdk.Error)
			if ok {
				return WriteJSON(w, r, msgListString, myError.Status)
			}
			return sdk.WrapError(globalError, "postWorkflowImportHandler> Unable import workflow %s", ew.Name)
		}

		if err := project.UpdateLastModified(tx, api.Cache, getUser(ctx), proj, sdk.ProjectPipelineLastModificationType); err != nil {
			return sdk.WrapError(err, "postWorkflowImportHandler> Unable to update project")
		}

		if err := tx.Commit(); err != nil {
			return sdk.WrapError(err, "postWorkflowImportHandler> Cannot commit transaction")
		}

		return WriteJSON(w, r, msgListString, http.StatusOK)
	}
}

func (api *API) postWorkflowPushHandler() Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		// Get project name in URL
		vars := mux.Vars(r)
		key := vars["permProjectKey"]

		//Load project
		proj, errp := project.Load(api.mustDB(), api.Cache, key, getUser(ctx),
			project.LoadOptions.WithGroups,
			project.LoadOptions.WithApplications,
			project.LoadOptions.WithEnvironments,
			project.LoadOptions.WithPipelines,
		)
		if errp != nil {
			log.Error("postWorkflowPushHandler> Unable to load project %s: %v", key, errp)
			return sdk.WrapError(errp, "postWorkflowPushHandler> Unable load project")
		}
		if r.Body == nil {
			return sdk.ErrWrongRequest
		}

		btes, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Error("postWorkflowPushHandler> Unable to read body: %v", err)
			return sdk.ErrWrongRequest
		}
		defer r.Body.Close()

		log.Debug("Read %d bytes from body", len(btes))

		tr := tar.NewReader(bytes.NewReader(btes))

		apps := make(map[string]exportentities.Application)
		pips := make(map[string]exportentities.PipelineV1)
		envs := make(map[string]exportentities.Environment)
		var wrkflw exportentities.Workflow

		mError := new(sdk.MultiError)
		for {
			hdr, err := tr.Next()
			if err == io.EOF {
				break
			}
			if err != nil {
				err = sdk.NewError(sdk.ErrWrongRequest, fmt.Errorf("Unable to read tar file"))
				return sdk.WrapError(err, "postWorkflowPushHandler>")
			}

			log.Debug("postWorkflowPushHandler> Reading %s", hdr.Name)

			buff := new(bytes.Buffer)
			if _, err := io.Copy(buff, tr); err != nil {
				err = sdk.NewError(sdk.ErrWrongRequest, fmt.Errorf("Unable to read tar file"))
				return sdk.WrapError(err, "postWorkflowPushHandler>")
			}

			b := buff.Bytes()
			switch {
			case strings.Contains(hdr.Name, ".app."):
				var app exportentities.Application
				if err := yaml.Unmarshal(b, &app); err != nil {
					log.Error("postWorkflowPushHandler> Unable to unmarshal application %s: %v", hdr.Name, err)
					mError.Append(fmt.Errorf("Unable to unmarshal application %s: %v", hdr.Name, err))
					continue
				}
				apps[hdr.Name] = app
			case strings.Contains(hdr.Name, ".pip."):
				var pip exportentities.PipelineV1
				if err := yaml.Unmarshal(b, &pip); err != nil {
					log.Error("postWorkflowPushHandler> Unable to unmarshal pipeline %s: %v", hdr.Name, err)
					mError.Append(fmt.Errorf("Unable to unmarshal pipeline %s: %v", hdr.Name, err))
					continue
				}
				pips[hdr.Name] = pip
			case strings.Contains(hdr.Name, ".env."):
				var env exportentities.Environment
				if err := yaml.Unmarshal(b, &env); err != nil {
					log.Error("postWorkflowPushHandler> Unable to unmarshal environment %s: %v", hdr.Name, err)
					mError.Append(fmt.Errorf("Unable to unmarshal environment %s: %v", hdr.Name, err))
					continue
				}
				envs[hdr.Name] = env
			default:
				if err := yaml.Unmarshal(b, &wrkflw); err != nil {
					log.Error("postWorkflowPushHandler> Unable to unmarshal workflow %s: %v", hdr.Name, err)
					mError.Append(fmt.Errorf("Unable to unmarshal workflow %s: %v", hdr.Name, err))
					continue
				}
			}
		}

		// We only use the multiError the une unmarshalling steps.
		// When a DB transaction has been started, just return at the first error
		// because transaction may have to be aborted
		if !mError.IsEmpty() {
			return mError
		}

		tx, err := api.mustDB().Begin()
		if err != nil {
			return sdk.WrapError(err, "postWorkflowPushHandler> Unable to start tx")
		}
		defer tx.Rollback()

		allMsg := []sdk.Message{}
		for filename, app := range apps {
			log.Debug("postWorkflowPushHandler> Parsing %s", filename)
			msgList, err := application.ParseAndImport(tx, api.Cache, proj, &app, true, project.DecryptWithBuiltinKey, getUser(ctx))
			if err != nil {
				err = sdk.SetError(err, "unable to import application %s", app.Name)
				return sdk.WrapError(err, "postWorkflowPushHandler> ", err)
			}
			allMsg = append(allMsg, msgList...)
			log.Debug("postWorkflowPushHandler> -- %s OK", filename)
		}

		for filename, env := range envs {
			log.Debug("postWorkflowPushHandler> Parsing %s", filename)
			msgList, err := environment.ParseAndImport(tx, api.Cache, proj, &env, true, project.DecryptWithBuiltinKey, getUser(ctx))
			if err != nil {
				err = sdk.SetError(err, "unable to import environment %s", env.Name)
				return sdk.WrapError(err, "postWorkflowPushHandler> ", err)
			}
			allMsg = append(allMsg, msgList...)
			log.Debug("postWorkflowPushHandler> -- %s OK", filename)
		}

		for filename, pip := range pips {
			log.Debug("postWorkflowPushHandler> Parsing %s", filename)
			msgList, err := pipeline.ParseAndImport(tx, api.Cache, proj, &pip, true, getUser(ctx))
			if err != nil {
				err = sdk.SetError(err, "unable to import pipeline %s", pip.Name)
				return sdk.WrapError(err, "postWorkflowPushHandler> ", err)
			}
			allMsg = append(allMsg, msgList...)
			log.Debug("postWorkflowPushHandler> -- %s OK", filename)
		}

		//Reload project to get apps, envs and pipelines updated
		proj, errp = project.Load(tx, api.Cache, key, getUser(ctx),
			project.LoadOptions.WithGroups,
			project.LoadOptions.WithApplications,
			project.LoadOptions.WithEnvironments,
			project.LoadOptions.WithPipelines,
		)
		if errp != nil {
			return sdk.WrapError(errp, "postWorkflowPushHandler> Unable reload project")
		}

		msgList, err := workflow.ParseAndImport(tx, api.Cache, proj, &wrkflw, true, getUser(ctx))
		if err != nil {
			err = sdk.SetError(err, "unable to import workflow %s", wrkflw.Name)
			return sdk.WrapError(err, "postWorkflowPushHandler> ", err)
		}

		allMsg = append(allMsg, msgList...)
		msgListString := translate(r, allMsg)

		if err := project.UpdateLastModified(tx, api.Cache, getUser(ctx), proj, sdk.ProjectPipelineLastModificationType); err != nil {
			return sdk.WrapError(err, "postWorkflowPushHandler> Unable to update project")
		}

		if err := tx.Commit(); err != nil {
			return sdk.WrapError(err, "postWorkflowPushHandler> Cannot commit transaction")
		}

		return WriteJSON(w, r, msgListString, http.StatusOK)
	}
}
<commit_msg>fix(api): fix mispell about content type yaml (#1994)<commit_after>package api

import (
	"archive/tar"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/gorilla/mux"
	"gopkg.in/yaml.v2"

	"github.com/ovh/cds/engine/api/application"
	"github.com/ovh/cds/engine/api/environment"
	"github.com/ovh/cds/engine/api/pipeline"
	"github.com/ovh/cds/engine/api/project"
	"github.com/ovh/cds/engine/api/workflow"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/exportentities"
	"github.com/ovh/cds/sdk/log"
)

func (api *API) postWorkflowImportHandler() Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		// Get project name in URL
		vars := mux.Vars(r)
		key := vars["permProjectKey"]
		force := FormBool(r, "force")

		//Load project
		proj, errp := project.Load(api.mustDB(), api.Cache, key, getUser(ctx),
			project.LoadOptions.WithGroups,
			project.LoadOptions.WithApplications,
			project.LoadOptions.WithEnvironments,
			project.LoadOptions.WithPipelines,
		)
		if errp != nil {
			return sdk.WrapError(errp, "postWorkflowImportHandler>> Unable load project")
		}

		body, errr := ioutil.ReadAll(r.Body)
		if errr != nil {
			return sdk.NewError(sdk.ErrWrongRequest, errr)
		}
		defer r.Body.Close()

		contentType := r.Header.Get("Content-Type")
		if contentType == "" {
			contentType = http.DetectContentType(body)
		}

		var ew = new(exportentities.Workflow)
		var errw error
		switch contentType {
		case "application/json":
			errw = json.Unmarshal(body, ew)
		case "application/x-yaml", "text/x-yaml":
			errw = yaml.Unmarshal(body, ew)
		default:
			return sdk.NewError(sdk.ErrWrongRequest, fmt.Errorf("unsupported content-type: %s", contentType))
		}

		if errw != nil {
			return sdk.NewError(sdk.ErrWrongRequest, errw)
		}

		tx, errtx := api.mustDB().Begin()
		if errtx != nil {
			return sdk.WrapError(errtx, "postWorkflowImportHandler> Unable to start tx")
		}
		defer tx.Rollback()

		msgList, globalError := workflow.ParseAndImport(tx, api.Cache, proj, ew, force, getUser(ctx))
		msgListString := translate(r, msgList)

		if globalError != nil {
			myError, ok := globalError.(sdk.Error)
			if ok {
				return WriteJSON(w, r, msgListString, myError.Status)
			}
			return sdk.WrapError(globalError, "postWorkflowImportHandler> Unable import workflow %s", ew.Name)
		}

		if err := project.UpdateLastModified(tx, api.Cache, getUser(ctx), proj, sdk.ProjectPipelineLastModificationType); err != nil {
			return sdk.WrapError(err, "postWorkflowImportHandler> Unable to update project")
		}

		if err := tx.Commit(); err != nil {
			return sdk.WrapError(err, "postWorkflowImportHandler> Cannot commit transaction")
		}

		return WriteJSON(w, r, msgListString, http.StatusOK)
	}
}

func (api *API) postWorkflowPushHandler() Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		// Get project name in URL
		vars := mux.Vars(r)
		key := vars["permProjectKey"]

		//Load project
		proj, errp := project.Load(api.mustDB(), api.Cache, key, getUser(ctx),
			project.LoadOptions.WithGroups,
			project.LoadOptions.WithApplications,
			project.LoadOptions.WithEnvironments,
			project.LoadOptions.WithPipelines,
		)
		if errp != nil {
			log.Error("postWorkflowPushHandler> Unable to load project %s: %v", key, errp)
			return sdk.WrapError(errp, "postWorkflowPushHandler> Unable load project")
		}
		if r.Body == nil {
			return sdk.ErrWrongRequest
		}

		btes, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Error("postWorkflowPushHandler> Unable to read body: %v", err)
			return sdk.ErrWrongRequest
		}
		defer r.Body.Close()

		log.Debug("Read %d bytes from body", len(btes))

		tr := tar.NewReader(bytes.NewReader(btes))

		apps := make(map[string]exportentities.Application)
		pips := make(map[string]exportentities.PipelineV1)
		envs := make(map[string]exportentities.Environment)
		var wrkflw exportentities.Workflow

		mError := new(sdk.MultiError)
		for {
			hdr, err := tr.Next()
			if err == io.EOF {
				break
			}
			if err != nil {
				err = sdk.NewError(sdk.ErrWrongRequest, fmt.Errorf("Unable to read tar file"))
				return sdk.WrapError(err, "postWorkflowPushHandler>")
			}

			log.Debug("postWorkflowPushHandler> Reading %s", hdr.Name)

			buff := new(bytes.Buffer)
			if _, err := io.Copy(buff, tr); err != nil {
				err = sdk.NewError(sdk.ErrWrongRequest, fmt.Errorf("Unable to read tar file"))
				return sdk.WrapError(err, "postWorkflowPushHandler>")
			}

			b := buff.Bytes()
			switch {
			case strings.Contains(hdr.Name, ".app."):
				var app exportentities.Application
				if err := yaml.Unmarshal(b, &app); err != nil {
					log.Error("postWorkflowPushHandler> Unable to unmarshal application %s: %v", hdr.Name, err)
					mError.Append(fmt.Errorf("Unable to unmarshal application %s: %v", hdr.Name, err))
					continue
				}
				apps[hdr.Name] = app
			case strings.Contains(hdr.Name, ".pip."):
				var pip exportentities.PipelineV1
				if err := yaml.Unmarshal(b, &pip); err != nil {
					log.Error("postWorkflowPushHandler> Unable to unmarshal pipeline %s: %v", hdr.Name, err)
					mError.Append(fmt.Errorf("Unable to unmarshal pipeline %s: %v", hdr.Name, err))
					continue
				}
				pips[hdr.Name] = pip
			case strings.Contains(hdr.Name, ".env."):
				var env exportentities.Environment
				if err := yaml.Unmarshal(b, &env); err != nil {
					log.Error("postWorkflowPushHandler> Unable to unmarshal environment %s: %v", hdr.Name, err)
					mError.Append(fmt.Errorf("Unable to unmarshal environment %s: %v", hdr.Name, err))
					continue
				}
				envs[hdr.Name] = env
			default:
				if err := yaml.Unmarshal(b, &wrkflw); err != nil {
					log.Error("postWorkflowPushHandler> Unable to unmarshal workflow %s: %v", hdr.Name, err)
					mError.Append(fmt.Errorf("Unable to unmarshal workflow %s: %v", hdr.Name, err))
					continue
				}
			}
		}

		// We only use the multiError the une unmarshalling steps.
		// When a DB transaction has been started, just return at the first error
		// because transaction may have to be aborted
		if !mError.IsEmpty() {
			return mError
		}

		tx, err := api.mustDB().Begin()
		if err != nil {
			return sdk.WrapError(err, "postWorkflowPushHandler> Unable to start tx")
		}
		defer tx.Rollback()

		allMsg := []sdk.Message{}
		for filename, app := range apps {
			log.Debug("postWorkflowPushHandler> Parsing %s", filename)
			msgList, err := application.ParseAndImport(tx, api.Cache, proj, &app, true, project.DecryptWithBuiltinKey, getUser(ctx))
			if err != nil {
				err = sdk.SetError(err, "unable to import application %s", app.Name)
				return sdk.WrapError(err, "postWorkflowPushHandler> ", err)
			}
			allMsg = append(allMsg, msgList...)
			log.Debug("postWorkflowPushHandler> -- %s OK", filename)
		}

		for filename, env := range envs {
			log.Debug("postWorkflowPushHandler> Parsing %s", filename)
			msgList, err := environment.ParseAndImport(tx, api.Cache, proj, &env, true, project.DecryptWithBuiltinKey, getUser(ctx))
			if err != nil {
				err = sdk.SetError(err, "unable to import environment %s", env.Name)
				return sdk.WrapError(err, "postWorkflowPushHandler> ", err)
			}
			allMsg = append(allMsg, msgList...)
			log.Debug("postWorkflowPushHandler> -- %s OK", filename)
		}

		for filename, pip := range pips {
			log.Debug("postWorkflowPushHandler> Parsing %s", filename)
			msgList, err := pipeline.ParseAndImport(tx, api.Cache, proj, &pip, true, getUser(ctx))
			if err != nil {
				err = sdk.SetError(err, "unable to import pipeline %s", pip.Name)
				return sdk.WrapError(err, "postWorkflowPushHandler> ", err)
			}
			allMsg = append(allMsg, msgList...)
			log.Debug("postWorkflowPushHandler> -- %s OK", filename)
		}

		//Reload project to get apps, envs and pipelines updated
		proj, errp = project.Load(tx, api.Cache, key, getUser(ctx),
			project.LoadOptions.WithGroups,
			project.LoadOptions.WithApplications,
			project.LoadOptions.WithEnvironments,
			project.LoadOptions.WithPipelines,
		)
		if errp != nil {
			return sdk.WrapError(errp, "postWorkflowPushHandler> Unable reload project")
		}

		msgList, err := workflow.ParseAndImport(tx, api.Cache, proj, &wrkflw, true, getUser(ctx))
		if err != nil {
			err = sdk.SetError(err, "unable to import workflow %s", wrkflw.Name)
			return sdk.WrapError(err, "postWorkflowPushHandler> ", err)
		}

		allMsg = append(allMsg, msgList...)
		msgListString := translate(r, allMsg)

		if err := project.UpdateLastModified(tx, api.Cache, getUser(ctx), proj, sdk.ProjectPipelineLastModificationType); err != nil {
			return sdk.WrapError(err, "postWorkflowPushHandler> Unable to update project")
		}

		if err := tx.Commit(); err != nil {
			return sdk.WrapError(err, "postWorkflowPushHandler> Cannot commit transaction")
		}

		return WriteJSON(w, r, msgListString, http.StatusOK)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"strings"

	"github.com/mackerelio/checkers"
)

type sshChecker struct {
	subcommand
}

func (c sshChecker) Execute(args []string) error {
	c.Executer = &c

	checker, err := c.executeAll()
	if err != nil {
		return err
	}

	checker.Name = "MasterHA"
	checker.Exit()
	return nil
}

func (c sshChecker) makeCommandName() string {
	return "masterha_check_ssh"
}

func (c sshChecker) makeCommandArgs() []string {
	return make([]string, 0, 2)
}

func (c sshChecker) parse(out string) (checkers.Status, string) {
	lines := extractNonEmptyLines(strings.Split(out, "\n"))
	lastLine := lines[len(lines)-1]
	if strings.Contains(lastLine, "All SSH connection tests passe") {
		return checkers.OK, lastLine
	} else if strings.Contains(lastLine, "SSH Configuration Check Failed!") {
		return checkers.OK, lastLine
	} else {
		msg := extractErrorMsg(out)
		return checkers.UNKNOWN, msg
	}
}
<commit_msg>should be critical status if ssh configuration check is failure<commit_after>package main

import (
	"strings"

	"github.com/mackerelio/checkers"
)

type sshChecker struct {
	subcommand
}

func (c sshChecker) Execute(args []string) error {
	c.Executer = &c

	checker, err := c.executeAll()
	if err != nil {
		return err
	}

	checker.Name = "MasterHA"
	checker.Exit()
	return nil
}

func (c sshChecker) makeCommandName() string {
	return "masterha_check_ssh"
}

func (c sshChecker) makeCommandArgs() []string {
	return make([]string, 0, 2)
}

func (c sshChecker) parse(out string) (checkers.Status, string) {
	lines := extractNonEmptyLines(strings.Split(out, "\n"))
	lastLine := lines[len(lines)-1]
	if strings.Contains(lastLine, "All SSH connection tests passe") {
		return checkers.OK, lastLine
	} else if strings.Contains(lastLine, "SSH Configuration Check Failed!") {
		return checkers.CRITICAL, lastLine
	} else {
		msg := extractErrorMsg(out)
		return checkers.UNKNOWN, msg
	}
}
<|endoftext|>


<commit_before>// Copyright 2014, The Serviced Authors. All rights reserved.
// Use of this source code is governed by a
// license that can be found in the LICENSE file.

package elasticsearch

import (
	"github.com/zenoss/glog"
	"github.com/zenoss/serviced/dao"
	"github.com/zenoss/serviced/domain/servicestate"
	"github.com/zenoss/serviced/rpc/agent"
)

func (this *ControlPlaneDao) GetServiceLogs(id string, logs *string) error {
	glog.V(3).Info("ControlPlaneDao.GetServiceLogs id=", id)
	var serviceStates []*servicestate.ServiceState
	err := this.zkDao.GetServiceStates(&serviceStates, id)
	if err != nil {
		return err
	}
	if len(serviceStates) == 0 {
		glog.V(1).Info("Unable to find any running services for ", id)
		return nil
	}

	serviceState := serviceStates[0]
	endpoint := serviceState.HostIP + ":4979"
	agentClient, err := agent.NewClient(endpoint)
	if err != nil {
		glog.Errorf("could not create client to %s", endpoint)
		return err
	}
	defer agentClient.Close()
	if mylogs, err := agentClient.GetDockerLogs(serviceState.DockerID); err != nil {
		glog.Errorf("could not get docker logs from agent client: %s", err)
		return err
	} else {
		*logs = mylogs
	}
	return nil
}

func (this *ControlPlaneDao) GetServiceStateLogs(request dao.ServiceStateRequest, logs *string) error {
	/* TODO: This command does not support logs on other hosts */
	glog.V(3).Info("ControlPlaneDao.GetServiceStateLogs id=", request)
	var serviceState servicestate.ServiceState
	err := this.zkDao.GetServiceState(&serviceState, request.ServiceID, request.ServiceStateID)
	if err != nil {
		glog.Errorf("ControlPlaneDao.GetServiceStateLogs servicestate=%+v err=%s", serviceState, err)
		return err
	}

	endpoint := serviceState.HostIP + ":4979"
	agentClient, err := agent.NewClient(endpoint)
	if err != nil {
		glog.Errorf("could not create client to %s", endpoint)
		return err
	}
	defer agentClient.Close()
	if mylogs, err := agentClient.GetDockerLogs(serviceState.DockerID); err != nil {
		glog.Errorf("could not get docker logs from agent client: %s", err)
		return err
	} else {
		*logs = mylogs
	}
	return nil
}
<commit_msg>add todo<commit_after>// Copyright 2014, The Serviced Authors. All rights reserved.
// Use of this source code is governed by a
// license that can be found in the LICENSE file.

package elasticsearch

import (
	"github.com/zenoss/glog"
	"github.com/zenoss/serviced/dao"
	"github.com/zenoss/serviced/domain/servicestate"
	"github.com/zenoss/serviced/rpc/agent"
)

func (this *ControlPlaneDao) GetServiceLogs(id string, logs *string) error {
	glog.V(3).Info("ControlPlaneDao.GetServiceLogs id=", id)
	var serviceStates []*servicestate.ServiceState
	err := this.zkDao.GetServiceStates(&serviceStates, id)
	if err != nil {
		return err
	}
	if len(serviceStates) == 0 {
		glog.V(1).Info("Unable to find any running services for ", id)
		return nil
	}

	serviceState := serviceStates[0]
	// FIXME: don't assume port is 4979
	endpoint := serviceState.HostIP + ":4979"
	agentClient, err := agent.NewClient(endpoint)
	if err != nil {
		glog.Errorf("could not create client to %s", endpoint)
		return err
	}
	defer agentClient.Close()
	if mylogs, err := agentClient.GetDockerLogs(serviceState.DockerID); err != nil {
		glog.Errorf("could not get docker logs from agent client: %s", err)
		return err
	} else {
		*logs = mylogs
	}
	return nil
}

func (this *ControlPlaneDao) GetServiceStateLogs(request dao.ServiceStateRequest, logs *string) error {
	/* TODO: This command does not support logs on other hosts */
	glog.V(3).Info("ControlPlaneDao.GetServiceStateLogs id=", request)
	var serviceState servicestate.ServiceState
	err := this.zkDao.GetServiceState(&serviceState, request.ServiceID, request.ServiceStateID)
	if err != nil {
		glog.Errorf("ControlPlaneDao.GetServiceStateLogs servicestate=%+v err=%s", serviceState, err)
		return err
	}

	// FIXME: don't assume port is 4979
	endpoint := serviceState.HostIP + ":4979"
	agentClient, err := agent.NewClient(endpoint)
	if err != nil {
		glog.Errorf("could not create client to %s", endpoint)
		return err
	}
	defer agentClient.Close()
	if mylogs, err := agentClient.GetDockerLogs(serviceState.DockerID); err != nil {
		glog.Errorf("could not get docker logs from agent client: %s", err)
		return err
	} else {
		*logs = mylogs
	}
	return nil
}
<|endoftext|>


<commit_before>// Application that captures webpage archives on a CT worker and uploads it to
// Google Storage.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"path/filepath"

	"github.com/skia-dev/glog"

	"strings"
	"time"

	"go.skia.org/infra/ct/go/util"
	"go.skia.org/infra/go/common"
	skutil "go.skia.org/infra/go/util"
)

var (
	workerNum     = flag.Int("worker_num", 1, "The number of this CT worker. It will be in the {1..100} range.")
	pagesetType   = flag.String("pageset_type", util.PAGESET_TYPE_MOBILE_10k, "The type of pagesets to create from the Alexa CSV list. Eg: 10k, Mobile10k, All.")
	chromiumBuild = flag.String("chromium_build", "", "The chromium build to use for this capture_archives run.")
)

func main() {
	common.Init()
	defer util.TimeTrack(time.Now(), "Capturing Archives")
	defer glog.Flush()

	// Create the task file so that the master knows this worker is still busy.
	skutil.LogErr(util.CreateTaskFile(util.ACTIVITY_CAPTURING_ARCHIVES))
	defer util.DeleteTaskFile(util.ACTIVITY_CAPTURING_ARCHIVES)

	if *chromiumBuild == "" {
		glog.Error("Must specify --chromium_build")
		return
	}

	// Reset the local chromium checkout.
	if err := util.ResetCheckout(util.ChromiumSrcDir); err != nil {
		glog.Errorf("Could not reset %s: %s", util.ChromiumSrcDir, err)
		return
	}
	// Sync the local chromium checkout.
	if err := util.SyncDir(util.ChromiumSrcDir); err != nil {
		glog.Errorf("Could not gclient sync %s: %s", util.ChromiumSrcDir, err)
		return
	}

	// Delete and remake the local webpage archives directory.
	pathToArchives := filepath.Join(util.WebArchivesDir, *pagesetType)
	skutil.RemoveAll(pathToArchives)
	skutil.MkdirAll(pathToArchives, 0700)

	// Instantiate GsUtil object.
	gs, err := util.NewGsUtil(nil)
	if err != nil {
		glog.Error(err)
		return
	}

	// Download the specified chromium build if it does not exist locally.
	if err := gs.DownloadChromiumBuild(*chromiumBuild); err != nil {
		glog.Error(err)
		return
	}

	// Download pagesets if they do not exist locally.
	if err := gs.DownloadWorkerArtifacts(util.PAGESETS_DIR_NAME, *pagesetType, *workerNum); err != nil {
		glog.Error(err)
		return
	}

	pathToPagesets := filepath.Join(util.PagesetsDir, *pagesetType)
	chromiumBinary := filepath.Join(util.ChromiumBuildsDir, *chromiumBuild, util.BINARY_CHROME)
	recordWprBinary := filepath.Join(util.TelemetryBinariesDir, util.BINARY_RECORD_WPR)
	timeoutSecs := util.PagesetTypeToInfo[*pagesetType].CaptureArchivesTimeoutSecs
	// Loop through all pagesets.
	fileInfos, err := ioutil.ReadDir(pathToPagesets)
	if err != nil {
		glog.Errorf("Unable to read the pagesets dir %s: %s", pathToPagesets, err)
		return
	}
	// TODO(rmistry): Remove this hack once the 1M webpage archives have been captured.
	glog.Infof("The length of fileInfos is: %s", len(fileInfos))
	fileInfos = fileInfos[11500:13000]
	glog.Infof("The fileInfos are: %s", fileInfos)
	for _, fileInfo := range fileInfos {
		pagesetBaseName := filepath.Base(fileInfo.Name())
		if pagesetBaseName == util.TIMESTAMP_FILE_NAME || filepath.Ext(pagesetBaseName) == ".pyc" {
			// Ignore timestamp files and .pyc files.
			continue
		}

		// Convert the filename into a format consumable by the record_wpr binary.
		pagesetArchiveName := strings.TrimSuffix(pagesetBaseName, filepath.Ext(pagesetBaseName))
		pagesetPath := filepath.Join(pathToPagesets, fileInfo.Name())

		glog.Infof("===== Processing %s =====", pagesetPath)
		args := []string{
			"--extra-browser-args=--disable-setuid-sandbox",
			"--browser=exact",
			"--browser-executable=" + chromiumBinary,
			fmt.Sprintf("%s_page_set", pagesetArchiveName),
			"--page-set-base-dir=" + pathToPagesets,
		}
		env := []string{
			fmt.Sprintf("PYTHONPATH=%s:$PYTHONPATH", pathToPagesets),
			"DISPLAY=:0",
		}
		skutil.LogErr(util.ExecuteCmd(recordWprBinary, args, env, time.Duration(timeoutSecs)*time.Second, nil, nil))
	}

	// Write timestamp to the webpage archives dir.
	skutil.LogErr(util.CreateTimestampFile(pathToArchives))

	// Upload webpage archives dir to Google Storage.
	if err := gs.UploadWorkerArtifacts(util.WEB_ARCHIVES_DIR_NAME, *pagesetType, *workerNum); err != nil {
		glog.Error(err)
		return
	}
}
<commit_msg>Capture CT archives from 13000-14499<commit_after>// Application that captures webpage archives on a CT worker and uploads it to
// Google Storage.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"path/filepath"

	"github.com/skia-dev/glog"

	"strings"
	"time"

	"go.skia.org/infra/ct/go/util"
	"go.skia.org/infra/go/common"
	skutil "go.skia.org/infra/go/util"
)

var (
	workerNum     = flag.Int("worker_num", 1, "The number of this CT worker. It will be in the {1..100} range.")
	pagesetType   = flag.String("pageset_type", util.PAGESET_TYPE_MOBILE_10k, "The type of pagesets to create from the Alexa CSV list. Eg: 10k, Mobile10k, All.")
	chromiumBuild = flag.String("chromium_build", "", "The chromium build to use for this capture_archives run.")
)

func main() {
	common.Init()
	defer util.TimeTrack(time.Now(), "Capturing Archives")
	defer glog.Flush()

	// Create the task file so that the master knows this worker is still busy.
	skutil.LogErr(util.CreateTaskFile(util.ACTIVITY_CAPTURING_ARCHIVES))
	defer util.DeleteTaskFile(util.ACTIVITY_CAPTURING_ARCHIVES)

	if *chromiumBuild == "" {
		glog.Error("Must specify --chromium_build")
		return
	}

	// Reset the local chromium checkout.
	if err := util.ResetCheckout(util.ChromiumSrcDir); err != nil {
		glog.Errorf("Could not reset %s: %s", util.ChromiumSrcDir, err)
		return
	}
	// Sync the local chromium checkout.
	if err := util.SyncDir(util.ChromiumSrcDir); err != nil {
		glog.Errorf("Could not gclient sync %s: %s", util.ChromiumSrcDir, err)
		return
	}

	// Delete and remake the local webpage archives directory.
	pathToArchives := filepath.Join(util.WebArchivesDir, *pagesetType)
	skutil.RemoveAll(pathToArchives)
	skutil.MkdirAll(pathToArchives, 0700)

	// Instantiate GsUtil object.
	gs, err := util.NewGsUtil(nil)
	if err != nil {
		glog.Error(err)
		return
	}

	// Download the specified chromium build if it does not exist locally.
	if err := gs.DownloadChromiumBuild(*chromiumBuild); err != nil {
		glog.Error(err)
		return
	}

	// Download pagesets if they do not exist locally.
	if err := gs.DownloadWorkerArtifacts(util.PAGESETS_DIR_NAME, *pagesetType, *workerNum); err != nil {
		glog.Error(err)
		return
	}

	pathToPagesets := filepath.Join(util.PagesetsDir, *pagesetType)
	chromiumBinary := filepath.Join(util.ChromiumBuildsDir, *chromiumBuild, util.BINARY_CHROME)
	recordWprBinary := filepath.Join(util.TelemetryBinariesDir, util.BINARY_RECORD_WPR)
	timeoutSecs := util.PagesetTypeToInfo[*pagesetType].CaptureArchivesTimeoutSecs
	// Loop through all pagesets.
	fileInfos, err := ioutil.ReadDir(pathToPagesets)
	if err != nil {
		glog.Errorf("Unable to read the pagesets dir %s: %s", pathToPagesets, err)
		return
	}
	// TODO(rmistry): Remove this hack once the 1M webpage archives have been captured.
	glog.Infof("The length of fileInfos is: %s", len(fileInfos))
	fileInfos = fileInfos[13000:14500]
	glog.Infof("The fileInfos are: %s", fileInfos)
	for _, fileInfo := range fileInfos {
		pagesetBaseName := filepath.Base(fileInfo.Name())
		if pagesetBaseName == util.TIMESTAMP_FILE_NAME || filepath.Ext(pagesetBaseName) == ".pyc" {
			// Ignore timestamp files and .pyc files.
			continue
		}

		// Convert the filename into a format consumable by the record_wpr binary.
		pagesetArchiveName := strings.TrimSuffix(pagesetBaseName, filepath.Ext(pagesetBaseName))
		pagesetPath := filepath.Join(pathToPagesets, fileInfo.Name())

		glog.Infof("===== Processing %s =====", pagesetPath)
		args := []string{
			"--extra-browser-args=--disable-setuid-sandbox",
			"--browser=exact",
			"--browser-executable=" + chromiumBinary,
			fmt.Sprintf("%s_page_set", pagesetArchiveName),
			"--page-set-base-dir=" + pathToPagesets,
		}
		env := []string{
			fmt.Sprintf("PYTHONPATH=%s:$PYTHONPATH", pathToPagesets),
			"DISPLAY=:0",
		}
		skutil.LogErr(util.ExecuteCmd(recordWprBinary, args, env, time.Duration(timeoutSecs)*time.Second, nil, nil))
	}

	// Write timestamp to the webpage archives dir.
	skutil.LogErr(util.CreateTimestampFile(pathToArchives))

	// Upload webpage archives dir to Google Storage.
	if err := gs.UploadWorkerArtifacts(util.WEB_ARCHIVES_DIR_NAME, *pagesetType, *workerNum); err != nil {
		glog.Error(err)
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/mdevilliers/redishappy/api"
	"github.com/mdevilliers/redishappy/configuration"
	"github.com/mdevilliers/redishappy/sentinel"
	"github.com/mdevilliers/redishappy/services/haproxy"
	"github.com/mdevilliers/redishappy/services/logger"
	"github.com/mdevilliers/redishappy/types"
	"github.com/mdevilliers/redishappy/util"
	"github.com/zenazn/goji"
)

func main() {

	//TODO : configure from command line
	logPath := "log" //var/log/redis-happy")
	configFile := "config.json"

	logger.InitLogging(logPath)

	logger.Info.Print("redis-happy started")

	configuration, err := configuration.LoadFromFile(configFile)

	if err != nil {
		logger.Error.Panic(err)
	}

	logger.Info.Printf("Parsed from config : %s\n", util.String(configuration))

	flipper := haproxy.NewFlipper(configuration)

	switchmasterchannel := make(chan types.MasterSwitchedEvent)

	go loopSentinelEvents(flipper, switchmasterchannel)

	sentinelManager := sentinel.NewManager(switchmasterchannel)

	go startMonitoring(flipper, sentinelManager, configuration)

	initApiServer(sentinelManager)
}

func initApiServer(manager *sentinel.SentinelManager) {

	logger.Info.Print("hosting json endpoint.")

	pongApi := api.PingApi{}
	sentinelApi := api.SentinelApi{Manager: manager}

	goji.Get("/api/ping", pongApi.Get)
	goji.Get("/api/sentinel", sentinelApi.Get)
	goji.Serve()
}

func startMonitoring(flipper types.FlipperClient, sentinelManager *sentinel.SentinelManager, configuration *configuration.Configuration) {

	started := 0
	detailcollection := types.NewMasterDetailsCollection()

	for _, configuredSentinel := range configuration.Sentinels {

		client, err := sentinelManager.NewSentinelClient(configuredSentinel)

		if err != nil {
			logger.Info.Printf("Error starting sentinel (%s) client : %s", configuredSentinel.GetLocation(), err.Error())
		} else {

			sentinelManager.NewSentinelMonitor(configuredSentinel)
			started++

			for _, clusterDetails := range configuration.Clusters {

				details, err := client.DiscoverMasterForCluster(clusterDetails.Name)

				if err == nil {
					details.ExternalPort = clusterDetails.MasterPort
					// TODO : last one wins?
					detailcollection.AddOrReplace(&details)
				}
			}
		}
	}

	flipper.InitialiseRunningState(detailcollection)

	if started == 0 {
		logger.Info.Printf("WARNING : no sentinels are currently being monitored.")
	}
}

func loopSentinelEvents(flipper types.FlipperClient, switchmasterchannel chan types.MasterSwitchedEvent) {

	for switchEvent := range switchmasterchannel {
		flipper.Orchestrate(switchEvent)
	}
}
<commit_msg>parse log file and config file from the command line<commit_after>package main

import (
	"flag"
	"github.com/mdevilliers/redishappy/api"
	"github.com/mdevilliers/redishappy/configuration"
	"github.com/mdevilliers/redishappy/sentinel"
	"github.com/mdevilliers/redishappy/services/haproxy"
	"github.com/mdevilliers/redishappy/services/logger"
	"github.com/mdevilliers/redishappy/types"
	"github.com/mdevilliers/redishappy/util"
	"github.com/zenazn/goji"
)

var configFile string
var logPath string

func init() {
	flag.StringVar(&configFile, "config", "config.json", "Full path of the configuration JSON file.")
	flag.StringVar(&logPath, "log", "log", "Folder for the logging folder.")
}

func main() {

	flag.Parse()
	logger.InitLogging(logPath)

	logger.Info.Print("redis-happy started")

	configuration, err := configuration.LoadFromFile(configFile)

	if err != nil {
		logger.Error.Panicf("Error opening config file : %s", err.Error())
	}

	logger.Info.Printf("Parsed from config : %s\n", util.String(configuration))

	flipper := haproxy.NewFlipper(configuration)

	switchmasterchannel := make(chan types.MasterSwitchedEvent)

	go loopSentinelEvents(flipper, switchmasterchannel)

	sentinelManager := sentinel.NewManager(switchmasterchannel)

	go startMonitoring(flipper, sentinelManager, configuration)

	initApiServer(sentinelManager)
}

func initApiServer(manager *sentinel.SentinelManager) {

	logger.Info.Print("hosting json endpoint.")

	pongApi := api.PingApi{}
	sentinelApi := api.SentinelApi{Manager: manager}

	goji.Get("/api/ping", pongApi.Get)
	goji.Get("/api/sentinel", sentinelApi.Get)
	goji.Serve()
}

func startMonitoring(flipper types.FlipperClient, sentinelManager *sentinel.SentinelManager, configuration *configuration.Configuration) {

	started := 0
	detailcollection := types.NewMasterDetailsCollection()

	for _, configuredSentinel := range configuration.Sentinels {

		client, err := sentinelManager.NewSentinelClient(configuredSentinel)

		if err != nil {
			logger.Info.Printf("Error starting sentinel (%s) client : %s", configuredSentinel.GetLocation(), err.Error())
		} else {

			sentinelManager.NewSentinelMonitor(configuredSentinel)
			started++

			for _, clusterDetails := range configuration.Clusters {

				details, err := client.DiscoverMasterForCluster(clusterDetails.Name)

				if err == nil {
					details.ExternalPort = clusterDetails.MasterPort
					// TODO : last one wins?
					detailcollection.AddOrReplace(&details)
				}
			}
		}
	}

	flipper.InitialiseRunningState(detailcollection)

	if started == 0 {
		logger.Info.Printf("WARNING : no sentinels are currently being monitored.")
	}
}

func loopSentinelEvents(flipper types.FlipperClient, switchmasterchannel chan types.MasterSwitchedEvent) {

	for switchEvent := range switchmasterchannel {
		flipper.Orchestrate(switchEvent)
	}
}
<|endoftext|>


<commit_before>package index

import (
	"bytes"
	"context"
	"io"
	"math"
	"strings"

	"github.com/gogo/protobuf/proto"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
	"github.com/pachyderm/pachyderm/src/server/pkg/storage/chunk"
	"github.com/pachyderm/pachyderm/src/server/pkg/storage/fileset/tar"
)

type levelReader struct {
	cr               *chunk.Reader
	tr               *tar.Reader
	currHdr, peekHdr *Header
}

// Reader is used for reading a multi-level index.
type Reader struct {
	ctx    context.Context
	objC   obj.Client
	chunks *chunk.Storage
	path   string
	prefix string
	levels []*levelReader
	done   bool
}

// NewReader create a new Reader.
func NewReader(ctx context.Context, objC obj.Client, chunks *chunk.Storage, path, prefix string) *Reader {
	return &Reader{
		ctx:    ctx,
		objC:   objC,
		chunks: chunks,
		path:   path,
		prefix: prefix,
	}
}

// Next gets the next header in the index.
func (r *Reader) Next() (*Header, error) {
	if err := r.setupLevels(); err != nil {
		return nil, err
	}
	if r.levels[0].peekHdr != nil {
		hdr := r.levels[0].peekHdr
		r.levels[0].peekHdr = nil
		return hdr, nil
	}
	if r.done {
		return nil, io.EOF
	}
	return r.next(len(r.levels) - 1)
}

func (r *Reader) setupLevels() error {
	if r.levels != nil {
		return nil
	}
	// Setup top level.
	objR, err := r.objC.Reader(r.ctx, r.path, 0, 0)
	if err != nil {
		return err
	}
	buf := &bytes.Buffer{}
	if _, err := io.Copy(buf, objR); err != nil {
		return err
	}
	if err := objR.Close(); err != nil {
		return err
	}
	r.levels = []*levelReader{&levelReader{tr: tar.NewReader(buf)}}
	// Traverse until we reach the first entry in the lowest level.
	for {
		hdr, err := r.peek(len(r.levels) - 1)
		if err != nil {
			return err
		}
		// Return when we are at the lowest level.
		if hdr.Hdr.Typeflag == indexType {
			return nil
		}
		_, err = r.next(len(r.levels) - 1)
		if err != nil {
			return err
		}
		// Setup next level.
		// (bryce) this whole process of reading at the offset is janky, needs to be re-thought.
		cr := r.chunks.NewReader(r.ctx, r.callback(len(r.levels)))
		dataRef := hdr.Idx.DataOp.DataRefs[0]
		dataRef.OffsetBytes = hdr.Idx.Range.Offset
		dataRef.SizeBytes -= hdr.Idx.Range.Offset
		cr.NextRange([]*chunk.DataRef{dataRef})
		r.levels = append(r.levels, &levelReader{
			cr: cr,
			tr: tar.NewReader(cr),
		})
	}
}

func (r *Reader) peek(level int) (*Header, error) {
	l := r.levels[level]
	if l.peekHdr != nil {
		return l.peekHdr, nil
	}
	var err error
	l.peekHdr, err = r.next(level)
	return l.peekHdr, err
}

func (r *Reader) next(level int) (*Header, error) {
	l := r.levels[level]
	if l.peekHdr != nil {
		hdr := l.peekHdr
		l.peekHdr = nil
		return hdr, nil
	}
	for {
		hdr, err := l.tr.Next()
		if err != nil {
			return nil, err
		}
		// Handle lowest level index.
		if hdr.Typeflag == indexType {
			cmpSize := int64(math.Min(float64(len(hdr.Name)), float64(len(r.prefix))))
			cmp := strings.Compare(hdr.Name[:cmpSize], r.prefix[:cmpSize])
			// If a header with the prefix cannot show up after the current header,
			// then we are done.
			if cmp > 0 {
				r.done = true
				return nil, io.EOF
			} else if cmp != 0 {
				continue
			}
			return deserialize(l.tr, hdr)
		}
		// Handle index level above lowest.
		fullHdr, err := deserialize(l.tr, hdr)
		if err != nil {
			return nil, err
		}
		// Skip to the starting header.
		if fullHdr.Idx.Range.LastPath < r.prefix {
			continue
		}
		return fullHdr, nil
	}
}

// Peek peeks ahead in the index.
func (r *Reader) Peek() (*Header, error) {
	if err := r.setupLevels(); err != nil {
		return nil, err
	}
	if r.levels[0].peekHdr != nil {
		return r.levels[0].peekHdr, nil
	}
	var err error
	r.levels[0].peekHdr, err = r.Next()
	return r.levels[0].peekHdr, err
}

func (r *Reader) callback(level int) chunk.ReaderFunc {
	return func() ([]*chunk.DataRef, error) {
		hdr, err := r.next(level - 1)
		if err != nil {
			return nil, err
		}
		r.levels[level-1].currHdr = hdr
		return hdr.Idx.DataOp.DataRefs, nil
	}
}

// Copy is the basic data structure to represent a copy of data from
// a reader to a writer.
type Copy struct {
	level int
	raw   *chunk.Copy
	hdrs  []*Header
}

// ReadCopyFunc returns a function for copying data from the reader.
func (r *Reader) ReadCopyFunc(pathBound ...string) func() (*Copy, error) {
	level := -1
	var offset int64
	var done bool
	return func() (*Copy, error) {
		if done {
			return nil, io.EOF
		}
		// Setup levels, initialize first level.
		if err := r.setupLevels(); err != nil {
			return nil, err
		}
		if level < 0 {
			level = len(r.levels) - 1
		}
		c := &Copy{level: r.wLevel(level)}
		cr := r.levels[level].cr
		// Handle index entries that span multiple chunks.
		if offset > 0 {
			raw, err := cr.ReadCopy(offset)
			if err != nil {
				return nil, err
			}
			c.raw = raw
		}
		// (bryce) this is janky, but we need the current header when copying a level above.
		if r.levels[level].currHdr != nil {
			r.levels[level].peekHdr, r.levels[level].currHdr = r.levels[level].currHdr, nil
		}
		// While not past a split point, get index entries to copy.
		pastSplit := false
		cr.OnSplit(func() { pastSplit = true })
		for !pastSplit {
			hdr, err := r.peek(level)
			if err != nil {
				if err == io.EOF {
					done = true
					return c, nil
				}
				return nil, err
			}
			// Stop copying when the last referenced (directly or indirectly) content
			// chunk has a path that is >= the path bound.
			if !BeforeBound(hdr.Idx.LastPathChunk, pathBound...) {
				if hdr.Idx.Range == nil {
					done = true
					return c, nil
				}
				level++
				offset = hdr.Idx.Range.Offset
				return c, nil
			}
			c.hdrs = append(c.hdrs, hdr)
			_, err = r.next(level)
			if err != nil {
				return nil, err
			}
		}
		level--
		return c, nil
	}
}

func (r *Reader) wLevel(rLevel int) int {
	return len(r.levels) - 1 - rLevel
}

// BeforeBound checks if the passed in string is before the string bound (exclusive).
// The string bound is optional, so if no string bound is passed then it returns true.
func BeforeBound(str string, strBound ...string) bool {
	return len(strBound) == 0 || strings.Compare(str, strBound[0]) < 0
}

// Close closes the reader.
func (r *Reader) Close() error {
	for i := 1; i < len(r.levels); i++ {
		if err := r.levels[i].cr.Close(); err != nil {
			return err
		}
	}
	return nil
}

func deserialize(tr *tar.Reader, hdr *tar.Header) (*Header, error) {
	data := &bytes.Buffer{}
	buf := grpcutil.GetBuffer()
	defer grpcutil.PutBuffer(buf)
	if _, err := io.CopyBuffer(data, tr, buf); err != nil {
		return nil, err
	}
	idx := &Index{}
	if err := proto.Unmarshal(data.Bytes(), idx); err != nil {
		return nil, err
	}
	return &Header{
		Hdr: hdr,
		Idx: idx,
	}, nil
}
<commit_msg>Adds documentation to wLevel function<commit_after>package index

import (
	"bytes"
	"context"
	"io"
	"math"
	"strings"

	"github.com/gogo/protobuf/proto"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
	"github.com/pachyderm/pachyderm/src/server/pkg/storage/chunk"
	"github.com/pachyderm/pachyderm/src/server/pkg/storage/fileset/tar"
)

type levelReader struct {
	cr               *chunk.Reader
	tr               *tar.Reader
	currHdr, peekHdr *Header
}

// Reader is used for reading a multi-level index.
type Reader struct {
	ctx    context.Context
	objC   obj.Client
	chunks *chunk.Storage
	path   string
	prefix string
	levels []*levelReader
	done   bool
}

// NewReader create a new Reader.
func NewReader(ctx context.Context, objC obj.Client, chunks *chunk.Storage, path, prefix string) *Reader {
	return &Reader{
		ctx:    ctx,
		objC:   objC,
		chunks: chunks,
		path:   path,
		prefix: prefix,
	}
}

// Next gets the next header in the index.
func (r *Reader) Next() (*Header, error) {
	if err := r.setupLevels(); err != nil {
		return nil, err
	}
	if r.levels[0].peekHdr != nil {
		hdr := r.levels[0].peekHdr
		r.levels[0].peekHdr = nil
		return hdr, nil
	}
	if r.done {
		return nil, io.EOF
	}
	return r.next(len(r.levels) - 1)
}

func (r *Reader) setupLevels() error {
	if r.levels != nil {
		return nil
	}
	// Setup top level.
	objR, err := r.objC.Reader(r.ctx, r.path, 0, 0)
	if err != nil {
		return err
	}
	buf := &bytes.Buffer{}
	if _, err := io.Copy(buf, objR); err != nil {
		return err
	}
	if err := objR.Close(); err != nil {
		return err
	}
	r.levels = []*levelReader{&levelReader{tr: tar.NewReader(buf)}}
	// Traverse until we reach the first entry in the lowest level.
	for {
		hdr, err := r.peek(len(r.levels) - 1)
		if err != nil {
			return err
		}
		// Return when we are at the lowest level.
		if hdr.Hdr.Typeflag == indexType {
			return nil
		}
		_, err = r.next(len(r.levels) - 1)
		if err != nil {
			return err
		}
		// Setup next level.
		// (bryce) this whole process of reading at the offset is janky, needs to be re-thought.
		cr := r.chunks.NewReader(r.ctx, r.callback(len(r.levels)))
		dataRef := hdr.Idx.DataOp.DataRefs[0]
		dataRef.OffsetBytes = hdr.Idx.Range.Offset
		dataRef.SizeBytes -= hdr.Idx.Range.Offset
		cr.NextRange([]*chunk.DataRef{dataRef})
		r.levels = append(r.levels, &levelReader{
			cr: cr,
			tr: tar.NewReader(cr),
		})
	}
}

func (r *Reader) peek(level int) (*Header, error) {
	l := r.levels[level]
	if l.peekHdr != nil {
		return l.peekHdr, nil
	}
	var err error
	l.peekHdr, err = r.next(level)
	return l.peekHdr, err
}

func (r *Reader) next(level int) (*Header, error) {
	l := r.levels[level]
	if l.peekHdr != nil {
		hdr := l.peekHdr
		l.peekHdr = nil
		return hdr, nil
	}
	for {
		hdr, err := l.tr.Next()
		if err != nil {
			return nil, err
		}
		// Handle lowest level index.
		if hdr.Typeflag == indexType {
			cmpSize := int64(math.Min(float64(len(hdr.Name)), float64(len(r.prefix))))
			cmp := strings.Compare(hdr.Name[:cmpSize], r.prefix[:cmpSize])
			// If a header with the prefix cannot show up after the current header,
			// then we are done.
			if cmp > 0 {
				r.done = true
				return nil, io.EOF
			} else if cmp != 0 {
				continue
			}
			return deserialize(l.tr, hdr)
		}
		// Handle index level above lowest.
		fullHdr, err := deserialize(l.tr, hdr)
		if err != nil {
			return nil, err
		}
		// Skip to the starting header.
		if fullHdr.Idx.Range.LastPath < r.prefix {
			continue
		}
		return fullHdr, nil
	}
}

// Peek peeks ahead in the index.
func (r *Reader) Peek() (*Header, error) {
	if err := r.setupLevels(); err != nil {
		return nil, err
	}
	if r.levels[0].peekHdr != nil {
		return r.levels[0].peekHdr, nil
	}
	var err error
	r.levels[0].peekHdr, err = r.Next()
	return r.levels[0].peekHdr, err
}

func (r *Reader) callback(level int) chunk.ReaderFunc {
	return func() ([]*chunk.DataRef, error) {
		hdr, err := r.next(level - 1)
		if err != nil {
			return nil, err
		}
		r.levels[level-1].currHdr = hdr
		return hdr.Idx.DataOp.DataRefs, nil
	}
}

// Copy is the basic data structure to represent a copy of data from
// a reader to a writer.
type Copy struct {
	level int
	raw   *chunk.Copy
	hdrs  []*Header
}

// ReadCopyFunc returns a function for copying data from the reader.
func (r *Reader) ReadCopyFunc(pathBound ...string) func() (*Copy, error) {
	level := -1
	var offset int64
	var done bool
	return func() (*Copy, error) {
		if done {
			return nil, io.EOF
		}
		// Setup levels, initialize first level.
		if err := r.setupLevels(); err != nil {
			return nil, err
		}
		if level < 0 {
			level = len(r.levels) - 1
		}
		c := &Copy{level: r.wLevel(level)}
		cr := r.levels[level].cr
		// Handle index entries that span multiple chunks.
		if offset > 0 {
			raw, err := cr.ReadCopy(offset)
			if err != nil {
				return nil, err
			}
			c.raw = raw
		}
		// (bryce) this is janky, but we need the current header when copying a level above.
		if r.levels[level].currHdr != nil {
			r.levels[level].peekHdr, r.levels[level].currHdr = r.levels[level].currHdr, nil
		}
		// While not past a split point, get index entries to copy.
		pastSplit := false
		cr.OnSplit(func() { pastSplit = true })
		for !pastSplit {
			hdr, err := r.peek(level)
			if err != nil {
				if err == io.EOF {
					done = true
					return c, nil
				}
				return nil, err
			}
			// Stop copying when the last referenced (directly or indirectly) content
			// chunk has a path that is >= the path bound.
			if !BeforeBound(hdr.Idx.LastPathChunk, pathBound...) {
				if hdr.Idx.Range == nil {
					done = true
					return c, nil
				}
				level++
				offset = hdr.Idx.Range.Offset
				return c, nil
			}
			c.hdrs = append(c.hdrs, hdr)
			_, err = r.next(level)
			if err != nil {
				return nil, err
			}
		}
		level--
		return c, nil
	}
}

// wLevel converts a reader level to the corresponding writer level.
// The reason this is necessary is because the corresponding writer level is the inverse of the reader level.
// The reader levels start at the top level index entry while the writer levels start at the content level.
func (r *Reader) wLevel(rLevel int) int {
	return len(r.levels) - 1 - rLevel
}

// BeforeBound checks if the passed in string is before the string bound (exclusive).
// The string bound is optional, so if no string bound is passed then it returns true.
func BeforeBound(str string, strBound ...string) bool {
	return len(strBound) == 0 || strings.Compare(str, strBound[0]) < 0
}

// Close closes the reader.
func (r *Reader) Close() error {
	for i := 1; i < len(r.levels); i++ {
		if err := r.levels[i].cr.Close(); err != nil {
			return err
		}
	}
	return nil
}

func deserialize(tr *tar.Reader, hdr *tar.Header) (*Header, error) {
	data := &bytes.Buffer{}
	buf := grpcutil.GetBuffer()
	defer grpcutil.PutBuffer(buf)
	if _, err := io.CopyBuffer(data, tr, buf); err != nil {
		return nil, err
	}
	idx := &Index{}
	if err := proto.Unmarshal(data.Bytes(), idx); err != nil {
		return nil, err
	}
	return &Header{
		Hdr: hdr,
		Idx: idx,
	}, nil
}
<|endoftext|>