<commit_before>package osc

import (
	"bufio"
	"bytes"
	"sync"
	"testing"
	"time"
)

func TestAppendArguments(t *testing.T) {
	oscAddress := "/address"
	message := NewOscMessage(oscAddress)
	if message.Address != oscAddress {
		t.Errorf("OSC address should be \"%s\" and is \"%s\"", oscAddress, message.Address)
	}

	message.Append("string argument")
	message.Append(123456789)
	message.Append(true)

	if message.CountArguments() != 3 {
		t.Errorf("Number of arguments should be %d and is %d", 3, message.CountArguments())
	}
}

func TestEqualMessage(t *testing.T) {
	msg1 := NewOscMessage("/address")
	msg2 := NewOscMessage("/address")

	msg1.Append(1234)
	msg2.Append(1234)
	msg1.Append("test string")
	msg2.Append("test string")

	if !msg1.Equals(msg2) {
		t.Error("Messages should be equal")
	}
}

func TestAddMsgHandler(t *testing.T) {
	server := NewOscServer("localhost", 6677)
	err := server.AddMsgHandler("/address/test", func(msg *OscMessage) {})
	if err != nil {
		t.Error("Expected that OSC address '/address/test' is valid")
	}
}

func TestAddMsgHandlerWithInvalidAddress(t *testing.T) {
	server := NewOscServer("localhost", 6677)
	err := server.AddMsgHandler("/address*/test", func(msg *OscMessage) {})
	if err == nil {
		t.Error("Expected error with '/address*/test'")
	}
}

func TestServerMessageDispatching(t *testing.T) {
	finish := make(chan bool)
	start := make(chan bool)
	var done sync.WaitGroup
	done.Add(2)

	// Start the OSC server in a new go-routine
	go func() {
		server := NewOscServer("localhost", 6677)
		err := server.AddMsgHandler("/address/test", func(msg *OscMessage) {
			if len(msg.Arguments) != 1 {
				t.Error("Argument length should be 1 and is: " + string(len(msg.Arguments)))
			}

			if msg.Arguments[0].(int32) != 1122 {
				t.Error("Argument should be 1122 and is: " + string(msg.Arguments[0].(int32)))
			}

			// Stop OSC server
			server.Close()
			finish <- true
		})

		if err != nil {
			t.Error("Error adding message handler")
		}

		start <- true
		server.ListenAndDispatch()
	}()

	go func() {
		timeout := time.After(5 * time.Second)
		select {
		case <-timeout:
		case <-start:
			time.Sleep(500 * time.Millisecond)
			client := NewOscClient("localhost", 6677)
			msg := NewOscMessage("/address/test")
			msg.Append(int32(1122))
			client.Send(msg)
		}

		done.Done()

		select {
		case <-timeout:
		case <-finish:
		}
		done.Done()
	}()

	done.Wait()
}

func TestServerMessageReceiving(t *testing.T) {
	finish := make(chan bool)
	start := make(chan bool)
	var done sync.WaitGroup
	done.Add(2)

	// Start the server in a go-routine
	go func() {
		server := NewOscServer("localhost", 6677)
		server.Listen()

		// Start the client
		start <- true

		for {
			packet, err := server.ReceivePacket()
			if err != nil {
				t.Error("Server error")
			}

			if packet != nil {
				msg := packet.(*OscMessage)
				if msg.CountArguments() != 2 {
					t.Errorf("Argument length should be 2 and is: %d\n", msg.CountArguments())
				}

				if msg.Arguments[0].(int32) != 1122 {
					t.Error("Argument should be 1122 and is: " + string(msg.Arguments[0].(int32)))
				}

				if msg.Arguments[1].(int32) != 3344 {
					t.Error("Argument should be 3344 and is: " + string(msg.Arguments[1].(int32)))
				}

				server.Close()
				finish <- true
			}
		}
	}()

	go func() {
		timeout := time.After(5 * time.Second)
		select {
		case <-timeout:
		case <-start:
			client := NewOscClient("localhost", 6677)
			msg := NewOscMessage("/address/test")
			msg.Append(int32(1122))
			msg.Append(int32(3344))
			time.Sleep(500 * time.Millisecond)
			client.Send(msg)
		}

		done.Done()

		select {
		case <-timeout:
		case <-finish:
		}
		done.Done()
	}()

	done.Wait()
}

func TestReadPaddedString(t *testing.T) {
	buf1 := []byte{'t', 'e', 's', 't', 's', 't', 'r', 'i', 'n', 'g', 0, 0}
	buf2 := []byte{'t', 'e', 's', 't', 0, 0, 0, 0}

	bytesBuffer := bytes.NewBuffer(buf1)
	st, n, err := readPaddedString(bufio.NewReader(bytesBuffer))
	if err != nil {
		t.Error("Error reading padded string: " + err.Error())
	}

	if n != 12 {
		t.Errorf("Number of bytes needs to be 12 and is: %d\n", n)
	}

	if st != "teststring" {
		t.Errorf("String should be \"teststring\" and is \"%s\"", st)
	}

	bytesBuffer = bytes.NewBuffer(buf2)
	st, n, err = readPaddedString(bufio.NewReader(bytesBuffer))
	if err != nil {
		t.Error("Error reading padded string: " + err.Error())
	}

	if n != 8 {
		t.Errorf("Number of bytes needs to be 8 and is: %d\n", n)
	}

	if st != "test" {
		t.Errorf("String should be \"test\" and is \"%s\"", st)
	}
}

func TestWritePaddedString(t *testing.T) {
	buf := []byte{}
	bytesBuffer := bytes.NewBuffer(buf)
	testString := "testString"
	expectedNumberOfWrittenBytes := len(testString) + padBytesNeeded(len(testString))

	n, err := writePaddedString(testString, bytesBuffer)
	if err != nil {
		t.Errorf(err.Error())
	}

	if n != expectedNumberOfWrittenBytes {
		t.Errorf("Expected number of written bytes should be \"%d\" and is \"%d\"", expectedNumberOfWrittenBytes, n)
	}
}

func TestPadBytesNeeded(t *testing.T) {
	var n int
	n = padBytesNeeded(4)
	if n != 4 {
		t.Errorf("Number of pad bytes should be 4 and is: %d", n)
	}

	n = padBytesNeeded(3)
	if n != 1 {
		t.Errorf("Number of pad bytes should be 1 and is: %d", n)
	}

	n = padBytesNeeded(1)
	if n != 3 {
		t.Errorf("Number of pad bytes should be 3 and is: %d", n)
	}

	n = padBytesNeeded(0)
	if n != 4 {
		t.Errorf("Number of pad bytes should be 4 and is: %d", n)
	}

	n = padBytesNeeded(32)
	if n != 4 {
		t.Errorf("Number of pad bytes should be 4 and is: %d", n)
	}

	n = padBytesNeeded(63)
	if n != 1 {
		t.Errorf("Number of pad bytes should be 1 and is: %d", n)
	}

	n = padBytesNeeded(10)
	if n != 2 {
		t.Errorf("Number of pad bytes should be 2 and is: %d", n)
	}
}

func TestTypeTagsString(t *testing.T) {
	msg := NewOscMessage("/some/address")
	msg.Append(int32(100))
	msg.Append(true)
	msg.Append(false)

	typeTags, err := msg.TypeTags()
	if err != nil {
		t.Error(err.Error())
	}

	if typeTags != ",iTF" {
		t.Errorf("Type tag string should be ',iTF' and is: %s", typeTags)
	}
}
<commit_msg>add test for closing a not running server<commit_after>package osc

import (
	"bufio"
	"bytes"
	"sync"
	"testing"
	"time"
)

func TestAppendArguments(t *testing.T) {
	oscAddress := "/address"
	message := NewOscMessage(oscAddress)
	if message.Address != oscAddress {
		t.Errorf("OSC address should be \"%s\" and is \"%s\"", oscAddress, message.Address)
	}

	message.Append("string argument")
	message.Append(123456789)
	message.Append(true)

	if message.CountArguments() != 3 {
		t.Errorf("Number of arguments should be %d and is %d", 3, message.CountArguments())
	}
}

func TestEqualMessage(t *testing.T) {
	msg1 := NewOscMessage("/address")
	msg2 := NewOscMessage("/address")

	msg1.Append(1234)
	msg2.Append(1234)
	msg1.Append("test string")
	msg2.Append("test string")

	if !msg1.Equals(msg2) {
		t.Error("Messages should be equal")
	}
}

func TestAddMsgHandler(t *testing.T) {
	server := NewOscServer("localhost", 6677)
	err := server.AddMsgHandler("/address/test", func(msg *OscMessage) {})
	if err != nil {
		t.Error("Expected that OSC address '/address/test' is valid")
	}
}

func TestAddMsgHandlerWithInvalidAddress(t *testing.T) {
	server := NewOscServer("localhost", 6677)
	err := server.AddMsgHandler("/address*/test", func(msg *OscMessage) {})
	if err == nil {
		t.Error("Expected error with '/address*/test'")
	}
}

func TestServerMessageDispatching(t *testing.T) {
	finish := make(chan bool)
	start := make(chan bool)
	var done sync.WaitGroup
	done.Add(2)

	// Start the OSC server in a new go-routine
	go func() {
		server := NewOscServer("localhost", 6677)
		err := server.AddMsgHandler("/address/test", func(msg *OscMessage) {
			if len(msg.Arguments) != 1 {
				t.Error("Argument length should be 1 and is: " + string(len(msg.Arguments)))
			}

			if msg.Arguments[0].(int32) != 1122 {
				t.Error("Argument should be 1122 and is: " + string(msg.Arguments[0].(int32)))
			}

			// Stop OSC server
			server.Close()
			finish <- true
		})

		if err != nil {
			t.Error("Error adding message handler")
		}

		start <- true
		server.ListenAndDispatch()
	}()

	go func() {
		timeout := time.After(5 * time.Second)
		select {
		case <-timeout:
		case <-start:
			time.Sleep(500 * time.Millisecond)
			client := NewOscClient("localhost", 6677)
			msg := NewOscMessage("/address/test")
			msg.Append(int32(1122))
			client.Send(msg)
		}

		done.Done()

		select {
		case <-timeout:
		case <-finish:
		}
		done.Done()
	}()

	done.Wait()
}

func TestServerMessageReceiving(t *testing.T) {
	finish := make(chan bool)
	start := make(chan bool)
	var done sync.WaitGroup
	done.Add(2)

	// Start the server in a go-routine
	go func() {
		server := NewOscServer("localhost", 6677)
		server.Listen()

		// Start the client
		start <- true

		for {
			packet, err := server.ReceivePacket()
			if err != nil {
				t.Error("Server error")
			}

			if packet != nil {
				msg := packet.(*OscMessage)
				if msg.CountArguments() != 2 {
					t.Errorf("Argument length should be 2 and is: %d\n", msg.CountArguments())
				}

				if msg.Arguments[0].(int32) != 1122 {
					t.Error("Argument should be 1122 and is: " + string(msg.Arguments[0].(int32)))
				}

				if msg.Arguments[1].(int32) != 3344 {
					t.Error("Argument should be 3344 and is: " + string(msg.Arguments[1].(int32)))
				}

				server.Close()
				finish <- true
			}
		}
	}()

	go func() {
		timeout := time.After(5 * time.Second)
		select {
		case <-timeout:
		case <-start:
			client := NewOscClient("localhost", 6677)
			msg := NewOscMessage("/address/test")
			msg.Append(int32(1122))
			msg.Append(int32(3344))
			time.Sleep(500 * time.Millisecond)
			client.Send(msg)
		}

		done.Done()

		select {
		case <-timeout:
		case <-finish:
		}
		done.Done()
	}()

	done.Wait()
}

func TestReadPaddedString(t *testing.T) {
	buf1 := []byte{'t', 'e', 's', 't', 's', 't', 'r', 'i', 'n', 'g', 0, 0}
	buf2 := []byte{'t', 'e', 's', 't', 0, 0, 0, 0}

	bytesBuffer := bytes.NewBuffer(buf1)
	st, n, err := readPaddedString(bufio.NewReader(bytesBuffer))
	if err != nil {
		t.Error("Error reading padded string: " + err.Error())
	}

	if n != 12 {
		t.Errorf("Number of bytes needs to be 12 and is: %d\n", n)
	}

	if st != "teststring" {
		t.Errorf("String should be \"teststring\" and is \"%s\"", st)
	}

	bytesBuffer = bytes.NewBuffer(buf2)
	st, n, err = readPaddedString(bufio.NewReader(bytesBuffer))
	if err != nil {
		t.Error("Error reading padded string: " + err.Error())
	}

	if n != 8 {
		t.Errorf("Number of bytes needs to be 8 and is: %d\n", n)
	}

	if st != "test" {
		t.Errorf("String should be \"test\" and is \"%s\"", st)
	}
}

func TestWritePaddedString(t *testing.T) {
	buf := []byte{}
	bytesBuffer := bytes.NewBuffer(buf)
	testString := "testString"
	expectedNumberOfWrittenBytes := len(testString) + padBytesNeeded(len(testString))

	n, err := writePaddedString(testString, bytesBuffer)
	if err != nil {
		t.Errorf(err.Error())
	}

	if n != expectedNumberOfWrittenBytes {
		t.Errorf("Expected number of written bytes should be \"%d\" and is \"%d\"", expectedNumberOfWrittenBytes, n)
	}
}

func TestPadBytesNeeded(t *testing.T) {
	var n int
	n = padBytesNeeded(4)
	if n != 4 {
		t.Errorf("Number of pad bytes should be 4 and is: %d", n)
	}

	n = padBytesNeeded(3)
	if n != 1 {
		t.Errorf("Number of pad bytes should be 1 and is: %d", n)
	}

	n = padBytesNeeded(1)
	if n != 3 {
		t.Errorf("Number of pad bytes should be 3 and is: %d", n)
	}

	n = padBytesNeeded(0)
	if n != 4 {
		t.Errorf("Number of pad bytes should be 4 and is: %d", n)
	}

	n = padBytesNeeded(32)
	if n != 4 {
		t.Errorf("Number of pad bytes should be 4 and is: %d", n)
	}

	n = padBytesNeeded(63)
	if n != 1 {
		t.Errorf("Number of pad bytes should be 1 and is: %d", n)
	}

	n = padBytesNeeded(10)
	if n != 2 {
		t.Errorf("Number of pad bytes should be 2 and is: %d", n)
	}
}

func TestTypeTagsString(t *testing.T) {
	msg := NewOscMessage("/some/address")
	msg.Append(int32(100))
	msg.Append(true)
	msg.Append(false)

	typeTags, err := msg.TypeTags()
	if err != nil {
		t.Error(err.Error())
	}

	if typeTags != ",iTF" {
		t.Errorf("Type tag string should be ',iTF' and is: %s", typeTags)
	}
}

func TestServerIsNotRunningAndGetsClosed(t *testing.T) {
	server := NewOscServer("127.0.0.1", 8000)
	err := server.Close()
	if err == nil {
		t.Errorf("Expected error if the the server is not running and it gets closed")
	}
}
<|endoftext|>


<commit_before>package s3

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

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"

	"github.com/docker/distribution/context"
	storagedriver "github.com/docker/distribution/registry/storage/driver"
	"github.com/docker/distribution/registry/storage/driver/testsuites"

	"gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { check.TestingT(t) }

var s3DriverConstructor func(rootDirectory, storageClass string) (*Driver, error)
var skipS3 func() string

func init() {
	accessKey := os.Getenv("AWS_ACCESS_KEY")
	secretKey := os.Getenv("AWS_SECRET_KEY")
	bucket := os.Getenv("S3_BUCKET")
	encrypt := os.Getenv("S3_ENCRYPT")
	secure := os.Getenv("S3_SECURE")
	region := os.Getenv("AWS_REGION")
	root, err := ioutil.TempDir("", "driver-")
	if err != nil {
		panic(err)
	}
	defer os.Remove(root)

	s3DriverConstructor = func(rootDirectory, storageClass string) (*Driver, error) {
		encryptBool := false
		if encrypt != "" {
			encryptBool, err = strconv.ParseBool(encrypt)
			if err != nil {
				return nil, err
			}
		}

		secureBool := true
		if secure != "" {
			secureBool, err = strconv.ParseBool(secure)
			if err != nil {
				return nil, err
			}
		}

		parameters := DriverParameters{
			accessKey,
			secretKey,
			bucket,
			region,
			encryptBool,
			secureBool,
			minChunkSize,
			rootDirectory,
			storageClass,
			driverName + "-test",
		}

		return New(parameters)
	}

	// Skip S3 storage driver tests if environment variable parameters are not provided
	skipS3 = func() string {
		if accessKey == "" || secretKey == "" || region == "" || bucket == "" || encrypt == "" {
			return "Must set AWS_ACCESS_KEY, AWS_SECRET_KEY, AWS_REGION, S3_BUCKET, and S3_ENCRYPT to run S3 tests"
		}
		return ""
	}

	testsuites.RegisterSuite(func() (storagedriver.StorageDriver, error) {
		return s3DriverConstructor(root, s3.StorageClassStandard)
	}, skipS3)
}

func TestEmptyRootList(t *testing.T) {
	if skipS3() != "" {
		t.Skip(skipS3())
	}

	validRoot, err := ioutil.TempDir("", "driver-")
	if err != nil {
		t.Fatalf("unexpected error creating temporary directory: %v", err)
	}
	defer os.Remove(validRoot)

	rootedDriver, err := s3DriverConstructor(validRoot, s3.StorageClassStandard)
	if err != nil {
		t.Fatalf("unexpected error creating rooted driver: %v", err)
	}

	emptyRootDriver, err := s3DriverConstructor("", s3.StorageClassStandard)
	if err != nil {
		t.Fatalf("unexpected error creating empty root driver: %v", err)
	}

	slashRootDriver, err := s3DriverConstructor("/", s3.StorageClassStandard)
	if err != nil {
		t.Fatalf("unexpected error creating slash root driver: %v", err)
	}

	filename := "/test"
	contents := []byte("contents")
	ctx := context.Background()
	err = rootedDriver.PutContent(ctx, filename, contents)
	if err != nil {
		t.Fatalf("unexpected error creating content: %v", err)
	}
	defer rootedDriver.Delete(ctx, filename)

	keys, err := emptyRootDriver.List(ctx, "/")
	for _, path := range keys {
		if !storagedriver.PathRegexp.MatchString(path) {
			t.Fatalf("unexpected string in path: %q != %q", path, storagedriver.PathRegexp)
		}
	}

	keys, err = slashRootDriver.List(ctx, "/")
	for _, path := range keys {
		if !storagedriver.PathRegexp.MatchString(path) {
			t.Fatalf("unexpected string in path: %q != %q", path, storagedriver.PathRegexp)
		}
	}
}

func TestStorageClass(t *testing.T) {
	if skipS3() != "" {
		t.Skip(skipS3())
	}

	rootDir, err := ioutil.TempDir("", "driver-")
	if err != nil {
		t.Fatalf("unexpected error creating temporary directory: %v", err)
	}
	defer os.Remove(rootDir)

	standardDriver, err := s3DriverConstructor(rootDir, s3.StorageClassStandard)
	if err != nil {
		t.Fatalf("unexpected error creating driver with standard storage: %v", err)
	}

	rrDriver, err := s3DriverConstructor(rootDir, s3.StorageClassReducedRedundancy)
	if err != nil {
		t.Fatalf("unexpected error creating driver with reduced redundancy storage: %v", err)
	}

	standardFilename := "/test-standard"
	rrFilename := "/test-rr"
	contents := []byte("contents")
	ctx := context.Background()

	err = standardDriver.PutContent(ctx, standardFilename, contents)
	if err != nil {
		t.Fatalf("unexpected error creating content: %v", err)
	}
	defer standardDriver.Delete(ctx, standardFilename)

	err = rrDriver.PutContent(ctx, rrFilename, contents)
	if err != nil {
		t.Fatalf("unexpected error creating content: %v", err)
	}
	defer rrDriver.Delete(ctx, rrFilename)

	standardDriverUnwrapped := standardDriver.Base.StorageDriver.(*driver)
	resp, err := standardDriverUnwrapped.S3.GetObject(&s3.GetObjectInput{
		Bucket: aws.String(standardDriverUnwrapped.Bucket),
		Key:    aws.String(standardDriverUnwrapped.s3Path(standardFilename)),
	})
	if err != nil {
		t.Fatalf("unexpected error retrieving standard storage file: %v", err)
	}
	defer resp.Body.Close()
	// Amazon only populates this header value for non-standard storage classes
	if resp.StorageClass != nil {
		t.Fatalf("unexpected storage class for standard file: %v", resp.StorageClass)
	}

	rrDriverUnwrapped := rrDriver.Base.StorageDriver.(*driver)
	resp, err = rrDriverUnwrapped.S3.GetObject(&s3.GetObjectInput{
		Bucket: aws.String(rrDriverUnwrapped.Bucket),
		Key:    aws.String(rrDriverUnwrapped.s3Path(standardFilename)),
	})
	if err != nil {
		t.Fatalf("unexpected error retrieving reduced-redundancy storage file: %v", err)
	}
	defer resp.Body.Close()
	if resp.StorageClass == nil {
		t.Fatalf("unexpected storage class for reduced-redundancy file: %v", s3.StorageClassStandard)
	} else if *resp.StorageClass != s3.StorageClassReducedRedundancy {
		t.Fatalf("unexpected storage class for reduced-redundancy file: %v", *resp.StorageClass)
	}

}
<commit_msg>[driver/s3aws] Fix TestStorageClass<commit_after>package s3

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

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"

	"github.com/docker/distribution/context"
	storagedriver "github.com/docker/distribution/registry/storage/driver"
	"github.com/docker/distribution/registry/storage/driver/testsuites"

	"gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { check.TestingT(t) }

var s3DriverConstructor func(rootDirectory, storageClass string) (*Driver, error)
var skipS3 func() string

func init() {
	accessKey := os.Getenv("AWS_ACCESS_KEY")
	secretKey := os.Getenv("AWS_SECRET_KEY")
	bucket := os.Getenv("S3_BUCKET")
	encrypt := os.Getenv("S3_ENCRYPT")
	secure := os.Getenv("S3_SECURE")
	region := os.Getenv("AWS_REGION")
	root, err := ioutil.TempDir("", "driver-")
	if err != nil {
		panic(err)
	}
	defer os.Remove(root)

	s3DriverConstructor = func(rootDirectory, storageClass string) (*Driver, error) {
		encryptBool := false
		if encrypt != "" {
			encryptBool, err = strconv.ParseBool(encrypt)
			if err != nil {
				return nil, err
			}
		}

		secureBool := true
		if secure != "" {
			secureBool, err = strconv.ParseBool(secure)
			if err != nil {
				return nil, err
			}
		}

		parameters := DriverParameters{
			accessKey,
			secretKey,
			bucket,
			region,
			encryptBool,
			secureBool,
			minChunkSize,
			rootDirectory,
			storageClass,
			driverName + "-test",
		}

		return New(parameters)
	}

	// Skip S3 storage driver tests if environment variable parameters are not provided
	skipS3 = func() string {
		if accessKey == "" || secretKey == "" || region == "" || bucket == "" || encrypt == "" {
			return "Must set AWS_ACCESS_KEY, AWS_SECRET_KEY, AWS_REGION, S3_BUCKET, and S3_ENCRYPT to run S3 tests"
		}
		return ""
	}

	testsuites.RegisterSuite(func() (storagedriver.StorageDriver, error) {
		return s3DriverConstructor(root, s3.StorageClassStandard)
	}, skipS3)
}

func TestEmptyRootList(t *testing.T) {
	if skipS3() != "" {
		t.Skip(skipS3())
	}

	validRoot, err := ioutil.TempDir("", "driver-")
	if err != nil {
		t.Fatalf("unexpected error creating temporary directory: %v", err)
	}
	defer os.Remove(validRoot)

	rootedDriver, err := s3DriverConstructor(validRoot, s3.StorageClassStandard)
	if err != nil {
		t.Fatalf("unexpected error creating rooted driver: %v", err)
	}

	emptyRootDriver, err := s3DriverConstructor("", s3.StorageClassStandard)
	if err != nil {
		t.Fatalf("unexpected error creating empty root driver: %v", err)
	}

	slashRootDriver, err := s3DriverConstructor("/", s3.StorageClassStandard)
	if err != nil {
		t.Fatalf("unexpected error creating slash root driver: %v", err)
	}

	filename := "/test"
	contents := []byte("contents")
	ctx := context.Background()
	err = rootedDriver.PutContent(ctx, filename, contents)
	if err != nil {
		t.Fatalf("unexpected error creating content: %v", err)
	}
	defer rootedDriver.Delete(ctx, filename)

	keys, err := emptyRootDriver.List(ctx, "/")
	for _, path := range keys {
		if !storagedriver.PathRegexp.MatchString(path) {
			t.Fatalf("unexpected string in path: %q != %q", path, storagedriver.PathRegexp)
		}
	}

	keys, err = slashRootDriver.List(ctx, "/")
	for _, path := range keys {
		if !storagedriver.PathRegexp.MatchString(path) {
			t.Fatalf("unexpected string in path: %q != %q", path, storagedriver.PathRegexp)
		}
	}
}

func TestStorageClass(t *testing.T) {
	if skipS3() != "" {
		t.Skip(skipS3())
	}

	rootDir, err := ioutil.TempDir("", "driver-")
	if err != nil {
		t.Fatalf("unexpected error creating temporary directory: %v", err)
	}
	defer os.Remove(rootDir)

	standardDriver, err := s3DriverConstructor(rootDir, s3.StorageClassStandard)
	if err != nil {
		t.Fatalf("unexpected error creating driver with standard storage: %v", err)
	}

	rrDriver, err := s3DriverConstructor(rootDir, s3.StorageClassReducedRedundancy)
	if err != nil {
		t.Fatalf("unexpected error creating driver with reduced redundancy storage: %v", err)
	}

	standardFilename := "/test-standard"
	rrFilename := "/test-rr"
	contents := []byte("contents")
	ctx := context.Background()

	err = standardDriver.PutContent(ctx, standardFilename, contents)
	if err != nil {
		t.Fatalf("unexpected error creating content: %v", err)
	}
	defer standardDriver.Delete(ctx, standardFilename)

	err = rrDriver.PutContent(ctx, rrFilename, contents)
	if err != nil {
		t.Fatalf("unexpected error creating content: %v", err)
	}
	defer rrDriver.Delete(ctx, rrFilename)

	standardDriverUnwrapped := standardDriver.Base.StorageDriver.(*driver)
	resp, err := standardDriverUnwrapped.S3.GetObject(&s3.GetObjectInput{
		Bucket: aws.String(standardDriverUnwrapped.Bucket),
		Key:    aws.String(standardDriverUnwrapped.s3Path(standardFilename)),
	})
	if err != nil {
		t.Fatalf("unexpected error retrieving standard storage file: %v", err)
	}
	defer resp.Body.Close()
	// Amazon only populates this header value for non-standard storage classes
	if resp.StorageClass != nil {
		t.Fatalf("unexpected storage class for standard file: %v", resp.StorageClass)
	}

	rrDriverUnwrapped := rrDriver.Base.StorageDriver.(*driver)
	resp, err = rrDriverUnwrapped.S3.GetObject(&s3.GetObjectInput{
		Bucket: aws.String(rrDriverUnwrapped.Bucket),
		Key:    aws.String(rrDriverUnwrapped.s3Path(rrFilename)),
	})
	if err != nil {
		t.Fatalf("unexpected error retrieving reduced-redundancy storage file: %v", err)
	}
	defer resp.Body.Close()
	if resp.StorageClass == nil {
		t.Fatalf("unexpected storage class for reduced-redundancy file: %v", s3.StorageClassStandard)
	} else if *resp.StorageClass != s3.StorageClassReducedRedundancy {
		t.Fatalf("unexpected storage class for reduced-redundancy file: %v", *resp.StorageClass)
	}

}
<|endoftext|>


<commit_before>package stopwatch

import (
	"net/http"
	"time"
)

// Measure latency in nanoseconds
func Measure(f func()) int64 {
	start := time.Now().UnixNano()
	f()
	return time.Now().UnixNano() - start
}

// LatencyHandler is a middleware that measures latency given http.Handler struct.
func LatencyHandler(resultChan chan int64, methods []string, next http.Handler) http.Handler {
	middle := func(w http.ResponseWriter, r *http.Request) {
		foundMethod := false
		for _, method := range methods {
			if r.Method == method {
				foundMethod = true
				break
			}
		}

		if foundMethod {
			start := time.Now().UnixNano()
			next.ServeHTTP(w, r)
			resultChan <- (time.Now().UnixNano() - start)
		} else {
			next.ServeHTTP(w, r)
		}
	}

	return http.HandlerFunc(middle)
}

// LatencyFuncHandler is a middleware that measures latency given request handler function.
func LatencyFuncHandler(resultChan chan int64, methods []string, nextFunc func(http.ResponseWriter, *http.Request)) http.Handler {
	return LatencyHandler(resultChan, methods, http.HandlerFunc(nextFunc))
}
<commit_msg>Use defer for cleaner code.<commit_after>package stopwatch

import (
	"net/http"
	"time"
)

// Measure latency in nanoseconds
func Measure(f func()) int64 {
	start := time.Now().UnixNano()
	f()
	return time.Now().UnixNano() - start
}

// LatencyHandler is a middleware that measures latency given http.Handler struct.
func LatencyHandler(resultChan chan int64, methods []string, next http.Handler) http.Handler {
	middle := func(w http.ResponseWriter, r *http.Request) {
		foundMethod := false
		for _, method := range methods {
			if r.Method == method {
				foundMethod = true
				break
			}
		}

		if foundMethod && resultChan != nil {
			start := time.Now().UnixNano()
			defer func() {
				resultChan <- (time.Now().UnixNano() - start)
			}()
		}

		next.ServeHTTP(w, r)
	}

	return http.HandlerFunc(middle)
}

// LatencyFuncHandler is a middleware that measures latency given request handler function.
func LatencyFuncHandler(resultChan chan int64, methods []string, nextFunc func(http.ResponseWriter, *http.Request)) http.Handler {
	return LatencyHandler(resultChan, methods, http.HandlerFunc(nextFunc))
}
<|endoftext|>


<commit_before>package codebase

import (
    "fmt"
    "log"
    "time"
)

type Event struct {
    Id        int    `xml:"id"`
    Title     string `xml:"title"`
    Date      string `xml:"timestamp"`
    Type      string `xml:"type"`
    HtmlTitle string `xml:"html-title"`
    HtmlText  string `xml:"html-text"`
    UserID    int    `xml:"user-id"`
    UserEmail string `xml:"actor-email"`
    ProjectID int    `xml:"project-id"`
    Deleted   bool   `xml:"deleted"`
    Avatar    string `xml:"avatar-url"`

    Raw struct {
        TicketID         int          `xml:"number"`
        Subject          string       `xml:"subject"`
        Content          string       `xml:"content"`
        Changes          EventChanges `xml:"changes"`
        ProjectPermalink string       `xml:"project-permalink"`
    }   `xml:"raw-properties"`
}

type EventChanges struct {
    Status   []string `xml:"status-id>status-id"`
    Assignee []string `xml:"assignee-id>assignee-id"`
    Subject  []string `xml:"subject>subject"`
    Priority []string `xml:"priority-id>priority-id"`
}

type eventQueryOptions struct {
    baseQueryOptions
    Raw   bool      `url:"raw,omitempty"`
    Since time.Time `url:"since,omitempty"`
}

func (c *CodeBaseAPI) Activities(since time.Time, user User) (events []Event) {
    type eventArray struct {
        Events []Event `xml:"event"`
    }

    queryOpts := eventQueryOptions{
        Raw:   true,
        Since: since,
    }
    queryOpts.Page = 1

    for {
        proxy := eventArray{}
        if err := c.fetchFromCodebase("activity", &proxy, queryOpts); err != nil {
            log.Fatalln("Could not fetch activities:", err)
        }

        if len(proxy.Events) == 0 {
            break
        }

        for _, event := range proxy.Events {
            if user.UserID != 0 && event.UserID != user.UserID {
                continue
            }

            if event.HasChanges() == false {
                continue
            }

            events = append(events, event)
        }
        queryOpts.Page++
    }

    return
}

func (e *Event) HasChanges() bool {
    return e.Raw.Changes.Changes() != ""
}

func (e *Event) TicketUrl(company string) string {
    return fmt.Sprintf("https://%s.codebasehq.com/projects/%s/tickets/%v", company, e.Raw.ProjectPermalink, e.Raw.TicketID)
}

func (e *Event) Day() time.Weekday {
    date, _ := time.Parse("2006-01-02 15:04:05 MST", e.Date)
    return date.Weekday()
}

func (e *EventChanges) Changes() string {
    changes := ""

    workType := map[string]string{
        "Blocked -> New":               "Unblocked",
        "Code Complete -> In Review":   "Reviewing",
        "In Progress -> Code Complete": "Finished",
        "In Progress -> New":           "Put back",
        "In Review -> Code Complete":   "Reviewed",
        "In Review -> Completed":       "Closed as completed",
        "New -> In Progress":           "Working on",
    }

    if len(e.Status) == 1 {
        changes += e.Status[0]
    } else if len(e.Status) == 2 {
        change := fmt.Sprintf("%s -> %s", e.Status[0], e.Status[1])
        if description, ok := workType[change]; ok {
            change = description
        }
        changes += change
    }

    return changes
}
<commit_msg>Remove hardcoded change discription, make it configurable instead<commit_after>package codebase

import (
    "fmt"
    "log"
    "time"
)

type Event struct {
    Id        int    `xml:"id"`
    Title     string `xml:"title"`
    Date      string `xml:"timestamp"`
    Type      string `xml:"type"`
    HtmlTitle string `xml:"html-title"`
    HtmlText  string `xml:"html-text"`
    UserID    int    `xml:"user-id"`
    UserEmail string `xml:"actor-email"`
    ProjectID int    `xml:"project-id"`
    Deleted   bool   `xml:"deleted"`
    Avatar    string `xml:"avatar-url"`

    Raw struct {
        TicketID         int          `xml:"number"`
        Subject          string       `xml:"subject"`
        Content          string       `xml:"content"`
        Changes          EventChanges `xml:"changes"`
        ProjectPermalink string       `xml:"project-permalink"`
    }   `xml:"raw-properties"`
}

type EventChanges struct {
    Status   []string `xml:"status-id>status-id"`
    Assignee []string `xml:"assignee-id>assignee-id"`
    Subject  []string `xml:"subject>subject"`
    Priority []string `xml:"priority-id>priority-id"`
}

type eventQueryOptions struct {
    baseQueryOptions
    Raw   bool      `url:"raw,omitempty"`
    Since time.Time `url:"since,omitempty"`
}

type ChangeMapping struct {
    Status map[string]string
}

func (c *CodeBaseAPI) Activities(since time.Time, user User) (events []Event) {
    type eventArray struct {
        Events []Event `xml:"event"`
    }

    queryOpts := eventQueryOptions{
        Raw:   true,
        Since: since,
    }
    queryOpts.Page = 1

    for {
        proxy := eventArray{}
        if err := c.fetchFromCodebase("activity", &proxy, queryOpts); err != nil {
            log.Fatalln("Could not fetch activities:", err)
        }

        if len(proxy.Events) == 0 {
            break
        }

        for _, event := range proxy.Events {
            if user.UserID != 0 && event.UserID != user.UserID {
                continue
            }

            if event.HasChanges(ChangeMapping{}) == false {
                continue
            }

            events = append(events, event)
        }
        queryOpts.Page++
    }

    return
}

func (e *Event) HasChanges(mapping ChangeMapping) bool {
    return e.Raw.Changes.Changes(mapping) != ""
}

func (e *Event) TicketUrl(company string) string {
    return fmt.Sprintf("https://%s.codebasehq.com/projects/%s/tickets/%v", company, e.Raw.ProjectPermalink, e.Raw.TicketID)
}

func (e *Event) Day() time.Weekday {
    date, _ := time.Parse("2006-01-02 15:04:05 MST", e.Date)
    return date.Weekday()
}

func (e *EventChanges) Changes(mapping ChangeMapping) string {
    changes := ""

    workType := mapping.Status

    if len(e.Status) == 1 {
        changes += e.Status[0]
    } else if len(e.Status) == 2 {
        change := fmt.Sprintf("%s -> %s", e.Status[0], e.Status[1])
        if description, ok := workType[change]; ok {
            change = description
        }

        if change != "" {
            changes += change
        }
    }

    return changes
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	HomeFlag     = "home"
	TraceFlag    = "trace"
	OutputFlag   = "output"
	EncodingFlag = "encoding"
)

// Executable is the minimal interface to *corba.Command, so we can
// wrap if desired before the test
type Executable interface {
	Execute() error
}

// PrepareBaseCmd is meant for tendermint and other servers
func PrepareBaseCmd(cmd *cobra.Command, envPrefix, defaultHome string) Executor {
	cobra.OnInitialize(func() { initEnv(envPrefix) })
	cmd.PersistentFlags().StringP(HomeFlag, "", defaultHome, "directory for config and data")
	cmd.PersistentFlags().Bool(TraceFlag, false, "print out full stack trace on errors")
	cmd.PersistentPreRunE = concatCobraCmdFuncs(bindFlagsLoadViper, cmd.PersistentPreRunE)
	return Executor{cmd, os.Exit}
}

// PrepareMainCmd is meant for client side libs that want some more flags
//
// This adds --encoding (hex, btc, base64) and --output (text, json) to
// the command.  These only really make sense in interactive commands.
func PrepareMainCmd(cmd *cobra.Command, envPrefix, defaultHome string) Executor {
	cmd.PersistentFlags().StringP(EncodingFlag, "e", "hex", "Binary encoding (hex|b64|btc)")
	cmd.PersistentFlags().StringP(OutputFlag, "o", "text", "Output format (text|json)")
	cmd.PersistentPreRunE = concatCobraCmdFuncs(validateOutput, cmd.PersistentPreRunE)
	return PrepareBaseCmd(cmd, envPrefix, defaultHome)
}

// initEnv sets to use ENV variables if set.
func initEnv(prefix string) {
	copyEnvVars(prefix)

	// env variables with TM prefix (eg. TM_ROOT)
	viper.SetEnvPrefix(prefix)
	viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_", "-", "_"))
	viper.AutomaticEnv()
}

// This copies all variables like TMROOT to TM_ROOT,
// so we can support both formats for the user
func copyEnvVars(prefix string) {
	prefix = strings.ToUpper(prefix)
	ps := prefix + "_"
	for _, e := range os.Environ() {
		kv := strings.SplitN(e, "=", 2)
		if len(kv) == 2 {
			k, v := kv[0], kv[1]
			if strings.HasPrefix(k, prefix) && !strings.HasPrefix(k, ps) {
				k2 := strings.Replace(k, prefix, ps, 1)
				os.Setenv(k2, v)
			}
		}
	}
}

// Executor wraps the cobra Command with a nicer Execute method
type Executor struct {
	*cobra.Command
	Exit func(int) // this is os.Exit by default, override in tests
}

type ExitCoder interface {
	ExitCode() int
}

// execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func (e Executor) Execute() error {
	e.SilenceUsage = true
	e.SilenceErrors = true
	err := e.Command.Execute()
	if err != nil {
		if viper.GetBool(TraceFlag) {
			fmt.Fprintf(os.Stderr, "ERROR: %+v\n", err)
		} else {
			fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		}

		// return error code 1 by default, can override it with a special error type
		exitCode := 1
		if ec, ok := err.(ExitCoder); ok {
			exitCode = ec.ExitCode()
		}
		e.Exit(exitCode)
	}
	return err
}

type cobraCmdFunc func(cmd *cobra.Command, args []string) error

// Returns a single function that calls each argument function in sequence
// RunE, PreRunE, PersistentPreRunE, etc. all have this same signature
func concatCobraCmdFuncs(fs ...cobraCmdFunc) cobraCmdFunc {
	return func(cmd *cobra.Command, args []string) error {
		for _, f := range fs {
			if f != nil {
				if err := f(cmd, args); err != nil {
					return err
				}
			}
		}
		return nil
	}
}

// Bind all flags and read the config into viper
func bindFlagsLoadViper(cmd *cobra.Command, args []string) error {
	// cmd.Flags() includes flags from this command and all persistent flags from the parent
	if err := viper.BindPFlags(cmd.Flags()); err != nil {
		return err
	}

	homeDir := viper.GetString(HomeFlag)
	viper.Set(HomeFlag, homeDir)
	viper.SetConfigName("config")                         // name of config file (without extension)
	viper.AddConfigPath(homeDir)                          // search root directory
	viper.AddConfigPath(filepath.Join(homeDir, "config")) // search root directory /config

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		// stderr, so if we redirect output to json file, this doesn't appear
		// fmt.Fprintln(os.Stderr, "Using config file:", viper.ConfigFileUsed())
	} else if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
		// we ignore not found error, only parse error
		// stderr, so if we redirect output to json file, this doesn't appear
		fmt.Fprintf(os.Stderr, "%#v", err)
	}
	return nil
}

func validateOutput(cmd *cobra.Command, args []string) error {
	// validate output format
	output := viper.GetString(OutputFlag)
	switch output {
	case "text", "json":
	default:
		return errors.Errorf("Unsupported output format: %s", output)
	}
	return nil
}
<commit_msg>Return config parse errors (#182)<commit_after>package cli

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	HomeFlag     = "home"
	TraceFlag    = "trace"
	OutputFlag   = "output"
	EncodingFlag = "encoding"
)

// Executable is the minimal interface to *corba.Command, so we can
// wrap if desired before the test
type Executable interface {
	Execute() error
}

// PrepareBaseCmd is meant for tendermint and other servers
func PrepareBaseCmd(cmd *cobra.Command, envPrefix, defaultHome string) Executor {
	cobra.OnInitialize(func() { initEnv(envPrefix) })
	cmd.PersistentFlags().StringP(HomeFlag, "", defaultHome, "directory for config and data")
	cmd.PersistentFlags().Bool(TraceFlag, false, "print out full stack trace on errors")
	cmd.PersistentPreRunE = concatCobraCmdFuncs(bindFlagsLoadViper, cmd.PersistentPreRunE)
	return Executor{cmd, os.Exit}
}

// PrepareMainCmd is meant for client side libs that want some more flags
//
// This adds --encoding (hex, btc, base64) and --output (text, json) to
// the command.  These only really make sense in interactive commands.
func PrepareMainCmd(cmd *cobra.Command, envPrefix, defaultHome string) Executor {
	cmd.PersistentFlags().StringP(EncodingFlag, "e", "hex", "Binary encoding (hex|b64|btc)")
	cmd.PersistentFlags().StringP(OutputFlag, "o", "text", "Output format (text|json)")
	cmd.PersistentPreRunE = concatCobraCmdFuncs(validateOutput, cmd.PersistentPreRunE)
	return PrepareBaseCmd(cmd, envPrefix, defaultHome)
}

// initEnv sets to use ENV variables if set.
func initEnv(prefix string) {
	copyEnvVars(prefix)

	// env variables with TM prefix (eg. TM_ROOT)
	viper.SetEnvPrefix(prefix)
	viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_", "-", "_"))
	viper.AutomaticEnv()
}

// This copies all variables like TMROOT to TM_ROOT,
// so we can support both formats for the user
func copyEnvVars(prefix string) {
	prefix = strings.ToUpper(prefix)
	ps := prefix + "_"
	for _, e := range os.Environ() {
		kv := strings.SplitN(e, "=", 2)
		if len(kv) == 2 {
			k, v := kv[0], kv[1]
			if strings.HasPrefix(k, prefix) && !strings.HasPrefix(k, ps) {
				k2 := strings.Replace(k, prefix, ps, 1)
				os.Setenv(k2, v)
			}
		}
	}
}

// Executor wraps the cobra Command with a nicer Execute method
type Executor struct {
	*cobra.Command
	Exit func(int) // this is os.Exit by default, override in tests
}

type ExitCoder interface {
	ExitCode() int
}

// execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func (e Executor) Execute() error {
	e.SilenceUsage = true
	e.SilenceErrors = true
	err := e.Command.Execute()
	if err != nil {
		if viper.GetBool(TraceFlag) {
			fmt.Fprintf(os.Stderr, "ERROR: %+v\n", err)
		} else {
			fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		}

		// return error code 1 by default, can override it with a special error type
		exitCode := 1
		if ec, ok := err.(ExitCoder); ok {
			exitCode = ec.ExitCode()
		}
		e.Exit(exitCode)
	}
	return err
}

type cobraCmdFunc func(cmd *cobra.Command, args []string) error

// Returns a single function that calls each argument function in sequence
// RunE, PreRunE, PersistentPreRunE, etc. all have this same signature
func concatCobraCmdFuncs(fs ...cobraCmdFunc) cobraCmdFunc {
	return func(cmd *cobra.Command, args []string) error {
		for _, f := range fs {
			if f != nil {
				if err := f(cmd, args); err != nil {
					return err
				}
			}
		}
		return nil
	}
}

// Bind all flags and read the config into viper
func bindFlagsLoadViper(cmd *cobra.Command, args []string) error {
	// cmd.Flags() includes flags from this command and all persistent flags from the parent
	if err := viper.BindPFlags(cmd.Flags()); err != nil {
		return err
	}

	homeDir := viper.GetString(HomeFlag)
	viper.Set(HomeFlag, homeDir)
	viper.SetConfigName("config")                         // name of config file (without extension)
	viper.AddConfigPath(homeDir)                          // search root directory
	viper.AddConfigPath(filepath.Join(homeDir, "config")) // search root directory /config

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		// stderr, so if we redirect output to json file, this doesn't appear
		// fmt.Fprintln(os.Stderr, "Using config file:", viper.ConfigFileUsed())
	} else if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
		// ignore not found error, return other errors
		return err
	}
	return nil
}

func validateOutput(cmd *cobra.Command, args []string) error {
	// validate output format
	output := viper.GetString(OutputFlag)
	switch output {
	case "text", "json":
	default:
		return errors.Errorf("Unsupported output format: %s", output)
	}
	return nil
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"strings"

	nomad "github.com/hashicorp/nomad/api"

	"github.com/jrasell/levant/helper"
	"github.com/jrasell/levant/levant"
	"github.com/jrasell/levant/levant/structs"
	"github.com/jrasell/levant/logging"
	"github.com/jrasell/levant/template"
)

// DeployCommand is the command implementation that allows users to deploy a
// Nomad job based on passed templates and variables.
type DeployCommand struct {
	Meta
}

// Help provides the help information for the deploy command.
func (c *DeployCommand) Help() string {
	helpText := `
Usage: levant deploy [options] [TEMPLATE]

  Deploy a Nomad job based on input templates and variable files. The deploy
  command supports passing variables individually on the command line. Multiple
  commands can be passed in the format of -var 'key=value'. Variables passed
  via the command line take precedence over the same variable declared within
  a passed variable file.

Arguments:

  TEMPLATE nomad job template
    If no argument is given we look for a single *.nomad file

General Options:

  -address=<http_address>
    The Nomad HTTP API address including port which Levant will use to make
    calls.

  -allow-stale
    Allow stale consistency mode for requests into nomad.

  -canary-auto-promote=<seconds>
    The time in seconds, after which Levant will auto-promote a canary job
    if all canaries within the deployment are healthy.

  -consul-address=<addr>
    The Consul host and port to use when making Consul KeyValue lookups for
    template rendering.

  -force-batch
    Forces a new instance of the periodic job. A new instance will be created
    even if it violates the job's prohibit_overlap settings.

  -force-count
    Use the taskgroup count from the Nomad jobfile instead of the count that
    is currently set in a running job.

  -ignore-no-changes
    By default if no changes are detected when running a deployment Levant will
    exit with a status 1 to indicate a deployment didn't happen. This behaviour
    can be changed using this flag so that Levant will exit cleanly ensuring CD
    pipelines don't fail when no changes are detected.

  -log-level=<level>
    Specify the verbosity level of Levant's logs. Valid values include DEBUG,
    INFO, and WARN, in decreasing order of verbosity. The default is INFO.

  -log-format=<format>
    Specify the format of Levant's logs. Valid values are HUMAN or JSON. The
    default is HUMAN.

  -var-file=<file>
    Used in conjunction with the -job-file will deploy a templated job to your
    Nomad cluster. You can repeat this flag multiple times to supply multiple var-files.
    [default: levant.(json|yaml|yml|tf)]
`
	return strings.TrimSpace(helpText)
}

// Synopsis is provides a brief summary of the deploy command.
func (c *DeployCommand) Synopsis() string {
	return "Render and deploy a Nomad job from a template"
}

// Run triggers a run of the Levant template and deploy functions.
func (c *DeployCommand) Run(args []string) int {

	var err error
	var addr string
	config := &structs.Config{}

	flags := c.Meta.FlagSet("deploy", FlagSetVars)
	flags.Usage = func() { c.UI.Output(c.Help()) }

	flags.StringVar(&config.Addr, "address", "", "")
	flags.IntVar(&config.Canary, "canary-auto-promote", 0, "")
	flags.StringVar(&addr, "consul-address", "", "")
	flags.BoolVar(&config.ForceBatch, "force-batch", false, "")
	flags.BoolVar(&config.ForceCount, "force-count", false, "")
	flags.BoolVar(&config.IgnoreNoChanges, "ignore-no-changes", false, "")
	flags.StringVar(&config.LogLevel, "log-level", "INFO", "")
	flags.StringVar(&config.LogFormat, "log-format", "HUMAN", "")
<<<<<<< HEAD
	flags.StringVar(&config.VaiableFile, "var-file", "", "")
	flags.BoolVar(&config.AllowStale, "allow-stale", false, "")
=======
	flags.Var((*helper.FlagStringSlice)(&config.VariableFiles), "var-file", "")
>>>>>>> b4741ff0afd8ccfaef7c23a1676e455eceeb77cb

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

	args = flags.Args()

	if err = logging.SetupLogger(config.LogLevel, config.LogFormat); err != nil {
		c.UI.Error(err.Error())
		return 1
	}

	if len(args) == 1 {
		config.TemplateFile = args[0]
	} else if len(args) == 0 {
		if config.TemplateFile = helper.GetDefaultTmplFile(); config.TemplateFile == "" {
			c.UI.Error(c.Help())
			c.UI.Error("\nERROR: Template arg missing and no default template found")
			return 1
		}
	} else {
		c.UI.Error(c.Help())
		return 1
	}

	config.Job, err = template.RenderJob(config.TemplateFile, config.VariableFiles, addr, &c.Meta.flagVars)
	if err != nil {
		c.UI.Error(fmt.Sprintf("[ERROR] levant/command: %v", err))
		return 1
	}

	if config.Canary > 0 {
		if err = c.checkCanaryAutoPromote(config.Job, config.Canary); err != nil {
			c.UI.Error(fmt.Sprintf("[ERROR] levant/command: %v", err))
			return 1
		}
	}

	if config.ForceBatch {
		if err = c.checkForceBatch(config.Job, config.ForceBatch); err != nil {
			c.UI.Error(fmt.Sprintf("[ERROR] levant/command: %v", err))
			return 1
		}
	}

	success := levant.TriggerDeployment(config, nil)
	if !success {
		return 1
	}

	return 0
}

func (c *DeployCommand) checkCanaryAutoPromote(job *nomad.Job, canaryAutoPromote int) error {
	if canaryAutoPromote == 0 {
		return nil
	}

	if job.Update != nil && job.Update.Canary != nil && *job.Update.Canary > 0 {
		return nil
	}

	for _, group := range job.TaskGroups {
		if group.Update != nil && group.Update.Canary != nil && *group.Update.Canary > 0 {
			return nil
		}
	}

	return fmt.Errorf("canary-auto-update of %v passed but job is not canary enabled", canaryAutoPromote)
}

// checkForceBatch ensures that if the force-batch flag is passed, the job is
// periodic.
func (c *DeployCommand) checkForceBatch(job *nomad.Job, forceBatch bool) error {

	if forceBatch && job.IsPeriodic() {
		return nil
	}

	return fmt.Errorf("force-batch passed but job is not periodic")
}
<commit_msg>Fix conflict<commit_after>package command

import (
	"fmt"
	"strings"

	nomad "github.com/hashicorp/nomad/api"

	"github.com/jrasell/levant/helper"
	"github.com/jrasell/levant/levant"
	"github.com/jrasell/levant/levant/structs"
	"github.com/jrasell/levant/logging"
	"github.com/jrasell/levant/template"
)

// DeployCommand is the command implementation that allows users to deploy a
// Nomad job based on passed templates and variables.
type DeployCommand struct {
	Meta
}

// Help provides the help information for the deploy command.
func (c *DeployCommand) Help() string {
	helpText := `
Usage: levant deploy [options] [TEMPLATE]

  Deploy a Nomad job based on input templates and variable files. The deploy
  command supports passing variables individually on the command line. Multiple
  commands can be passed in the format of -var 'key=value'. Variables passed
  via the command line take precedence over the same variable declared within
  a passed variable file.

Arguments:

  TEMPLATE nomad job template
    If no argument is given we look for a single *.nomad file

General Options:

  -address=<http_address>
    The Nomad HTTP API address including port which Levant will use to make
    calls.

  -allow-stale
    Allow stale consistency mode for requests into nomad.

  -canary-auto-promote=<seconds>
    The time in seconds, after which Levant will auto-promote a canary job
    if all canaries within the deployment are healthy.

  -consul-address=<addr>
    The Consul host and port to use when making Consul KeyValue lookups for
    template rendering.

  -force-batch
    Forces a new instance of the periodic job. A new instance will be created
    even if it violates the job's prohibit_overlap settings.

  -force-count
    Use the taskgroup count from the Nomad jobfile instead of the count that
    is currently set in a running job.

  -ignore-no-changes
    By default if no changes are detected when running a deployment Levant will
    exit with a status 1 to indicate a deployment didn't happen. This behaviour
    can be changed using this flag so that Levant will exit cleanly ensuring CD
    pipelines don't fail when no changes are detected.

  -log-level=<level>
    Specify the verbosity level of Levant's logs. Valid values include DEBUG,
    INFO, and WARN, in decreasing order of verbosity. The default is INFO.

  -log-format=<format>
    Specify the format of Levant's logs. Valid values are HUMAN or JSON. The
    default is HUMAN.

  -var-file=<file>
    Used in conjunction with the -job-file will deploy a templated job to your
    Nomad cluster. You can repeat this flag multiple times to supply multiple var-files.
    [default: levant.(json|yaml|yml|tf)]
`
	return strings.TrimSpace(helpText)
}

// Synopsis is provides a brief summary of the deploy command.
func (c *DeployCommand) Synopsis() string {
	return "Render and deploy a Nomad job from a template"
}

// Run triggers a run of the Levant template and deploy functions.
func (c *DeployCommand) Run(args []string) int {

	var err error
	var addr string
	config := &structs.Config{}

	flags := c.Meta.FlagSet("deploy", FlagSetVars)
	flags.Usage = func() { c.UI.Output(c.Help()) }

	flags.StringVar(&config.Addr, "address", "", "")
	flags.BoolVar(&config.AllowStale, "allow-stale", false, "")d
	flags.IntVar(&config.Canary, "canary-auto-promote", 0, "")
	flags.StringVar(&addr, "consul-address", "", "")
	flags.BoolVar(&config.ForceBatch, "force-batch", false, "")
	flags.BoolVar(&config.ForceCount, "force-count", false, "")
	flags.BoolVar(&config.IgnoreNoChanges, "ignore-no-changes", false, "")
	flags.StringVar(&config.LogLevel, "log-level", "INFO", "")
	flags.StringVar(&config.LogFormat, "log-format", "HUMAN", "")
	flags.Var((*helper.FlagStringSlice)(&config.VariableFiles), "var-file", "")

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

	args = flags.Args()

	if err = logging.SetupLogger(config.LogLevel, config.LogFormat); err != nil {
		c.UI.Error(err.Error())
		return 1
	}

	if len(args) == 1 {
		config.TemplateFile = args[0]
	} else if len(args) == 0 {
		if config.TemplateFile = helper.GetDefaultTmplFile(); config.TemplateFile == "" {
			c.UI.Error(c.Help())
			c.UI.Error("\nERROR: Template arg missing and no default template found")
			return 1
		}
	} else {
		c.UI.Error(c.Help())
		return 1
	}

	config.Job, err = template.RenderJob(config.TemplateFile, config.VariableFiles, addr, &c.Meta.flagVars)
	if err != nil {
		c.UI.Error(fmt.Sprintf("[ERROR] levant/command: %v", err))
		return 1
	}

	if config.Canary > 0 {
		if err = c.checkCanaryAutoPromote(config.Job, config.Canary); err != nil {
			c.UI.Error(fmt.Sprintf("[ERROR] levant/command: %v", err))
			return 1
		}
	}

	if config.ForceBatch {
		if err = c.checkForceBatch(config.Job, config.ForceBatch); err != nil {
			c.UI.Error(fmt.Sprintf("[ERROR] levant/command: %v", err))
			return 1
		}
	}

	success := levant.TriggerDeployment(config, nil)
	if !success {
		return 1
	}

	return 0
}

func (c *DeployCommand) checkCanaryAutoPromote(job *nomad.Job, canaryAutoPromote int) error {
	if canaryAutoPromote == 0 {
		return nil
	}

	if job.Update != nil && job.Update.Canary != nil && *job.Update.Canary > 0 {
		return nil
	}

	for _, group := range job.TaskGroups {
		if group.Update != nil && group.Update.Canary != nil && *group.Update.Canary > 0 {
			return nil
		}
	}

	return fmt.Errorf("canary-auto-update of %v passed but job is not canary enabled", canaryAutoPromote)
}

// checkForceBatch ensures that if the force-batch flag is passed, the job is
// periodic.
func (c *DeployCommand) checkForceBatch(job *nomad.Job, forceBatch bool) error {

	if forceBatch && job.IsPeriodic() {
		return nil
	}

	return fmt.Errorf("force-batch passed but job is not periodic")
}
<|endoftext|>


<commit_before>//
// command/docker.go
//
// Copyright (c) 2016-2017 Junpei Kawamoto
//
// This software is released under the MIT License.
//
// http://opensource.org/licenses/mit-license.php
//

package command

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"text/template"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"
)

// DockerfileAsset defines a asset name for Dockerfile.
const DockerfileAsset = "assets/Dockerfile"

// DockerfileOpt defines option variables used in Dockerfile templates.
type DockerfileOpt struct {
	// Customize FROM clause.
	BaseImage string
	// Git repository.
	Repository string
	// URL for an Apt proxy.
	AptProxy string
	// URL for a pypi proxy.
	PypiProxy string
	// URL for a HTTP proxy.
	HTTPProxy string
	// URL for a HTTPS proxy.
	HTTPSProxy string
	// Comma separated URL lists.
	NoProxy string
}

type travisExt struct {
	*Travis
	*DockerfileOpt
	Archive string
}

// Dockerfile creates a Dockerfile from an instance of Travis.
func Dockerfile(travis *Travis, opt *DockerfileOpt, archive string) (res []byte, err error) {

	var data []byte

	// Loading the base template.
	data, err = Asset(DockerfileAsset)
	if err != nil {
		return
	}
	base, err := template.New("").Parse(string(data))
	if err != nil {
		return
	}

	// Loading a child template.
	name := fmt.Sprintf("%s-%s", DockerfileAsset, travis.Language)
	data, err = Asset(name)
	if err != nil {
		data, err = Asset(fmt.Sprintf("%s-python", DockerfileAsset))
		if err != nil {
			return
		}
	}
	temp, err := base.Parse(string(data))
	if err != nil {
		return
	}

	// Checking optional parameters.
	opt.PypiProxy = strings.TrimSuffix(opt.PypiProxy, "/")

	// Creating Dockerfile.
	param := travisExt{
		Travis:        travis,
		DockerfileOpt: opt,
		Archive:       archive,
	}
	buf := bytes.Buffer{}
	if err = temp.ExecuteTemplate(&buf, "base", &param); err != nil {
		return
	}
	res = buf.Bytes()

	return

}

// Build builds a docker image from a directory. The built image named tag.
// The directory must have Dockerfile.
func Build(ctx context.Context, dir, tag string) (err error) {

	cd, err := os.Getwd()
	if err != nil {
		return
	}
	if err = os.Chdir(dir); err != nil {
		return
	}
	defer os.Chdir(cd)

	cmd := exec.CommandContext(ctx, "docker", "build", "-t", tag, ".")
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return
	}

	go io.Copy(os.Stdout, stdout)
	go io.Copy(os.Stderr, stderr)

	return cmd.Run()

}

// Start runs a container to run tests.
func Start(ctx context.Context, tag, name string, args ...string) (err error) {

	// Create a docker client.
	cli, err := client.NewClient(client.DefaultDockerHost, "", nil, nil)
	if err != nil {
		return
	}
	defer cli.Close()

	// Create a docker container.
	config := container.Config{
		Image: tag,
		Cmd:   args,
	}
	host := container.HostConfig{}
	nw := network.NetworkingConfig{}
	container, err := cli.ContainerCreate(ctx, &config, &host, &nw, name)
	if err != nil {
		return
	}
	if name != "" {
		// If any container name isn't given, remove the container.
		defer cli.ContainerRemove(ctx, container.ID, types.ContainerRemoveOptions{})
	}

	// Attach stdout of the container.
	stdout, err := cli.ContainerAttach(ctx, container.ID, types.ContainerAttachOptions{
		Stream: true,
		Stdout: true,
	})
	if err != nil {
		return
	}
	defer stdout.Close()
	go io.Copy(os.Stdout, stdout.Reader)

	// Attach stderr of the container.
	stderr, err := cli.ContainerAttach(ctx, container.ID, types.ContainerAttachOptions{
		Stream: true,
		Stderr: true,
	})
	if err != nil {
		return
	}
	defer stderr.Close()
	go io.Copy(os.Stderr, stderr.Reader)

	// Start the container.
	options := types.ContainerStartOptions{}
	if err = cli.ContainerStart(ctx, container.ID, options); err != nil {
		return
	}

	// Wait until the container ends.
	exit, err := cli.ContainerWait(ctx, container.ID)
	if err != nil {
		return
	} else if exit != 0 {
		return fmt.Errorf("Testing container returns an error:", exit)
	}

	return

}
<commit_msg>Use docker APIs to build images instead of docker command.<commit_after>//
// command/docker.go
//
// Copyright (c) 2016-2017 Junpei Kawamoto
//
// This software is released under the MIT License.
//
// http://opensource.org/licenses/mit-license.php
//

package command

import (
	"archive/tar"
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/client"
)

// DockerfileAsset defines a asset name for Dockerfile.
const DockerfileAsset = "assets/Dockerfile"

// DockerfileOpt defines option variables used in Dockerfile templates.
type DockerfileOpt struct {
	// Customize FROM clause.
	BaseImage string
	// Git repository.
	Repository string
	// URL for an Apt proxy.
	AptProxy string
	// URL for a pypi proxy.
	PypiProxy string
	// URL for a HTTP proxy.
	HTTPProxy string
	// URL for a HTTPS proxy.
	HTTPSProxy string
	// Comma separated URL lists.
	NoProxy string
}

type travisExt struct {
	*Travis
	*DockerfileOpt
	Archive string
}

// Dockerfile creates a Dockerfile from an instance of Travis.
func Dockerfile(travis *Travis, opt *DockerfileOpt, archive string) (res []byte, err error) {

	var data []byte

	// Loading the base template.
	data, err = Asset(DockerfileAsset)
	if err != nil {
		return
	}
	base, err := template.New("").Parse(string(data))
	if err != nil {
		return
	}

	// Loading a child template.
	name := fmt.Sprintf("%s-%s", DockerfileAsset, travis.Language)
	data, err = Asset(name)
	if err != nil {
		data, err = Asset(fmt.Sprintf("%s-python", DockerfileAsset))
		if err != nil {
			return
		}
	}
	temp, err := base.Parse(string(data))
	if err != nil {
		return
	}

	// Checking optional parameters.
	opt.PypiProxy = strings.TrimSuffix(opt.PypiProxy, "/")

	// Creating Dockerfile.
	param := travisExt{
		Travis:        travis,
		DockerfileOpt: opt,
		Archive:       archive,
	}
	buf := bytes.Buffer{}
	if err = temp.ExecuteTemplate(&buf, "base", &param); err != nil {
		return
	}
	res = buf.Bytes()

	return

}

// Build builds a docker image from a directory. The built image named tag.
// The directory must have Dockerfile.
func Build(ctx context.Context, dir, tag string) (err error) {

	// Create a docker client.
	cli, err := client.NewClient(client.DefaultDockerHost, "", nil, nil)
	if err != nil {
		return
	}
	defer cli.Close()

	// Create a pipe.
	reader, writer := io.Pipe()

	// Send the build context.
	go func() {
		archiveContext(ctx, dir, writer)
		writer.Close()
	}()

	// Start to build an image.
	res, err := cli.ImageBuild(ctx, reader, types.ImageBuildOptions{
		Tags: []string{tag},
	})
	if err != nil {
		return
	}
	defer res.Body.Close()

	// Wait untile the copy ends or the context will be canceled.
	done := make(chan struct{})
	go func() {
		io.Copy(os.Stderr, res.Body)
		close(done)
	}()

	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-done:
		return
	}

}

// Start runs a container to run tests.
func Start(ctx context.Context, tag, name string, args ...string) (err error) {

	// Create a docker client.
	cli, err := client.NewClient(client.DefaultDockerHost, "", nil, nil)
	if err != nil {
		return
	}
	defer cli.Close()

	// Create a docker container.
	config := container.Config{
		Image: tag,
		Cmd:   args,
	}
	container, err := cli.ContainerCreate(ctx, &config, nil, nil, name)
	if err != nil {
		return
	}
	if name != "" {
		// If any container name isn't given, remove the container.
		defer cli.ContainerRemove(ctx, container.ID, types.ContainerRemoveOptions{})
	}

	// Attach stdout of the container.
	stdout, err := cli.ContainerAttach(ctx, container.ID, types.ContainerAttachOptions{
		Stream: true,
		Stdout: true,
	})
	if err != nil {
		return
	}
	defer stdout.Close()
	go io.Copy(os.Stdout, stdout.Reader)

	// Attach stderr of the container.
	stderr, err := cli.ContainerAttach(ctx, container.ID, types.ContainerAttachOptions{
		Stream: true,
		Stderr: true,
	})
	if err != nil {
		return
	}
	defer stderr.Close()
	go io.Copy(os.Stderr, stderr.Reader)

	// Start the container.
	options := types.ContainerStartOptions{}
	if err = cli.ContainerStart(ctx, container.ID, options); err != nil {
		return
	}

	// Wait until the container ends.
	exit, err := cli.ContainerWait(ctx, container.ID)
	if err != nil {
		return
	} else if exit != 0 {
		return fmt.Errorf("Testing container returns an error:", exit)
	}

	return

}

// archiveContext makes a tar.gz stream consists of files.
func archiveContext(ctx context.Context, root string, writer io.Writer) (err error) {

	// Create a buffered writer.
	bufWriter := bufio.NewWriter(writer)
	defer bufWriter.Flush()

	// Create a zipped writer on the bufferd writer.
	zipWriter, err := gzip.NewWriterLevel(bufWriter, gzip.BestCompression)
	if err != nil {
		return
	}
	defer zipWriter.Close()

	// Create a tarball writer on the zipped writer.
	tarWriter := tar.NewWriter(zipWriter)
	defer tarWriter.Close()

	// Create a tarball.
	// TODO: Use ioutil.ReadDir
	sources, err := filepath.Glob(filepath.Join(root, "*"))
	if err != nil {
		return
	}
	for _, path := range sources {

		// Write a file header.
		info, err := os.Stat(path)
		if err != nil {
			fmt.Printf("Cannot find %s (%s)", path, err.Error())
			break
		}

		header, err := tar.FileInfoHeader(info, path)
		if err != nil {
			fmt.Println(err.Error())
			break
		}
		tarWriter.WriteHeader(header)

		// Write the body.
		if err = copyFile(path, tarWriter); err != nil {
			break
		}

	}

	return

}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"strings"

	"github.com/mitchellh/cli"
	"github.com/posener/complete"
)

var _ cli.Command = (*KVGetCommand)(nil)
var _ cli.CommandAutocomplete = (*KVGetCommand)(nil)

type KVGetCommand struct {
	*BaseCommand

	flagVersion int
}

func (c *KVGetCommand) Synopsis() string {
	return "Retrieves data from the KV store"
}

func (c *KVGetCommand) Help() string {
	helpText := `
Usage: vault kv get [options] KEY

  Retrieves the value from Vault's key-value store at the given key name. If no
  key exists with that name, an error is returned. If a key exists with that
  name but has no data, nothing is returned.

      $ vault kv get secret/foo

  To view the given key name at a specific version in time, specify the "-version"
  flag:

      $ vault kv get -version=1 secret/foo

  Additional flags and more advanced use cases are detailed below.

` + c.Flags().Help()
	return strings.TrimSpace(helpText)
}

func (c *KVGetCommand) Flags() *FlagSets {
	set := c.flagSet(FlagSetHTTP | FlagSetOutputField | FlagSetOutputFormat)

	// Common Options
	f := set.NewFlagSet("Common Options")

	f.IntVar(&IntVar{
		Name:    "version",
		Target:  &c.flagVersion,
		Default: 0,
		Usage:   `If passed, the value at the version number will be returned.`,
	})

	return set
}

func (c *KVGetCommand) AutocompleteArgs() complete.Predictor {
	return nil
}

func (c *KVGetCommand) AutocompleteFlags() complete.Flags {
	return c.Flags().Completions()
}

func (c *KVGetCommand) Run(args []string) int {
	f := c.Flags()

	if err := f.Parse(args); err != nil {
		c.UI.Error(err.Error())
		return 1
	}

	args = f.Args()
	switch {
	case len(args) < 1:
		c.UI.Error(fmt.Sprintf("Not enough arguments (expected 1, got %d)", len(args)))
		return 1
	case len(args) > 1:
		c.UI.Error(fmt.Sprintf("Too many arguments (expected 1, got %d)", len(args)))
		return 1
	}

	client, err := c.Client()
	if err != nil {
		c.UI.Error(err.Error())
		return 2
	}

	path := sanitizePath(args[0])
	mountPath, v2, err := isKVv2(path, client)
	if err != nil {
		c.UI.Error(err.Error())
		return 2
	}

	var versionParam map[string]string

	if v2 {
		path = addPrefixToVKVPath(path, mountPath, "data")
		if err != nil {
			c.UI.Error(err.Error())
			return 2
		}

		if c.flagVersion > 0 {
			versionParam = map[string]string{
				"version": fmt.Sprintf("%d", c.flagVersion),
			}
		}
	}

	secret, err := kvReadRequest(client, path, versionParam)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error reading %s: %s", path, err))
		if secret != nil {
			OutputSecret(c.UI, secret)
		}
		return 2
	}
	if secret == nil {
		c.UI.Error(fmt.Sprintf("No value found at %s", path))
		return 2
	}

	if c.flagField != "" {
		if v2 {
			// This is a v2, pass in the data field
			if data, ok := secret.Data["data"]; ok && data != nil {
				return PrintRawField(c.UI, data, c.flagField)
			} else {
				c.UI.Error(fmt.Sprintf("No data found at %s", path))
				return 2
			}
		} else {
			return PrintRawField(c.UI, secret, c.flagField)
		}
	}

	// If we have wrap info print the secret normally.
	if secret.WrapInfo != nil || c.flagFormat != "table" {
		return OutputSecret(c.UI, secret)
	}

	if len(secret.Warnings) > 0 {
		tf := TableFormatter{}
		tf.printWarnings(c.UI, secret)
	}

	if metadata, ok := secret.Data["metadata"]; ok && metadata != nil {
		c.UI.Info(getHeaderForMap("Metadata", metadata.(map[string]interface{})))
		OutputData(c.UI, metadata)
		c.UI.Info("")
	}

	data := secret.Data
	if v2 && data != nil {
		data = nil
		dataRaw := secret.Data["data"]
		if dataRaw != nil {
			data = dataRaw.(map[string]interface{})
		}
	}

	if data != nil {
		c.UI.Info(getHeaderForMap("Data", data))
		OutputData(c.UI, data)
	}

	return 0
}
<commit_msg>Add in logic for handling a field of 'data' to `kv get` (#4895)<commit_after>package command

import (
	"fmt"
	"strings"

	"github.com/mitchellh/cli"
	"github.com/posener/complete"
)

var _ cli.Command = (*KVGetCommand)(nil)
var _ cli.CommandAutocomplete = (*KVGetCommand)(nil)

type KVGetCommand struct {
	*BaseCommand

	flagVersion int
}

func (c *KVGetCommand) Synopsis() string {
	return "Retrieves data from the KV store"
}

func (c *KVGetCommand) Help() string {
	helpText := `
Usage: vault kv get [options] KEY

  Retrieves the value from Vault's key-value store at the given key name. If no
  key exists with that name, an error is returned. If a key exists with that
  name but has no data, nothing is returned.

      $ vault kv get secret/foo

  To view the given key name at a specific version in time, specify the "-version"
  flag:

      $ vault kv get -version=1 secret/foo

  Additional flags and more advanced use cases are detailed below.

` + c.Flags().Help()
	return strings.TrimSpace(helpText)
}

func (c *KVGetCommand) Flags() *FlagSets {
	set := c.flagSet(FlagSetHTTP | FlagSetOutputField | FlagSetOutputFormat)

	// Common Options
	f := set.NewFlagSet("Common Options")

	f.IntVar(&IntVar{
		Name:    "version",
		Target:  &c.flagVersion,
		Default: 0,
		Usage:   `If passed, the value at the version number will be returned.`,
	})

	return set
}

func (c *KVGetCommand) AutocompleteArgs() complete.Predictor {
	return nil
}

func (c *KVGetCommand) AutocompleteFlags() complete.Flags {
	return c.Flags().Completions()
}

func (c *KVGetCommand) Run(args []string) int {
	f := c.Flags()

	if err := f.Parse(args); err != nil {
		c.UI.Error(err.Error())
		return 1
	}

	args = f.Args()
	switch {
	case len(args) < 1:
		c.UI.Error(fmt.Sprintf("Not enough arguments (expected 1, got %d)", len(args)))
		return 1
	case len(args) > 1:
		c.UI.Error(fmt.Sprintf("Too many arguments (expected 1, got %d)", len(args)))
		return 1
	}

	client, err := c.Client()
	if err != nil {
		c.UI.Error(err.Error())
		return 2
	}

	path := sanitizePath(args[0])
	mountPath, v2, err := isKVv2(path, client)
	if err != nil {
		c.UI.Error(err.Error())
		return 2
	}

	var versionParam map[string]string

	if v2 {
		path = addPrefixToVKVPath(path, mountPath, "data")
		if err != nil {
			c.UI.Error(err.Error())
			return 2
		}

		if c.flagVersion > 0 {
			versionParam = map[string]string{
				"version": fmt.Sprintf("%d", c.flagVersion),
			}
		}
	}

	secret, err := kvReadRequest(client, path, versionParam)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error reading %s: %s", path, err))
		if secret != nil {
			OutputSecret(c.UI, secret)
		}
		return 2
	}
	if secret == nil {
		c.UI.Error(fmt.Sprintf("No value found at %s", path))
		return 2
	}

	if c.flagField != "" {
		if v2 {
			// This is a v2, pass in the data field
			if data, ok := secret.Data["data"]; ok && data != nil {
				// If they requested a literal "data" see if they meant actual
				// value or the data block itself
				if c.flagField == "data" {
					if dataMap, ok := data.(map[string]interface{}); ok {
						if _, ok := dataMap["data"]; ok {
							return PrintRawField(c.UI, dataMap, c.flagField)
						}
					}
					return PrintRawField(c.UI, secret, c.flagField)
				}
				return PrintRawField(c.UI, data, c.flagField)
			} else {
				c.UI.Error(fmt.Sprintf("No data found at %s", path))
				return 2
			}
		} else {
			return PrintRawField(c.UI, secret, c.flagField)
		}
	}

	// If we have wrap info print the secret normally.
	if secret.WrapInfo != nil || c.flagFormat != "table" {
		return OutputSecret(c.UI, secret)
	}

	if len(secret.Warnings) > 0 {
		tf := TableFormatter{}
		tf.printWarnings(c.UI, secret)
	}

	if metadata, ok := secret.Data["metadata"]; ok && metadata != nil {
		c.UI.Info(getHeaderForMap("Metadata", metadata.(map[string]interface{})))
		OutputData(c.UI, metadata)
		c.UI.Info("")
	}

	data := secret.Data
	if v2 && data != nil {
		data = nil
		dataRaw := secret.Data["data"]
		if dataRaw != nil {
			data = dataRaw.(map[string]interface{})
		}
	}

	if data != nil {
		c.UI.Info(getHeaderForMap("Data", data))
		OutputData(c.UI, data)
	}

	return 0
}
<|endoftext|>


<commit_before>package commands

import (
	"errors"
	"strings"

	"github.com/emersion/go-imap"
)

// Fetch is a FETCH command, as defined in RFC 3501 section 6.4.5.
type Fetch struct {
	SeqSet *imap.SeqSet
	Items  []imap.FetchItem
}

func (cmd *Fetch) Command() *imap.Command {
	items := make([]interface{}, len(cmd.Items))
	for i, item := range cmd.Items {
		items[i] = imap.RawString(item)
	}

	return &imap.Command{
		Name:      "FETCH",
		Arguments: []interface{}{cmd.SeqSet, items},
	}
}

func (cmd *Fetch) Parse(fields []interface{}) error {
	if len(fields) < 2 {
		return errors.New("No enough arguments")
	}

	var err error
	if seqset, ok := fields[0].(string); !ok {
		return errors.New("Sequence set must be an atom")
	} else if cmd.SeqSet, err = imap.ParseSeqSet(seqset); err != nil {
		return err
	}

	switch items := fields[1].(type) {
	case string: // A macro or a single item
		cmd.Items = imap.FetchItem(strings.ToUpper(items)).Expand()
	case []interface{}: // A list of items
		cmd.Items = make([]imap.FetchItem, 0, len(items))
		for _, v := range items {
			itemStr, _ := v.(string)
			item := imap.FetchItem(strings.ToUpper(itemStr))
			cmd.Items = append(cmd.Items, item.Expand()...)
		}
	default:
		return errors.New("Items must be either a string or a list")
	}

	return nil
}
<commit_msg>Fix FETCH command with ALL, FAST and FULL macro as fetch items<commit_after>package commands

import (
	"errors"
	"strings"

	"github.com/emersion/go-imap"
)

// Fetch is a FETCH command, as defined in RFC 3501 section 6.4.5.
type Fetch struct {
	SeqSet *imap.SeqSet
	Items  []imap.FetchItem
}

func (cmd *Fetch) Command() *imap.Command {
	// Handle FETCH macros separately as they should not be serialized within parentheses
	if len(cmd.Items) == 1 && (cmd.Items[0] == imap.FetchAll || cmd.Items[0] == imap.FetchFast || cmd.Items[0] == imap.FetchFull) {
		return &imap.Command{
			Name:      "FETCH",
			Arguments: []interface{}{cmd.SeqSet, imap.RawString(cmd.Items[0])},
		}
	} else {
		items := make([]interface{}, len(cmd.Items))
		for i, item := range cmd.Items {
			items[i] = imap.RawString(item)
		}

		return &imap.Command{
			Name:      "FETCH",
			Arguments: []interface{}{cmd.SeqSet, items},
		}
	}
}

func (cmd *Fetch) Parse(fields []interface{}) error {
	if len(fields) < 2 {
		return errors.New("No enough arguments")
	}

	var err error
	if seqset, ok := fields[0].(string); !ok {
		return errors.New("Sequence set must be an atom")
	} else if cmd.SeqSet, err = imap.ParseSeqSet(seqset); err != nil {
		return err
	}

	switch items := fields[1].(type) {
	case string: // A macro or a single item
		cmd.Items = imap.FetchItem(strings.ToUpper(items)).Expand()
	case []interface{}: // A list of items
		cmd.Items = make([]imap.FetchItem, 0, len(items))
		for _, v := range items {
			itemStr, _ := v.(string)
			item := imap.FetchItem(strings.ToUpper(itemStr))
			cmd.Items = append(cmd.Items, item.Expand()...)
		}
	default:
		return errors.New("Items must be either a string or a list")
	}

	return nil
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"log"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/spf13/cobra"
	api "github.com/yungsang/dockerclient"
	"github.com/yungsang/tablewriter"
	"github.com/yungsang/talk2docker/client"
)

var cmdIs = &cobra.Command{
	Use:     "ls [NAME[:TAG]]",
	Aliases: []string{"images"},
	Short:   "List images",
	Long:    appName + " ls - List images",
	Run:     listImages,
}

var cmdImage = &cobra.Command{
	Use:     "image [command]",
	Aliases: []string{"img"},
	Short:   "Manage images",
	Long:    appName + " image - Manage images",
	Run: func(ctx *cobra.Command, args []string) {
		ctx.Usage()
	},
}

var cmdListImages = &cobra.Command{
	Use:     "list [NAME[:TAG]]",
	Aliases: []string{"ls"},
	Short:   "List images",
	Long:    appName + " image list - List images",
	Run:     listImages,
}

var cmdPullImage = &cobra.Command{
	Use:   "pull <NAME[:TAG]>",
	Short: "Pull an image",
	Long:  appName + " image pull - Pull an image",
	Run:   pullImage,
}

// Define at ps.go
// var boolAll, boolQuiet, boolNoHeader bool

func init() {
	cmdIs.Flags().BoolVarP(&boolAll, "all", "a", false, "Show all images. Only named/taged and leaf images are shown by default.")
	cmdIs.Flags().BoolVarP(&boolQuiet, "quiet", "q", false, "Only display numeric IDs")
	cmdIs.Flags().BoolVarP(&boolNoHeader, "no-header", "n", false, "Omit the header")

	cmdListImages.Flags().BoolVarP(&boolAll, "all", "a", false, "Show all images. Only named/taged and leaf images are shown by default.")
	cmdListImages.Flags().BoolVarP(&boolQuiet, "quiet", "q", false, "Only display numeric IDs")
	cmdListImages.Flags().BoolVarP(&boolNoHeader, "no-header", "n", false, "Omit the header")

	cmdPullImage.Flags().BoolVarP(&boolAll, "all", "a", false, "Pull all tagged images in the repository. Only the \"latest\" tagged image is pulled by default.")

	cmdImage.AddCommand(cmdListImages)
	cmdImage.AddCommand(cmdPullImage)
}

func listImages(ctx *cobra.Command, args []string) {
	docker, err := client.GetDockerClient(configPath, hostName)
	if err != nil {
		log.Fatal(err)
	}

	images, err := docker.ListImages(boolAll)
	if err != nil {
		log.Fatal(err)
	}

	var matchName = ""
	if len(args) > 0 {
		matchName = args[0]
	}

	matchImageByName := func(tags []string, name string) bool {
		_name := strings.Split(name, ":")

		for _, tag := range tags {
			_tag := strings.Split(tag, ":")
			if _tag[0] == _name[0] {
				if (len(_name) < 2) || (_tag[1] == _name[1]) {
					return true
				}
			}
		}

		return false
	}

	if boolQuiet {
		for _, image := range images {
			if (matchName == "") || matchImageByName(image.RepoTags, matchName) {
				fmt.Println(Truncate(image.Id, 12))
			}
		}
		return
	}

	var items [][]string

	if boolAll {
		roots := make([]*api.Image, 0)
		parents := make(map[string][]*api.Image)
		for _, image := range images {
			if image.ParentId == "" {
				roots = append(roots, image)
			} else {
				if children, exists := parents[image.ParentId]; exists {
					parents[image.ParentId] = append(children, image)
				} else {
					children := make([]*api.Image, 0)
					parents[image.ParentId] = append(children, image)
				}
			}
		}

		items = walkTree(roots, parents, "", items)
	} else {
		for _, image := range images {
			if (matchName == "") || matchImageByName(image.RepoTags, matchName) {
				name := strings.Join(image.RepoTags, ", ")
				if name == "<none>:<none>" {
					name = "<none>"
				}
				out := []string{
					Truncate(image.Id, 12),
					FormatNonBreakingString(name),
					fmt.Sprintf("%.3f", float64(image.VirtualSize)/1000000.0),
					FormatDateTime(time.Unix(image.Created, 0)),
				}
				items = append(items, out)
			}
		}
	}

	var header = []string{
		"ID",
		"Name:Tags",
		"Size(MB)",
	}
	if !boolAll {
		header = append(header, "Created at")
	}

	table := tablewriter.NewWriter(os.Stdout)
	if !boolNoHeader {
		table.SetHeader(header)
	}
	table.SetBorder(false)
	table.AppendBulk(items)
	table.Render()
}

func walkTree(images []*api.Image, parents map[string][]*api.Image, prefix string, items [][]string) [][]string {
	printImage := func(prefix string, image *api.Image, isLeaf bool) {
		name := strings.Join(image.RepoTags, ", ")
		if name == "<none>:<none>" {
			if isLeaf {
				name = "<none>"
			} else {
				name = ""
			}
		}
		out := []string{
			FormatNonBreakingString(fmt.Sprintf("%s %s", prefix, Truncate(image.Id, 12))),
			FormatNonBreakingString(name),
			fmt.Sprintf("%.3f", float64(image.VirtualSize)/1000000.0),
		}
		items = append(items, out)
	}

	length := len(images)
	if length > 1 {
		for index, image := range images {
			if (index + 1) == length {
				subimages, exists := parents[image.Id]
				printImage(prefix+"└", image, !exists)
				if exists {
					items = walkTree(subimages, parents, prefix+" ", items)
				}
			} else {
				subimages, exists := parents[image.Id]
				printImage(prefix+"├", image, !exists)
				if exists {
					items = walkTree(subimages, parents, prefix+"│", items)
				}
			}
		}
	} else {
		for _, image := range images {
			subimages, exists := parents[image.Id]
			printImage(prefix+"└", image, !exists)
			if exists {
				items = walkTree(subimages, parents, prefix+" ", items)
			}
		}
	}
	return items
}

func pullImage(ctx *cobra.Command, args []string) {
	if len(args) < 1 {
		fmt.Println("Needs an argument <NAME> to pull")
		ctx.Usage()
		return
	}

	var name = args[0]
	var tag = ""
	n := strings.LastIndex(name, ":")
	if n >= 0 {
		if !strings.Contains(name[n+1:], "/") {
			tag = name[n+1:]
			name = name[:n]
		}
	}

	if tag == "" {
		tag = "latest"
	}

	var repository = name + ":" + tag

	if boolAll {
		repository = name
	}

	path := os.ExpandEnv(configPath)

	config, err := client.LoadConfig(path)
	if err != nil {
		log.Fatal(err)
	}

	host, err := config.GetHost(hostName)
	if err != nil {
		log.Fatal(err)
	}

	docker, err := client.GetDockerClient(configPath, host.Name)
	if err != nil {
		log.Fatal(err)
	}

	info, err := docker.Info()
	if err != nil {
		log.Fatal(err)
	}

	server, err := config.GetIndexServer(info.IndexServerAddress)
	if err != nil {
		log.Fatal(err)
	}

	v := url.Values{}
	v.Set("fromImage", repository)
	uri := fmt.Sprintf("/%s/images/create?%s", api.APIVersion, v.Encode())

	headers := make(map[string]string)
	headers["X-Registry-Auth"] = server.Auth

	err = client.DoStreamRequest(docker, "POST", uri, nil, headers)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Set the appropriate error message<commit_after>package commands

import (
	"fmt"
	"log"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/spf13/cobra"
	api "github.com/yungsang/dockerclient"
	"github.com/yungsang/tablewriter"
	"github.com/yungsang/talk2docker/client"
)

var cmdIs = &cobra.Command{
	Use:     "ls [NAME[:TAG]]",
	Aliases: []string{"images"},
	Short:   "List images",
	Long:    appName + " ls - List images",
	Run:     listImages,
}

var cmdImage = &cobra.Command{
	Use:     "image [command]",
	Aliases: []string{"img"},
	Short:   "Manage images",
	Long:    appName + " image - Manage images",
	Run: func(ctx *cobra.Command, args []string) {
		ctx.Usage()
	},
}

var cmdListImages = &cobra.Command{
	Use:     "list [NAME[:TAG]]",
	Aliases: []string{"ls"},
	Short:   "List images",
	Long:    appName + " image list - List images",
	Run:     listImages,
}

var cmdPullImage = &cobra.Command{
	Use:   "pull <NAME[:TAG]>",
	Short: "Pull an image",
	Long:  appName + " image pull - Pull an image",
	Run:   pullImage,
}

// Define at ps.go
// var boolAll, boolQuiet, boolNoHeader bool

func init() {
	cmdIs.Flags().BoolVarP(&boolAll, "all", "a", false, "Show all images. Only named/taged and leaf images are shown by default.")
	cmdIs.Flags().BoolVarP(&boolQuiet, "quiet", "q", false, "Only display numeric IDs")
	cmdIs.Flags().BoolVarP(&boolNoHeader, "no-header", "n", false, "Omit the header")

	cmdListImages.Flags().BoolVarP(&boolAll, "all", "a", false, "Show all images. Only named/taged and leaf images are shown by default.")
	cmdListImages.Flags().BoolVarP(&boolQuiet, "quiet", "q", false, "Only display numeric IDs")
	cmdListImages.Flags().BoolVarP(&boolNoHeader, "no-header", "n", false, "Omit the header")

	cmdPullImage.Flags().BoolVarP(&boolAll, "all", "a", false, "Pull all tagged images in the repository. Only the \"latest\" tagged image is pulled by default.")

	cmdImage.AddCommand(cmdListImages)
	cmdImage.AddCommand(cmdPullImage)
}

func listImages(ctx *cobra.Command, args []string) {
	docker, err := client.GetDockerClient(configPath, hostName)
	if err != nil {
		log.Fatal(err)
	}

	images, err := docker.ListImages(boolAll)
	if err != nil {
		log.Fatal(err)
	}

	var matchName = ""
	if len(args) > 0 {
		matchName = args[0]
	}

	matchImageByName := func(tags []string, name string) bool {
		_name := strings.Split(name, ":")

		for _, tag := range tags {
			_tag := strings.Split(tag, ":")
			if _tag[0] == _name[0] {
				if (len(_name) < 2) || (_tag[1] == _name[1]) {
					return true
				}
			}
		}

		return false
	}

	if boolQuiet {
		for _, image := range images {
			if (matchName == "") || matchImageByName(image.RepoTags, matchName) {
				fmt.Println(Truncate(image.Id, 12))
			}
		}
		return
	}

	var items [][]string

	if boolAll {
		roots := make([]*api.Image, 0)
		parents := make(map[string][]*api.Image)
		for _, image := range images {
			if image.ParentId == "" {
				roots = append(roots, image)
			} else {
				if children, exists := parents[image.ParentId]; exists {
					parents[image.ParentId] = append(children, image)
				} else {
					children := make([]*api.Image, 0)
					parents[image.ParentId] = append(children, image)
				}
			}
		}

		items = walkTree(roots, parents, "", items)
	} else {
		for _, image := range images {
			if (matchName == "") || matchImageByName(image.RepoTags, matchName) {
				name := strings.Join(image.RepoTags, ", ")
				if name == "<none>:<none>" {
					name = "<none>"
				}
				out := []string{
					Truncate(image.Id, 12),
					FormatNonBreakingString(name),
					fmt.Sprintf("%.3f", float64(image.VirtualSize)/1000000.0),
					FormatDateTime(time.Unix(image.Created, 0)),
				}
				items = append(items, out)
			}
		}
	}

	var header = []string{
		"ID",
		"Name:Tags",
		"Size(MB)",
	}
	if !boolAll {
		header = append(header, "Created at")
	}

	table := tablewriter.NewWriter(os.Stdout)
	if !boolNoHeader {
		table.SetHeader(header)
	}
	table.SetBorder(false)
	table.AppendBulk(items)
	table.Render()
}

func walkTree(images []*api.Image, parents map[string][]*api.Image, prefix string, items [][]string) [][]string {
	printImage := func(prefix string, image *api.Image, isLeaf bool) {
		name := strings.Join(image.RepoTags, ", ")
		if name == "<none>:<none>" {
			if isLeaf {
				name = "<none>"
			} else {
				name = ""
			}
		}
		out := []string{
			FormatNonBreakingString(fmt.Sprintf("%s %s", prefix, Truncate(image.Id, 12))),
			FormatNonBreakingString(name),
			fmt.Sprintf("%.3f", float64(image.VirtualSize)/1000000.0),
		}
		items = append(items, out)
	}

	length := len(images)
	if length > 1 {
		for index, image := range images {
			if (index + 1) == length {
				subimages, exists := parents[image.Id]
				printImage(prefix+"└", image, !exists)
				if exists {
					items = walkTree(subimages, parents, prefix+" ", items)
				}
			} else {
				subimages, exists := parents[image.Id]
				printImage(prefix+"├", image, !exists)
				if exists {
					items = walkTree(subimages, parents, prefix+"│", items)
				}
			}
		}
	} else {
		for _, image := range images {
			subimages, exists := parents[image.Id]
			printImage(prefix+"└", image, !exists)
			if exists {
				items = walkTree(subimages, parents, prefix+" ", items)
			}
		}
	}
	return items
}

func pullImage(ctx *cobra.Command, args []string) {
	if len(args) < 1 {
		fmt.Println("Needs an argument <NAME> to pull")
		ctx.Usage()
		return
	}

	var name = args[0]
	var tag = ""
	n := strings.LastIndex(name, ":")
	if n >= 0 {
		if !strings.Contains(name[n+1:], "/") {
			tag = name[n+1:]
			name = name[:n]
		}
	}

	if tag == "" {
		tag = "latest"
	}

	var repository = name + ":" + tag

	if boolAll {
		repository = name
	}

	path := os.ExpandEnv(configPath)

	config, err := client.LoadConfig(path)
	if err != nil {
		log.Fatal(err)
	}

	host, err := config.GetHost(hostName)
	if err != nil {
		log.Fatal(err)
	}

	docker, err := client.GetDockerClient(configPath, host.Name)
	if err != nil {
		log.Fatal(err)
	}

	info, err := docker.Info()
	if err != nil {
		log.Fatal(err)
	}

	server, err := config.GetIndexServer(info.IndexServerAddress)
	if err != nil {
		log.Fatal("Please login prior to pulling an image.")
	}

	v := url.Values{}
	v.Set("fromImage", repository)
	uri := fmt.Sprintf("/%s/images/create?%s", api.APIVersion, v.Encode())

	headers := make(map[string]string)
	headers["X-Registry-Auth"] = server.Auth

	err = client.DoStreamRequest(docker, "POST", uri, nil, headers)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package cmap

import (
	"encoding/json"
	"hash/fnv"
	"sync"
)

var SHARD_COUNT = 32

// TODO: Add Keys function which returns an array of keys for the map.

// A "thread" safe map of type string:Anything.
// To avoid lock bottlenecks this map is dived to several (SHARD_COUNT) map shards.
type ConcurrentMap []*ConcurrentMapShared

// A "thread" safe string to anything map.
type ConcurrentMapShared struct {
	items        map[string]interface{}
	sync.RWMutex // Read Write mutex, guards access to internal map.
}

// Creates a new concurrent map.
func New() ConcurrentMap {
	m := make(ConcurrentMap, SHARD_COUNT)
	for i := 0; i < SHARD_COUNT; i++ {
		m[i] = &ConcurrentMapShared{items: make(map[string]interface{})}
	}
	return m
}

// Returns shard under given key
func (m ConcurrentMap) GetShard(key string) *ConcurrentMapShared {
	hasher := fnv.New32()
	hasher.Write([]byte(key))
	return m[uint(hasher.Sum32())%uint(SHARD_COUNT)]
}

// Sets the given value under the specified key.
func (m *ConcurrentMap) Set(key string, value interface{}) {
	// Get map shard.
	shard := m.GetShard(key)
	shard.Lock()
	defer shard.Unlock()
	shard.items[key] = value
}

// Sets the given value under the specified key if no value was associated with it.
func (m *ConcurrentMap) SetIfAbsent(key string, value interface{}) bool {
	// Get map shard.
	shard := m.GetShard(key)
	shard.Lock()
	defer shard.Unlock()
	_, ok := shard.items[key]
	if !ok {
		shard.items[key] = value
	}
	return !ok
}

// Retrieves an element from map under given key.
func (m ConcurrentMap) Get(key string) (interface{}, bool) {
	// Get shard
	shard := m.GetShard(key)
	shard.RLock()
	defer shard.RUnlock()

	// Get item from shard.
	val, ok := shard.items[key]
	return val, ok
}

// Returns the number of elements within the map.
func (m ConcurrentMap) Count() int {
	count := 0
	for i := 0; i < SHARD_COUNT; i++ {
		shard := m[i]
		shard.RLock()
		count += len(shard.items)
		shard.RUnlock()
	}
	return count
}

// Looks up an item under specified key
func (m *ConcurrentMap) Has(key string) bool {
	// Get shard
	shard := m.GetShard(key)
	shard.RLock()
	defer shard.RUnlock()

	// See if element is within shard.
	_, ok := shard.items[key]
	return ok
}

// Removes an element from the map.
func (m *ConcurrentMap) Remove(key string) {
	// Try to get shard.
	shard := m.GetShard(key)
	shard.Lock()
	defer shard.Unlock()
	delete(shard.items, key)
}

// Checks if map is empty.
func (m *ConcurrentMap) IsEmpty() bool {
	return m.Count() == 0
}

// Used by the Iter & IterBuffered functions to wrap two variables together over a channel,
type Tuple struct {
	Key string
	Val interface{}
}

// Returns an iterator which could be used in a for range loop.
func (m ConcurrentMap) Iter() <-chan Tuple {
	ch := make(chan Tuple)
	go func() {
		// Foreach shard.
		for _, shard := range m {
			// Foreach key, value pair.
			shard.RLock()
			for key, val := range shard.items {
				ch <- Tuple{key, val}
			}
			shard.RUnlock()
		}
		close(ch)
	}()
	return ch
}

// Returns a buffered iterator which could be used in a for range loop.
func (m ConcurrentMap) IterBuffered() <-chan Tuple {
	ch := make(chan Tuple, m.Count())
	go func() {
		// Foreach shard.
		for _, shard := range m {
			// Foreach key, value pair.
			shard.RLock()
			for key, val := range shard.items {
				ch <- Tuple{key, val}
			}
			shard.RUnlock()
		}
		close(ch)
	}()
	return ch
}

// Returns all items as map[string]interface{}
func (m ConcurrentMap) Items() map[string]interface{} {
	tmp := make(map[string]interface{})

	// Insert items to temporary map.
	for item := range m.Iter() {
		tmp[item.Key] = item.Val
	}

	return tmp
}

//Reviles ConcurrentMap "private" variables to json marshal.
func (m ConcurrentMap) MarshalJSON() ([]byte, error) {
	// Create a temporary map, which will hold all item spread across shards.
	tmp := make(map[string]interface{})

	// Insert items to temporary map.
	for item := range m.Iter() {
		tmp[item.Key] = item.Val
	}
	return json.Marshal(tmp)
}

// Concurrent map uses Interface{} as its value, therefor JSON Unmarshal
// will probably won't know which to type to unmarshal into, in such case
// we'll end up with a value of type map[string]interface{}, In most cases this isn't
// out value type, this is why we've decided to remove this functionality.

// func (m *ConcurrentMap) UnmarshalJSON(b []byte) (err error) {
// 	// Reverse process of Marshal.

// 	tmp := make(map[string]interface{})

// 	// Unmarshal into a single map.
// 	if err := json.Unmarshal(b, &tmp); err != nil {
// 		return nil
// 	}

// 	// foreach key,value pair in temporary map insert into our concurrent map.
// 	for key, val := range tmp {
// 		m.Set(key, val)
// 	}
// 	return nil
// }
<commit_msg>change defer unlock to immediatly unlock to avoid delay call which can speed up all func call<commit_after>package cmap

import (
	"encoding/json"
	"hash/fnv"
	"sync"
)

var SHARD_COUNT = 32

// TODO: Add Keys function which returns an array of keys for the map.

// A "thread" safe map of type string:Anything.
// To avoid lock bottlenecks this map is dived to several (SHARD_COUNT) map shards.
type ConcurrentMap []*ConcurrentMapShared

// A "thread" safe string to anything map.
type ConcurrentMapShared struct {
	items        map[string]interface{}
	sync.RWMutex // Read Write mutex, guards access to internal map.
}

// Creates a new concurrent map.
func New() ConcurrentMap {
	m := make(ConcurrentMap, SHARD_COUNT)
	for i := 0; i < SHARD_COUNT; i++ {
		m[i] = &ConcurrentMapShared{items: make(map[string]interface{})}
	}
	return m
}

// Returns shard under given key
func (m ConcurrentMap) GetShard(key string) *ConcurrentMapShared {
	hasher := fnv.New32()
	hasher.Write([]byte(key))
	return m[uint(hasher.Sum32())%uint(SHARD_COUNT)]
}

// Sets the given value under the specified key.
func (m *ConcurrentMap) Set(key string, value interface{}) {
	// Get map shard.
	shard := m.GetShard(key)
	shard.Lock()
	shard.items[key] = value
	shard.Unlock()
}

// Sets the given value under the specified key if no value was associated with it.
func (m *ConcurrentMap) SetIfAbsent(key string, value interface{}) bool {
	// Get map shard.
	shard := m.GetShard(key)
	shard.Lock()
	_, ok := shard.items[key]
	if !ok {
		shard.items[key] = value
	}
	shard.Unlock()
	return !ok
}

// Retrieves an element from map under given key.
func (m ConcurrentMap) Get(key string) (interface{}, bool) {
	// Get shard
	shard := m.GetShard(key)
	shard.RLock()
	// Get item from shard.
	val, ok := shard.items[key]
	shard.RUnlock()
	return val, ok
}

// Returns the number of elements within the map.
func (m ConcurrentMap) Count() int {
	count := 0
	for i := 0; i < SHARD_COUNT; i++ {
		shard := m[i]
		shard.RLock()
		count += len(shard.items)
		shard.RUnlock()
	}
	return count
}

// Looks up an item under specified key
func (m *ConcurrentMap) Has(key string) bool {
	// Get shard
	shard := m.GetShard(key)
	shard.RLock()
	// See if element is within shard.
	_, ok := shard.items[key]
	shard.RUnlock()
	return ok
}

// Removes an element from the map.
func (m *ConcurrentMap) Remove(key string) {
	// Try to get shard.
	shard := m.GetShard(key)
	shard.Lock()
	delete(shard.items, key)
	shard.Unlock()
}

// Checks if map is empty.
func (m *ConcurrentMap) IsEmpty() bool {
	return m.Count() == 0
}

// Used by the Iter & IterBuffered functions to wrap two variables together over a channel,
type Tuple struct {
	Key string
	Val interface{}
}

// Returns an iterator which could be used in a for range loop.
func (m ConcurrentMap) Iter() <-chan Tuple {
	ch := make(chan Tuple)
	go func() {
		// Foreach shard.
		for _, shard := range m {
			// Foreach key, value pair.
			shard.RLock()
			for key, val := range shard.items {
				ch <- Tuple{key, val}
			}
			shard.RUnlock()
		}
		close(ch)
	}()
	return ch
}

// Returns a buffered iterator which could be used in a for range loop.
func (m ConcurrentMap) IterBuffered() <-chan Tuple {
	ch := make(chan Tuple, m.Count())
	go func() {
		// Foreach shard.
		for _, shard := range m {
			// Foreach key, value pair.
			shard.RLock()
			for key, val := range shard.items {
				ch <- Tuple{key, val}
			}
			shard.RUnlock()
		}
		close(ch)
	}()
	return ch
}

// Returns all items as map[string]interface{}
func (m ConcurrentMap) Items() map[string]interface{} {
	tmp := make(map[string]interface{})

	// Insert items to temporary map.
	for item := range m.Iter() {
		tmp[item.Key] = item.Val
	}

	return tmp
}

//Reviles ConcurrentMap "private" variables to json marshal.
func (m ConcurrentMap) MarshalJSON() ([]byte, error) {
	// Create a temporary map, which will hold all item spread across shards.
	tmp := make(map[string]interface{})

	// Insert items to temporary map.
	for item := range m.Iter() {
		tmp[item.Key] = item.Val
	}
	return json.Marshal(tmp)
}

// Concurrent map uses Interface{} as its value, therefor JSON Unmarshal
// will probably won't know which to type to unmarshal into, in such case
// we'll end up with a value of type map[string]interface{}, In most cases this isn't
// out value type, this is why we've decided to remove this functionality.

// func (m *ConcurrentMap) UnmarshalJSON(b []byte) (err error) {
// 	// Reverse process of Marshal.

// 	tmp := make(map[string]interface{})

// 	// Unmarshal into a single map.
// 	if err := json.Unmarshal(b, &tmp); err != nil {
// 		return nil
// 	}

// 	// foreach key,value pair in temporary map insert into our concurrent map.
// 	for key, val := range tmp {
// 		m.Set(key, val)
// 	}
// 	return nil
// }
<|endoftext|>