<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package gen

import (
	"encoding/json"
	"math/rand"
	"reflect"
	"testing"
	"testing/quick"
	"time"

	tc "go.uber.org/thriftrw/gen/internal/tests/containers"
	te "go.uber.org/thriftrw/gen/internal/tests/enums"
	tx "go.uber.org/thriftrw/gen/internal/tests/exceptions"
	tf "go.uber.org/thriftrw/gen/internal/tests/services"
	ts "go.uber.org/thriftrw/gen/internal/tests/structs"
	td "go.uber.org/thriftrw/gen/internal/tests/typedefs"
	tu "go.uber.org/thriftrw/gen/internal/tests/unions"
	"go.uber.org/thriftrw/wire"
	"go.uber.org/zap/zapcore"

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

func defaultValueGenerator(typ reflect.Type) func(*testing.T, *rand.Rand) thriftType {
	return func(t *testing.T, rand *rand.Rand) thriftType {
		for {
			// We will keep trying to generate a value until a valid one
			// is found.

			v, ok := quick.Value(typ, rand)
			require.True(t, ok, "failed to generate a value")

			tval := v.Addr().Interface().(thriftType)

			// TODO(abg): ToWire + EvaluateValue to validate here means we end
			// up serializing this value twice. We may want to include a
			// Validate method on generated types.

			w, err := tval.ToWire()
			if err != nil {
				// Value fails validity check. Try again.
				continue
			}

			// Because we evaluate collections lazily, validation issues
			// with items in them won't be known until we try to serialize
			// it or explicitly evaluate the lazy lists with
			// wire.EvaluateValue.
			if err := wire.EvaluateValue(w); err != nil {
				// Value fails validity check. Try again.
				continue
			}

			return tval
		}
	}
}

// enumValueGenerator builds a generator for random enum values given the
// `*_Values` function for that enum.
func enumValueGenerator(valuesFunc interface{}) func(*testing.T, *rand.Rand) thriftType {
	vfunc := reflect.ValueOf(valuesFunc)
	typ := vfunc.Type().Out(0).Elem() // Foo_Values() []Foo -> Foo
	return func(t *testing.T, rand *rand.Rand) thriftType {
		knownValues := vfunc.Call(nil)[0]

		var giveV reflect.Value
		// Flip a coin to decide whether we're evaluating a known or
		// unknown value.
		if rand.Int()%2 == 0 && knownValues.Len() > 0 {
			// Pick a known value at random
			giveV = knownValues.Index(rand.Intn(knownValues.Len()))
		} else {
			// give = MyEnum($randomValue)
			giveV = reflect.New(typ).Elem()
			giveV.Set(reflect.ValueOf(rand.Int31()).Convert(typ))
		}

		return giveV.Addr().Interface().(thriftType)
	}
}

func TestQuickRoundTrip(t *testing.T) {
	type testCase struct {
		// Sample value of the type to be tested.
		Sample interface{}

		// Specifies how we generate valid values of this type. Defaults to
		// defaultValueGenerator(Type) if unspecified.
		Generator func(*testing.T, *rand.Rand) thriftType

		// If set, logging for this type will not be tested. This is needed
		// for typedefs of primitives which can't implement ArrayMarshaler or
		// ObjectMarshaler.
		NoLog bool

		// Whether we should evaluate JSON round-tripping. This is opt-in
		// rather than opt-out because struct types that use collections won't
		// round-trip with JSON successfully due to nil versus empty
		// collection differences.
		JSON bool
	}

	// The following types from our tests have been skipped.
	// - unions.ArbitraryValue: Self-reference causes testing/quick to loop
	//   for too long
	// - services.KeyValue_SetValue_Args{}: Accepts an ArbitraryValue
	// - services.KeyValue_SetValueV2_Args: Accepts an ArbitraryValue
	// - services.KeyValue_GetManyValues_Result{}: Produces an ArbitraryValue
	// - services.KeyValue_GetValue_Result{}: Produces an ArbitraryValue

	// TODO(abg): ^Use custom generators to make this not-a-problem.

	tests := []testCase{
		// structs, unions, and exceptions
		{Sample: tc.ContainersOfContainers{}},
		{Sample: tc.EnumContainers{}},
		{Sample: tc.ListOfConflictingEnums{}},
		{Sample: tc.ListOfConflictingUUIDs{}},
		{Sample: tc.MapOfBinaryAndString{}},
		{Sample: tc.PrimitiveContainersRequired{}},
		{Sample: tc.PrimitiveContainers{}},
		{Sample: td.DefaultPrimitiveTypedef{}},
		{Sample: td.Event{}},
		{Sample: td.I128{}},
		{Sample: td.Transition{}},
		{Sample: te.StructWithOptionalEnum{}},
		{Sample: tf.Cache_Clear_Args{}},
		{Sample: tf.Cache_ClearAfter_Args{}},
		{Sample: tf.ConflictingNames_SetValue_Args{}},
		{Sample: tf.ConflictingNames_SetValue_Result{}},
		{Sample: tf.ConflictingNamesSetValueArgs{}},
		{Sample: tf.InternalError{}},
		{Sample: tf.KeyValue_DeleteValue_Args{}},
		{Sample: tf.KeyValue_DeleteValue_Result{}},
		{Sample: tf.KeyValue_GetManyValues_Args{}},
		{Sample: tf.KeyValue_GetValue_Args{}},
		{Sample: tf.KeyValue_SetValue_Result{}},
		{Sample: tf.KeyValue_SetValueV2_Result{}},
		{Sample: tf.KeyValue_Size_Args{}},
		{Sample: tf.KeyValue_Size_Result{}},
		{Sample: tf.NonStandardServiceName_NonStandardFunctionName_Args{}},
		{Sample: tf.NonStandardServiceName_NonStandardFunctionName_Result{}},
		{Sample: ts.ContactInfo{}},
		{Sample: ts.DefaultsStruct{}},
		{Sample: ts.Edge{}},
		{Sample: ts.EmptyStruct{}},
		{Sample: ts.Frame{}},
		{Sample: ts.GoTags{}},
		{Sample: ts.Graph{}},
		{Sample: ts.Node{}},
		{Sample: ts.Omit{}},
		{Sample: ts.Point{}},
		{Sample: ts.PrimitiveOptionalStruct{}},
		{Sample: ts.PrimitiveRequiredStruct{}},
		{Sample: ts.Rename{}},
		{Sample: ts.Size{}},
		{Sample: ts.StructLabels{}},
		{Sample: ts.User{}},
		{Sample: ts.ZapOptOutStruct{}},
		{Sample: tu.Document{}},
		{Sample: tu.EmptyUnion{}},
		{Sample: tx.DoesNotExistException{}},
		{Sample: tx.EmptyException{}},

		// typedefs
		{Sample: td.BinarySet{}},
		{Sample: td.EdgeMap{}},
		{Sample: td.FrameGroup{}},
		{Sample: td.MyEnum(0)},
		{Sample: td.PDF{}, NoLog: true},
		{Sample: td.PointMap{}},
		{Sample: td.State(""), NoLog: true},
		{Sample: td.StateMap{}},
		{Sample: td.Timestamp(0), NoLog: true},
		{Sample: td.UUID{}},

		// enums
		{
			Sample:    te.EmptyEnum(0),
			Generator: enumValueGenerator(te.EmptyEnum_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumDefault(0),
			Generator: enumValueGenerator(te.EnumDefault_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumWithDuplicateName(0),
			Generator: enumValueGenerator(te.EnumWithDuplicateName_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumWithDuplicateValues(0),
			Generator: enumValueGenerator(te.EnumWithDuplicateValues_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumWithLabel(0),
			Generator: enumValueGenerator(te.EnumWithLabel_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumWithValues(0),
			Generator: enumValueGenerator(te.EnumWithValues_Values),
			JSON:      true,
		},
		{
			Sample:    te.LowerCaseEnum(0),
			Generator: enumValueGenerator(te.LowerCaseEnum_Values),
			JSON:      true,
		},
		{
			Sample:    te.RecordType(0),
			Generator: enumValueGenerator(te.RecordType_Values),
			JSON:      true,
		},
		{
			Sample:    te.RecordTypeValues(0),
			Generator: enumValueGenerator(te.RecordTypeValues_Values),
			JSON:      true,
		},
	}

	// Log the seed so that we can reproduce this if it ever fails.
	seed := time.Now().UnixNano()
	rand := rand.New(rand.NewSource(seed))
	t.Logf("Using seed %v for testing/quick", seed)

	const numValues = 1000 // number of values to test against
	for _, tt := range tests {
		typ := reflect.TypeOf(tt.Sample)
		t.Run(typ.Name(), func(t *testing.T) {
			generator := tt.Generator
			if generator == nil {
				generator = defaultValueGenerator(typ)
			}

			values := make([]thriftType, numValues)
			for i := range values {
				values[i] = generator(t, rand)
			}

			t.Run("Thrift", func(t *testing.T) {
				for _, give := range values {
					w, err := give.ToWire()
					require.NoError(t, err, "failed to Thrift encode %v", give)

					got := reflect.New(typ).Interface().(thriftType)
					require.NoError(t, got.FromWire(w), "failed to Thrift decode from %v", w)

					assert.Equal(t, got, give)
				}
			})

			t.Run("String", func(t *testing.T) {
				for _, give := range values {
					assert.NotPanics(t, func() {
						_ = give.String()
					}, "failed to String %#v", give)
				}
			})

			if tt.JSON {
				t.Run("JSON", func(t *testing.T) {
					for _, giveValue := range values {
						give, ok := giveValue.(json.Marshaler)
						require.True(t, ok, "Type does not implement json.Marshaler")

						bs, err := give.MarshalJSON()
						require.NoError(t, err, "failed to encode %v", give)

						got, ok := reflect.New(typ).Interface().(json.Unmarshaler)
						require.True(t, ok, "Type does not implement json.Unmarshaler")

						require.NoError(t, got.UnmarshalJSON(bs), "failed to decode from %q", bs)
						assert.Equal(t, got, give, "could not round-trip")
					}
				})
			}

			if !tt.NoLog {
				t.Run("Zap", func(t *testing.T) {
					for _, give := range values {
						assert.NotPanics(t, func() {
							enc := zapcore.NewMapObjectEncoder()

							if obj, ok := give.(zapcore.ObjectMarshaler); ok {
								assert.NoErrorf(t, obj.MarshalLogObject(enc), "failed to log %v", give)
								return
							}

							if arr, ok := give.(zapcore.ArrayMarshaler); ok {
								assert.NoErrorf(t, enc.AddArray("values", arr), "failed to log %v", give)
								return
							}

							t.Fatal(
								"Type does not implement zapcore.ObjectMarshaler or zapcore.ArrayMarshaler. "+
									"Did you mean to add NoLog?", typ)
						}, "failed to log %v", give)
					}
				})
			}
		})
	}
}
<commit_msg>gen/quick/test: Text round-tripping<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package gen

import (
	"encoding"
	"encoding/json"
	"math/rand"
	"reflect"
	"testing"
	"testing/quick"
	"time"

	tc "go.uber.org/thriftrw/gen/internal/tests/containers"
	te "go.uber.org/thriftrw/gen/internal/tests/enums"
	tx "go.uber.org/thriftrw/gen/internal/tests/exceptions"
	tf "go.uber.org/thriftrw/gen/internal/tests/services"
	ts "go.uber.org/thriftrw/gen/internal/tests/structs"
	td "go.uber.org/thriftrw/gen/internal/tests/typedefs"
	tu "go.uber.org/thriftrw/gen/internal/tests/unions"
	"go.uber.org/thriftrw/wire"
	"go.uber.org/zap/zapcore"

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

func defaultValueGenerator(typ reflect.Type) func(*testing.T, *rand.Rand) thriftType {
	return func(t *testing.T, rand *rand.Rand) thriftType {
		for {
			// We will keep trying to generate a value until a valid one
			// is found.

			v, ok := quick.Value(typ, rand)
			require.True(t, ok, "failed to generate a value")

			tval := v.Addr().Interface().(thriftType)

			// TODO(abg): ToWire + EvaluateValue to validate here means we end
			// up serializing this value twice. We may want to include a
			// Validate method on generated types.

			w, err := tval.ToWire()
			if err != nil {
				// Value fails validity check. Try again.
				continue
			}

			// Because we evaluate collections lazily, validation issues
			// with items in them won't be known until we try to serialize
			// it or explicitly evaluate the lazy lists with
			// wire.EvaluateValue.
			if err := wire.EvaluateValue(w); err != nil {
				// Value fails validity check. Try again.
				continue
			}

			return tval
		}
	}
}

// enumValueGenerator builds a generator for random enum values given the
// `*_Values` function for that enum.
func enumValueGenerator(valuesFunc interface{}) func(*testing.T, *rand.Rand) thriftType {
	vfunc := reflect.ValueOf(valuesFunc)
	typ := vfunc.Type().Out(0).Elem() // Foo_Values() []Foo -> Foo
	return func(t *testing.T, rand *rand.Rand) thriftType {
		knownValues := vfunc.Call(nil)[0]

		var giveV reflect.Value
		// Flip a coin to decide whether we're evaluating a known or
		// unknown value.
		if rand.Int()%2 == 0 && knownValues.Len() > 0 {
			// Pick a known value at random
			giveV = knownValues.Index(rand.Intn(knownValues.Len()))
		} else {
			// give = MyEnum($randomValue)
			giveV = reflect.New(typ).Elem()
			giveV.Set(reflect.ValueOf(rand.Int31()).Convert(typ))
		}

		return giveV.Addr().Interface().(thriftType)
	}
}

func TestQuickRoundTrip(t *testing.T) {
	type testCase struct {
		// Sample value of the type to be tested.
		Sample interface{}

		// Specifies how we generate valid values of this type. Defaults to
		// defaultValueGenerator(Type) if unspecified.
		Generator func(*testing.T, *rand.Rand) thriftType

		// If set, logging for this type will not be tested. This is needed
		// for typedefs of primitives which can't implement ArrayMarshaler or
		// ObjectMarshaler.
		NoLog bool

		// Whether we should evaluate JSON round-tripping. This is opt-in
		// rather than opt-out because struct types that use collections won't
		// round-trip with JSON successfully due to nil versus empty
		// collection differences.
		JSON bool

		// Whether we should evaluate encoding.TextMarshaler round-tripping.
		// This is only suported on enums.
		Text bool
	}

	// The following types from our tests have been skipped.
	// - unions.ArbitraryValue: Self-reference causes testing/quick to loop
	//   for too long
	// - services.KeyValue_SetValue_Args{}: Accepts an ArbitraryValue
	// - services.KeyValue_SetValueV2_Args: Accepts an ArbitraryValue
	// - services.KeyValue_GetManyValues_Result{}: Produces an ArbitraryValue
	// - services.KeyValue_GetValue_Result{}: Produces an ArbitraryValue

	// TODO(abg): ^Use custom generators to make this not-a-problem.

	tests := []testCase{
		// structs, unions, and exceptions
		{Sample: tc.ContainersOfContainers{}},
		{Sample: tc.EnumContainers{}},
		{Sample: tc.ListOfConflictingEnums{}},
		{Sample: tc.ListOfConflictingUUIDs{}},
		{Sample: tc.MapOfBinaryAndString{}},
		{Sample: tc.PrimitiveContainersRequired{}},
		{Sample: tc.PrimitiveContainers{}},
		{Sample: td.DefaultPrimitiveTypedef{}},
		{Sample: td.Event{}},
		{Sample: td.I128{}},
		{Sample: td.Transition{}},
		{Sample: te.StructWithOptionalEnum{}},
		{Sample: tf.Cache_Clear_Args{}},
		{Sample: tf.Cache_ClearAfter_Args{}},
		{Sample: tf.ConflictingNames_SetValue_Args{}},
		{Sample: tf.ConflictingNames_SetValue_Result{}},
		{Sample: tf.ConflictingNamesSetValueArgs{}},
		{Sample: tf.InternalError{}},
		{Sample: tf.KeyValue_DeleteValue_Args{}},
		{Sample: tf.KeyValue_DeleteValue_Result{}},
		{Sample: tf.KeyValue_GetManyValues_Args{}},
		{Sample: tf.KeyValue_GetValue_Args{}},
		{Sample: tf.KeyValue_SetValue_Result{}},
		{Sample: tf.KeyValue_SetValueV2_Result{}},
		{Sample: tf.KeyValue_Size_Args{}},
		{Sample: tf.KeyValue_Size_Result{}},
		{Sample: tf.NonStandardServiceName_NonStandardFunctionName_Args{}},
		{Sample: tf.NonStandardServiceName_NonStandardFunctionName_Result{}},
		{Sample: ts.ContactInfo{}},
		{Sample: ts.DefaultsStruct{}},
		{Sample: ts.Edge{}},
		{Sample: ts.EmptyStruct{}},
		{Sample: ts.Frame{}},
		{Sample: ts.GoTags{}},
		{Sample: ts.Graph{}},
		{Sample: ts.Node{}},
		{Sample: ts.Omit{}},
		{Sample: ts.Point{}},
		{Sample: ts.PrimitiveOptionalStruct{}},
		{Sample: ts.PrimitiveRequiredStruct{}},
		{Sample: ts.Rename{}},
		{Sample: ts.Size{}},
		{Sample: ts.StructLabels{}},
		{Sample: ts.User{}},
		{Sample: ts.ZapOptOutStruct{}},
		{Sample: tu.Document{}},
		{Sample: tu.EmptyUnion{}},
		{Sample: tx.DoesNotExistException{}},
		{Sample: tx.EmptyException{}},

		// typedefs
		{Sample: td.BinarySet{}},
		{Sample: td.EdgeMap{}},
		{Sample: td.FrameGroup{}},
		{Sample: td.MyEnum(0)},
		{Sample: td.PDF{}, NoLog: true},
		{Sample: td.PointMap{}},
		{Sample: td.State(""), NoLog: true},
		{Sample: td.StateMap{}},
		{Sample: td.Timestamp(0), NoLog: true},
		{Sample: td.UUID{}},

		// enums
		{
			Sample:    te.EmptyEnum(0),
			Generator: enumValueGenerator(te.EmptyEnum_Values),
			JSON:      true,
			Text:      true,
		},
		{
			Sample:    te.EnumDefault(0),
			Generator: enumValueGenerator(te.EnumDefault_Values),
			JSON:      true,
			Text:      true,
		},
		{
			Sample:    te.EnumWithDuplicateName(0),
			Generator: enumValueGenerator(te.EnumWithDuplicateName_Values),
			JSON:      true,
			Text:      true,
		},
		{
			Sample:    te.EnumWithDuplicateValues(0),
			Generator: enumValueGenerator(te.EnumWithDuplicateValues_Values),
			JSON:      true,
			Text:      true,
		},
		{
			Sample:    te.EnumWithLabel(0),
			Generator: enumValueGenerator(te.EnumWithLabel_Values),
			JSON:      true,
			Text:      true,
		},
		{
			Sample:    te.EnumWithValues(0),
			Generator: enumValueGenerator(te.EnumWithValues_Values),
			JSON:      true,
			Text:      true,
		},
		{
			Sample:    te.LowerCaseEnum(0),
			Generator: enumValueGenerator(te.LowerCaseEnum_Values),
			JSON:      true,
			Text:      true,
		},
		{
			Sample:    te.RecordType(0),
			Generator: enumValueGenerator(te.RecordType_Values),
			JSON:      true,
			Text:      true,
		},
		{
			Sample:    te.RecordTypeValues(0),
			Generator: enumValueGenerator(te.RecordTypeValues_Values),
			JSON:      true,
			Text:      true,
		},
	}

	// Log the seed so that we can reproduce this if it ever fails.
	seed := time.Now().UnixNano()
	rand := rand.New(rand.NewSource(seed))
	t.Logf("Using seed %v for testing/quick", seed)

	const numValues = 1000 // number of values to test against
	for _, tt := range tests {
		typ := reflect.TypeOf(tt.Sample)
		t.Run(typ.Name(), func(t *testing.T) {
			generator := tt.Generator
			if generator == nil {
				generator = defaultValueGenerator(typ)
			}

			values := make([]thriftType, numValues)
			for i := range values {
				values[i] = generator(t, rand)
			}

			t.Run("Thrift", func(t *testing.T) {
				for _, give := range values {
					w, err := give.ToWire()
					require.NoError(t, err, "failed to Thrift encode %v", give)

					got := reflect.New(typ).Interface().(thriftType)
					require.NoError(t, got.FromWire(w), "failed to Thrift decode from %v", w)

					assert.Equal(t, got, give)
				}
			})

			t.Run("String", func(t *testing.T) {
				for _, give := range values {
					assert.NotPanics(t, func() {
						_ = give.String()
					}, "failed to String %#v", give)
				}
			})

			if tt.JSON {
				t.Run("JSON", func(t *testing.T) {
					for _, giveValue := range values {
						give, ok := giveValue.(json.Marshaler)
						require.True(t, ok, "Type does not implement json.Marshaler")

						bs, err := give.MarshalJSON()
						require.NoError(t, err, "failed to encode %v", give)

						got, ok := reflect.New(typ).Interface().(json.Unmarshaler)
						require.True(t, ok, "Type does not implement json.Unmarshaler")

						require.NoError(t, got.UnmarshalJSON(bs), "failed to decode from %q", bs)
						assert.Equal(t, got, give, "could not round-trip")
					}
				})
			}

			if tt.Text {
				t.Run("Text", func(t *testing.T) {
					for _, giveValue := range values {
						give, ok := giveValue.(encoding.TextMarshaler)
						require.True(t, ok, "Type does not implement encoding.TextMarshaler")

						bs, err := give.MarshalText()
						require.NoError(t, err, "failed to encode %v", give)

						got, ok := reflect.New(typ).Interface().(encoding.TextUnmarshaler)
						require.True(t, ok, "Type does not implement encoding.TextUnmarshaler")

						require.NoError(t, got.UnmarshalText(bs), "failed to decode from %q", bs)
						assert.Equal(t, got, give, "could not round-trip")
					}
				})
			}

			if !tt.NoLog {
				t.Run("Zap", func(t *testing.T) {
					for _, give := range values {
						assert.NotPanics(t, func() {
							enc := zapcore.NewMapObjectEncoder()

							if obj, ok := give.(zapcore.ObjectMarshaler); ok {
								assert.NoErrorf(t, obj.MarshalLogObject(enc), "failed to log %v", give)
								return
							}

							if arr, ok := give.(zapcore.ArrayMarshaler); ok {
								assert.NoErrorf(t, enc.AddArray("values", arr), "failed to log %v", give)
								return
							}

							t.Fatal(
								"Type does not implement zapcore.ObjectMarshaler or zapcore.ArrayMarshaler. "+
									"Did you mean to add NoLog?", typ)
						}, "failed to log %v", give)
					}
				})
			}
		})
	}
}
<|endoftext|>


<commit_before>package websocket

import (
	"net"

	"github.com/gorilla/websocket"
	"github.com/spring1843/chat-server/src/shared/errs"
	"github.com/spring1843/chat-server/src/shared/logs"
)

// ChatConnection is an middleman between the WebSocket connection and Chat Server
type ChatConnection struct {
	Connection *websocket.Conn
	Incoming   chan []byte
}

// NewChatConnection returns a new ChatConnection
func NewChatConnection() *ChatConnection {
	return &ChatConnection{
		Incoming: make(chan []byte),
	}
}

// Write to a ChatConnection
func (c *ChatConnection) Write(p []byte) (int, error) {
	if err := handleOutgoing(1, c, p); err != nil {
		return -1, errs.Wrap(err, "Error while trying to write to connection")
	}
	return len(p) - 1, nil
}

// Close a ChatConnection
func (c *ChatConnection) Close() error {
	if err := c.Connection.Close(); err != nil {
		return errs.Wrap(err, "Error closing WebSocket connection")
	}
	return nil
}

// RemoteAddr returns the remote address of the connected user
func (c *ChatConnection) RemoteAddr() net.Addr {
	return c.Connection.RemoteAddr()
}

// Read from a ChatConnection
// P is a buffered, write only from the start and maintain the size
func (c *ChatConnection) Read(p []byte) (int, error) {
	i := 0
	message := <-c.Incoming
	message = append(message, byte('\n'))

	if len(p) < len(message) {
		p = make([]byte, len(message), len(message))
	}

	for i, bit := range message {
		p[i] = bit
	}
	return i, nil
}

func handleOutgoing(msgType int, c *ChatConnection, message []byte) error {
	err := c.Connection.WriteMessage(msgType, message)
	if err != nil {
		return errs.Wrap(err, "Error handling Websocket Outgoging")
	}
	return nil
}

func listen(c *ChatConnection) {
	for {
		msgType, message, err := c.Connection.ReadMessage()
		if err != nil {
			logs.ErrIfErrf(err, "Error reading from WebSocket connection")
			break
		}
		if msgType == 1 {
			c.Incoming <- message
		}
	}
	logs.Infof("No longer listening to %s", c.RemoteAddr())
}
<commit_msg>Use next writer instead<commit_after>package websocket

import (
	"net"

	"github.com/gorilla/websocket"
	"github.com/spring1843/chat-server/src/shared/errs"
	"github.com/spring1843/chat-server/src/shared/logs"
)

// ChatConnection is an middleman between the WebSocket connection and Chat Server
type ChatConnection struct {
	Connection *websocket.Conn
	Incoming   chan []byte
}

// NewChatConnection returns a new ChatConnection
func NewChatConnection() *ChatConnection {
	return &ChatConnection{
		Incoming: make(chan []byte),
	}
}

// Write to a ChatConnection
func (c *ChatConnection) Write(p []byte) (int, error) {
	w, err := c.Connection.NextWriter(websocket.TextMessage)
	if err != nil {
		return 0, errs.Wrap(err, "Error getting nextwriter from WebSocket connection.")
	}
	defer w.Close()
	return w.Write(p)
}

// Close a ChatConnection
func (c *ChatConnection) Close() error {
	if err := c.Connection.Close(); err != nil {
		return errs.Wrap(err, "Error closing WebSocket connection")
	}
	return nil
}

// RemoteAddr returns the remote address of the connected user
func (c *ChatConnection) RemoteAddr() net.Addr {
	return c.Connection.RemoteAddr()
}

// Read from a ChatConnection
// P is a buffered, write only from the start and maintain the size
func (c *ChatConnection) Read(p []byte) (int, error) {
	i := 0
	message := <-c.Incoming
	message = append(message, byte('\n'))

	if len(p) < len(message) {
		p = make([]byte, len(message), len(message))
	}

	for i, bit := range message {
		p[i] = bit
	}
	return i, nil
}

func listen(c *ChatConnection) {
	for {
		msgType, message, err := c.Connection.ReadMessage()
		if err != nil {
			logs.ErrIfErrf(err, "Error reading from WebSocket connection")
			break
		}
		if msgType == 1 {
			c.Incoming <- message
		}
	}
	logs.Infof("No longer listening to %s", c.RemoteAddr())
}
<|endoftext|>


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

package openstack

import (
	"fmt"
	"net/url"

	"github.com/juju/schema"
	"gopkg.in/goose.v1/identity"
	"gopkg.in/juju/environschema.v1"

	"github.com/juju/juju/environs/config"
)

var configSchema = environschema.Fields{
	"username": {
		Description: "The user name  to use when auth-mode is userpass.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvUser,
		Group:       environschema.AccountGroup,
	},
	"password": {
		Description: "The password to use when auth-mode is userpass.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvSecrets,
		Group:       environschema.AccountGroup,
	},
	"tenant-name": {
		Description: "The openstack tenant name.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvTenantName,
		Group:       environschema.AccountGroup,
	},
	"auth-url": {
		Description: "The keystone URL for authentication.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvAuthURL,
		Example:     "https://yourkeystoneurl:443/v2.0/",
		Group:       environschema.AccountGroup,
	},
	"auth-mode": {
		Description: "The authentication mode to use. When set to keypair, the access-key and secret-key parameters should be set; when set to userpass or legacy, the username and password parameters should be set.",
		Type:        environschema.Tstring,
		Values:      []interface{}{AuthKeyPair, AuthLegacy, AuthUserPass},
		Group:       environschema.AccountGroup,
	},
	"access-key": {
		Description: "The access key to use when auth-mode is set to keypair.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvUser,
		Group:       environschema.AccountGroup,
	},
	"secret-key": {
		Description: "The secret key to use when auth-mode is set to keypair.",
		EnvVars:     identity.CredEnvSecrets,
		Group:       environschema.AccountGroup,
		Type:        environschema.Tstring,
	},
	"region": {
		Description: "The openstack region.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvRegion,
	},
	"control-bucket": {
		Description: "The name to use for the control bucket (do not set unless you know what you are doing!).",
		Type:        environschema.Tstring,
	},
	"use-floating-ip": {
		Description: "Whether a floating IP address is required to give the nodes a public IP address. Some installations assign public IP addresses by default without requiring a floating IP address.",
		Type:        environschema.Tbool,
	},
	"use-default-secgroup": {
		Description: `Whether new machine instances should have the "default" Openstack security group assigned.`,
		Type:        environschema.Tbool,
	},
	"network": {
		Description: "The network label or UUID to bring machines up on when multiple networks exist.",
		Type:        environschema.Tstring,
	},
}

var configFields = func() schema.Fields {
	fs, _, err := configSchema.ValidationSchema()
	if err != nil {
		panic(err)
	}
	return fs
}()

var configDefaults = schema.Defaults{
	"username":             "",
	"password":             "",
	"tenant-name":          "",
	"auth-url":             "",
	"auth-mode":            string(AuthUserPass),
	"access-key":           "",
	"secret-key":           "",
	"region":               "",
	"control-bucket":       "",
	"use-floating-ip":      false,
	"use-default-secgroup": false,
	"network":              "",
}

type environConfig struct {
	*config.Config
	attrs map[string]interface{}
}

func (c *environConfig) region() string {
	return c.attrs["region"].(string)
}

func (c *environConfig) username() string {
	return c.attrs["username"].(string)
}

func (c *environConfig) password() string {
	return c.attrs["password"].(string)
}

func (c *environConfig) tenantName() string {
	return c.attrs["tenant-name"].(string)
}

func (c *environConfig) authURL() string {
	return c.attrs["auth-url"].(string)
}

func (c *environConfig) authMode() AuthMode {
	return AuthMode(c.attrs["auth-mode"].(string))
}

func (c *environConfig) accessKey() string {
	return c.attrs["access-key"].(string)
}

func (c *environConfig) secretKey() string {
	return c.attrs["secret-key"].(string)
}

func (c *environConfig) controlBucket() string {
	return c.attrs["control-bucket"].(string)
}

func (c *environConfig) useFloatingIP() bool {
	return c.attrs["use-floating-ip"].(bool)
}

func (c *environConfig) useDefaultSecurityGroup() bool {
	return c.attrs["use-default-secgroup"].(bool)
}

func (c *environConfig) network() string {
	return c.attrs["network"].(string)
}

func (p environProvider) newConfig(cfg *config.Config) (*environConfig, error) {
	valid, err := p.Validate(cfg, nil)
	if err != nil {
		return nil, err
	}
	return &environConfig{valid, valid.UnknownAttrs()}, nil
}

type AuthMode string

const (
	AuthKeyPair  AuthMode = "keypair"
	AuthLegacy   AuthMode = "legacy"
	AuthUserPass AuthMode = "userpass"
)

// Schema returns the configuration schema for an environment.
func (environProvider) Schema() environschema.Fields {
	fields, err := config.Schema(configSchema)
	if err != nil {
		panic(err)
	}
	return fields
}

func (p environProvider) Validate(cfg, old *config.Config) (valid *config.Config, err error) {
	// Check for valid changes for the base config values.
	if err := config.Validate(cfg, old); err != nil {
		return nil, err
	}

	validated, err := cfg.ValidateUnknownAttrs(configFields, configDefaults)
	if err != nil {
		return nil, err
	}

	// Add Openstack specific defaults.
	providerDefaults := make(map[string]interface{})

	// Storage.
	if _, ok := cfg.StorageDefaultBlockSource(); !ok {
		providerDefaults[config.StorageDefaultBlockSourceKey] = CinderProviderType
	}
	if len(providerDefaults) > 0 {
		if cfg, err = cfg.Apply(providerDefaults); err != nil {
			return nil, err
		}
	}

	ecfg := &environConfig{cfg, validated}

	if ecfg.authURL() != "" {
		parts, err := url.Parse(ecfg.authURL())
		if err != nil || parts.Host == "" || parts.Scheme == "" {
			return nil, fmt.Errorf("invalid auth-url value %q", ecfg.authURL())
		}
	}
	cred := identity.CredentialsFromEnv()
	format := "required environment variable not set for credentials attribute: %s"
	switch ecfg.authMode() {
	case AuthUserPass, AuthLegacy:
		if ecfg.username() == "" {
			if cred.User == "" {
				return nil, fmt.Errorf(format, "User")
			}
			ecfg.attrs["username"] = cred.User
		}
		if ecfg.password() == "" {
			if cred.Secrets == "" {
				return nil, fmt.Errorf(format, "Secrets")
			}
			ecfg.attrs["password"] = cred.Secrets
		}
	case AuthKeyPair:
		if ecfg.accessKey() == "" {
			if cred.User == "" {
				return nil, fmt.Errorf(format, "User")
			}
			ecfg.attrs["access-key"] = cred.User
		}
		if ecfg.secretKey() == "" {
			if cred.Secrets == "" {
				return nil, fmt.Errorf(format, "Secrets")
			}
			ecfg.attrs["secret-key"] = cred.Secrets
		}
	default:
		return nil, fmt.Errorf("unexpected authentication mode %q", ecfg.authMode())
	}
	if ecfg.authURL() == "" {
		if cred.URL == "" {
			return nil, fmt.Errorf(format, "URL")
		}
		ecfg.attrs["auth-url"] = cred.URL
	}
	if ecfg.tenantName() == "" {
		if cred.TenantName == "" {
			return nil, fmt.Errorf(format, "TenantName")
		}
		ecfg.attrs["tenant-name"] = cred.TenantName
	}
	if ecfg.region() == "" {
		if cred.Region == "" {
			return nil, fmt.Errorf(format, "Region")
		}
		ecfg.attrs["region"] = cred.Region
	}

	if old != nil {
		attrs := old.UnknownAttrs()
		if region, _ := attrs["region"].(string); ecfg.region() != region {
			return nil, fmt.Errorf("cannot change region from %q to %q", region, ecfg.region())
		}
		if controlBucket, _ := attrs["control-bucket"].(string); ecfg.controlBucket() != controlBucket {
			return nil, fmt.Errorf("cannot change control-bucket from %q to %q", controlBucket, ecfg.controlBucket())
		}
	}

	// Check for deprecated fields and log a warning. We also print to stderr to ensure the user sees the message
	// even if they are not running with --debug.
	cfgAttrs := cfg.AllAttrs()
	if defaultImageId := cfgAttrs["default-image-id"]; defaultImageId != nil && defaultImageId.(string) != "" {
		msg := fmt.Sprintf(
			"Config attribute %q (%v) is deprecated and ignored.\n"+
				"Your cloud provider should have set up image metadata to provide the correct image id\n"+
				"for your chosen series and archietcure. If this is a private Openstack deployment without\n"+
				"existing image metadata, please run 'juju-metadata help' to see how suitable image"+
				"metadata can be generated.",
			"default-image-id", defaultImageId)
		logger.Warningf(msg)
	}
	if defaultInstanceType := cfgAttrs["default-instance-type"]; defaultInstanceType != nil && defaultInstanceType.(string) != "" {
		msg := fmt.Sprintf(
			"Config attribute %q (%v) is deprecated and ignored.\n"+
				"The correct instance flavor is determined using constraints, globally specified\n"+
				"when an environment is bootstrapped, or individually when a charm is deployed.\n"+
				"See 'juju help bootstrap' or 'juju help deploy'.",
			"default-instance-type", defaultInstanceType)
		logger.Warningf(msg)
	}
	// Construct a new config with the deprecated attributes removed.
	for _, attr := range []string{"default-image-id", "default-instance-type"} {
		delete(cfgAttrs, attr)
		delete(ecfg.attrs, attr)
	}
	for k, v := range ecfg.attrs {
		cfgAttrs[k] = v
	}
	return config.New(config.NoDefaults, cfgAttrs)
}
<commit_msg>provider/openstack: add appropriate secret attributes to config<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package openstack

import (
	"fmt"
	"net/url"

	"github.com/juju/schema"
	"gopkg.in/goose.v1/identity"
	"gopkg.in/juju/environschema.v1"

	"github.com/juju/juju/environs/config"
)

var configSchema = environschema.Fields{
	"username": {
		Description: "The user name  to use when auth-mode is userpass.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvUser,
		Group:       environschema.AccountGroup,
	},
	"password": {
		Description: "The password to use when auth-mode is userpass.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvSecrets,
		Group:       environschema.AccountGroup,
		Secret:      true,
	},
	"tenant-name": {
		Description: "The openstack tenant name.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvTenantName,
		Group:       environschema.AccountGroup,
	},
	"auth-url": {
		Description: "The keystone URL for authentication.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvAuthURL,
		Example:     "https://yourkeystoneurl:443/v2.0/",
		Group:       environschema.AccountGroup,
	},
	"auth-mode": {
		Description: "The authentication mode to use. When set to keypair, the access-key and secret-key parameters should be set; when set to userpass or legacy, the username and password parameters should be set.",
		Type:        environschema.Tstring,
		Values:      []interface{}{AuthKeyPair, AuthLegacy, AuthUserPass},
		Group:       environschema.AccountGroup,
	},
	"access-key": {
		Description: "The access key to use when auth-mode is set to keypair.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvUser,
		Group:       environschema.AccountGroup,
		Secret:      true,
	},
	"secret-key": {
		Description: "The secret key to use when auth-mode is set to keypair.",
		EnvVars:     identity.CredEnvSecrets,
		Group:       environschema.AccountGroup,
		Type:        environschema.Tstring,
		Secret:      true,
	},
	"region": {
		Description: "The openstack region.",
		Type:        environschema.Tstring,
		EnvVars:     identity.CredEnvRegion,
	},
	"control-bucket": {
		Description: "The name to use for the control bucket (do not set unless you know what you are doing!).",
		Type:        environschema.Tstring,
	},
	"use-floating-ip": {
		Description: "Whether a floating IP address is required to give the nodes a public IP address. Some installations assign public IP addresses by default without requiring a floating IP address.",
		Type:        environschema.Tbool,
	},
	"use-default-secgroup": {
		Description: `Whether new machine instances should have the "default" Openstack security group assigned.`,
		Type:        environschema.Tbool,
	},
	"network": {
		Description: "The network label or UUID to bring machines up on when multiple networks exist.",
		Type:        environschema.Tstring,
	},
}

var configFields = func() schema.Fields {
	fs, _, err := configSchema.ValidationSchema()
	if err != nil {
		panic(err)
	}
	return fs
}()

var configDefaults = schema.Defaults{
	"username":             "",
	"password":             "",
	"tenant-name":          "",
	"auth-url":             "",
	"auth-mode":            string(AuthUserPass),
	"access-key":           "",
	"secret-key":           "",
	"region":               "",
	"control-bucket":       "",
	"use-floating-ip":      false,
	"use-default-secgroup": false,
	"network":              "",
}

type environConfig struct {
	*config.Config
	attrs map[string]interface{}
}

func (c *environConfig) region() string {
	return c.attrs["region"].(string)
}

func (c *environConfig) username() string {
	return c.attrs["username"].(string)
}

func (c *environConfig) password() string {
	return c.attrs["password"].(string)
}

func (c *environConfig) tenantName() string {
	return c.attrs["tenant-name"].(string)
}

func (c *environConfig) authURL() string {
	return c.attrs["auth-url"].(string)
}

func (c *environConfig) authMode() AuthMode {
	return AuthMode(c.attrs["auth-mode"].(string))
}

func (c *environConfig) accessKey() string {
	return c.attrs["access-key"].(string)
}

func (c *environConfig) secretKey() string {
	return c.attrs["secret-key"].(string)
}

func (c *environConfig) controlBucket() string {
	return c.attrs["control-bucket"].(string)
}

func (c *environConfig) useFloatingIP() bool {
	return c.attrs["use-floating-ip"].(bool)
}

func (c *environConfig) useDefaultSecurityGroup() bool {
	return c.attrs["use-default-secgroup"].(bool)
}

func (c *environConfig) network() string {
	return c.attrs["network"].(string)
}

func (p environProvider) newConfig(cfg *config.Config) (*environConfig, error) {
	valid, err := p.Validate(cfg, nil)
	if err != nil {
		return nil, err
	}
	return &environConfig{valid, valid.UnknownAttrs()}, nil
}

type AuthMode string

const (
	AuthKeyPair  AuthMode = "keypair"
	AuthLegacy   AuthMode = "legacy"
	AuthUserPass AuthMode = "userpass"
)

// Schema returns the configuration schema for an environment.
func (environProvider) Schema() environschema.Fields {
	fields, err := config.Schema(configSchema)
	if err != nil {
		panic(err)
	}
	return fields
}

func (p environProvider) Validate(cfg, old *config.Config) (valid *config.Config, err error) {
	// Check for valid changes for the base config values.
	if err := config.Validate(cfg, old); err != nil {
		return nil, err
	}

	validated, err := cfg.ValidateUnknownAttrs(configFields, configDefaults)
	if err != nil {
		return nil, err
	}

	// Add Openstack specific defaults.
	providerDefaults := make(map[string]interface{})

	// Storage.
	if _, ok := cfg.StorageDefaultBlockSource(); !ok {
		providerDefaults[config.StorageDefaultBlockSourceKey] = CinderProviderType
	}
	if len(providerDefaults) > 0 {
		if cfg, err = cfg.Apply(providerDefaults); err != nil {
			return nil, err
		}
	}

	ecfg := &environConfig{cfg, validated}

	if ecfg.authURL() != "" {
		parts, err := url.Parse(ecfg.authURL())
		if err != nil || parts.Host == "" || parts.Scheme == "" {
			return nil, fmt.Errorf("invalid auth-url value %q", ecfg.authURL())
		}
	}
	cred := identity.CredentialsFromEnv()
	format := "required environment variable not set for credentials attribute: %s"
	switch ecfg.authMode() {
	case AuthUserPass, AuthLegacy:
		if ecfg.username() == "" {
			if cred.User == "" {
				return nil, fmt.Errorf(format, "User")
			}
			ecfg.attrs["username"] = cred.User
		}
		if ecfg.password() == "" {
			if cred.Secrets == "" {
				return nil, fmt.Errorf(format, "Secrets")
			}
			ecfg.attrs["password"] = cred.Secrets
		}
	case AuthKeyPair:
		if ecfg.accessKey() == "" {
			if cred.User == "" {
				return nil, fmt.Errorf(format, "User")
			}
			ecfg.attrs["access-key"] = cred.User
		}
		if ecfg.secretKey() == "" {
			if cred.Secrets == "" {
				return nil, fmt.Errorf(format, "Secrets")
			}
			ecfg.attrs["secret-key"] = cred.Secrets
		}
	default:
		return nil, fmt.Errorf("unexpected authentication mode %q", ecfg.authMode())
	}
	if ecfg.authURL() == "" {
		if cred.URL == "" {
			return nil, fmt.Errorf(format, "URL")
		}
		ecfg.attrs["auth-url"] = cred.URL
	}
	if ecfg.tenantName() == "" {
		if cred.TenantName == "" {
			return nil, fmt.Errorf(format, "TenantName")
		}
		ecfg.attrs["tenant-name"] = cred.TenantName
	}
	if ecfg.region() == "" {
		if cred.Region == "" {
			return nil, fmt.Errorf(format, "Region")
		}
		ecfg.attrs["region"] = cred.Region
	}

	if old != nil {
		attrs := old.UnknownAttrs()
		if region, _ := attrs["region"].(string); ecfg.region() != region {
			return nil, fmt.Errorf("cannot change region from %q to %q", region, ecfg.region())
		}
		if controlBucket, _ := attrs["control-bucket"].(string); ecfg.controlBucket() != controlBucket {
			return nil, fmt.Errorf("cannot change control-bucket from %q to %q", controlBucket, ecfg.controlBucket())
		}
	}

	// Check for deprecated fields and log a warning. We also print to stderr to ensure the user sees the message
	// even if they are not running with --debug.
	cfgAttrs := cfg.AllAttrs()
	if defaultImageId := cfgAttrs["default-image-id"]; defaultImageId != nil && defaultImageId.(string) != "" {
		msg := fmt.Sprintf(
			"Config attribute %q (%v) is deprecated and ignored.\n"+
				"Your cloud provider should have set up image metadata to provide the correct image id\n"+
				"for your chosen series and archietcure. If this is a private Openstack deployment without\n"+
				"existing image metadata, please run 'juju-metadata help' to see how suitable image"+
				"metadata can be generated.",
			"default-image-id", defaultImageId)
		logger.Warningf(msg)
	}
	if defaultInstanceType := cfgAttrs["default-instance-type"]; defaultInstanceType != nil && defaultInstanceType.(string) != "" {
		msg := fmt.Sprintf(
			"Config attribute %q (%v) is deprecated and ignored.\n"+
				"The correct instance flavor is determined using constraints, globally specified\n"+
				"when an environment is bootstrapped, or individually when a charm is deployed.\n"+
				"See 'juju help bootstrap' or 'juju help deploy'.",
			"default-instance-type", defaultInstanceType)
		logger.Warningf(msg)
	}
	// Construct a new config with the deprecated attributes removed.
	for _, attr := range []string{"default-image-id", "default-instance-type"} {
		delete(cfgAttrs, attr)
		delete(ecfg.attrs, attr)
	}
	for k, v := range ecfg.attrs {
		cfgAttrs[k] = v
	}
	return config.New(config.NoDefaults, cfgAttrs)
}
<|endoftext|>


<commit_before>package smart

import (
	"io"
	"net/url"
)

type TransportProgressCallback func(bytesDone, totalBytes int64)

// The transport interface abstracts away how the smart provider talks to the server
// It might do this over a persistent SSH connection, sending data across in/out streams,
// or it might process each request as a discrete request/response pair over REST
// Note each transport instance is stateful and associated with a server/connection, see SmartTransportFactory for how they are created
type Transport interface {
	// Release any resources associated with this transport (including any persostent connections)
	Release()
	// Ask the server for a list of capabilities
	QueryCaps() ([]string, error)
	// Request that the server enable capabilities for this exchange (note, non-persistent transports can store & send this with every request)
	SetEnabledCaps(caps []string) error

	// Return whether LOB metadata exists on the server (also returns size)
	MetadataExists(lobsha string) (ex bool, sz int64, e error)
	// Return whether LOB chunk content exists on the server
	ChunkExists(lobsha string, chunk int) (ex bool, sz int64, e error)
	// Return whether LOB chunk content exists on the server, and is of a specific size
	ChunkExistsAndIsOfSize(lobsha string, chunk int, sz int64) (bool, error)
	// Entire LOB exists? Also returns entire content size
	LOBExists(lobsha string) (ex bool, sz int64, e error)

	// Upload metadata for a LOB (from a stream); no progress callback as very small
	UploadMetadata(lobsha string, sz int64, data io.Reader) error
	// Upload chunk content for a LOB (from a stream); must call back progress
	UploadChunk(lobsha string, chunk int, sz int64, data io.Reader, callback TransportProgressCallback) error
	// Download metadata for a LOB (to a stream); no progress callback as very small
	DownloadMetadata(lobsha string, out io.Writer) error
	// Download chunk content for a LOB (from a stream); must call back progress
	// This is a non-delta download operation, just provide entire chunk content
	DownloadChunk(lobsha string, chunk int, out io.Writer, callback TransportProgressCallback) error

	// Return the LOB which the server has a complete copy of, from a list of candidates
	// Server must test in the order provided & return the earliest one which is complete on the server
	// Server doesn't have to test full integrity of LOB, just completeness (check size against meta)
	// Return a blank string if none are available
	GetFirstCompleteLOBFromList(candidateSHAs []string) (string, error)
	// Upload a binary delta to apply against a LOB the server already has, to generate a new LOB
	// Deltas apply to whole LOB content and are not per-chunk
	// Returns a boolean to determine whether the upload was accepted or not (server may prefer not to accept, not an error)
	// In the case of false return, client will fall back to non-delta upload.
	// On true, server must return nil error only after data is fully received, applied, saved as targetSHA and the
	// integrity confirmed by recalculating the SHA of the final patched data.
	UploadDelta(baseSHA, targetSHA string, deltaSize int64, data io.Reader, callback TransportProgressCallback) (bool, error)
	// Prepare a binary delta between 2 LOBs and report the size
	DownloadDeltaPrepare(baseSHA, targetSHA string) (int64, error)
	// Generate (if not already cached) and download a binary delta that the client can apply locally to generate a new LOB
	// Deltas apply to whole LOB content and are not per-chunk
	// The server should respect sizeLimit and if the delta is larger than that, abandon the process
	// Return a bool to indicate whether the delta went ahead or not (client will fall back to non-delta on false)
	DownloadDelta(baseSHA, targetSHA string, sizeLimit int64, out io.Writer, callback TransportProgressCallback) (bool, error)
}

// Interface for a factory which creates persistent transports for use by SmartSyncProvider
type TransportFactory interface {
	// Does this factory want to handle the URL passed in?
	WillHandleUrl(u *url.URL) bool
	// Provide a new, connected (may not be persistent, but if not test connection/auth) transport for given URL
	Connect(u *url.URL) (Transport, error)
}

var (
	transportFactories []TransportFactory
)

// Registers an instance of a SmartTransportFactory for creating connections
// Must only be called from the main thread, not thread safe
// Later factories registered will take precedence over earlier ones (including core)
func RegisterTransportFactory(f TransportFactory) {
	transportFactories = append(transportFactories, f)
}

// Retrieve the best ConnectionFactory for a given URL (or nil)
func GetTransportFactory(u *url.URL) TransportFactory {
	// Iterate in reverse order
	for i := len(transportFactories) - 1; i > 0; i-- {
		if transportFactories[i].WillHandleUrl(u) {
			return transportFactories[i]
		}
	}
	return nil
}
<commit_msg>Fix off by one error which stopped one of the smart transports being found for a URL<commit_after>package smart

import (
	"io"
	"net/url"
)

type TransportProgressCallback func(bytesDone, totalBytes int64)

// The transport interface abstracts away how the smart provider talks to the server
// It might do this over a persistent SSH connection, sending data across in/out streams,
// or it might process each request as a discrete request/response pair over REST
// Note each transport instance is stateful and associated with a server/connection, see SmartTransportFactory for how they are created
type Transport interface {
	// Release any resources associated with this transport (including any persostent connections)
	Release()
	// Ask the server for a list of capabilities
	QueryCaps() ([]string, error)
	// Request that the server enable capabilities for this exchange (note, non-persistent transports can store & send this with every request)
	SetEnabledCaps(caps []string) error

	// Return whether LOB metadata exists on the server (also returns size)
	MetadataExists(lobsha string) (ex bool, sz int64, e error)
	// Return whether LOB chunk content exists on the server
	ChunkExists(lobsha string, chunk int) (ex bool, sz int64, e error)
	// Return whether LOB chunk content exists on the server, and is of a specific size
	ChunkExistsAndIsOfSize(lobsha string, chunk int, sz int64) (bool, error)
	// Entire LOB exists? Also returns entire content size
	LOBExists(lobsha string) (ex bool, sz int64, e error)

	// Upload metadata for a LOB (from a stream); no progress callback as very small
	UploadMetadata(lobsha string, sz int64, data io.Reader) error
	// Upload chunk content for a LOB (from a stream); must call back progress
	UploadChunk(lobsha string, chunk int, sz int64, data io.Reader, callback TransportProgressCallback) error
	// Download metadata for a LOB (to a stream); no progress callback as very small
	DownloadMetadata(lobsha string, out io.Writer) error
	// Download chunk content for a LOB (from a stream); must call back progress
	// This is a non-delta download operation, just provide entire chunk content
	DownloadChunk(lobsha string, chunk int, out io.Writer, callback TransportProgressCallback) error

	// Return the LOB which the server has a complete copy of, from a list of candidates
	// Server must test in the order provided & return the earliest one which is complete on the server
	// Server doesn't have to test full integrity of LOB, just completeness (check size against meta)
	// Return a blank string if none are available
	GetFirstCompleteLOBFromList(candidateSHAs []string) (string, error)
	// Upload a binary delta to apply against a LOB the server already has, to generate a new LOB
	// Deltas apply to whole LOB content and are not per-chunk
	// Returns a boolean to determine whether the upload was accepted or not (server may prefer not to accept, not an error)
	// In the case of false return, client will fall back to non-delta upload.
	// On true, server must return nil error only after data is fully received, applied, saved as targetSHA and the
	// integrity confirmed by recalculating the SHA of the final patched data.
	UploadDelta(baseSHA, targetSHA string, deltaSize int64, data io.Reader, callback TransportProgressCallback) (bool, error)
	// Prepare a binary delta between 2 LOBs and report the size
	DownloadDeltaPrepare(baseSHA, targetSHA string) (int64, error)
	// Generate (if not already cached) and download a binary delta that the client can apply locally to generate a new LOB
	// Deltas apply to whole LOB content and are not per-chunk
	// The server should respect sizeLimit and if the delta is larger than that, abandon the process
	// Return a bool to indicate whether the delta went ahead or not (client will fall back to non-delta on false)
	DownloadDelta(baseSHA, targetSHA string, sizeLimit int64, out io.Writer, callback TransportProgressCallback) (bool, error)
}

// Interface for a factory which creates persistent transports for use by SmartSyncProvider
type TransportFactory interface {
	// Does this factory want to handle the URL passed in?
	WillHandleUrl(u *url.URL) bool
	// Provide a new, connected (may not be persistent, but if not test connection/auth) transport for given URL
	Connect(u *url.URL) (Transport, error)
}

var (
	transportFactories []TransportFactory
)

// Registers an instance of a SmartTransportFactory for creating connections
// Must only be called from the main thread, not thread safe
// Later factories registered will take precedence over earlier ones (including core)
func RegisterTransportFactory(f TransportFactory) {
	transportFactories = append(transportFactories, f)
}

// Retrieve the best ConnectionFactory for a given URL (or nil)
func GetTransportFactory(u *url.URL) TransportFactory {
	// Iterate in reverse order
	for i := len(transportFactories) - 1; i >= 0; i-- {
		if transportFactories[i].WillHandleUrl(u) {
			return transportFactories[i]
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package blob

import (
)

// Return a blob store that wraps the supplied one, confirming that the blob
// contents and scores it returns are correct, guarding against silent data
// corruption.
func NewCheckingStore(wrapped Store) Store {
	// TODO
	return nil
}
<commit_msg>Added a stub implementation.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package blob

import (
	"fmt"
)

// Return a blob store that wraps the supplied one, confirming that the blob
// contents and scores it returns are correct, guarding against silent data
// corruption.
func NewCheckingStore(wrapped Store) Store {
	return &checkingStore{wrapped}
}

type checkingStore struct {
	wrapped Store
}

func (s *checkingStore) Store(blob []byte) (score Score, err error) {
	return nil, fmt.Errorf("TODO")
}

func (s *checkingStore) Load(score Score) (blob []byte, err error) {
	return nil, fmt.Errorf("TODO")
}
<|endoftext|>


<commit_before>// Use and distribution licensed under the Apache license version 2.
//
// See the COPYING file in the root project directory for full text.
//

package ghw

import (
	"strings"

	"github.com/StackExchange/wmi"
)

const wqlDiskDrive = "SELECT Caption, CreationClassName, Description, DeviceID, Index, InterfaceType, Manufacturer, MediaType, Model, Name, Partitions, SerialNumber, Size, TotalCylinders, TotalHeads, TotalSectors, TotalTracks, TracksPerCylinder FROM Win32_DiskDrive"

type win32DiskDrive struct {
	Caption           string
	CreationClassName string
	Description       string
	DeviceID          string
	Index             uint32 // Used to link with partition
	InterfaceType     string
	Manufacturer      string
	MediaType         string
	Model             string
	Name              string
	Partitions        int32
	SerialNumber      string
	Size              uint64
	TotalCylinders    int64
	TotalHeads        int32
	TotalSectors      int64
	TotalTracks       int64
	TracksPerCylinder int32
}

const wqlDiskPartition = "SELECT Access, BlockSize, Caption, CreationClassName, Description, DeviceID, DiskIndex, Index, Name, Size, SystemName, Type FROM Win32_DiskPartition"

type win32DiskPartition struct {
	Access            uint16
	BlockSize         uint64
	Caption           string
	CreationClassName string
	Description       string
	DeviceID          string
	DiskIndex         uint32 // Used to link with Disk Drive
	Index             uint32
	Name              string
	Size              int64
	SystemName        string
	Type              string
}

const wqlLogicalDiskToPartition = "SELECT Antecedent, Dependent FROM Win32_LogicalDiskToPartition"

type win32LogicalDiskToPartition struct {
	Antecedent string
	Dependent  string
}

const wqlLogicalDisk = "SELECT Caption, CreationClassName, Description, DeviceID, FileSystem, FreeSpace, Name, Size, SystemName FROM Win32_LogicalDisk"

type win32LogicalDisk struct {
	Caption           string
	CreationClassName string
	Description       string
	DeviceID          string
	FileSystem        string
	FreeSpace         uint64
	Name              string
	Size              uint64
	SystemName        string
}

func (ctx *context) blockFillInfo(info *BlockInfo) error {
	win32DiskDriveDescriptions, err := getDiskDrives()
	if err != nil {
		return err
	}

	win32DiskPartitionDescriptions, err := getDiskPartitions()
	if err != nil {
		return err
	}

	win32LogicalDiskToPartitionDescriptions, err := getLogicalDisksToPartitions()
	if err != nil {
		return err
	}

	win32LogicalDiskDescriptions, err := getLogicalDisks()
	if err != nil {
		return err
	}

	// Converting into standard structures
	disks := make([]*Disk, 0)
	for _, diskdrive := range win32DiskDriveDescriptions {
		disk := &Disk{
			Name:                   diskdrive.Name,
			SizeBytes:              diskdrive.Size,
			PhysicalBlockSizeBytes: 0,
			DriveType:              toDriveType(diskdrive.MediaType),
			StorageController:      toStorageController(diskdrive.InterfaceType),
			BusType:                toBusType(diskdrive.InterfaceType),
			BusPath:                UNKNOWN, // TODO: add information
			Vendor:                 UNKNOWN, // TODO: add information
			Model:                  diskdrive.Caption,
			SerialNumber:           diskdrive.SerialNumber,
			WWN:                    UNKNOWN, // TODO: add information
			Partitions:             make([]*Partition, 0),
		}
		for _, diskpartition := range win32DiskPartitionDescriptions {
			// Finding disk partition linked to current disk drive
			if diskdrive.Index == diskpartition.DiskIndex {
				disk.PhysicalBlockSizeBytes = diskpartition.BlockSize
				// Finding logical partition linked to current disk partition
				for _, logicaldisk := range win32LogicalDiskDescriptions {
					for _, logicaldisktodiskpartition := range win32LogicalDiskToPartitionDescriptions {
						var desiredAntecedent = "\\\\" + diskpartition.SystemName + "\\root\\cimv2:" + diskpartition.CreationClassName + ".DeviceID=\"" + diskpartition.DeviceID + "\""
						var desiredDependent = "\\\\" + logicaldisk.SystemName + "\\root\\cimv2:" + logicaldisk.CreationClassName + ".DeviceID=\"" + logicaldisk.DeviceID + "\""
						if logicaldisktodiskpartition.Antecedent == desiredAntecedent && logicaldisktodiskpartition.Dependent == desiredDependent {
							// Appending Partition
							p := &Partition{
								Name:       logicaldisk.Caption,
								Label:      logicaldisk.Caption,
								SizeBytes:  logicaldisk.Size,
								MountPoint: logicaldisk.DeviceID,
								Type:       diskpartition.Type,
								IsReadOnly: toReadOnly(diskpartition.Access), // TODO: add information
							}
							disk.Partitions = append(disk.Partitions, p)
							break
						}
					}
				}
			}
		}
		disks = append(disks, disk)
	}

	info.Disks = disks
	var tpb uint64
	for _, d := range info.Disks {
		tpb += d.SizeBytes
	}
	info.TotalPhysicalBytes = tpb
	return nil
}

func getDiskDrives() ([]win32DiskDrive, error) {
	// Getting disks drives data from WMI
	var win3232DiskDriveDescriptions []win32DiskDrive
	if err := wmi.Query(wqlDiskDrive, &win3232DiskDriveDescriptions); err != nil {
		return nil, err
	}
	return win3232DiskDriveDescriptions, nil
}

func getDiskPartitions() ([]win32DiskPartition, error) {
	// Getting disk partitions from WMI
	var win32DiskPartitionDescriptions []win32DiskPartition
	if err := wmi.Query(wqlDiskPartition, &win32DiskPartitionDescriptions); err != nil {
		return nil, err
	}
	return win32DiskPartitionDescriptions, nil
}

func getLogicalDisksToPartitions() ([]win32LogicalDiskToPartition, error) {
	// Getting links between logical disks and partitions from WMI
	var win32LogicalDiskToPartitionDescriptions []win32LogicalDiskToPartition
	if err := wmi.Query(wqlLogicalDiskToPartition, &win32LogicalDiskToPartitionDescriptions); err != nil {
		return nil, err
	}
	return win32LogicalDiskToPartitionDescriptions, nil
}

func getLogicalDisks() ([]win32LogicalDisk, error) {
	// Getting logical disks from WMI
	var win32LogicalDiskDescriptions []win32LogicalDisk
	if err := wmi.Query(wqlLogicalDisk, &win32LogicalDiskDescriptions); err != nil {
		return nil, err
	}
	return win32LogicalDiskDescriptions, nil
}

// TODO: improve
func toDriveType(mediaType string) DriveType {
	var driveType DriveType
	mediaType = strings.ToLower(mediaType)
	if strings.Contains(mediaType, "fixed") || strings.Contains(mediaType, "ssd") {
		driveType = DRIVE_TYPE_SSD
	} else if strings.ContainsAny(mediaType, "hdd") {
		driveType = DRIVE_TYPE_HDD
	} else {
		driveType = DRIVE_TYPE_UNKNOWN
	}
	return driveType
}

// TODO: improve
func toStorageController(interfaceType string) StorageController {
	var storageController StorageController
	switch interfaceType {
	case "SCSI":
		storageController = STORAGE_CONTROLLER_SCSI
	case "IDE":
		storageController = STORAGE_CONTROLLER_IDE
	default:
		storageController = STORAGE_CONTROLLER_UNKNOWN
	}
	return storageController
}

// TODO: improve
func toBusType(interfaceType string) BusType {
	var busType BusType
	switch interfaceType {
	case "SCSI":
		busType = BUS_TYPE_SCSI
	case "IDE":
		busType = BUS_TYPE_IDE
	default:
		busType = BUS_TYPE_UNKNOWN
	}
	return busType
}

// TODO: improve
func toReadOnly(access uint16) bool {
	// See Access property from: https://docs.microsoft.com/en-us/windows/win32/cimwin32prov/win32-diskpartition
	return access == 0x1
}
<commit_msg>[windows] look in diskdrive.Caption for "ssd"<commit_after>// Use and distribution licensed under the Apache license version 2.
//
// See the COPYING file in the root project directory for full text.
//

package ghw

import (
	"strings"

	"github.com/StackExchange/wmi"
)

const wqlDiskDrive = "SELECT Caption, CreationClassName, Description, DeviceID, Index, InterfaceType, Manufacturer, MediaType, Model, Name, Partitions, SerialNumber, Size, TotalCylinders, TotalHeads, TotalSectors, TotalTracks, TracksPerCylinder FROM Win32_DiskDrive"

type win32DiskDrive struct {
	Caption           string
	CreationClassName string
	Description       string
	DeviceID          string
	Index             uint32 // Used to link with partition
	InterfaceType     string
	Manufacturer      string
	MediaType         string
	Model             string
	Name              string
	Partitions        int32
	SerialNumber      string
	Size              uint64
	TotalCylinders    int64
	TotalHeads        int32
	TotalSectors      int64
	TotalTracks       int64
	TracksPerCylinder int32
}

const wqlDiskPartition = "SELECT Access, BlockSize, Caption, CreationClassName, Description, DeviceID, DiskIndex, Index, Name, Size, SystemName, Type FROM Win32_DiskPartition"

type win32DiskPartition struct {
	Access            uint16
	BlockSize         uint64
	Caption           string
	CreationClassName string
	Description       string
	DeviceID          string
	DiskIndex         uint32 // Used to link with Disk Drive
	Index             uint32
	Name              string
	Size              int64
	SystemName        string
	Type              string
}

const wqlLogicalDiskToPartition = "SELECT Antecedent, Dependent FROM Win32_LogicalDiskToPartition"

type win32LogicalDiskToPartition struct {
	Antecedent string
	Dependent  string
}

const wqlLogicalDisk = "SELECT Caption, CreationClassName, Description, DeviceID, FileSystem, FreeSpace, Name, Size, SystemName FROM Win32_LogicalDisk"

type win32LogicalDisk struct {
	Caption           string
	CreationClassName string
	Description       string
	DeviceID          string
	FileSystem        string
	FreeSpace         uint64
	Name              string
	Size              uint64
	SystemName        string
}

func (ctx *context) blockFillInfo(info *BlockInfo) error {
	win32DiskDriveDescriptions, err := getDiskDrives()
	if err != nil {
		return err
	}

	win32DiskPartitionDescriptions, err := getDiskPartitions()
	if err != nil {
		return err
	}

	win32LogicalDiskToPartitionDescriptions, err := getLogicalDisksToPartitions()
	if err != nil {
		return err
	}

	win32LogicalDiskDescriptions, err := getLogicalDisks()
	if err != nil {
		return err
	}

	// Converting into standard structures
	disks := make([]*Disk, 0)
	for _, diskdrive := range win32DiskDriveDescriptions {
		disk := &Disk{
			Name:                   diskdrive.Name,
			SizeBytes:              diskdrive.Size,
			PhysicalBlockSizeBytes: 0,
			DriveType:              toDriveType(diskdrive.MediaType, diskdrive.Caption),
			StorageController:      toStorageController(diskdrive.InterfaceType),
			BusType:                toBusType(diskdrive.InterfaceType),
			BusPath:                UNKNOWN, // TODO: add information
			Vendor:                 UNKNOWN, // TODO: add information
			Model:                  strings.TrimSpace(diskdrive.Caption),
			SerialNumber:           strings.TrimSpace(diskdrive.SerialNumber),
			WWN:                    UNKNOWN, // TODO: add information
			Partitions:             make([]*Partition, 0),
		}
		for _, diskpartition := range win32DiskPartitionDescriptions {
			// Finding disk partition linked to current disk drive
			if diskdrive.Index == diskpartition.DiskIndex {
				disk.PhysicalBlockSizeBytes = diskpartition.BlockSize
				// Finding logical partition linked to current disk partition
				for _, logicaldisk := range win32LogicalDiskDescriptions {
					for _, logicaldisktodiskpartition := range win32LogicalDiskToPartitionDescriptions {
						var desiredAntecedent = "\\\\" + diskpartition.SystemName + "\\root\\cimv2:" + diskpartition.CreationClassName + ".DeviceID=\"" + diskpartition.DeviceID + "\""
						var desiredDependent = "\\\\" + logicaldisk.SystemName + "\\root\\cimv2:" + logicaldisk.CreationClassName + ".DeviceID=\"" + logicaldisk.DeviceID + "\""
						if logicaldisktodiskpartition.Antecedent == desiredAntecedent && logicaldisktodiskpartition.Dependent == desiredDependent {
							// Appending Partition
							p := &Partition{
								Name:       strings.TrimSpace(logicaldisk.Caption),
								Label:      strings.TrimSpace(logicaldisk.Caption),
								SizeBytes:  logicaldisk.Size,
								MountPoint: logicaldisk.DeviceID,
								Type:       diskpartition.Type,
								IsReadOnly: toReadOnly(diskpartition.Access), // TODO: add information
							}
							disk.Partitions = append(disk.Partitions, p)
							break
						}
					}
				}
			}
		}
		disks = append(disks, disk)
	}

	info.Disks = disks
	var tpb uint64
	for _, d := range info.Disks {
		tpb += d.SizeBytes
	}
	info.TotalPhysicalBytes = tpb
	return nil
}

func getDiskDrives() ([]win32DiskDrive, error) {
	// Getting disks drives data from WMI
	var win3232DiskDriveDescriptions []win32DiskDrive
	if err := wmi.Query(wqlDiskDrive, &win3232DiskDriveDescriptions); err != nil {
		return nil, err
	}
	return win3232DiskDriveDescriptions, nil
}

func getDiskPartitions() ([]win32DiskPartition, error) {
	// Getting disk partitions from WMI
	var win32DiskPartitionDescriptions []win32DiskPartition
	if err := wmi.Query(wqlDiskPartition, &win32DiskPartitionDescriptions); err != nil {
		return nil, err
	}
	return win32DiskPartitionDescriptions, nil
}

func getLogicalDisksToPartitions() ([]win32LogicalDiskToPartition, error) {
	// Getting links between logical disks and partitions from WMI
	var win32LogicalDiskToPartitionDescriptions []win32LogicalDiskToPartition
	if err := wmi.Query(wqlLogicalDiskToPartition, &win32LogicalDiskToPartitionDescriptions); err != nil {
		return nil, err
	}
	return win32LogicalDiskToPartitionDescriptions, nil
}

func getLogicalDisks() ([]win32LogicalDisk, error) {
	// Getting logical disks from WMI
	var win32LogicalDiskDescriptions []win32LogicalDisk
	if err := wmi.Query(wqlLogicalDisk, &win32LogicalDiskDescriptions); err != nil {
		return nil, err
	}
	return win32LogicalDiskDescriptions, nil
}

func toDriveType(mediaType string, caption string) DriveType {
	mediaType = strings.ToLower(mediaType)
	caption = strings.ToLower(caption)
	if strings.Contains(mediaType, "fixed") || strings.Contains(mediaType, "ssd") || strings.Contains(caption, "ssd") {
		return DRIVE_TYPE_SSD
	} else if strings.ContainsAny(mediaType, "hdd") {
		return DRIVE_TYPE_HDD
	}
	return DRIVE_TYPE_UNKNOWN
}

// TODO: improve
func toStorageController(interfaceType string) StorageController {
	var storageController StorageController
	switch interfaceType {
	case "SCSI":
		storageController = STORAGE_CONTROLLER_SCSI
	case "IDE":
		storageController = STORAGE_CONTROLLER_IDE
	default:
		storageController = STORAGE_CONTROLLER_UNKNOWN
	}
	return storageController
}

// TODO: improve
func toBusType(interfaceType string) BusType {
	var busType BusType
	switch interfaceType {
	case "SCSI":
		busType = BUS_TYPE_SCSI
	case "IDE":
		busType = BUS_TYPE_IDE
	default:
		busType = BUS_TYPE_UNKNOWN
	}
	return busType
}

// TODO: improve
func toReadOnly(access uint16) bool {
	// See Access property from: https://docs.microsoft.com/en-us/windows/win32/cimwin32prov/win32-diskpartition
	return access == 0x1
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"regexp"
	"strconv"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSLambdaEventSourceMapping_basic(t *testing.T) {
	var conf lambda.EventSourceMappingConfiguration

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckLambdaEventSourceMappingDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSLambdaEventSourceMappingConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaEventSourceMappingExists("aws_lambda_event_source_mapping.lambda_event_source_mapping_test", &conf),
					testAccCheckAWSLambdaEventSourceMappingAttributes(&conf),
				),
			},
			resource.TestStep{
				Config: testAccAWSLambdaEventSourceMappingConfigUpdate,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaEventSourceMappingExists("aws_lambda_event_source_mapping.lambda_event_source_mapping_test", &conf),
					resource.TestCheckResourceAttr("aws_lambda_event_source_mapping.lambda_event_source_mapping_test",
						"batch_size",
						strconv.Itoa(200)),
					resource.TestCheckResourceAttr("aws_lambda_event_source_mapping.lambda_event_source_mapping_test",
						"enabled",
						strconv.FormatBool(false)),
					resource.TestMatchResourceAttr(
						"aws_lambda_event_source_mapping.lambda_event_source_mapping_test",
						"function_arn",
						regexp.MustCompile("example_lambda_name_update$"),
					),
				),
			},
		},
	})
}

func TestAccAWSLambdaEventSourceMapping_disappears(t *testing.T) {
	var conf lambda.EventSourceMappingConfiguration

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckLambdaEventSourceMappingDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSLambdaEventSourceMappingConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaEventSourceMappingExists("aws_lambda_event_source_mapping.lambda_event_source_mapping_test", &conf),
					testAccCheckAWSLambdaEventSourceMappingDisappears(&conf),
				),
				ExpectNonEmptyPlan: true,
			},
		},
	})
}

func testAccCheckAWSLambdaEventSourceMappingDisappears(conf *lambda.EventSourceMappingConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.Meta().(*AWSClient).lambdaconn

		params := &lambda.DeleteEventSourceMappingInput{
			UUID: conf.UUID,
		}

		_, err := conn.DeleteEventSourceMapping(params)
		if err != nil {
			if err != nil {
				return err
			}
		}

		return resource.Retry(10*time.Minute, func() *resource.RetryError {
			params := &lambda.GetEventSourceMappingInput{
				UUID: conf.UUID,
			}
			_, err := conn.GetEventSourceMapping(params)
			if err != nil {
				cgw, ok := err.(awserr.Error)
				if ok && cgw.Code() == "ResourceNotFoundException" {
					return nil
				}
				return resource.NonRetryableError(
					fmt.Errorf("Error retrieving Lambda Event Source Mapping: %s", err))
			}
			return resource.RetryableError(fmt.Errorf(
				"Waiting for Lambda Event Source Mapping: %v", conf.UUID))
		})
	}
}

func testAccCheckLambdaEventSourceMappingDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).lambdaconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_lambda_event_source_mapping" {
			continue
		}

		_, err := conn.GetEventSourceMapping(&lambda.GetEventSourceMappingInput{
			UUID: aws.String(rs.Primary.ID),
		})

		if err == nil {
			return fmt.Errorf("Lambda event source mapping was not deleted")
		}

	}

	return nil

}

func testAccCheckAwsLambdaEventSourceMappingExists(n string, mapping *lambda.EventSourceMappingConfiguration) resource.TestCheckFunc {
	// Wait for IAM role
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Lambda event source mapping not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("Lambda event source mapping ID not set")
		}

		conn := testAccProvider.Meta().(*AWSClient).lambdaconn

		params := &lambda.GetEventSourceMappingInput{
			UUID: aws.String(rs.Primary.ID),
		}

		getSourceMappingConfiguration, err := conn.GetEventSourceMapping(params)
		if err != nil {
			return err
		}

		*mapping = *getSourceMappingConfiguration

		return nil
	}
}

func testAccCheckAWSLambdaEventSourceMappingAttributes(mapping *lambda.EventSourceMappingConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		uuid := *mapping.UUID
		if uuid == "" {
			return fmt.Errorf("Could not read Lambda event source mapping's UUID")
		}

		return nil
	}
}

const testAccAWSLambdaEventSourceMappingConfig = `
resource "aws_iam_role" "iam_for_lambda" {
    name = "iam_for_lambda"
    assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "lambda.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_policy" "policy_for_role" {
    name = "policy_for_role"
    path = "/"
    description = "IAM policy for for Lamda event mapping testing"
    policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
      {
          "Effect": "Allow",
          "Action": [
            "kinesis:GetRecords",
            "kinesis:GetShardIterator",
            "kinesis:DescribeStream"
          ],
          "Resource": "*"
      },
      {
          "Effect": "Allow",
          "Action": [
            "kinesis:ListStreams"
          ],
          "Resource": "*"
      }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "policy_attachment_for_role" {
    name = "policy_attachment_for_role"
    roles = ["${aws_iam_role.iam_for_lambda.name}"]
    policy_arn = "${aws_iam_policy.policy_for_role.arn}"
}

resource "aws_kinesis_stream" "kinesis_stream_test" {
    name = "kinesis_stream_test"
    shard_count = 1
}

resource "aws_lambda_function" "lambda_function_test_create" {
    filename = "test-fixtures/lambdatest.zip"
    function_name = "example_lambda_name_create"
    role = "${aws_iam_role.iam_for_lambda.arn}"
    handler = "exports.example"
    runtime = "nodejs4.3"
}

resource "aws_lambda_function" "lambda_function_test_update" {
    filename = "test-fixtures/lambdatest.zip"
    function_name = "example_lambda_name_update"
    role = "${aws_iam_role.iam_for_lambda.arn}"
    handler = "exports.example"
    runtime = "nodejs4.3"
}

resource "aws_lambda_event_source_mapping" "lambda_event_source_mapping_test" {
		batch_size = 100
		event_source_arn = "${aws_kinesis_stream.kinesis_stream_test.arn}"
		enabled = true
		depends_on = ["aws_iam_policy_attachment.policy_attachment_for_role"]
		function_name = "${aws_lambda_function.lambda_function_test_create.arn}"
		starting_position = "TRIM_HORIZON"
}
`

const testAccAWSLambdaEventSourceMappingConfigUpdate = `
resource "aws_iam_role" "iam_for_lambda" {
    name = "iam_for_lambda"
    assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "lambda.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_policy" "policy_for_role" {
    name = "policy_for_role"
    path = "/"
    description = "IAM policy for for Lamda event mapping testing"
    policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
      {
          "Effect": "Allow",
          "Action": [
            "kinesis:GetRecords",
            "kinesis:GetShardIterator",
            "kinesis:DescribeStream"
          ],
          "Resource": "*"
      },
      {
          "Effect": "Allow",
          "Action": [
            "kinesis:ListStreams"
          ],
          "Resource": "*"
      }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "policy_attachment_for_role" {
    name = "policy_attachment_for_role"
    roles = ["${aws_iam_role.iam_for_lambda.name}"]
    policy_arn = "${aws_iam_policy.policy_for_role.arn}"
}

resource "aws_kinesis_stream" "kinesis_stream_test" {
    name = "kinesis_stream_test"
    shard_count = 1
}

resource "aws_lambda_function" "lambda_function_test_create" {
    filename = "test-fixtures/lambdatest.zip"
    function_name = "example_lambda_name_create"
    role = "${aws_iam_role.iam_for_lambda.arn}"
    handler = "exports.example"
    runtime = "nodejs4.3"
}

resource "aws_lambda_function" "lambda_function_test_update" {
    filename = "test-fixtures/lambdatest.zip"
    function_name = "example_lambda_name_update"
    role = "${aws_iam_role.iam_for_lambda.arn}"
    handler = "exports.example"
    runtime = "nodejs4.3"
}

resource "aws_lambda_event_source_mapping" "lambda_event_source_mapping_test" {
		batch_size = 200
		event_source_arn = "${aws_kinesis_stream.kinesis_stream_test.arn}"
		enabled = false
		depends_on = ["aws_iam_policy_attachment.policy_attachment_for_role"]
		function_name = "${aws_lambda_function.lambda_function_test_update.arn}"
		starting_position = "TRIM_HORIZON"
}
`
<commit_msg>random some lambda tests<commit_after>package aws

import (
	"fmt"
	"regexp"
	"strconv"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSLambdaEventSourceMapping_basic(t *testing.T) {
	var conf lambda.EventSourceMappingConfiguration
	rInt := acctest.RandInt()

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckLambdaEventSourceMappingDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSLambdaEventSourceMappingConfig(rInt),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaEventSourceMappingExists("aws_lambda_event_source_mapping.lambda_event_source_mapping_test", &conf),
					testAccCheckAWSLambdaEventSourceMappingAttributes(&conf),
				),
			},
			resource.TestStep{
				Config: testAccAWSLambdaEventSourceMappingConfigUpdate(rInt),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaEventSourceMappingExists("aws_lambda_event_source_mapping.lambda_event_source_mapping_test", &conf),
					resource.TestCheckResourceAttr("aws_lambda_event_source_mapping.lambda_event_source_mapping_test",
						"batch_size",
						strconv.Itoa(200)),
					resource.TestCheckResourceAttr("aws_lambda_event_source_mapping.lambda_event_source_mapping_test",
						"enabled",
						strconv.FormatBool(false)),
					resource.TestMatchResourceAttr(
						"aws_lambda_event_source_mapping.lambda_event_source_mapping_test",
						"function_arn",
						regexp.MustCompile("example_lambda_name_update$"),
					),
				),
			},
		},
	})
}

func TestAccAWSLambdaEventSourceMapping_disappears(t *testing.T) {
	var conf lambda.EventSourceMappingConfiguration

	rInt := acctest.RandInt()

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckLambdaEventSourceMappingDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSLambdaEventSourceMappingConfig(rInt),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAwsLambdaEventSourceMappingExists("aws_lambda_event_source_mapping.lambda_event_source_mapping_test", &conf),
					testAccCheckAWSLambdaEventSourceMappingDisappears(&conf),
				),
				ExpectNonEmptyPlan: true,
			},
		},
	})
}

func testAccCheckAWSLambdaEventSourceMappingDisappears(conf *lambda.EventSourceMappingConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.Meta().(*AWSClient).lambdaconn

		params := &lambda.DeleteEventSourceMappingInput{
			UUID: conf.UUID,
		}

		_, err := conn.DeleteEventSourceMapping(params)
		if err != nil {
			if err != nil {
				return err
			}
		}

		return resource.Retry(10*time.Minute, func() *resource.RetryError {
			params := &lambda.GetEventSourceMappingInput{
				UUID: conf.UUID,
			}
			_, err := conn.GetEventSourceMapping(params)
			if err != nil {
				cgw, ok := err.(awserr.Error)
				if ok && cgw.Code() == "ResourceNotFoundException" {
					return nil
				}
				return resource.NonRetryableError(
					fmt.Errorf("Error retrieving Lambda Event Source Mapping: %s", err))
			}
			return resource.RetryableError(fmt.Errorf(
				"Waiting for Lambda Event Source Mapping: %v", conf.UUID))
		})
	}
}

func testAccCheckLambdaEventSourceMappingDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).lambdaconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_lambda_event_source_mapping" {
			continue
		}

		_, err := conn.GetEventSourceMapping(&lambda.GetEventSourceMappingInput{
			UUID: aws.String(rs.Primary.ID),
		})

		if err == nil {
			return fmt.Errorf("Lambda event source mapping was not deleted")
		}

	}

	return nil

}

func testAccCheckAwsLambdaEventSourceMappingExists(n string, mapping *lambda.EventSourceMappingConfiguration) resource.TestCheckFunc {
	// Wait for IAM role
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Lambda event source mapping not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("Lambda event source mapping ID not set")
		}

		conn := testAccProvider.Meta().(*AWSClient).lambdaconn

		params := &lambda.GetEventSourceMappingInput{
			UUID: aws.String(rs.Primary.ID),
		}

		getSourceMappingConfiguration, err := conn.GetEventSourceMapping(params)
		if err != nil {
			return err
		}

		*mapping = *getSourceMappingConfiguration

		return nil
	}
}

func testAccCheckAWSLambdaEventSourceMappingAttributes(mapping *lambda.EventSourceMappingConfiguration) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		uuid := *mapping.UUID
		if uuid == "" {
			return fmt.Errorf("Could not read Lambda event source mapping's UUID")
		}

		return nil
	}
}

func testAccAWSLambdaEventSourceMappingConfig(rInt int) string {
	return fmt.Sprintf(`
resource "aws_iam_role" "iam_for_lambda" {
    name = "iam_for_lambda_%d"
    assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "lambda.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_policy" "policy_for_role" {
    name = "policy_for_role_%d"
    path = "/"
    description = "IAM policy for for Lamda event mapping testing"
    policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
      {
          "Effect": "Allow",
          "Action": [
            "kinesis:GetRecords",
            "kinesis:GetShardIterator",
            "kinesis:DescribeStream"
          ],
          "Resource": "*"
      },
      {
          "Effect": "Allow",
          "Action": [
            "kinesis:ListStreams"
          ],
          "Resource": "*"
      }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "policy_attachment_for_role" {
    name = "policy_attachment_for_role_%d"
    roles = ["${aws_iam_role.iam_for_lambda.name}"]
    policy_arn = "${aws_iam_policy.policy_for_role.arn}"
}

resource "aws_kinesis_stream" "kinesis_stream_test" {
    name = "kinesis_stream_test_%d"
    shard_count = 1
}

resource "aws_lambda_function" "lambda_function_test_create" {
    filename = "test-fixtures/lambdatest.zip"
    function_name = "%d_example_lambda_name_create"
    role = "${aws_iam_role.iam_for_lambda.arn}"
    handler = "exports.example"
    runtime = "nodejs4.3"
}

resource "aws_lambda_function" "lambda_function_test_update" {
    filename = "test-fixtures/lambdatest.zip"
    function_name = "%d_example_lambda_name_update"
    role = "${aws_iam_role.iam_for_lambda.arn}"
    handler = "exports.example"
    runtime = "nodejs4.3"
}

resource "aws_lambda_event_source_mapping" "lambda_event_source_mapping_test" {
		batch_size = 100
		event_source_arn = "${aws_kinesis_stream.kinesis_stream_test.arn}"
		enabled = true
		depends_on = ["aws_iam_policy_attachment.policy_attachment_for_role"]
		function_name = "${aws_lambda_function.lambda_function_test_create.arn}"
		starting_position = "TRIM_HORIZON"
}`, rInt, rInt, rInt, rInt, rInt, rInt)
}

func testAccAWSLambdaEventSourceMappingConfigUpdate(rInt int) string {
	return fmt.Sprintf(`
resource "aws_iam_role" "iam_for_lambda" {
    name = "iam_for_lambda_%d"
    assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "lambda.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_policy" "policy_for_role" {
    name = "policy_for_role_%d"
    path = "/"
    description = "IAM policy for for Lamda event mapping testing"
    policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
      {
          "Effect": "Allow",
          "Action": [
            "kinesis:GetRecords",
            "kinesis:GetShardIterator",
            "kinesis:DescribeStream"
          ],
          "Resource": "*"
      },
      {
          "Effect": "Allow",
          "Action": [
            "kinesis:ListStreams"
          ],
          "Resource": "*"
      }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "policy_attachment_for_role" {
    name = "policy_attachment_for_role_%d"
    roles = ["${aws_iam_role.iam_for_lambda.name}"]
    policy_arn = "${aws_iam_policy.policy_for_role.arn}"
}

resource "aws_kinesis_stream" "kinesis_stream_test" {
    name = "kinesis_stream_test_%d"
    shard_count = 1
}

resource "aws_lambda_function" "lambda_function_test_create" {
    filename = "test-fixtures/lambdatest.zip"
    function_name = "%d_example_lambda_name_create"
    role = "${aws_iam_role.iam_for_lambda.arn}"
    handler = "exports.example"
    runtime = "nodejs4.3"
}

resource "aws_lambda_function" "lambda_function_test_update" {
    filename = "test-fixtures/lambdatest.zip"
    function_name = "%d_example_lambda_name_update"
    role = "${aws_iam_role.iam_for_lambda.arn}"
    handler = "exports.example"
    runtime = "nodejs4.3"
}

resource "aws_lambda_event_source_mapping" "lambda_event_source_mapping_test" {
		batch_size = 200
		event_source_arn = "${aws_kinesis_stream.kinesis_stream_test.arn}"
		enabled = false
		depends_on = ["aws_iam_policy_attachment.policy_attachment_for_role"]
		function_name = "${aws_lambda_function.lambda_function_test_update.arn}"
		starting_position = "TRIM_HORIZON"
}`, rInt, rInt, rInt, rInt, rInt, rInt)
}
<|endoftext|>