<commit_before>// Copyright 2019 The Kubernetes Authors.
// SPDX-License-Identifier: Apache-2.0

package builtinconfig

import (
	"log"
	"sort"

	"sigs.k8s.io/kustomize/api/ifc"
	"sigs.k8s.io/kustomize/api/konfig/builtinpluginconsts"
	"sigs.k8s.io/kustomize/api/types"
)

// TransformerConfig holds the data needed to perform transformations.
type TransformerConfig struct {
	NamePrefix        types.FsSlice `json:"namePrefix,omitempty" yaml:"namePrefix,omitempty"`
	NameSuffix        types.FsSlice `json:"nameSuffix,omitempty" yaml:"nameSuffix,omitempty"`
	NameSpace         types.FsSlice `json:"namespace,omitempty" yaml:"namespace,omitempty"`
	CommonLabels      types.FsSlice `json:"commonLabels,omitempty" yaml:"commonLabels,omitempty"`
	TemplateLabels    types.FsSlice `json:"templateLabels,omitempty" yaml:"templateLabels,omitempty"`
	CommonAnnotations types.FsSlice `json:"commonAnnotations,omitempty" yaml:"commonAnnotations,omitempty"`
	NameReference     nbrSlice      `json:"nameReference,omitempty" yaml:"nameReference,omitempty"`
	VarReference      types.FsSlice `json:"varReference,omitempty" yaml:"varReference,omitempty"`
	Images            types.FsSlice `json:"images,omitempty" yaml:"images,omitempty"`
	Replicas          types.FsSlice `json:"replicas,omitempty" yaml:"replicas,omitempty"`
}

// MakeEmptyConfig returns an empty TransformerConfig object
func MakeEmptyConfig() *TransformerConfig {
	return &TransformerConfig{}
}

// MakeDefaultConfig returns a default TransformerConfig.
func MakeDefaultConfig() *TransformerConfig {
	c, err := makeTransformerConfigFromBytes(
		builtinpluginconsts.GetDefaultFieldSpecs())
	if err != nil {
		log.Fatalf("Unable to make default transformconfig: %v", err)
	}
	return c
}

// MakeTransformerConfig returns a merger of custom config,
// if any, with default config.
func MakeTransformerConfig(
	ldr ifc.Loader, paths []string) (*TransformerConfig, error) {
	t1 := MakeDefaultConfig()
	if len(paths) == 0 {
		return t1, nil
	}
	t2, err := loadDefaultConfig(ldr, paths)
	if err != nil {
		return nil, err
	}
	return t1.Merge(t2)
}

// sortFields provides determinism in logging, tests, etc.
func (t *TransformerConfig) sortFields() {
	sort.Sort(t.NamePrefix)
	sort.Sort(t.NameSpace)
	sort.Sort(t.CommonLabels)
	sort.Sort(t.TemplateLabels)
	sort.Sort(t.CommonAnnotations)
	sort.Sort(t.NameReference)
	sort.Sort(t.VarReference)
	sort.Sort(t.Images)
	sort.Sort(t.Replicas)
}

// AddPrefixFieldSpec adds a FieldSpec to NamePrefix
func (t *TransformerConfig) AddPrefixFieldSpec(fs types.FieldSpec) (err error) {
	t.NamePrefix, err = t.NamePrefix.MergeOne(fs)
	return err
}

// AddSuffixFieldSpec adds a FieldSpec to NameSuffix
func (t *TransformerConfig) AddSuffixFieldSpec(fs types.FieldSpec) (err error) {
	t.NameSuffix, err = t.NameSuffix.MergeOne(fs)
	return err
}

// AddLabelFieldSpec adds a FieldSpec to CommonLabels
func (t *TransformerConfig) AddLabelFieldSpec(fs types.FieldSpec) (err error) {
	t.CommonLabels, err = t.CommonLabels.MergeOne(fs)
	return err
}

// AddAnnotationFieldSpec adds a FieldSpec to CommonAnnotations
func (t *TransformerConfig) AddAnnotationFieldSpec(fs types.FieldSpec) (err error) {
	t.CommonAnnotations, err = t.CommonAnnotations.MergeOne(fs)
	return err
}

// AddNamereferenceFieldSpec adds a NameBackReferences to NameReference
func (t *TransformerConfig) AddNamereferenceFieldSpec(
	nbrs NameBackReferences) (err error) {
	t.NameReference, err = t.NameReference.mergeOne(nbrs)
	return err
}

// Merge merges two TransformerConfigs objects into
// a new TransformerConfig object
func (t *TransformerConfig) Merge(input *TransformerConfig) (
	merged *TransformerConfig, err error) {
	if input == nil {
		return t, nil
	}
	merged = &TransformerConfig{}
	merged.NamePrefix, err = t.NamePrefix.MergeAll(input.NamePrefix)
	if err != nil {
		return nil, err
	}
	merged.NameSuffix, err = t.NameSuffix.MergeAll(input.NameSuffix)
	if err != nil {
		return nil, err
	}
	merged.NameSpace, err = t.NameSpace.MergeAll(input.NameSpace)
	if err != nil {
		return nil, err
	}
	merged.CommonAnnotations, err = t.CommonAnnotations.MergeAll(
		input.CommonAnnotations)
	if err != nil {
		return nil, err
	}
	merged.CommonLabels, err = t.CommonLabels.MergeAll(input.CommonLabels)
	if err != nil {
		return nil, err
	}
	merged.TemplateLabels, err = t.TemplateLabels.MergeAll(input.TemplateLabels)
	if err != nil {
		return nil, err
	}
	merged.VarReference, err = t.VarReference.MergeAll(input.VarReference)
	if err != nil {
		return nil, err
	}
	merged.NameReference, err = t.NameReference.mergeAll(input.NameReference)
	if err != nil {
		return nil, err
	}
	merged.Images, err = t.Images.MergeAll(input.Images)
	if err != nil {
		return nil, err
	}
	merged.Replicas, err = t.Replicas.MergeAll(input.Replicas)
	if err != nil {
		return nil, err
	}
	merged.sortFields()
	return merged, nil
}
<commit_msg>error wraps for linter<commit_after>// Copyright 2019 The Kubernetes Authors.
// SPDX-License-Identifier: Apache-2.0

package builtinconfig

import (
	"log"
	"sort"

	"github.com/pkg/errors"
	"sigs.k8s.io/kustomize/api/ifc"
	"sigs.k8s.io/kustomize/api/konfig/builtinpluginconsts"
	"sigs.k8s.io/kustomize/api/types"
)

// TransformerConfig holds the data needed to perform transformations.
type TransformerConfig struct {
	NamePrefix        types.FsSlice `json:"namePrefix,omitempty" yaml:"namePrefix,omitempty"`
	NameSuffix        types.FsSlice `json:"nameSuffix,omitempty" yaml:"nameSuffix,omitempty"`
	NameSpace         types.FsSlice `json:"namespace,omitempty" yaml:"namespace,omitempty"`
	CommonLabels      types.FsSlice `json:"commonLabels,omitempty" yaml:"commonLabels,omitempty"`
	TemplateLabels    types.FsSlice `json:"templateLabels,omitempty" yaml:"templateLabels,omitempty"`
	CommonAnnotations types.FsSlice `json:"commonAnnotations,omitempty" yaml:"commonAnnotations,omitempty"`
	NameReference     nbrSlice      `json:"nameReference,omitempty" yaml:"nameReference,omitempty"`
	VarReference      types.FsSlice `json:"varReference,omitempty" yaml:"varReference,omitempty"`
	Images            types.FsSlice `json:"images,omitempty" yaml:"images,omitempty"`
	Replicas          types.FsSlice `json:"replicas,omitempty" yaml:"replicas,omitempty"`
}

// MakeEmptyConfig returns an empty TransformerConfig object
func MakeEmptyConfig() *TransformerConfig {
	return &TransformerConfig{}
}

// MakeDefaultConfig returns a default TransformerConfig.
func MakeDefaultConfig() *TransformerConfig {
	c, err := makeTransformerConfigFromBytes(
		builtinpluginconsts.GetDefaultFieldSpecs())
	if err != nil {
		log.Fatalf("Unable to make default transformconfig: %v", err)
	}
	return c
}

// MakeTransformerConfig returns a merger of custom config,
// if any, with default config.
func MakeTransformerConfig(
	ldr ifc.Loader, paths []string) (*TransformerConfig, error) {
	t1 := MakeDefaultConfig()
	if len(paths) == 0 {
		return t1, nil
	}
	t2, err := loadDefaultConfig(ldr, paths)
	if err != nil {
		return nil, err
	}
	return t1.Merge(t2)
}

// sortFields provides determinism in logging, tests, etc.
func (t *TransformerConfig) sortFields() {
	sort.Sort(t.NamePrefix)
	sort.Sort(t.NameSpace)
	sort.Sort(t.CommonLabels)
	sort.Sort(t.TemplateLabels)
	sort.Sort(t.CommonAnnotations)
	sort.Sort(t.NameReference)
	sort.Sort(t.VarReference)
	sort.Sort(t.Images)
	sort.Sort(t.Replicas)
}

// AddPrefixFieldSpec adds a FieldSpec to NamePrefix
func (t *TransformerConfig) AddPrefixFieldSpec(fs types.FieldSpec) (err error) {
	t.NamePrefix, err = t.NamePrefix.MergeOne(fs)
	return err
}

// AddSuffixFieldSpec adds a FieldSpec to NameSuffix
func (t *TransformerConfig) AddSuffixFieldSpec(fs types.FieldSpec) (err error) {
	t.NameSuffix, err = t.NameSuffix.MergeOne(fs)
	return err
}

// AddLabelFieldSpec adds a FieldSpec to CommonLabels
func (t *TransformerConfig) AddLabelFieldSpec(fs types.FieldSpec) (err error) {
	t.CommonLabels, err = t.CommonLabels.MergeOne(fs)
	return err
}

// AddAnnotationFieldSpec adds a FieldSpec to CommonAnnotations
func (t *TransformerConfig) AddAnnotationFieldSpec(fs types.FieldSpec) (err error) {
	t.CommonAnnotations, err = t.CommonAnnotations.MergeOne(fs)
	return err
}

// AddNamereferenceFieldSpec adds a NameBackReferences to NameReference
func (t *TransformerConfig) AddNamereferenceFieldSpec(
	nbrs NameBackReferences) (err error) {
	t.NameReference, err = t.NameReference.mergeOne(nbrs)
	return err
}

// Merge merges two TransformerConfigs objects into
// a new TransformerConfig object
func (t *TransformerConfig) Merge(input *TransformerConfig) (
	merged *TransformerConfig, err error) {
	if input == nil {
		return t, nil
	}
	merged = &TransformerConfig{}
	merged.NamePrefix, err = t.NamePrefix.MergeAll(input.NamePrefix)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge NamePrefix fieldSpec")
	}
	merged.NameSuffix, err = t.NameSuffix.MergeAll(input.NameSuffix)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge NameSuffix fieldSpec")
	}
	merged.NameSpace, err = t.NameSpace.MergeAll(input.NameSpace)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge NameSpace fieldSpec")
	}
	merged.CommonAnnotations, err = t.CommonAnnotations.MergeAll(
		input.CommonAnnotations)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge CommonAnnotations fieldSpec")
	}
	merged.CommonLabels, err = t.CommonLabels.MergeAll(input.CommonLabels)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge CommonLabels fieldSpec")
	}
	merged.TemplateLabels, err = t.TemplateLabels.MergeAll(input.TemplateLabels)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge TemplateLabels fieldSpec")
	}
	merged.VarReference, err = t.VarReference.MergeAll(input.VarReference)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge VarReference fieldSpec")
	}
	merged.NameReference, err = t.NameReference.mergeAll(input.NameReference)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge NameReference fieldSpec")
	}
	merged.Images, err = t.Images.MergeAll(input.Images)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge Images fieldSpec")
	}
	merged.Replicas, err = t.Replicas.MergeAll(input.Replicas)
	if err != nil {
		return nil, errors.Wrap(err, "failed to merge Replicas fieldSpec")
	}
	merged.sortFields()
	return merged, nil
}
<|endoftext|>


<commit_before>package mesh

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	wt "github.com/weaveworks/weave/testing"
)

// TODO test gossip unicast; atm we only test topology gossip and
// surrogates, neither of which employ unicast.

type mockChannelConnection struct {
	RemoteConnection
	dest *Router
}

func NewTestRouter(name string) *Router {
	peerName, _ := PeerNameFromString(name)
	router := NewRouter(Config{}, peerName, "nick", nil)
	router.Start()
	return router
}

func (conn *mockChannelConnection) SendProtocolMsg(protocolMsg ProtocolMsg) {
	if err := conn.dest.handleGossip(protocolMsg.tag, protocolMsg.msg); err != nil {
		panic(err)
	}
}

func sendPendingGossip(routers ...*Router) {
	// Loop until all routers report they didn't send anything
	for sentSomething := true; sentSomething; {
		sentSomething = false
		for _, router := range routers {
			sentSomething = router.sendPendingGossip() || sentSomething
		}
	}
}

func (router *Router) AddTestChannelConnection(r *Router) {
	fromPeer := NewPeerFrom(router.Ourself.Peer)
	toPeer := NewPeerFrom(r.Ourself.Peer)

	r.Peers.FetchWithDefault(fromPeer)             // Has side-effect of incrementing refcount
	toPeer = router.Peers.FetchWithDefault(toPeer) //

	conn := &mockChannelConnection{RemoteConnection{router.Ourself.Peer, toPeer, "", false, true}, r}
	router.Ourself.handleAddConnection(conn)
	router.Ourself.handleConnectionEstablished(conn)
}

func (router *Router) DeleteTestChannelConnection(r *Router) {
	fromName := router.Ourself.Peer.Name
	toName := r.Ourself.Peer.Name

	fromPeer := r.Peers.Fetch(fromName)
	toPeer := router.Peers.Fetch(toName)

	r.Peers.Dereference(fromPeer)
	router.Peers.Dereference(toPeer)

	conn, _ := router.Ourself.ConnectionTo(toName)
	router.Ourself.handleDeleteConnection(conn)
}

func TestGossipTopology(t *testing.T) {
	wt.RunWithTimeout(t, 5*time.Second, func() {
		implTestGossipTopology(t)
	})
}

// Create a Peer representing the receiver router, with connections to
// the routers supplied as arguments, carrying across all UID and
// version information.
func (router *Router) tp(routers ...*Router) *Peer {
	peer := NewPeerFrom(router.Ourself.Peer)
	connections := make(map[PeerName]Connection)
	for _, r := range routers {
		p := NewPeerFrom(r.Ourself.Peer)
		connections[r.Ourself.Peer.Name] = newMockConnection(peer, p)
	}
	peer.Version = router.Ourself.Peer.Version
	peer.connections = connections
	return peer
}

// Check that the topology of router matches the peers and all of their connections
func checkTopology(t *testing.T, router *Router, wantedPeers ...*Peer) {
	router.Peers.RLock()
	checkTopologyPeers(t, true, router.Peers.allPeers(), wantedPeers...)
	router.Peers.RUnlock()
}

func implTestGossipTopology(t *testing.T) {
	// Create some peers that will talk to each other
	r1 := NewTestRouter("01:00:00:01:00:00")
	r2 := NewTestRouter("02:00:00:02:00:00")
	r3 := NewTestRouter("03:00:00:03:00:00")

	// Check state when they have no connections
	checkTopology(t, r1, r1.tp())
	checkTopology(t, r2, r2.tp())

	// Now try adding some connections
	r1.AddTestChannelConnection(r2)
	sendPendingGossip(r1, r2)
	checkTopology(t, r1, r1.tp(r2), r2.tp())
	checkTopology(t, r2, r1.tp(r2), r2.tp())
	r2.AddTestChannelConnection(r1)
	sendPendingGossip(r1, r2)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1))

	// Currently, the connection from 2 to 3 is one-way only
	r2.AddTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp())
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp())
	// When r2 gossiped to r3, 1 was unreachable from r3 so it got removed from the
	// list of peers, but remains referenced in the connection from 1 to 3.
	checkTopology(t, r3, r2.tp(r1, r3), r3.tp())

	// Add a connection from 3 to 1 and now r1 is reachable.
	r3.AddTestChannelConnection(r1)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r3, r1.tp(), r2.tp(r1, r3), r3.tp(r1))

	r1.AddTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))

	// Drop the connection from 2 to 3
	r2.DeleteTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))

	// Drop the connection from 1 to 3
	r1.DeleteTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))

	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1))
	// r3 still thinks r1 has a connection to it
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))

	// On a timer, r3 will gossip to r1
	r3.SendAllGossip()
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))
}

func TestGossipSurrogate(t *testing.T) {
	// create the topology r1 <-> r2 <-> r3
	r1 := NewTestRouter("01:00:00:01:00:00")
	r2 := NewTestRouter("02:00:00:02:00:00")
	r3 := NewTestRouter("03:00:00:03:00:00")
	r1.AddTestChannelConnection(r2)
	r2.AddTestChannelConnection(r1)
	r3.AddTestChannelConnection(r2)
	r2.AddTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
	checkTopology(t, r3, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))

	// create a gossiper at either end, but not the middle
	g1 := newTestGossiper()
	g3 := newTestGossiper()
	s1 := r1.NewGossip("Test", g1)
	s3 := r3.NewGossip("Test", g3)

	// broadcast a message from each end, check it reaches the other
	broadcast(s1, 1)
	broadcast(s3, 2)
	sendPendingGossip(r1, r2, r3)
	g1.checkHas(t, 2)
	g3.checkHas(t, 1)

	// check that each end gets their message back through periodic
	// gossip
	r1.SendAllGossip()
	r3.SendAllGossip()
	sendPendingGossip(r1, r2, r3)
	g1.checkHas(t, 1, 2)
	g3.checkHas(t, 1, 2)
}

type testGossiper struct {
	sync.RWMutex
	state map[byte]struct{}
}

func newTestGossiper() *testGossiper {
	return &testGossiper{state: make(map[byte]struct{})}
}

func (g *testGossiper) OnGossipUnicast(sender PeerName, msg []byte) error {
	return nil
}

func (g *testGossiper) OnGossipBroadcast(_ PeerName, update []byte) (GossipData, error) {
	g.Lock()
	defer g.Unlock()
	for _, v := range update {
		g.state[v] = void
	}
	return NewSurrogateGossipData(update), nil
}

func (g *testGossiper) Gossip() GossipData {
	g.RLock()
	defer g.RUnlock()
	state := make([]byte, len(g.state))
	for v := range g.state {
		state = append(state, v)
	}
	return NewSurrogateGossipData(state)
}

func (g *testGossiper) OnGossip(update []byte) (GossipData, error) {
	g.Lock()
	defer g.Unlock()
	var delta []byte
	for _, v := range update {
		if _, found := g.state[v]; !found {
			delta = append(delta, v)
			g.state[v] = void
		}
	}
	if len(delta) > 0 {
		return NewSurrogateGossipData(delta), nil
	}
	return nil, nil
}

func (g *testGossiper) checkHas(t *testing.T, vs ...byte) {
	g.RLock()
	defer g.RUnlock()
	for _, v := range vs {
		if _, found := g.state[v]; !found {
			require.FailNow(t, fmt.Sprintf("%d is missing", v))
		}
	}
}

func broadcast(s Gossip, v byte) {
	s.GossipBroadcast(NewSurrogateGossipData([]byte{v}))
}
<commit_msg>remove redundant code in test<commit_after>package mesh

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	wt "github.com/weaveworks/weave/testing"
)

// TODO test gossip unicast; atm we only test topology gossip and
// surrogates, neither of which employ unicast.

type mockChannelConnection struct {
	RemoteConnection
	dest *Router
}

func NewTestRouter(name string) *Router {
	peerName, _ := PeerNameFromString(name)
	router := NewRouter(Config{}, peerName, "nick", nil)
	router.Start()
	return router
}

func (conn *mockChannelConnection) SendProtocolMsg(protocolMsg ProtocolMsg) {
	if err := conn.dest.handleGossip(protocolMsg.tag, protocolMsg.msg); err != nil {
		panic(err)
	}
}

func sendPendingGossip(routers ...*Router) {
	// Loop until all routers report they didn't send anything
	for sentSomething := true; sentSomething; {
		sentSomething = false
		for _, router := range routers {
			sentSomething = router.sendPendingGossip() || sentSomething
		}
	}
}

func (router *Router) AddTestChannelConnection(r *Router) {
	fromPeer := NewPeerFrom(router.Ourself.Peer)
	toPeer := NewPeerFrom(r.Ourself.Peer)

	r.Peers.FetchWithDefault(fromPeer)             // Has side-effect of incrementing refcount
	toPeer = router.Peers.FetchWithDefault(toPeer) //

	conn := &mockChannelConnection{RemoteConnection{router.Ourself.Peer, toPeer, "", false, true}, r}
	router.Ourself.handleAddConnection(conn)
	router.Ourself.handleConnectionEstablished(conn)
}

func (router *Router) DeleteTestChannelConnection(r *Router) {
	fromName := router.Ourself.Peer.Name
	toName := r.Ourself.Peer.Name

	fromPeer := r.Peers.Fetch(fromName)
	toPeer := router.Peers.Fetch(toName)

	r.Peers.Dereference(fromPeer)
	router.Peers.Dereference(toPeer)

	conn, _ := router.Ourself.ConnectionTo(toName)
	router.Ourself.handleDeleteConnection(conn)
}

func TestGossipTopology(t *testing.T) {
	wt.RunWithTimeout(t, 5*time.Second, func() {
		implTestGossipTopology(t)
	})
}

// Create a Peer representing the receiver router, with connections to
// the routers supplied as arguments, carrying across all UID and
// version information.
func (router *Router) tp(routers ...*Router) *Peer {
	peer := NewPeerFrom(router.Ourself.Peer)
	connections := make(map[PeerName]Connection)
	for _, r := range routers {
		p := NewPeerFrom(r.Ourself.Peer)
		connections[r.Ourself.Peer.Name] = newMockConnection(peer, p)
	}
	peer.Version = router.Ourself.Peer.Version
	peer.connections = connections
	return peer
}

// Check that the topology of router matches the peers and all of their connections
func checkTopology(t *testing.T, router *Router, wantedPeers ...*Peer) {
	router.Peers.RLock()
	checkTopologyPeers(t, true, router.Peers.allPeers(), wantedPeers...)
	router.Peers.RUnlock()
}

func implTestGossipTopology(t *testing.T) {
	// Create some peers that will talk to each other
	r1 := NewTestRouter("01:00:00:01:00:00")
	r2 := NewTestRouter("02:00:00:02:00:00")
	r3 := NewTestRouter("03:00:00:03:00:00")

	// Check state when they have no connections
	checkTopology(t, r1, r1.tp())
	checkTopology(t, r2, r2.tp())

	// Now try adding some connections
	r1.AddTestChannelConnection(r2)
	sendPendingGossip(r1, r2)
	checkTopology(t, r1, r1.tp(r2), r2.tp())
	checkTopology(t, r2, r1.tp(r2), r2.tp())
	r2.AddTestChannelConnection(r1)
	sendPendingGossip(r1, r2)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1))

	// Currently, the connection from 2 to 3 is one-way only
	r2.AddTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp())
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp())
	// When r2 gossiped to r3, 1 was unreachable from r3 so it got removed from the
	// list of peers, but remains referenced in the connection from 1 to 3.
	checkTopology(t, r3, r2.tp(r1, r3), r3.tp())

	// Add a connection from 3 to 1 and now r1 is reachable.
	r3.AddTestChannelConnection(r1)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r3, r1.tp(), r2.tp(r1, r3), r3.tp(r1))

	r1.AddTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))

	// Drop the connection from 2 to 3
	r2.DeleteTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))

	// Drop the connection from 1 to 3
	r1.DeleteTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1))
	// r3 still thinks r1 has a connection to it
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))

	// On a timer, r3 will gossip to r1
	r3.SendAllGossip()
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))
}

func TestGossipSurrogate(t *testing.T) {
	// create the topology r1 <-> r2 <-> r3
	r1 := NewTestRouter("01:00:00:01:00:00")
	r2 := NewTestRouter("02:00:00:02:00:00")
	r3 := NewTestRouter("03:00:00:03:00:00")
	r1.AddTestChannelConnection(r2)
	r2.AddTestChannelConnection(r1)
	r3.AddTestChannelConnection(r2)
	r2.AddTestChannelConnection(r3)
	sendPendingGossip(r1, r2, r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
	checkTopology(t, r3, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))

	// create a gossiper at either end, but not the middle
	g1 := newTestGossiper()
	g3 := newTestGossiper()
	s1 := r1.NewGossip("Test", g1)
	s3 := r3.NewGossip("Test", g3)

	// broadcast a message from each end, check it reaches the other
	broadcast(s1, 1)
	broadcast(s3, 2)
	sendPendingGossip(r1, r2, r3)
	g1.checkHas(t, 2)
	g3.checkHas(t, 1)

	// check that each end gets their message back through periodic
	// gossip
	r1.SendAllGossip()
	r3.SendAllGossip()
	sendPendingGossip(r1, r2, r3)
	g1.checkHas(t, 1, 2)
	g3.checkHas(t, 1, 2)
}

type testGossiper struct {
	sync.RWMutex
	state map[byte]struct{}
}

func newTestGossiper() *testGossiper {
	return &testGossiper{state: make(map[byte]struct{})}
}

func (g *testGossiper) OnGossipUnicast(sender PeerName, msg []byte) error {
	return nil
}

func (g *testGossiper) OnGossipBroadcast(_ PeerName, update []byte) (GossipData, error) {
	g.Lock()
	defer g.Unlock()
	for _, v := range update {
		g.state[v] = void
	}
	return NewSurrogateGossipData(update), nil
}

func (g *testGossiper) Gossip() GossipData {
	g.RLock()
	defer g.RUnlock()
	state := make([]byte, len(g.state))
	for v := range g.state {
		state = append(state, v)
	}
	return NewSurrogateGossipData(state)
}

func (g *testGossiper) OnGossip(update []byte) (GossipData, error) {
	g.Lock()
	defer g.Unlock()
	var delta []byte
	for _, v := range update {
		if _, found := g.state[v]; !found {
			delta = append(delta, v)
			g.state[v] = void
		}
	}
	if len(delta) > 0 {
		return NewSurrogateGossipData(delta), nil
	}
	return nil, nil
}

func (g *testGossiper) checkHas(t *testing.T, vs ...byte) {
	g.RLock()
	defer g.RUnlock()
	for _, v := range vs {
		if _, found := g.state[v]; !found {
			require.FailNow(t, fmt.Sprintf("%d is missing", v))
		}
	}
}

func broadcast(s Gossip, v byte) {
	s.GossipBroadcast(NewSurrogateGossipData([]byte{v}))
}
<|endoftext|>


<commit_before>package operator

import (
	"io/ioutil"
	"path"

	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/app_helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/assets"

	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/ginkgo"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gbytes"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gexec"

	archive_helpers "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/pivotal-golang/archiver/extractor/test_helper"
)

var _ = Describe("Environment Variables Groups", func() {
	var createBuildpack = func() string {
		tmpPath, err := ioutil.TempDir("", "env-group-staging")
		Expect(err).ToNot(HaveOccurred())

		buildpackArchivePath := path.Join(tmpPath, "buildpack.zip")

		archive_helpers.CreateZipArchive(buildpackArchivePath, []archive_helpers.ArchiveFile{
			{
				Name: "bin/compile",
				Body: `#!/usr/bin/env bash
sleep 5
echo $CATS_STAGING_TEST_VAR
exit 1
`,
			},
			{
				Name: "bin/detect",
				Body: `#!/bin/bash
exit 1
`,
			},
			{
				Name: "bin/release",
				Body: `#!/usr/bin/env bash
exit 1
`,
			},
		})

		return buildpackArchivePath
	}

	Context("Staging environment variable groups", func() {
		var originalStagingEnv string
		var appName string
		var buildpackName string

		BeforeEach(func() {
			appName = generator.PrefixedRandomName("CATS-APP-")
			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				session := cf.Cf("curl", "/v2/config/environment_variable_groups/staging").Wait(DEFAULT_TIMEOUT)
				Expect(session).To(Exit(0))
				originalStagingEnv = string(session.Out.Contents())
			})
		})

		AfterEach(func() {
			app_helpers.AppReport(appName, DEFAULT_TIMEOUT)

			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				Expect(cf.Cf("curl", "/v2/config/environment_variable_groups/staging", "-X", "PUT", "-d", originalStagingEnv).Wait(DEFAULT_TIMEOUT)).To(Exit(0))

				if buildpackName != "" {
					Expect(cf.Cf("delete-buildpack", buildpackName, "-f").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				}
			})

			Expect(cf.Cf("delete", appName, "-f", "-r").Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
		})

		It("Applies environment variables while staging apps", func() {
			buildpackName = generator.RandomName()
			buildpackZip := createBuildpack()

			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				Expect(cf.Cf("set-staging-environment-variable-group", `{"CATS_STAGING_TEST_VAR":"staging_env_value"}`).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				Expect(cf.Cf("create-buildpack", buildpackName, buildpackZip, "999").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			})

			Expect(cf.Cf("push", appName, "--no-start", "-m", DEFAULT_MEMORY_LIMIT, "-b", buildpackName, "-p", assets.NewAssets().HelloWorld, "-d", config.AppsDomain).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			app_helpers.SetBackend(appName)
			Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(Exit(1))

			Eventually(func() *Session {
				appLogsSession := cf.Cf("logs", "--recent", appName)
				Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				return appLogsSession
			}, DEFAULT_TIMEOUT).Should(Say("staging_env_value"))
		})
	})
	Context("Running environment variable groups", func() {
		var originalRunningEnv string
		var appName string
		var buildpackName string

		BeforeEach(func() {
			appName = generator.PrefixedRandomName("CATS-APP-")
			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				session := cf.Cf("curl", "/v2/config/environment_variable_groups/running").Wait(DEFAULT_TIMEOUT)
				Expect(session).To(Exit(0))
				originalRunningEnv = string(session.Out.Contents())
			})
		})

		AfterEach(func() {
			app_helpers.AppReport(appName, DEFAULT_TIMEOUT)

			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				Expect(cf.Cf("curl", "/v2/config/environment_variable_groups/running", "-X", "PUT", "-d", originalRunningEnv).Wait(DEFAULT_TIMEOUT)).To(Exit(0))

				if buildpackName != "" {
					Expect(cf.Cf("delete-buildpack", buildpackName, "-f").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				}
			})

			Expect(cf.Cf("delete", appName, "-f", "-r").Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
		})

		It("Applies correct environment variables while running apps", func() {
			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				Expect(cf.Cf("set-running-environment-variable-group", `{"CATS_RUNNING_TEST_VAR":"running_env_value"}`).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			})

			Expect(cf.Cf("push", appName, "--no-start", "-b", config.RubyBuildpackName, "-m", DEFAULT_MEMORY_LIMIT, "-p", assets.NewAssets().Dora, "-d", config.AppsDomain).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			app_helpers.SetBackend(appName)
			Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))

			env := helpers.CurlApp(appName, "/env")

			Expect(env).To(ContainSubstring("running_env_value"))
		})
	})
})
<commit_msg>Remove unused buildpack from operator test<commit_after>package operator

import (
	"io/ioutil"
	"path"

	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/app_helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/assets"

	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/ginkgo"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gbytes"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gexec"

	archive_helpers "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/pivotal-golang/archiver/extractor/test_helper"
)

var _ = Describe("Environment Variables Groups", func() {
	var createBuildpack = func() string {
		tmpPath, err := ioutil.TempDir("", "env-group-staging")
		Expect(err).ToNot(HaveOccurred())

		buildpackArchivePath := path.Join(tmpPath, "buildpack.zip")

		archive_helpers.CreateZipArchive(buildpackArchivePath, []archive_helpers.ArchiveFile{
			{
				Name: "bin/compile",
				Body: `#!/usr/bin/env bash
sleep 5
echo $CATS_STAGING_TEST_VAR
exit 1
`,
			},
			{
				Name: "bin/detect",
				Body: `#!/bin/bash
exit 1
`,
			},
			{
				Name: "bin/release",
				Body: `#!/usr/bin/env bash
exit 1
`,
			},
		})

		return buildpackArchivePath
	}

	Context("Staging environment variable groups", func() {
		var originalStagingEnv string
		var appName string
		var buildpackName string

		BeforeEach(func() {
			appName = generator.PrefixedRandomName("CATS-APP-")
			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				session := cf.Cf("curl", "/v2/config/environment_variable_groups/staging").Wait(DEFAULT_TIMEOUT)
				Expect(session).To(Exit(0))
				originalStagingEnv = string(session.Out.Contents())
			})
		})

		AfterEach(func() {
			app_helpers.AppReport(appName, DEFAULT_TIMEOUT)

			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				Expect(cf.Cf("curl", "/v2/config/environment_variable_groups/staging", "-X", "PUT", "-d", originalStagingEnv).Wait(DEFAULT_TIMEOUT)).To(Exit(0))

				if buildpackName != "" {
					Expect(cf.Cf("delete-buildpack", buildpackName, "-f").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				}
			})

			Expect(cf.Cf("delete", appName, "-f", "-r").Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
		})

		It("Applies environment variables while staging apps", func() {
			buildpackName = generator.RandomName()
			buildpackZip := createBuildpack()

			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				Expect(cf.Cf("set-staging-environment-variable-group", `{"CATS_STAGING_TEST_VAR":"staging_env_value"}`).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				Expect(cf.Cf("create-buildpack", buildpackName, buildpackZip, "999").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			})

			Expect(cf.Cf("push", appName, "--no-start", "-m", DEFAULT_MEMORY_LIMIT, "-b", buildpackName, "-p", assets.NewAssets().HelloWorld, "-d", config.AppsDomain).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			app_helpers.SetBackend(appName)
			Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(Exit(1))

			Eventually(func() *Session {
				appLogsSession := cf.Cf("logs", "--recent", appName)
				Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				return appLogsSession
			}, DEFAULT_TIMEOUT).Should(Say("staging_env_value"))
		})
	})
	Context("Running environment variable groups", func() {
		var originalRunningEnv string
		var appName string

		BeforeEach(func() {
			appName = generator.PrefixedRandomName("CATS-APP-")
			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				session := cf.Cf("curl", "/v2/config/environment_variable_groups/running").Wait(DEFAULT_TIMEOUT)
				Expect(session).To(Exit(0))
				originalRunningEnv = string(session.Out.Contents())
			})
		})

		AfterEach(func() {
			app_helpers.AppReport(appName, DEFAULT_TIMEOUT)

			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				Expect(cf.Cf("curl", "/v2/config/environment_variable_groups/running", "-X", "PUT", "-d", originalRunningEnv).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			})

			Expect(cf.Cf("delete", appName, "-f", "-r").Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
		})

		It("Applies correct environment variables while running apps", func() {
			cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
				Expect(cf.Cf("set-running-environment-variable-group", `{"CATS_RUNNING_TEST_VAR":"running_env_value"}`).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			})

			Expect(cf.Cf("push", appName, "--no-start", "-b", config.RubyBuildpackName, "-m", DEFAULT_MEMORY_LIMIT, "-p", assets.NewAssets().Dora, "-d", config.AppsDomain).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			app_helpers.SetBackend(appName)
			Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))

			env := helpers.CurlApp(appName, "/env")

			Expect(env).To(ContainSubstring("running_env_value"))
		})
	})
})
<|endoftext|>


<commit_before>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// go test *.go -bench=".*" -benchmem

package gqueue_test

import (
	"github.com/gogf/gf/g/container/gqueue"
	"github.com/gogf/gf/g/test/gtest"
	"testing"
)

func TestQueue_Len(t *testing.T) {

	max := 100
	for n := 10; n < max; n++ {
		q1 := gqueue.New(max)
		for i := 0; i < max; i++ {
			q1.Push(i)
		}
		gtest.Assert(q1.Len(), max)
		gtest.Assert(q1.Size(), max)
	}
}

func TestQueue_Basic(t *testing.T) {
	q := gqueue.New()
	for i := 0; i < 100; i++ {
		q.Push(i)
	}
	gtest.Assert(q.Pop(), 0)
	gtest.Assert(q.Pop(), 1)
}

func TestQueue_Pop(t *testing.T) {
	q1 := gqueue.New()
	q1.Push(1)
	q1.Push(2)
	q1.Push(3)
	q1.Push(4)
	i1 := q1.Pop()
	gtest.Assert(i1, 1)

}

func TestQueue_Close(t *testing.T) {
	q1 := gqueue.New()
	q1.Push(1)
	q1.Push(2)
	gtest.Assert(q1.Len(), 2)
	q1.Close()
}
<commit_msg>Update gqueue_unit_test.go<commit_after>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// go test *.go -bench=".*" -benchmem

package gqueue_test

import (
	"github.com/gogf/gf/g/container/gqueue"
	"github.com/gogf/gf/g/test/gtest"
	"testing"
)

func TestQueue_Len(t *testing.T) {

	max := 100
	for n := 10; n < max; n++ {
		q1 := gqueue.New(max)
		for i := 0; i < max; i++ {
			q1.Push(i)
		}
		gtest.Assert(q1.Len(), max)
		gtest.Assert(q1.Size(), max)
	}
}

func TestQueue_Basic(t *testing.T) {
	q := gqueue.New()
	for i := 0; i < 100; i++ {
		q.Push(i)
	}
	gtest.Assert(q.Pop(), 0)
	gtest.Assert(q.Pop(), 1)
}

func TestQueue_Pop(t *testing.T) {
	q1 := gqueue.New()
	q1.Push(1)
	q1.Push(2)
	q1.Push(3)
	q1.Push(4)
	i1 := q1.Pop()
	gtest.Assert(i1, 1)

}

func TestQueue_Close(t *testing.T) {
	q1 := gqueue.New()
	q1.Push(1)
	q1.Push(2)
	gtest.Assert(q1.Len(), 2)
	q1.Close()
	q1.Pop()
	gtest.Assert(q1.Len(), 2)
}
<|endoftext|>


<commit_before>package aws

import (
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
	"testing"
)

func TestAccDataSourceAwsEc2InstanceType_basic(t *testing.T) {
	resourceBasic := "data.aws_ec2_instance_type.basic"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccDataSourceEc2InstanceTypeBasic,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceBasic, "auto_recovery_supported", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "bare_metal", "false"),
					resource.TestCheckResourceAttr(resourceBasic, "burstable_performance_supported", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "current_generation", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "dedicated_hosts_supported", "false"),
					resource.TestCheckResourceAttr(resourceBasic, "default_cores", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "default_threads_per_core", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "default_vcpus", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "ebs_encryption_support", "supported"),
					resource.TestCheckResourceAttr(resourceBasic, "ebs_nvme_support", "unsupported"),
					resource.TestCheckResourceAttr(resourceBasic, "ebs_optimized_support", "unsupported"),
					resource.TestCheckResourceAttr(resourceBasic, "efa_supported", "false"),
					resource.TestCheckResourceAttr(resourceBasic, "ena_support", "unsupported"),
					resource.TestCheckResourceAttr(resourceBasic, "free_tier_eligible", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "hibernation_supported", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "hypervisor", "xen"),
					resource.TestCheckResourceAttr(resourceBasic, "instance_storage_supported", "false"),
					resource.TestCheckResourceAttr(resourceBasic, "instance_type", "t2.micro"),
					resource.TestCheckResourceAttr(resourceBasic, "ipv6_supported", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "maximum_ipv4_addresses_per_interface", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "maximum_ipv6_addresses_per_interface", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "maximum_network_interfaces", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "memory_size", "1024"),
					resource.TestCheckResourceAttr(resourceBasic, "network_performance", "Low to Moderate"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_architectures.#", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_architectures.0", "i386"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_architectures.1", "x86_64"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_placement_strategies.#", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_placement_strategies.0", "partition"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_placement_strategies.1", "spread"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_root_device_types.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_root_device_types.0", "ebs"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_usages_classes.#", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_usages_classes.0", "on-demand"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_usages_classes.1", "spot"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_virtualization_types.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_virtualization_types.0", "hvm"),
					resource.TestCheckResourceAttr(resourceBasic, "sustained_clock_speed", "2.5"),
					resource.TestCheckResourceAttr(resourceBasic, "valid_cores.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "valid_cores.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "valid_threads_per_core.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "valid_threads_per_core.0", "1"),
				),
			},
		},
	})
}

func TestAccDataSourceAwsEc2InstanceType_metal(t *testing.T) {
	resourceMetal := "data.aws_ec2_instance_type.metal"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccDataSourceEc2InstanceTypeMetal,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_baseline_bandwidth", "19000"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_baseline_throughput", "2375"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_baseline_iops", "80000"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_maximum_bandwidth", "19000"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_maximum_throughput", "2375"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_maximum_iops", "80000"),
					resource.TestCheckResourceAttr(resourceMetal, "instance_disks.#", "1"),
					resource.TestCheckResourceAttr(resourceMetal, "instance_disks.0.count", "8"),
					resource.TestCheckResourceAttr(resourceMetal, "instance_disks.0.size", "7500"),
					resource.TestCheckResourceAttr(resourceMetal, "instance_disks.0.type", "ssd"),
					resource.TestCheckResourceAttr(resourceMetal, "total_instance_storage", "60000"),
				),
			},
		},
	})
}

func TestAccDataSourceAwsEc2InstanceType_gpu(t *testing.T) {
	resourceGpu := "data.aws_ec2_instance_type.gpu"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccDataSourceEc2InstanceTypeGpu,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceGpu, "gpus.#", "1"),
					resource.TestCheckResourceAttr(resourceGpu, "gpus.0.count", "1"),
					resource.TestCheckResourceAttr(resourceGpu, "gpus.0.manufacturer", "NVIDIA"),
					resource.TestCheckResourceAttr(resourceGpu, "gpus.0.memory_size", "4096"),
					resource.TestCheckResourceAttr(resourceGpu, "gpus.0.name", "K520"),
				),
			},
		},
	})
}

func TestAccDataSourceAwsEc2InstanceType_fpga(t *testing.T) {
	resourceFpga := "data.aws_ec2_instance_type.fpga"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccDataSourceEc2InstanceTypeFgpa,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.#", "1"),
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.0.count", "1"),
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.0.manufacturer", "Xilinx"),
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.0.memory_size", "65536"),
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.0.name", "Virtex UltraScale (VU9P)"),
					resource.TestCheckResourceAttr(resourceFpga, "total_fpga_memory", "65536"),
				),
			},
		},
	})
}

func TestAccDataSourceAwsEc2InstanceType_accelerator(t *testing.T) {
	resourceAccelerator := "data.aws_ec2_instance_type.accelerator"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccDataSourceEc2InstanceTypeAccelerator,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceAccelerator, "inference_accelerators.#", "1"),
					resource.TestCheckResourceAttr(resourceAccelerator, "inference_accelerators.0.count", "1"),
					resource.TestCheckResourceAttr(resourceAccelerator, "inference_accelerators.0.manufacturer", "AWS"),
					resource.TestCheckResourceAttr(resourceAccelerator, "inference_accelerators.0.name", "Inferentia"),
				),
			},
		},
	})
}

const testAccDataSourceEc2InstanceTypeBasic = `
data "aws_ec2_instance_type" "basic" {
  instance_type = "t2.micro"
}
`

const testAccDataSourceEc2InstanceTypeMetal = `
data "aws_ec2_instance_type" "metal" {
  instance_type = "i3en.metal"
}
`

const testAccDataSourceEc2InstanceTypeGpu = `
data "aws_ec2_instance_type" "gpu" {
  instance_type = "g2.2xlarge"
}
`

const testAccDataSourceEc2InstanceTypeFgpa = `
data "aws_ec2_instance_type" "fpga" {
  instance_type = "f1.2xlarge"
}
`

const testAccDataSourceEc2InstanceTypeAccelerator = `
data "aws_ec2_instance_type" "accelerator" {
  instance_type = "inf1.xlarge"
}
`
<commit_msg>Adapt test for cross regions<commit_after>package aws

import (
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
	"testing"
)

func TestAccDataSourceAwsEc2InstanceType_basic(t *testing.T) {
	resourceBasic := "data.aws_ec2_instance_type.basic"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccDataSourceEc2InstanceTypeBasic,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceBasic, "auto_recovery_supported", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "bare_metal", "false"),
					resource.TestCheckResourceAttr(resourceBasic, "burstable_performance_supported", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "current_generation", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "dedicated_hosts_supported", "false"),
					resource.TestCheckResourceAttr(resourceBasic, "default_cores", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "default_threads_per_core", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "default_vcpus", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "ebs_encryption_support", "supported"),
					resource.TestCheckResourceAttr(resourceBasic, "ebs_nvme_support", "unsupported"),
					resource.TestCheckResourceAttr(resourceBasic, "ebs_optimized_support", "unsupported"),
					resource.TestCheckResourceAttr(resourceBasic, "efa_supported", "false"),
					resource.TestCheckResourceAttr(resourceBasic, "ena_support", "unsupported"),
					resource.TestCheckResourceAttr(resourceBasic, "free_tier_eligible", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "hibernation_supported", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "hypervisor", "xen"),
					resource.TestCheckResourceAttr(resourceBasic, "instance_storage_supported", "false"),
					resource.TestCheckResourceAttr(resourceBasic, "instance_type", "t2.micro"),
					resource.TestCheckResourceAttr(resourceBasic, "ipv6_supported", "true"),
					resource.TestCheckResourceAttr(resourceBasic, "maximum_ipv4_addresses_per_interface", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "maximum_ipv6_addresses_per_interface", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "maximum_network_interfaces", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "memory_size", "1024"),
					resource.TestCheckResourceAttr(resourceBasic, "network_performance", "Low to Moderate"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_architectures.#", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_architectures.0", "i386"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_architectures.1", "x86_64"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_placement_strategies.#", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_placement_strategies.0", "partition"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_placement_strategies.1", "spread"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_root_device_types.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_root_device_types.0", "ebs"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_usages_classes.#", "2"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_usages_classes.0", "on-demand"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_usages_classes.1", "spot"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_virtualization_types.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "supported_virtualization_types.0", "hvm"),
					resource.TestCheckResourceAttr(resourceBasic, "sustained_clock_speed", "2.5"),
					resource.TestCheckResourceAttr(resourceBasic, "valid_cores.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "valid_cores.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "valid_threads_per_core.#", "1"),
					resource.TestCheckResourceAttr(resourceBasic, "valid_threads_per_core.0", "1"),
				),
			},
		},
	})
}

func TestAccDataSourceAwsEc2InstanceType_metal(t *testing.T) {
	resourceMetal := "data.aws_ec2_instance_type.metal"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccDataSourceEc2InstanceTypeMetal,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_baseline_bandwidth", "19000"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_baseline_throughput", "2375"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_baseline_iops", "80000"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_maximum_bandwidth", "19000"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_maximum_throughput", "2375"),
					resource.TestCheckResourceAttr(resourceMetal, "ebs_performance_maximum_iops", "80000"),
					resource.TestCheckResourceAttr(resourceMetal, "instance_disks.#", "1"),
					resource.TestCheckResourceAttr(resourceMetal, "instance_disks.0.count", "8"),
					resource.TestCheckResourceAttr(resourceMetal, "instance_disks.0.size", "7500"),
					resource.TestCheckResourceAttr(resourceMetal, "instance_disks.0.type", "ssd"),
					resource.TestCheckResourceAttr(resourceMetal, "total_instance_storage", "60000"),
				),
			},
		},
	})
}

func TestAccDataSourceAwsEc2InstanceType_gpu(t *testing.T) {
	resourceGpu := "data.aws_ec2_instance_type.gpu"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccDataSourceEc2InstanceTypeGpu,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceGpu, "gpus.#", "1"),
					resource.TestCheckResourceAttr(resourceGpu, "gpus.0.count", "1"),
					resource.TestCheckResourceAttr(resourceGpu, "gpus.0.manufacturer", "NVIDIA"),
					resource.TestCheckResourceAttr(resourceGpu, "gpus.0.memory_size", "8192"),
					resource.TestCheckResourceAttr(resourceGpu, "gpus.0.name", "M60"),
				),
			},
		},
	})
}

func TestAccDataSourceAwsEc2InstanceType_fpga(t *testing.T) {
	resourceFpga := "data.aws_ec2_instance_type.fpga"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccDataSourceEc2InstanceTypeFgpa,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.#", "1"),
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.0.count", "1"),
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.0.manufacturer", "Xilinx"),
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.0.memory_size", "65536"),
					resource.TestCheckResourceAttr(resourceFpga, "fpgas.0.name", "Virtex UltraScale (VU9P)"),
					resource.TestCheckResourceAttr(resourceFpga, "total_fpga_memory", "65536"),
				),
			},
		},
	})
}

const testAccDataSourceEc2InstanceTypeBasic = `
data "aws_ec2_instance_type" "basic" {
  instance_type = "t2.micro"
}
`

const testAccDataSourceEc2InstanceTypeMetal = `
data "aws_ec2_instance_type" "metal" {
  instance_type = "i3en.metal"
}
`

const testAccDataSourceEc2InstanceTypeGpu = `
data "aws_ec2_instance_type" "gpu" {
  instance_type = "g3.4xlarge"
}
`

const testAccDataSourceEc2InstanceTypeFgpa = `
data "aws_ec2_instance_type" "fpga" {
  instance_type = "f1.2xlarge"
}
`
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/lakeformation"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"
)

func AwsLakeFormationPermissions() []string {
	return []string{
		"ALL",
		"SELECT",
		"ALTER",
		"DROP",
		"DELETE",
		"INSERT",
		"CREATE_DATABASE",
		"CREATE_TABLE",
		"DATA_LOCATION_ACCESS",
	}
}

func resourceAwsLakeFormationPermissions() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLakeFormationPermissionsGrant,
		Read:   resourceAwsLakeFormationPermissionsList,
		Delete: resourceAwsLakeFormationPermissionsRevoke,

		Schema: map[string]*schema.Schema{
			"catalog_id": {
				Type:         schema.TypeString,
				ForceNew:     true,
				Optional:     true,
				Computed:     true,
				ValidateFunc: validateAwsAccountId,
			},
			"permissions": {
				Type:     schema.TypeList,
				Required: true,
				ForceNew: true,
				MinItems: 1,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validation.StringInSlice(AwsLakeFormationPermissions(), false),
				},
			},
			"permissions_with_grant_option": {
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				Computed: true,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validation.StringInSlice(AwsLakeFormationPermissions(), false),
				},
			},
			"principal": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateArn,
			},
			"database": {
				Type:          schema.TypeString,
				Optional:      true,
				ForceNew:      true,
				ConflictsWith: []string{"location", "table"},
			},
			"location": {
				Type:          schema.TypeString,
				Optional:      true,
				ForceNew:      true,
				ValidateFunc:  validateArn,
				ConflictsWith: []string{"database", "table"},
			},
			"table": {
				Type:          schema.TypeList,
				Optional:      true,
				ForceNew:      true,
				ConflictsWith: []string{"database", "location"},
				MinItems:      0,
				MaxItems:      1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"database": {
							Type:         schema.TypeString,
							Required:     true,
							ForceNew:     true,
							ValidateFunc: validation.NoZeroValues,
						},
						"name": {
							Type:         schema.TypeString,
							Required:     true,
							ForceNew:     true,
							ValidateFunc: validation.NoZeroValues,
						},
						"column_names": {
							Type:     schema.TypeList,
							Optional: true,
							ForceNew: true,
							Elem: &schema.Schema{
								Type:         schema.TypeString,
								ValidateFunc: validation.NoZeroValues,
							},
						},
						"excluded_column_names": {
							Type:     schema.TypeList,
							Optional: true,
							ForceNew: true,
							Elem: &schema.Schema{
								Type:         schema.TypeString,
								ValidateFunc: validation.NoZeroValues,
							},
						},
					},
				},
			},
		},
	}
}

func resourceAwsLakeFormationPermissionsGrant(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	catalogId := createAwsDataCatalogId(d, meta.(*AWSClient).accountid)
	resource := expandAwsLakeFormationResource(d)

	input := &lakeformation.GrantPermissionsInput{
		CatalogId:   aws.String(catalogId),
		Permissions: expandStringList(d.Get("permissions").([]interface{})),
		Principal:   expandAwsLakeFormationPrincipal(d),
		Resource:    resource,
	}
	if vs, ok := d.GetOk("permissions_with_grant_option"); ok {
		input.PermissionsWithGrantOption = expandStringList(vs.([]interface{}))
	}

	// Catalog: CREATE_DATABASE
	// Location: DATA_LOCATION_ACCESS
	// Database: ALTER, CREATE_TABLE, DROP, (ALL ~ Super)
	// Table: ALTER, INSERT, DELETE, DROP, SELECT, (ALL ~ Super)
	// TableWithColumns: SELECT

	_, err := conn.GrantPermissions(input)
	if err != nil {
		return fmt.Errorf("Error granting LakeFormation Permissions: %s", err)
	}

	d.SetId(fmt.Sprintf("lakeformation:resource:%s:%s", catalogId, time.Now().UTC().String()))
	d.Set("catalog_id", catalogId)

	return resourceAwsLakeFormationPermissionsList(d, meta)
}

func resourceAwsLakeFormationPermissionsList(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	catalogId := d.Get("catalog_id").(string)

	input := &lakeformation.ListPermissionsInput{
		CatalogId: aws.String(catalogId),
		Principal: expandAwsLakeFormationPrincipal(d),
		Resource:  expandAwsLakeFormationResource(d),
	}

	out, err := conn.ListPermissions(input)
	if err != nil {
		return fmt.Errorf("Error listing LakeFormation Permissions: %s", err)
	}

	permissions := out.PrincipalResourcePermissions
	if len(permissions) == 0 {
		return fmt.Errorf("Error no LakeFormation Permissions found: %s", input)
	}

	permissionsHead := permissions[0]
	d.Set("catalog_id", catalogId)
	d.Set("principal", permissionsHead.Principal.DataLakePrincipalIdentifier)
	if dataLocation := permissionsHead.Resource.DataLocation; dataLocation != nil {
		d.Set("location", dataLocation.ResourceArn)
	}
	if database := permissionsHead.Resource.Database; database != nil {
		d.Set("database", database.Name)
	}
	if table := permissionsHead.Resource.Table; table != nil {
		d.Set("table", flattenAWSLakeFormationTable(table))
	}
	if table := permissionsHead.Resource.TableWithColumns; table != nil {
		d.Set("table", flattenAWSLakeFormationTableWithColumns(table))
	}
	var allPermissions, allPermissionsWithGrant []*string
	for _, p := range permissions {
		allPermissions = append(allPermissions, p.Permissions...)
		allPermissionsWithGrant = append(allPermissionsWithGrant, p.PermissionsWithGrantOption...)
	}
	d.Set("permissions", allPermissions)
	d.Set("permissions_with_grant_option", allPermissionsWithGrant)

	return nil
}

func resourceAwsLakeFormationPermissionsRevoke(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	catalogId := d.Get("catalog_id").(string)

	input := &lakeformation.RevokePermissionsInput{
		CatalogId:   aws.String(catalogId),
		Permissions: expandStringList(d.Get("permissions").([]interface{})),
		Principal:   expandAwsLakeFormationPrincipal(d),
		Resource:    expandAwsLakeFormationResource(d),
	}
	if vs, ok := d.GetOk("permissions_with_grant_option"); ok {
		input.PermissionsWithGrantOption = expandStringList(vs.([]interface{}))
	}

	_, err := conn.RevokePermissions(input)
	if err != nil {
		return fmt.Errorf("Error revoking LakeFormation Permissions: %s", err)
	}

	return nil
}

func expandAwsLakeFormationPrincipal(d *schema.ResourceData) *lakeformation.DataLakePrincipal {
	return &lakeformation.DataLakePrincipal{
		DataLakePrincipalIdentifier: aws.String(d.Get("principal").(string)),
	}
}

func expandAwsLakeFormationResource(d *schema.ResourceData) *lakeformation.Resource {
	if v, ok := d.GetOk("database"); ok {
		databaseName := v.(string)
		if len(databaseName) > 0 {
			return &lakeformation.Resource{
				Database: &lakeformation.DatabaseResource{
					Name: aws.String(databaseName),
				},
			}
		}
	}
	if v, ok := d.GetOk("location"); ok {
		location := v.(string)
		if len(location) > 0 {
			return &lakeformation.Resource{
				DataLocation: &lakeformation.DataLocationResource{
					ResourceArn: aws.String(v.(string)),
				},
			}
		}
	}
	if vs, ok := d.GetOk("table"); ok {
		tables := vs.([]interface{})
		if len(tables) > 0 {
			table := tables[0].(map[string]interface{})

			resource := &lakeformation.Resource{}
			var databaseName, tableName string
			var columnNames, excludedColumnNames []interface{}
			if x, ok := table["database"]; ok {
				databaseName = x.(string)
			}
			if x, ok := table["name"]; ok {
				tableName = x.(string)
			}
			if xs, ok := table["column_names"]; ok {
				columnNames = xs.([]interface{})
			}
			if xs, ok := table["excluded_column_names"]; ok {
				excludedColumnNames = xs.([]interface{})
			}

			if len(columnNames) > 0 || len(excludedColumnNames) > 0 {
				tableWithColumns := &lakeformation.TableWithColumnsResource{
					DatabaseName: aws.String(databaseName),
					Name:         aws.String(tableName),
				}
				if len(columnNames) > 0 {
					tableWithColumns.ColumnNames = expandStringList(columnNames)
				}
				if len(excludedColumnNames) > 0 {
					tableWithColumns.ColumnWildcard = &lakeformation.ColumnWildcard{
						ExcludedColumnNames: expandStringList(excludedColumnNames),
					}
				}
				resource.TableWithColumns = tableWithColumns
			} else {
				resource.Table = &lakeformation.TableResource{
					DatabaseName: aws.String(databaseName),
					Name:         aws.String(tableName),
				}
			}
			return resource
		}
	}
	return &lakeformation.Resource{
		Catalog: &lakeformation.CatalogResource{},
	}
}

func flattenAWSLakeFormationTable(tb *lakeformation.TableResource) map[string]interface{} {
	m := make(map[string]interface{})

	m["database"] = tb.DatabaseName
	m["name"] = tb.Name

	return m
}

func flattenAWSLakeFormationTableWithColumns(tb *lakeformation.TableWithColumnsResource) map[string]interface{} {
	m := make(map[string]interface{})

	m["database"] = tb.DatabaseName
	m["name"] = tb.Name
	if columnNames := tb.ColumnNames; columnNames != nil {
		m["column_names"] = columnNames
	}
	if columnWildcard := tb.ColumnWildcard; columnWildcard != nil {
		m["excluded_column_names"] = columnWildcard.ExcludedColumnNames
	}

	return m
}
<commit_msg>Use const defined in API for permissions<commit_after>package aws

import (
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/lakeformation"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"
)

func AwsLakeFormationPermissions() []string {
	return []string{
		lakeformation.PermissionAll,
		lakeformation.PermissionSelect,
		lakeformation.PermissionAlter,
		lakeformation.PermissionDrop,
		lakeformation.PermissionDelete,
		lakeformation.PermissionInsert,
		lakeformation.PermissionCreateDatabase,
		lakeformation.PermissionCreateTable,
		lakeformation.PermissionDataLocationAccess,
	}
}

func resourceAwsLakeFormationPermissions() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLakeFormationPermissionsGrant,
		Read:   resourceAwsLakeFormationPermissionsList,
		Delete: resourceAwsLakeFormationPermissionsRevoke,

		Schema: map[string]*schema.Schema{
			"catalog_id": {
				Type:         schema.TypeString,
				ForceNew:     true,
				Optional:     true,
				Computed:     true,
				ValidateFunc: validateAwsAccountId,
			},
			"permissions": {
				Type:     schema.TypeList,
				Required: true,
				ForceNew: true,
				MinItems: 1,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validation.StringInSlice(AwsLakeFormationPermissions(), false),
				},
			},
			"permissions_with_grant_option": {
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				Computed: true,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validation.StringInSlice(AwsLakeFormationPermissions(), false),
				},
			},
			"principal": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateArn,
			},
			"database": {
				Type:          schema.TypeString,
				Optional:      true,
				ForceNew:      true,
				ConflictsWith: []string{"location", "table"},
			},
			"location": {
				Type:          schema.TypeString,
				Optional:      true,
				ForceNew:      true,
				ValidateFunc:  validateArn,
				ConflictsWith: []string{"database", "table"},
			},
			"table": {
				Type:          schema.TypeList,
				Optional:      true,
				ForceNew:      true,
				ConflictsWith: []string{"database", "location"},
				MinItems:      0,
				MaxItems:      1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"database": {
							Type:         schema.TypeString,
							Required:     true,
							ForceNew:     true,
							ValidateFunc: validation.NoZeroValues,
						},
						"name": {
							Type:         schema.TypeString,
							Required:     true,
							ForceNew:     true,
							ValidateFunc: validation.NoZeroValues,
						},
						"column_names": {
							Type:     schema.TypeList,
							Optional: true,
							ForceNew: true,
							Elem: &schema.Schema{
								Type:         schema.TypeString,
								ValidateFunc: validation.NoZeroValues,
							},
						},
						"excluded_column_names": {
							Type:     schema.TypeList,
							Optional: true,
							ForceNew: true,
							Elem: &schema.Schema{
								Type:         schema.TypeString,
								ValidateFunc: validation.NoZeroValues,
							},
						},
					},
				},
			},
		},
	}
}

func resourceAwsLakeFormationPermissionsGrant(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	catalogId := createAwsDataCatalogId(d, meta.(*AWSClient).accountid)
	resource := expandAwsLakeFormationResource(d)

	input := &lakeformation.GrantPermissionsInput{
		CatalogId:   aws.String(catalogId),
		Permissions: expandStringList(d.Get("permissions").([]interface{})),
		Principal:   expandAwsLakeFormationPrincipal(d),
		Resource:    resource,
	}
	if vs, ok := d.GetOk("permissions_with_grant_option"); ok {
		input.PermissionsWithGrantOption = expandStringList(vs.([]interface{}))
	}

	// Catalog: CREATE_DATABASE
	// Location: DATA_LOCATION_ACCESS
	// Database: ALTER, CREATE_TABLE, DROP, (ALL ~ Super)
	// Table: ALTER, INSERT, DELETE, DROP, SELECT, (ALL ~ Super)
	// TableWithColumns: SELECT

	_, err := conn.GrantPermissions(input)
	if err != nil {
		return fmt.Errorf("Error granting LakeFormation Permissions: %s", err)
	}

	d.SetId(fmt.Sprintf("lakeformation:resource:%s:%s", catalogId, time.Now().UTC().String()))
	d.Set("catalog_id", catalogId)

	return resourceAwsLakeFormationPermissionsList(d, meta)
}

func resourceAwsLakeFormationPermissionsList(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	catalogId := d.Get("catalog_id").(string)

	input := &lakeformation.ListPermissionsInput{
		CatalogId: aws.String(catalogId),
		Principal: expandAwsLakeFormationPrincipal(d),
		Resource:  expandAwsLakeFormationResource(d),
	}

	out, err := conn.ListPermissions(input)
	if err != nil {
		return fmt.Errorf("Error listing LakeFormation Permissions: %s", err)
	}

	permissions := out.PrincipalResourcePermissions
	if len(permissions) == 0 {
		return fmt.Errorf("Error no LakeFormation Permissions found: %s", input)
	}

	permissionsHead := permissions[0]
	d.Set("catalog_id", catalogId)
	d.Set("principal", permissionsHead.Principal.DataLakePrincipalIdentifier)
	if dataLocation := permissionsHead.Resource.DataLocation; dataLocation != nil {
		d.Set("location", dataLocation.ResourceArn)
	}
	if database := permissionsHead.Resource.Database; database != nil {
		d.Set("database", database.Name)
	}
	if table := permissionsHead.Resource.Table; table != nil {
		d.Set("table", flattenAWSLakeFormationTable(table))
	}
	if table := permissionsHead.Resource.TableWithColumns; table != nil {
		d.Set("table", flattenAWSLakeFormationTableWithColumns(table))
	}
	var allPermissions, allPermissionsWithGrant []*string
	for _, p := range permissions {
		allPermissions = append(allPermissions, p.Permissions...)
		allPermissionsWithGrant = append(allPermissionsWithGrant, p.PermissionsWithGrantOption...)
	}
	d.Set("permissions", allPermissions)
	d.Set("permissions_with_grant_option", allPermissionsWithGrant)

	return nil
}

func resourceAwsLakeFormationPermissionsRevoke(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	catalogId := d.Get("catalog_id").(string)

	input := &lakeformation.RevokePermissionsInput{
		CatalogId:   aws.String(catalogId),
		Permissions: expandStringList(d.Get("permissions").([]interface{})),
		Principal:   expandAwsLakeFormationPrincipal(d),
		Resource:    expandAwsLakeFormationResource(d),
	}
	if vs, ok := d.GetOk("permissions_with_grant_option"); ok {
		input.PermissionsWithGrantOption = expandStringList(vs.([]interface{}))
	}

	_, err := conn.RevokePermissions(input)
	if err != nil {
		return fmt.Errorf("Error revoking LakeFormation Permissions: %s", err)
	}

	return nil
}

func expandAwsLakeFormationPrincipal(d *schema.ResourceData) *lakeformation.DataLakePrincipal {
	return &lakeformation.DataLakePrincipal{
		DataLakePrincipalIdentifier: aws.String(d.Get("principal").(string)),
	}
}

func expandAwsLakeFormationResource(d *schema.ResourceData) *lakeformation.Resource {
	if v, ok := d.GetOk("database"); ok {
		databaseName := v.(string)
		if len(databaseName) > 0 {
			return &lakeformation.Resource{
				Database: &lakeformation.DatabaseResource{
					Name: aws.String(databaseName),
				},
			}
		}
	}
	if v, ok := d.GetOk("location"); ok {
		location := v.(string)
		if len(location) > 0 {
			return &lakeformation.Resource{
				DataLocation: &lakeformation.DataLocationResource{
					ResourceArn: aws.String(v.(string)),
				},
			}
		}
	}
	if vs, ok := d.GetOk("table"); ok {
		tables := vs.([]interface{})
		if len(tables) > 0 {
			table := tables[0].(map[string]interface{})

			resource := &lakeformation.Resource{}
			var databaseName, tableName string
			var columnNames, excludedColumnNames []interface{}
			if x, ok := table["database"]; ok {
				databaseName = x.(string)
			}
			if x, ok := table["name"]; ok {
				tableName = x.(string)
			}
			if xs, ok := table["column_names"]; ok {
				columnNames = xs.([]interface{})
			}
			if xs, ok := table["excluded_column_names"]; ok {
				excludedColumnNames = xs.([]interface{})
			}

			if len(columnNames) > 0 || len(excludedColumnNames) > 0 {
				tableWithColumns := &lakeformation.TableWithColumnsResource{
					DatabaseName: aws.String(databaseName),
					Name:         aws.String(tableName),
				}
				if len(columnNames) > 0 {
					tableWithColumns.ColumnNames = expandStringList(columnNames)
				}
				if len(excludedColumnNames) > 0 {
					tableWithColumns.ColumnWildcard = &lakeformation.ColumnWildcard{
						ExcludedColumnNames: expandStringList(excludedColumnNames),
					}
				}
				resource.TableWithColumns = tableWithColumns
			} else {
				resource.Table = &lakeformation.TableResource{
					DatabaseName: aws.String(databaseName),
					Name:         aws.String(tableName),
				}
			}
			return resource
		}
	}
	return &lakeformation.Resource{
		Catalog: &lakeformation.CatalogResource{},
	}
}

func flattenAWSLakeFormationTable(tb *lakeformation.TableResource) map[string]interface{} {
	m := make(map[string]interface{})

	m["database"] = tb.DatabaseName
	m["name"] = tb.Name

	return m
}

func flattenAWSLakeFormationTableWithColumns(tb *lakeformation.TableWithColumnsResource) map[string]interface{} {
	m := make(map[string]interface{})

	m["database"] = tb.DatabaseName
	m["name"] = tb.Name
	if columnNames := tb.ColumnNames; columnNames != nil {
		m["column_names"] = columnNames
	}
	if columnWildcard := tb.ColumnWildcard; columnWildcard != nil {
		m["excluded_column_names"] = columnWildcard.ExcludedColumnNames
	}

	return m
}
<|endoftext|>


<commit_before>package mocks

import (
	"sync"
	"sync/atomic"

	"github.com/Shopify/sarama"
)

// Consumer implements sarama's Consumer interface for testing purposes.
// Before you can start consuming from this consumer, you have to register
// topic/partitions using ExpectConsumePartition, and set expectations on them.
type Consumer struct {
	l                  sync.Mutex
	t                  ErrorReporter
	config             *sarama.Config
	partitionConsumers map[string]map[int32]*PartitionConsumer
	metadata           map[string][]int32
}

// NewConsumer returns a new mock Consumer instance. The t argument should
// be the *testing.T instance of your test method. An error will be written to it if
// an expectation is violated. The config argument is currently unused and can be set to nil.
func NewConsumer(t ErrorReporter, config *sarama.Config) *Consumer {
	if config == nil {
		config = sarama.NewConfig()
	}

	c := &Consumer{
		t:                  t,
		config:             config,
		partitionConsumers: make(map[string]map[int32]*PartitionConsumer),
	}
	return c
}

///////////////////////////////////////////////////
// Consumer interface implementation
///////////////////////////////////////////////////

// ConsumePartition implements the ConsumePartition method from the sarama.Consumer interface.
// Before you can start consuming a partition, you have to set expectations on it using
// ExpectConsumePartition. You can only consume a partition once per consumer.
func (c *Consumer) ConsumePartition(topic string, partition int32, offset int64) (sarama.PartitionConsumer, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.partitionConsumers[topic] == nil || c.partitionConsumers[topic][partition] == nil {
		c.t.Errorf("No expectations set for %s/%d", topic, partition)
		return nil, errOutOfExpectations
	}

	pc := c.partitionConsumers[topic][partition]
	if pc.consumed {
		return nil, sarama.ConfigurationError("The topic/partition is already being consumed")
	}

	if pc.offset != AnyOffset && pc.offset != offset {
		c.t.Errorf("Unexpected offset when calling ConsumePartition for %s/%d. Expected %d, got %d.", topic, partition, pc.offset, offset)
	}

	pc.consumed = true
	return pc, nil
}

// Topics returns a list of topics, as registered with SetMetadata
func (c *Consumer) Topics() ([]string, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.metadata == nil {
		c.t.Errorf("Unexpected call to Topics. Initialize the mock's topic metadata with SetMetadata.")
		return nil, sarama.ErrOutOfBrokers
	}

	var result []string
	for topic := range c.metadata {
		result = append(result, topic)
	}
	return result, nil
}

// Partitions returns the list of parititons for the given topic, as registered with SetMetadata
func (c *Consumer) Partitions(topic string) ([]int32, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.metadata == nil {
		c.t.Errorf("Unexpected call to Partitions. Initialize the mock's topic metadata with SetMetadata.")
		return nil, sarama.ErrOutOfBrokers
	}
	if c.metadata[topic] == nil {
		return nil, sarama.ErrUnknownTopicOrPartition
	}

	return c.metadata[topic], nil
}

// Close implements the Close method from the sarama.Consumer interface. It will close
// all registered PartitionConsumer instances.
func (c *Consumer) Close() error {
	c.l.Lock()
	defer c.l.Unlock()

	for _, partitions := range c.partitionConsumers {
		for _, partitionConsumer := range partitions {
			_ = partitionConsumer.Close()
		}
	}

	return nil
}

///////////////////////////////////////////////////
// Expectation API
///////////////////////////////////////////////////

// SetTopicMetadata sets the clusters topic/partition metadata,
// which will be returned by Topics() and Partitions().
func (c *Consumer) SetTopicMetadata(metadata map[string][]int32) {
	c.l.Lock()
	defer c.l.Unlock()

	c.metadata = metadata
}

// ExpectConsumePartition will register a topic/partition, so you can set expectations on it.
// The registered PartitionConsumer will be returned, so you can set expectations
// on it using method chaining. Once a topic/partition is registered, you are
// expected to start consuming it using ConsumePartition. If that doesn't happen,
// an error will be written to the error reporter once the mock consumer is closed. It will
// also expect that the
func (c *Consumer) ExpectConsumePartition(topic string, partition int32, offset int64) *PartitionConsumer {
	c.l.Lock()
	defer c.l.Unlock()

	if c.partitionConsumers[topic] == nil {
		c.partitionConsumers[topic] = make(map[int32]*PartitionConsumer)
	}

	if c.partitionConsumers[topic][partition] == nil {
		c.partitionConsumers[topic][partition] = &PartitionConsumer{
			t:         c.t,
			topic:     topic,
			partition: partition,
			offset:    offset,
			messages:  make(chan *sarama.ConsumerMessage, c.config.ChannelBufferSize),
			errors:    make(chan *sarama.ConsumerError, c.config.ChannelBufferSize),
		}
	}

	return c.partitionConsumers[topic][partition]
}

///////////////////////////////////////////////////
// PartitionConsumer mock type
///////////////////////////////////////////////////

// PartitionConsumer implements sarama's PartitionConsumer interface for testing purposes.
// It is returned by the mock Consumers ConsumePartitionMethod, but only if it is
// registered first using the Consumer's ExpectConsumePartition method. Before consuming the
// Errors and Messages channel, you should specify what values will be provided on these
// channels using YieldMessage and YieldError.
type PartitionConsumer struct {
	l                       sync.Mutex
	t                       ErrorReporter
	topic                   string
	partition               int32
	offset                  int64
	messages                chan *sarama.ConsumerMessage
	errors                  chan *sarama.ConsumerError
	singleClose             sync.Once
	consumed                bool
	errorsShouldBeDrained   bool
	messagesShouldBeDrained bool
	highWaterMarkOffset     int64
}

///////////////////////////////////////////////////
// PartitionConsumer interface implementation
///////////////////////////////////////////////////

// AsyncClose implements the AsyncClose method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) AsyncClose() {
	pc.singleClose.Do(func() {
		close(pc.messages)
		close(pc.errors)
	})
}

// Close implements the Close method from the sarama.PartitionConsumer interface. It will
// verify whether the partition consumer was actually started.
func (pc *PartitionConsumer) Close() error {
	if !pc.consumed {
		pc.t.Errorf("Expectations set on %s/%d, but no partition consumer was started.", pc.topic, pc.partition)
		return errPartitionConsumerNotStarted
	}

	if pc.errorsShouldBeDrained && len(pc.errors) > 0 {
		pc.t.Errorf("Expected the errors channel for %s/%d to be drained on close, but found %d errors.", pc.topic, pc.partition, len(pc.errors))
	}

	if pc.messagesShouldBeDrained && len(pc.messages) > 0 {
		pc.t.Errorf("Expected the messages channel for %s/%d to be drained on close, but found %d messages.", pc.topic, pc.partition, len(pc.messages))
	}

	pc.AsyncClose()

	var (
		closeErr error
		wg       sync.WaitGroup
	)

	wg.Add(1)
	go func() {
		defer wg.Done()

		var errs = make(sarama.ConsumerErrors, 0)
		for err := range pc.errors {
			errs = append(errs, err)
		}

		if len(errs) > 0 {
			closeErr = errs
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		for _ = range pc.messages {
			// drain
		}
	}()

	wg.Wait()
	return closeErr
}

// Errors implements the Errors method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) Errors() <-chan *sarama.ConsumerError {
	return pc.errors
}

// Messages implements the Messages method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) Messages() <-chan *sarama.ConsumerMessage {
	return pc.messages
}

func (pc *PartitionConsumer) HighWaterMarkOffset() int64 {
	return atomic.LoadInt64(&pc.highWaterMarkOffset) + 1
}

///////////////////////////////////////////////////
// Expectation API
///////////////////////////////////////////////////

// YieldMessage will yield a messages Messages channel of this partition consumer
// when it is consumed. By default, the mock consumer will not verify whether this
// message was consumed from the Messages channel, because there are legitimate
// reasons forthis not to happen. ou can call ExpectMessagesDrainedOnClose so it will
// verify that the channel is empty on close.
func (pc *PartitionConsumer) YieldMessage(msg *sarama.ConsumerMessage) {
	pc.l.Lock()
	defer pc.l.Unlock()

	msg.Topic = pc.topic
	msg.Partition = pc.partition
	msg.Offset = atomic.AddInt64(&pc.highWaterMarkOffset, 1)

	pc.messages <- msg
}

// YieldError will yield an error on the Errors channel of this partition consumer
// when it is consumed. By default, the mock consumer will not verify whether this error was
// consumed from the Errors channel, because there are legitimate reasons for this
// not to happen. You can call ExpectErrorsDrainedOnClose so it will verify that
// the channel is empty on close.
func (pc *PartitionConsumer) YieldError(err error) {
	pc.errors <- &sarama.ConsumerError{
		Topic:     pc.topic,
		Partition: pc.partition,
		Err:       err,
	}
}

// ExpectMessagesDrainedOnClose sets an expectation on the partition consumer
// that the messages channel will be fully drained when Close is called. If this
// expectation is not met, an error is reported to the error reporter.
func (pc *PartitionConsumer) ExpectMessagesDrainedOnClose() {
	pc.messagesShouldBeDrained = true
}

// ExpectErrorsDrainedOnClose sets an expectation on the partition consumer
// that the errors channel will be fully drained when Close is called. If this
// expectation is not met, an error is reported to the error reporter.
func (pc *PartitionConsumer) ExpectErrorsDrainedOnClose() {
	pc.errorsShouldBeDrained = true
}
<commit_msg>fix mocks<commit_after>package mocks

import (
	"sync"
	"sync/atomic"

	"github.com/Shopify/sarama"
)

// Consumer implements sarama's Consumer interface for testing purposes.
// Before you can start consuming from this consumer, you have to register
// topic/partitions using ExpectConsumePartition, and set expectations on them.
type Consumer struct {
	l                  sync.Mutex
	t                  ErrorReporter
	config             *sarama.Config
	partitionConsumers map[string]map[int32]*PartitionConsumer
	metadata           map[string][]int32
}

// NewConsumer returns a new mock Consumer instance. The t argument should
// be the *testing.T instance of your test method. An error will be written to it if
// an expectation is violated. The config argument is currently unused and can be set to nil.
func NewConsumer(t ErrorReporter, config *sarama.Config) *Consumer {
	if config == nil {
		config = sarama.NewConfig()
	}

	c := &Consumer{
		t:                  t,
		config:             config,
		partitionConsumers: make(map[string]map[int32]*PartitionConsumer),
	}
	return c
}

///////////////////////////////////////////////////
// Consumer interface implementation
///////////////////////////////////////////////////

// ConsumePartition implements the ConsumePartition method from the sarama.Consumer interface.
// Before you can start consuming a partition, you have to set expectations on it using
// ExpectConsumePartition. You can only consume a partition once per consumer.
func (c *Consumer) ConsumePartition(topic string, partition int32, offset int64) (sarama.PartitionConsumer, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.partitionConsumers[topic] == nil || c.partitionConsumers[topic][partition] == nil {
		c.t.Errorf("No expectations set for %s/%d", topic, partition)
		return nil, errOutOfExpectations
	}

	pc := c.partitionConsumers[topic][partition]
	if pc.consumed {
		return nil, sarama.ConfigurationError("The topic/partition is already being consumed")
	}

	if pc.offset != AnyOffset && pc.offset != offset {
		c.t.Errorf("Unexpected offset when calling ConsumePartition for %s/%d. Expected %d, got %d.", topic, partition, pc.offset, offset)
	}

	pc.consumed = true
	return pc, nil
}

// Topics returns a list of topics, as registered with SetMetadata
func (c *Consumer) Topics() ([]string, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.metadata == nil {
		c.t.Errorf("Unexpected call to Topics. Initialize the mock's topic metadata with SetMetadata.")
		return nil, sarama.ErrOutOfBrokers
	}

	var result []string
	for topic := range c.metadata {
		result = append(result, topic)
	}
	return result, nil
}

// Partitions returns the list of parititons for the given topic, as registered with SetMetadata
func (c *Consumer) Partitions(topic string) ([]int32, error) {
	c.l.Lock()
	defer c.l.Unlock()

	if c.metadata == nil {
		c.t.Errorf("Unexpected call to Partitions. Initialize the mock's topic metadata with SetMetadata.")
		return nil, sarama.ErrOutOfBrokers
	}
	if c.metadata[topic] == nil {
		return nil, sarama.ErrUnknownTopicOrPartition
	}

	return c.metadata[topic], nil
}

func (c *Consumer) HighWaterMarks() map[string]map[int32]int64 {
	c.l.Lock()
	defer c.l.Unlock()

	hwms := make(map[string]map[int32]int64, len(c.partitionConsumers))
	for topic, partitionConsumers := range c.partitionConsumers {
		hwm := make(map[int32]int64, len(partitionConsumers))
		for partition, pc := range partitionConsumers {
			hwm[partition] = pc.HighWaterMarkOffset()
		}
		hwms[topic] = hwm
	}

	return hwms
}

// Close implements the Close method from the sarama.Consumer interface. It will close
// all registered PartitionConsumer instances.
func (c *Consumer) Close() error {
	c.l.Lock()
	defer c.l.Unlock()

	for _, partitions := range c.partitionConsumers {
		for _, partitionConsumer := range partitions {
			_ = partitionConsumer.Close()
		}
	}

	return nil
}

///////////////////////////////////////////////////
// Expectation API
///////////////////////////////////////////////////

// SetTopicMetadata sets the clusters topic/partition metadata,
// which will be returned by Topics() and Partitions().
func (c *Consumer) SetTopicMetadata(metadata map[string][]int32) {
	c.l.Lock()
	defer c.l.Unlock()

	c.metadata = metadata
}

// ExpectConsumePartition will register a topic/partition, so you can set expectations on it.
// The registered PartitionConsumer will be returned, so you can set expectations
// on it using method chaining. Once a topic/partition is registered, you are
// expected to start consuming it using ConsumePartition. If that doesn't happen,
// an error will be written to the error reporter once the mock consumer is closed. It will
// also expect that the
func (c *Consumer) ExpectConsumePartition(topic string, partition int32, offset int64) *PartitionConsumer {
	c.l.Lock()
	defer c.l.Unlock()

	if c.partitionConsumers[topic] == nil {
		c.partitionConsumers[topic] = make(map[int32]*PartitionConsumer)
	}

	if c.partitionConsumers[topic][partition] == nil {
		c.partitionConsumers[topic][partition] = &PartitionConsumer{
			t:         c.t,
			topic:     topic,
			partition: partition,
			offset:    offset,
			messages:  make(chan *sarama.ConsumerMessage, c.config.ChannelBufferSize),
			errors:    make(chan *sarama.ConsumerError, c.config.ChannelBufferSize),
		}
	}

	return c.partitionConsumers[topic][partition]
}

///////////////////////////////////////////////////
// PartitionConsumer mock type
///////////////////////////////////////////////////

// PartitionConsumer implements sarama's PartitionConsumer interface for testing purposes.
// It is returned by the mock Consumers ConsumePartitionMethod, but only if it is
// registered first using the Consumer's ExpectConsumePartition method. Before consuming the
// Errors and Messages channel, you should specify what values will be provided on these
// channels using YieldMessage and YieldError.
type PartitionConsumer struct {
	l                       sync.Mutex
	t                       ErrorReporter
	topic                   string
	partition               int32
	offset                  int64
	messages                chan *sarama.ConsumerMessage
	errors                  chan *sarama.ConsumerError
	singleClose             sync.Once
	consumed                bool
	errorsShouldBeDrained   bool
	messagesShouldBeDrained bool
	highWaterMarkOffset     int64
}

///////////////////////////////////////////////////
// PartitionConsumer interface implementation
///////////////////////////////////////////////////

// AsyncClose implements the AsyncClose method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) AsyncClose() {
	pc.singleClose.Do(func() {
		close(pc.messages)
		close(pc.errors)
	})
}

// Close implements the Close method from the sarama.PartitionConsumer interface. It will
// verify whether the partition consumer was actually started.
func (pc *PartitionConsumer) Close() error {
	if !pc.consumed {
		pc.t.Errorf("Expectations set on %s/%d, but no partition consumer was started.", pc.topic, pc.partition)
		return errPartitionConsumerNotStarted
	}

	if pc.errorsShouldBeDrained && len(pc.errors) > 0 {
		pc.t.Errorf("Expected the errors channel for %s/%d to be drained on close, but found %d errors.", pc.topic, pc.partition, len(pc.errors))
	}

	if pc.messagesShouldBeDrained && len(pc.messages) > 0 {
		pc.t.Errorf("Expected the messages channel for %s/%d to be drained on close, but found %d messages.", pc.topic, pc.partition, len(pc.messages))
	}

	pc.AsyncClose()

	var (
		closeErr error
		wg       sync.WaitGroup
	)

	wg.Add(1)
	go func() {
		defer wg.Done()

		var errs = make(sarama.ConsumerErrors, 0)
		for err := range pc.errors {
			errs = append(errs, err)
		}

		if len(errs) > 0 {
			closeErr = errs
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		for _ = range pc.messages {
			// drain
		}
	}()

	wg.Wait()
	return closeErr
}

// Errors implements the Errors method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) Errors() <-chan *sarama.ConsumerError {
	return pc.errors
}

// Messages implements the Messages method from the sarama.PartitionConsumer interface.
func (pc *PartitionConsumer) Messages() <-chan *sarama.ConsumerMessage {
	return pc.messages
}

func (pc *PartitionConsumer) HighWaterMarkOffset() int64 {
	return atomic.LoadInt64(&pc.highWaterMarkOffset) + 1
}

///////////////////////////////////////////////////
// Expectation API
///////////////////////////////////////////////////

// YieldMessage will yield a messages Messages channel of this partition consumer
// when it is consumed. By default, the mock consumer will not verify whether this
// message was consumed from the Messages channel, because there are legitimate
// reasons forthis not to happen. ou can call ExpectMessagesDrainedOnClose so it will
// verify that the channel is empty on close.
func (pc *PartitionConsumer) YieldMessage(msg *sarama.ConsumerMessage) {
	pc.l.Lock()
	defer pc.l.Unlock()

	msg.Topic = pc.topic
	msg.Partition = pc.partition
	msg.Offset = atomic.AddInt64(&pc.highWaterMarkOffset, 1)

	pc.messages <- msg
}

// YieldError will yield an error on the Errors channel of this partition consumer
// when it is consumed. By default, the mock consumer will not verify whether this error was
// consumed from the Errors channel, because there are legitimate reasons for this
// not to happen. You can call ExpectErrorsDrainedOnClose so it will verify that
// the channel is empty on close.
func (pc *PartitionConsumer) YieldError(err error) {
	pc.errors <- &sarama.ConsumerError{
		Topic:     pc.topic,
		Partition: pc.partition,
		Err:       err,
	}
}

// ExpectMessagesDrainedOnClose sets an expectation on the partition consumer
// that the messages channel will be fully drained when Close is called. If this
// expectation is not met, an error is reported to the error reporter.
func (pc *PartitionConsumer) ExpectMessagesDrainedOnClose() {
	pc.messagesShouldBeDrained = true
}

// ExpectErrorsDrainedOnClose sets an expectation on the partition consumer
// that the errors channel will be fully drained when Close is called. If this
// expectation is not met, an error is reported to the error reporter.
func (pc *PartitionConsumer) ExpectErrorsDrainedOnClose() {
	pc.errorsShouldBeDrained = true
}
<|endoftext|>


<commit_before>package countdown

import (
	"fmt"
	"math"
	"time"
)

// Countdown contains the target date and a ticker to compare it to
type Countdown struct {
	To time.Time
}

// New creates a new countdown object
func New(to time.Time) (*Countdown, error) {
	if !happened(to) {
		return &Countdown{To: to}, nil
	}
	return &Countdown{}, fmt.Errorf("%v is already in the past", to.Format(time.RFC822))
}

// DHMS returns "<days>d <hours>h <mins>m <sec>s"
func (c *Countdown) DHMS() string {
	return fmt.Sprintf("%vd %vh %vm %vs", c.TotalDays(), c.RemainingHours(), c.RemainingMinutes(), c.RemainingSeconds())
}

// DHMSColonSeparated returns "<days>:<hours>:<mins>:<secs>"
func (c *Countdown) DHMSColonSeparated() string {
	return fmt.Sprintf("%v:%v:%v:%v", digital(c.TotalDays()), digital(c.RemainingHours()), digital(c.RemainingMinutes()), digital(c.RemainingSeconds()))
}

// HMS returns "<hours>:<mins>:<sec>"
func (c *Countdown) HMS() string {
	return fmt.Sprintf("%v:%v:%v", digital(c.TotalHours()), digital(c.RemainingMinutes()), digital(c.RemainingSeconds()))
}

// MS returns "<mins>:<sec>"
func (c *Countdown) MS() string {
	return fmt.Sprintf("%v:%v", digital(c.TotalMinutes()), digital(c.RemainingSeconds()))
}

// TotalDays returns the countdown in days
func (c *Countdown) TotalDays() int {
	if !happened(c.To) {
		return int(math.Floor(c.To.Sub(time.Now()).Seconds() / 3600 / 24))
	}
	return 0
}

// TotalHours return the countdown in hours
func (c *Countdown) TotalHours() int {
	if !happened(c.To) {
		return int(math.Floor(c.To.Sub(time.Now()).Seconds() / 3600))
	}
	return 0
}

// TotalMinutes returns the countdown in minutes
func (c *Countdown) TotalMinutes() int {
	if !happened(c.To) {
		return int(math.Floor(c.To.Sub(time.Now()).Seconds() / 3600 / 60))
	}
	return 0
}

// TotalSeconds returns the conutdown in seconds
func (c *Countdown) TotalSeconds() int {
	if !happened(c.To) {
		return int(c.To.Sub(time.Now()).Seconds())
	}
	return 0
}

// RemainingHours returns the remaining hours
func (c *Countdown) RemainingHours() int {
	if !happened(c.To) {
		return int(math.Floor(math.Mod(c.To.Sub(time.Now()).Seconds()/3600, 24)))
	}
	return 0
}

// RemainingMinutes returns the remaining minutes
func (c *Countdown) RemainingMinutes() int {
	if !happened(c.To) {
		return int(math.Floor(math.Mod(c.To.Sub(time.Now()).Seconds()/60, 60)))
	}
	return 0
}

// RemainingSeconds returns the remaining seconds
func (c *Countdown) RemainingSeconds() int {
	if !happened(c.To) {
		return int(math.Floor(math.Mod(c.To.Sub(time.Now()).Seconds(), 60)))
	}
	return 0
}

func happened(to time.Time) bool {
	if time.Now().Before(to) {
		return false
	}
	return true
}

func digital(i int) string {
	if i < 10 {
		return fmt.Sprintf("0%v", i)
	}
	return fmt.Sprintf("%v", i)
}
<commit_msg>fixed redundant comment<commit_after>package countdown

import (
	"fmt"
	"math"
	"time"
)

// Countdown contains the target date
type Countdown struct {
	To time.Time
}

// New creates a new countdown object
func New(to time.Time) (*Countdown, error) {
	if !happened(to) {
		return &Countdown{To: to}, nil
	}
	return &Countdown{}, fmt.Errorf("%v is already in the past", to.Format(time.RFC822))
}

// DHMS returns "<days>d <hours>h <mins>m <sec>s"
func (c *Countdown) DHMS() string {
	return fmt.Sprintf("%vd %vh %vm %vs", c.TotalDays(), c.RemainingHours(), c.RemainingMinutes(), c.RemainingSeconds())
}

// DHMSColonSeparated returns "<days>:<hours>:<mins>:<secs>"
func (c *Countdown) DHMSColonSeparated() string {
	return fmt.Sprintf("%v:%v:%v:%v", digital(c.TotalDays()), digital(c.RemainingHours()), digital(c.RemainingMinutes()), digital(c.RemainingSeconds()))
}

// HMS returns "<hours>:<mins>:<sec>"
func (c *Countdown) HMS() string {
	return fmt.Sprintf("%v:%v:%v", digital(c.TotalHours()), digital(c.RemainingMinutes()), digital(c.RemainingSeconds()))
}

// MS returns "<mins>:<sec>"
func (c *Countdown) MS() string {
	return fmt.Sprintf("%v:%v", digital(c.TotalMinutes()), digital(c.RemainingSeconds()))
}

// TotalDays returns the countdown in days
func (c *Countdown) TotalDays() int {
	if !happened(c.To) {
		return int(math.Floor(c.To.Sub(time.Now()).Seconds() / 3600 / 24))
	}
	return 0
}

// TotalHours return the countdown in hours
func (c *Countdown) TotalHours() int {
	if !happened(c.To) {
		return int(math.Floor(c.To.Sub(time.Now()).Seconds() / 3600))
	}
	return 0
}

// TotalMinutes returns the countdown in minutes
func (c *Countdown) TotalMinutes() int {
	if !happened(c.To) {
		return int(math.Floor(c.To.Sub(time.Now()).Seconds() / 3600 / 60))
	}
	return 0
}

// TotalSeconds returns the conutdown in seconds
func (c *Countdown) TotalSeconds() int {
	if !happened(c.To) {
		return int(c.To.Sub(time.Now()).Seconds())
	}
	return 0
}

// RemainingHours returns the remaining hours
func (c *Countdown) RemainingHours() int {
	if !happened(c.To) {
		return int(math.Floor(math.Mod(c.To.Sub(time.Now()).Seconds()/3600, 24)))
	}
	return 0
}

// RemainingMinutes returns the remaining minutes
func (c *Countdown) RemainingMinutes() int {
	if !happened(c.To) {
		return int(math.Floor(math.Mod(c.To.Sub(time.Now()).Seconds()/60, 60)))
	}
	return 0
}

// RemainingSeconds returns the remaining seconds
func (c *Countdown) RemainingSeconds() int {
	if !happened(c.To) {
		return int(math.Floor(math.Mod(c.To.Sub(time.Now()).Seconds(), 60)))
	}
	return 0
}

func happened(to time.Time) bool {
	if time.Now().Before(to) {
		return false
	}
	return true
}

func digital(i int) string {
	if i < 10 {
		return fmt.Sprintf("0%v", i)
	}
	return fmt.Sprintf("%v", i)
}
<|endoftext|>