<commit_before>package window

// An array-based moving window; a moving window
// is a queue with a maximum size and the property
// that when the size is reached, pushing a new
// element into the queue causes the head to be
// popped.
//
// You can optimize the amount of slice copying
// that the MovingWindow will be doing by trading
// off with space complexity. Namely, the underlying
// array is allocated with a size that is the
// multiple of the intended capacity of the queue
// so that copying is less frequent.
type MovingWindow struct {
    arr  []float64
    size int
    head int
    tail int
}

// PushBack will push a new piece of data into
// the moving window
func (m *MovingWindow) PushBack(v float64) {
    // if the array is full, rewind
    if m.tail == len(m.arr) {
        m.rewind()
    }
    // push the value
    m.arr[m.tail] = v
    // check if the window is full,
    // and move head pointer appropriately
    if m.tail-m.head >= m.size {
        m.head++
    }
    m.tail++
}

// rewind will copy the last size-1 elements
// from the end of the underlying array to
// the front, starting at index 0.
func (m *MovingWindow) rewind() {
    l := len(m.arr)
    for i := 0; i < m.size-1; i++ {
        m.arr[i] = m.arr[l-m.size+i+1]
    }
    m.head, m.tail = 0, m.size-1
}

// Slice will present the MovingWindow in
// the form of a slice. This operation never
// requires array copying of any kind.
func (m *MovingWindow) Slice() []float64 {
    return m.arr[m.head:m.tail]
}

// Size returns the size of the moving window,
// which is set at initialization
func (m *MovingWindow) Size() int {
    return m.size
}

// New creates a new moving window,
// with the size and multiple specified.
//
// This data structures trades off space
// and copying complexity; more precisely,
// the number of moving windows that can
// be displayed without having to do any
// array copying is proportional to approx 1/M,
// where M is the multiple.
func New(size, multiple int) *MovingWindow {
    if size < 1 {
        panic("Must have positive size")
    }
    capacity := size*multiple
    return &MovingWindow{
        arr:  make([]float64, capacity, capacity),
        size: size,
    }
}
<commit_msg>Added extra checks.<commit_after>package window

// An array-based moving window; a moving window
// is a queue with a maximum size and the property
// that when the size is reached, pushing a new
// element into the queue causes the head to be
// popped.
//
// You can optimize the amount of slice copying
// that the MovingWindow will be doing by trading
// off with space complexity. Namely, the underlying
// array is allocated with a size that is the
// multiple of the intended capacity of the queue
// so that copying is less frequent.
type MovingWindow struct {
    arr  []float64
    size int
    head int
    tail int
}

// PushBack will push a new piece of data into
// the moving window
func (m *MovingWindow) PushBack(v float64) {
    // if the array is full, rewind
    if m.tail == len(m.arr) {
        m.rewind()
    }
    // push the value
    m.arr[m.tail] = v
    // check if the window is full,
    // and move head pointer appropriately
    if m.tail-m.head >= m.size {
        m.head++
    }
    m.tail++
}

// rewind will copy the last size-1 elements
// from the end of the underlying array to
// the front, starting at index 0.
func (m *MovingWindow) rewind() {
    l := len(m.arr)
    for i := 0; i < m.size-1; i++ {
        m.arr[i] = m.arr[l-m.size+i+1]
    }
    m.head, m.tail = 0, m.size-1
}

// Slice will present the MovingWindow in
// the form of a slice. This operation never
// requires array copying of any kind.
func (m *MovingWindow) Slice() []float64 {
    return m.arr[m.head:m.tail]
}

// Size returns the size of the moving window,
// which is set at initialization
func (m *MovingWindow) Size() int {
    return m.size
}

// New creates a new moving window,
// with the size and multiple specified.
//
// This data structures trades off space
// and copying complexity; more precisely,
// the number of moving windows that can
// be displayed without having to do any
// array copying is proportional to approx 1/M,
// where M is the multiple.
func New(size, multiple int) *MovingWindow {
    if size < 1 || multiple < 1{
        panic("Must have positive size and multiple")
    }
    capacity := size*multiple
    return &MovingWindow{
        arr:  make([]float64, capacity, capacity),
        size: size,
    }
}
<|endoftext|>


<commit_before>package upgrader

import (
	"context"
	"errors"
	"fmt"
	"net"
	"reflect"
	"time"

	"github.com/libp2p/go-libp2p-core/connmgr"
	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	ipnet "github.com/libp2p/go-libp2p-core/pnet"
	"github.com/libp2p/go-libp2p-core/sec"
	"github.com/libp2p/go-libp2p-core/transport"

	pnet "github.com/libp2p/go-libp2p-pnet"
	manet "github.com/multiformats/go-multiaddr/net"
)

// ErrNilPeer is returned when attempting to upgrade an outbound connection
// without specifying a peer ID.
var ErrNilPeer = errors.New("nil peer")

// AcceptQueueLength is the number of connections to fully setup before not accepting any new connections
var AcceptQueueLength = 16

const defaultAcceptTimeout = 15 * time.Second

type Option func(*upgrader) error

func WithPSK(psk ipnet.PSK) Option {
	return func(u *upgrader) error {
		u.psk = psk
		return nil
	}
}

func WithAcceptTimeout(t time.Duration) Option {
	return func(u *upgrader) error {
		u.acceptTimeout = t
		return nil
	}
}

func WithConnectionGater(g connmgr.ConnectionGater) Option {
	return func(u *upgrader) error {
		u.connGater = g
		return nil
	}
}

func WithResourceManager(m network.ResourceManager) Option {
	return func(u *upgrader) error {
		u.rcmgr = m
		return nil
	}
}

// Upgrader is a multistream upgrader that can upgrade an underlying connection
// to a full transport connection (secure and multiplexed).
type upgrader struct {
	secure sec.SecureMuxer
	muxer  network.Multiplexer

	psk       ipnet.PSK
	connGater connmgr.ConnectionGater
	rcmgr     network.ResourceManager

	// AcceptTimeout is the maximum duration an Accept is allowed to take.
	// This includes the time between accepting the raw network connection,
	// protocol selection as well as the handshake, if applicable.
	//
	// If unset, the default value (15s) is used.
	acceptTimeout time.Duration
}

var _ transport.Upgrader = &upgrader{}

func New(secureMuxer sec.SecureMuxer, muxer network.Multiplexer, opts ...Option) (transport.Upgrader, error) {
	u := &upgrader{
		secure:        secureMuxer,
		muxer:         muxer,
		acceptTimeout: defaultAcceptTimeout,
	}
	for _, opt := range opts {
		if err := opt(u); err != nil {
			return nil, err
		}
	}
	if u.rcmgr == nil {
		u.rcmgr = network.NullResourceManager
	}
	return u, nil
}

// UpgradeListener upgrades the passed multiaddr-net listener into a full libp2p-transport listener.
func (u *upgrader) UpgradeListener(t transport.Transport, list manet.Listener) transport.Listener {
	ctx, cancel := context.WithCancel(context.Background())
	l := &listener{
		Listener:  list,
		upgrader:  u,
		transport: t,
		rcmgr:     u.rcmgr,
		threshold: newThreshold(AcceptQueueLength),
		incoming:  make(chan transport.CapableConn),
		cancel:    cancel,
		ctx:       ctx,
	}
	go l.handleIncoming()
	return l
}

// Upgrade upgrades the multiaddr/net connection into a full libp2p-transport connection.
func (u *upgrader) Upgrade(ctx context.Context, t transport.Transport, maconn manet.Conn, dir network.Direction, p peer.ID, connScope network.ConnManagementScope) (transport.CapableConn, error) {
	c, err := u.upgrade(ctx, t, maconn, dir, p, connScope)
	if err != nil {
		connScope.Done()
		return nil, err
	}
	return c, nil
}

func (u *upgrader) upgrade(ctx context.Context, t transport.Transport, maconn manet.Conn, dir network.Direction, p peer.ID, connScope network.ConnManagementScope) (transport.CapableConn, error) {
	if dir == network.DirOutbound && p == "" {
		return nil, ErrNilPeer
	}
	var stat network.ConnStats
	if cs, ok := maconn.(network.ConnStat); ok {
		stat = cs.Stat()
	}

	var conn net.Conn = maconn
	if u.psk != nil {
		pconn, err := pnet.NewProtectedConn(u.psk, conn)
		if err != nil {
			conn.Close()
			return nil, fmt.Errorf("failed to setup private network protector: %s", err)
		}
		conn = pconn
	} else if ipnet.ForcePrivateNetwork {
		log.Error("tried to dial with no Private Network Protector but usage of Private Networks is forced by the environment")
		return nil, ipnet.ErrNotInPrivateNetwork
	}

	sconn, server, err := u.setupSecurity(ctx, conn, p, dir)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to negotiate security protocol: %s", err)
	}

	// call the connection gater, if one is registered.
	if u.connGater != nil && !u.connGater.InterceptSecured(dir, sconn.RemotePeer(), maconn) {
		if err := maconn.Close(); err != nil {
			log.Errorw("failed to close connection", "peer", p, "addr", maconn.RemoteMultiaddr(), "error", err)
		}
		return nil, fmt.Errorf("gater rejected connection with peer %s and addr %s with direction %d",
			sconn.RemotePeer().Pretty(), maconn.RemoteMultiaddr(), dir)
	}
	// Only call SetPeer if it hasn't already been set -- this can happen when we don't know
	// the peer in advance and in some bug scenarios.
	if ps := connScope.PeerScope(); ps == nil || reflect.ValueOf(ps).IsNil() {
		if err := connScope.SetPeer(sconn.RemotePeer()); err != nil {
			log.Debugw("resource manager blocked connection for peer", "peer", sconn.RemotePeer(), "addr", conn.RemoteAddr(), "error", err)
			if err := maconn.Close(); err != nil {
				log.Errorw("failed to close connection", "peer", p, "addr", maconn.RemoteMultiaddr(), "error", err)
			}
			return nil, fmt.Errorf("resource manager connection with peer %s and addr %s with direction %d",
				sconn.RemotePeer().Pretty(), maconn.RemoteMultiaddr(), dir)
		}
	}

	smconn, err := u.setupMuxer(ctx, sconn, server, connScope.PeerScope())
	if err != nil {
		sconn.Close()
		return nil, fmt.Errorf("failed to negotiate stream multiplexer: %s", err)
	}

	tc := &transportConn{
		MuxedConn:      smconn,
		ConnMultiaddrs: maconn,
		ConnSecurity:   sconn,
		transport:      t,
		stat:           stat,
		scope:          connScope,
	}
	return tc, nil
}

func (u *upgrader) setupSecurity(ctx context.Context, conn net.Conn, p peer.ID, dir network.Direction) (sec.SecureConn, bool, error) {
	if dir == network.DirInbound {
		return u.secure.SecureInbound(ctx, conn, p)
	}
	return u.secure.SecureOutbound(ctx, conn, p)
}

func (u *upgrader) setupMuxer(ctx context.Context, conn net.Conn, server bool, scope network.PeerScope) (network.MuxedConn, error) {
	// TODO: The muxer should take a context.
	done := make(chan struct{})

	var smconn network.MuxedConn
	var err error
	go func() {
		defer close(done)
		smconn, err = u.muxer.NewConn(conn, server, scope)
	}()

	select {
	case <-done:
		return smconn, err
	case <-ctx.Done():
		// interrupt this process
		conn.Close()
		// wait to finish
		<-done
		return nil, ctx.Err()
	}
}
<commit_msg>avoid reflection<commit_after>package upgrader

import (
	"context"
	"errors"
	"fmt"
	"net"
	"time"

	"github.com/libp2p/go-libp2p-core/connmgr"
	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	ipnet "github.com/libp2p/go-libp2p-core/pnet"
	"github.com/libp2p/go-libp2p-core/sec"
	"github.com/libp2p/go-libp2p-core/transport"

	pnet "github.com/libp2p/go-libp2p-pnet"
	manet "github.com/multiformats/go-multiaddr/net"
)

// ErrNilPeer is returned when attempting to upgrade an outbound connection
// without specifying a peer ID.
var ErrNilPeer = errors.New("nil peer")

// AcceptQueueLength is the number of connections to fully setup before not accepting any new connections
var AcceptQueueLength = 16

const defaultAcceptTimeout = 15 * time.Second

type Option func(*upgrader) error

func WithPSK(psk ipnet.PSK) Option {
	return func(u *upgrader) error {
		u.psk = psk
		return nil
	}
}

func WithAcceptTimeout(t time.Duration) Option {
	return func(u *upgrader) error {
		u.acceptTimeout = t
		return nil
	}
}

func WithConnectionGater(g connmgr.ConnectionGater) Option {
	return func(u *upgrader) error {
		u.connGater = g
		return nil
	}
}

func WithResourceManager(m network.ResourceManager) Option {
	return func(u *upgrader) error {
		u.rcmgr = m
		return nil
	}
}

// Upgrader is a multistream upgrader that can upgrade an underlying connection
// to a full transport connection (secure and multiplexed).
type upgrader struct {
	secure sec.SecureMuxer
	muxer  network.Multiplexer

	psk       ipnet.PSK
	connGater connmgr.ConnectionGater
	rcmgr     network.ResourceManager

	// AcceptTimeout is the maximum duration an Accept is allowed to take.
	// This includes the time between accepting the raw network connection,
	// protocol selection as well as the handshake, if applicable.
	//
	// If unset, the default value (15s) is used.
	acceptTimeout time.Duration
}

var _ transport.Upgrader = &upgrader{}

func New(secureMuxer sec.SecureMuxer, muxer network.Multiplexer, opts ...Option) (transport.Upgrader, error) {
	u := &upgrader{
		secure:        secureMuxer,
		muxer:         muxer,
		acceptTimeout: defaultAcceptTimeout,
	}
	for _, opt := range opts {
		if err := opt(u); err != nil {
			return nil, err
		}
	}
	if u.rcmgr == nil {
		u.rcmgr = network.NullResourceManager
	}
	return u, nil
}

// UpgradeListener upgrades the passed multiaddr-net listener into a full libp2p-transport listener.
func (u *upgrader) UpgradeListener(t transport.Transport, list manet.Listener) transport.Listener {
	ctx, cancel := context.WithCancel(context.Background())
	l := &listener{
		Listener:  list,
		upgrader:  u,
		transport: t,
		rcmgr:     u.rcmgr,
		threshold: newThreshold(AcceptQueueLength),
		incoming:  make(chan transport.CapableConn),
		cancel:    cancel,
		ctx:       ctx,
	}
	go l.handleIncoming()
	return l
}

// Upgrade upgrades the multiaddr/net connection into a full libp2p-transport connection.
func (u *upgrader) Upgrade(ctx context.Context, t transport.Transport, maconn manet.Conn, dir network.Direction, p peer.ID, connScope network.ConnManagementScope) (transport.CapableConn, error) {
	c, err := u.upgrade(ctx, t, maconn, dir, p, connScope)
	if err != nil {
		connScope.Done()
		return nil, err
	}
	return c, nil
}

func (u *upgrader) upgrade(ctx context.Context, t transport.Transport, maconn manet.Conn, dir network.Direction, p peer.ID, connScope network.ConnManagementScope) (transport.CapableConn, error) {
	if dir == network.DirOutbound && p == "" {
		return nil, ErrNilPeer
	}
	var stat network.ConnStats
	if cs, ok := maconn.(network.ConnStat); ok {
		stat = cs.Stat()
	}

	var conn net.Conn = maconn
	if u.psk != nil {
		pconn, err := pnet.NewProtectedConn(u.psk, conn)
		if err != nil {
			conn.Close()
			return nil, fmt.Errorf("failed to setup private network protector: %s", err)
		}
		conn = pconn
	} else if ipnet.ForcePrivateNetwork {
		log.Error("tried to dial with no Private Network Protector but usage of Private Networks is forced by the environment")
		return nil, ipnet.ErrNotInPrivateNetwork
	}

	sconn, server, err := u.setupSecurity(ctx, conn, p, dir)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to negotiate security protocol: %s", err)
	}

	// call the connection gater, if one is registered.
	if u.connGater != nil && !u.connGater.InterceptSecured(dir, sconn.RemotePeer(), maconn) {
		if err := maconn.Close(); err != nil {
			log.Errorw("failed to close connection", "peer", p, "addr", maconn.RemoteMultiaddr(), "error", err)
		}
		return nil, fmt.Errorf("gater rejected connection with peer %s and addr %s with direction %d",
			sconn.RemotePeer().Pretty(), maconn.RemoteMultiaddr(), dir)
	}
	// Only call SetPeer if it hasn't already been set -- this can happen when we don't know
	// the peer in advance and in some bug scenarios.
	if connScope.PeerScope() == nil {
		if err := connScope.SetPeer(sconn.RemotePeer()); err != nil {
			log.Debugw("resource manager blocked connection for peer", "peer", sconn.RemotePeer(), "addr", conn.RemoteAddr(), "error", err)
			if err := maconn.Close(); err != nil {
				log.Errorw("failed to close connection", "peer", p, "addr", maconn.RemoteMultiaddr(), "error", err)
			}
			return nil, fmt.Errorf("resource manager connection with peer %s and addr %s with direction %d",
				sconn.RemotePeer().Pretty(), maconn.RemoteMultiaddr(), dir)
		}
	}

	smconn, err := u.setupMuxer(ctx, sconn, server, connScope.PeerScope())
	if err != nil {
		sconn.Close()
		return nil, fmt.Errorf("failed to negotiate stream multiplexer: %s", err)
	}

	tc := &transportConn{
		MuxedConn:      smconn,
		ConnMultiaddrs: maconn,
		ConnSecurity:   sconn,
		transport:      t,
		stat:           stat,
		scope:          connScope,
	}
	return tc, nil
}

func (u *upgrader) setupSecurity(ctx context.Context, conn net.Conn, p peer.ID, dir network.Direction) (sec.SecureConn, bool, error) {
	if dir == network.DirInbound {
		return u.secure.SecureInbound(ctx, conn, p)
	}
	return u.secure.SecureOutbound(ctx, conn, p)
}

func (u *upgrader) setupMuxer(ctx context.Context, conn net.Conn, server bool, scope network.PeerScope) (network.MuxedConn, error) {
	// TODO: The muxer should take a context.
	done := make(chan struct{})

	var smconn network.MuxedConn
	var err error
	go func() {
		defer close(done)
		smconn, err = u.muxer.NewConn(conn, server, scope)
	}()

	select {
	case <-done:
		return smconn, err
	case <-ctx.Done():
		// interrupt this process
		conn.Close()
		// wait to finish
		<-done
		return nil, ctx.Err()
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package connpool

import (
	"runtime"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/mysql/fakesqldb"
	"vitess.io/vitess/go/vt/callerid"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"

	"golang.org/x/net/context"
)

func TestConnPoolGet(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	dbConn, err := connPool.Get(context.Background())
	if err != nil {
		t.Fatalf("should not get an error, but got: %v", err)
	}
	if dbConn == nil {
		t.Fatalf("db conn should not be nil")
	}
	// There is no context, it should not use appdebug connection
	if dbConn.pool == nil {
		t.Fatalf("db conn pool should not be nil")
	}
	dbConn.Recycle()
}

func TestConnPoolTimeout(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := New(tabletenv.NewTestEnv(nil, nil, "PoolTest"), "TestPool", tabletenv.ConnPoolConfig{
		Size:               1,
		TimeoutSeconds:     1,
		IdleTimeoutSeconds: 10,
	})
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	dbConn, err := connPool.Get(context.Background())
	require.NoError(t, err)
	defer dbConn.Recycle()
	_, err = connPool.Get(context.Background())
	assert.EqualError(t, err, "resource pool timed out")
}

func TestConnPoolMaxWaiters(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := New(tabletenv.NewTestEnv(nil, nil, "PoolTest"), "TestPool", tabletenv.ConnPoolConfig{
		Size:       1,
		MaxWaiters: 1,
	})
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	dbConn, err := connPool.Get(context.Background())
	require.NoError(t, err)

	// waiter 1
	go func() {
		c1, err := connPool.Get(context.Background())
		assert.NoError(t, err)
		c1.Recycle()
	}()
	// Wait for the first waiter to increment count.
	for {
		runtime.Gosched()
		if connPool.waiterCount.Get() == 1 {
			break
		}
	}
	_, err = connPool.Get(context.Background())
	assert.EqualError(t, err, "pool TestPool waiter count exceeded")

	// This recycle will make waiter1 succeed.
	dbConn.Recycle()
}

func TestConnPoolGetEmptyDebugConfig(t *testing.T) {
	db := fakesqldb.New(t)
	debugConn := db.ConnParamsWithUname("")
	defer db.Close()
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), debugConn)
	im := callerid.NewImmediateCallerID("")
	ecid := callerid.NewEffectiveCallerID("p", "c", "sc")
	ctx := context.Background()
	ctx = callerid.NewContext(ctx, ecid, im)
	defer connPool.Close()
	dbConn, err := connPool.Get(ctx)
	if err != nil {
		t.Fatalf("should not get an error, but got: %v", err)
	}
	if dbConn == nil {
		t.Fatalf("db conn should not be nil")
	}
	// Context is empty, it should not use appdebug connection
	if dbConn.pool == nil {
		t.Fatalf("db conn pool should not be nil")
	}
	dbConn.Recycle()
}

func TestConnPoolGetAppDebug(t *testing.T) {
	db := fakesqldb.New(t)
	debugConn := db.ConnParamsWithUname("debugUsername")
	ctx := context.Background()
	im := callerid.NewImmediateCallerID("debugUsername")
	ecid := callerid.NewEffectiveCallerID("p", "c", "sc")
	ctx = callerid.NewContext(ctx, ecid, im)
	defer db.Close()
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), debugConn)
	defer connPool.Close()
	dbConn, err := connPool.Get(ctx)
	if err != nil {
		t.Fatalf("should not get an error, but got: %v", err)
	}
	if dbConn == nil {
		t.Fatalf("db conn should not be nil")
	}
	if dbConn.pool != nil {
		t.Fatalf("db conn pool should be nil for appDebug")
	}
	dbConn.Recycle()
	if !dbConn.IsClosed() {
		t.Fatalf("db conn should be closed after recycle")
	}
}

func TestConnPoolPutWhilePoolIsClosed(t *testing.T) {
	connPool := newPool()
	defer func() {
		if recover() == nil {
			t.Fatalf("pool is closed, should get an error")
		}
	}()
	connPool.Put(nil)
}

func TestConnPoolSetCapacity(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	err := connPool.SetCapacity(-10)
	if err == nil {
		t.Fatalf("set capacity should return error for negative capacity")
	}
	err = connPool.SetCapacity(10)
	if err != nil {
		t.Fatalf("set capacity should succeed")
	}
	if connPool.Capacity() != 10 {
		t.Fatalf("capacity should be 10")
	}
}

func TestConnPoolStatJSON(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := newPool()
	if connPool.StatsJSON() != "{}" {
		t.Fatalf("pool is closed, stats json should be empty: {}")
	}
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	statsJSON := connPool.StatsJSON()
	if statsJSON == "" || statsJSON == "{}" {
		t.Fatalf("stats json should not be empty")
	}
}

func TestConnPoolStateWhilePoolIsClosed(t *testing.T) {
	connPool := newPool()
	if connPool.Capacity() != 0 {
		t.Fatalf("pool capacity should be 0 because it is still closed")
	}
	if connPool.Available() != 0 {
		t.Fatalf("pool available connections should be 0 because it is still closed")
	}
	if connPool.MaxCap() != 0 {
		t.Fatalf("pool max capacity should be 0 because it is still closed")
	}
	if connPool.WaitCount() != 0 {
		t.Fatalf("pool wait count should be 0 because it is still closed")
	}
	if connPool.WaitTime() != 0 {
		t.Fatalf("pool wait time should be 0 because it is still closed")
	}
	if connPool.IdleTimeout() != 0 {
		t.Fatalf("pool idle timeout should be 0 because it is still closed")
	}
}

func TestConnPoolStateWhilePoolIsOpen(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	idleTimeout := 10 * time.Second
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	if connPool.Capacity() != 100 {
		t.Fatalf("pool capacity should be 100")
	}
	if connPool.MaxCap() != 100 {
		t.Fatalf("pool max capacity should be 100")
	}
	if connPool.WaitTime() != 0 {
		t.Fatalf("pool wait time should be 0")
	}
	if connPool.WaitCount() != 0 {
		t.Fatalf("pool wait count should be 0")
	}
	if connPool.IdleTimeout() != idleTimeout {
		t.Fatalf("pool idle timeout should be 0")
	}
	if connPool.Available() != 100 {
		t.Fatalf("pool available connections should be 100")
	}
	if connPool.Active() != 0 {
		t.Fatalf("pool active connections should be 0")
	}
	if connPool.InUse() != 0 {
		t.Fatalf("pool inUse connections should be 0")
	}
	dbConn, _ := connPool.Get(context.Background())
	if connPool.Available() != 99 {
		t.Fatalf("pool available connections should be 99")
	}
	if connPool.Active() != 1 {
		t.Fatalf("pool active connections should be 1")
	}
	if connPool.InUse() != 1 {
		t.Fatalf("pool inUse connections should be 1")
	}

	dbConn.Recycle()
	if connPool.Available() != 100 {
		t.Fatalf("pool available connections should be 100")
	}
	if connPool.Active() != 1 {
		t.Fatalf("pool active connections should be 1")
	}
	if connPool.InUse() != 0 {
		t.Fatalf("pool inUse connections should be 0")
	}
}

func newPool() *Pool {
	return New(tabletenv.NewTestEnv(nil, nil, "PoolTest"), "TestPool", tabletenv.ConnPoolConfig{
		Size:               100,
		IdleTimeoutSeconds: 10,
	})
}
<commit_msg>yaml: deflake pool_test<commit_after>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package connpool

import (
	"runtime"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/mysql/fakesqldb"
	"vitess.io/vitess/go/vt/callerid"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"

	"golang.org/x/net/context"
)

func TestConnPoolGet(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	dbConn, err := connPool.Get(context.Background())
	if err != nil {
		t.Fatalf("should not get an error, but got: %v", err)
	}
	if dbConn == nil {
		t.Fatalf("db conn should not be nil")
	}
	// There is no context, it should not use appdebug connection
	if dbConn.pool == nil {
		t.Fatalf("db conn pool should not be nil")
	}
	dbConn.Recycle()
}

func TestConnPoolTimeout(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := New(tabletenv.NewTestEnv(nil, nil, "PoolTest"), "TestPool", tabletenv.ConnPoolConfig{
		Size:               1,
		TimeoutSeconds:     1,
		IdleTimeoutSeconds: 10,
	})
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	dbConn, err := connPool.Get(context.Background())
	require.NoError(t, err)
	defer dbConn.Recycle()
	_, err = connPool.Get(context.Background())
	assert.EqualError(t, err, "resource pool timed out")
}

func TestConnPoolMaxWaiters(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := New(tabletenv.NewTestEnv(nil, nil, "PoolTest"), "TestPool", tabletenv.ConnPoolConfig{
		Size:       1,
		MaxWaiters: 1,
	})
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	dbConn, err := connPool.Get(context.Background())
	require.NoError(t, err)

	// waiter 1
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		c1, err := connPool.Get(context.Background())
		assert.NoError(t, err)
		c1.Recycle()
	}()
	// Wait for the first waiter to increment count.
	for {
		runtime.Gosched()
		if connPool.waiterCount.Get() == 1 {
			break
		}
	}

	// waiter 2
	_, err = connPool.Get(context.Background())
	assert.EqualError(t, err, "pool TestPool waiter count exceeded")

	// This recycle will make waiter1 succeed.
	dbConn.Recycle()
	wg.Wait()
}

func TestConnPoolGetEmptyDebugConfig(t *testing.T) {
	db := fakesqldb.New(t)
	debugConn := db.ConnParamsWithUname("")
	defer db.Close()
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), debugConn)
	im := callerid.NewImmediateCallerID("")
	ecid := callerid.NewEffectiveCallerID("p", "c", "sc")
	ctx := context.Background()
	ctx = callerid.NewContext(ctx, ecid, im)
	defer connPool.Close()
	dbConn, err := connPool.Get(ctx)
	if err != nil {
		t.Fatalf("should not get an error, but got: %v", err)
	}
	if dbConn == nil {
		t.Fatalf("db conn should not be nil")
	}
	// Context is empty, it should not use appdebug connection
	if dbConn.pool == nil {
		t.Fatalf("db conn pool should not be nil")
	}
	dbConn.Recycle()
}

func TestConnPoolGetAppDebug(t *testing.T) {
	db := fakesqldb.New(t)
	debugConn := db.ConnParamsWithUname("debugUsername")
	ctx := context.Background()
	im := callerid.NewImmediateCallerID("debugUsername")
	ecid := callerid.NewEffectiveCallerID("p", "c", "sc")
	ctx = callerid.NewContext(ctx, ecid, im)
	defer db.Close()
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), debugConn)
	defer connPool.Close()
	dbConn, err := connPool.Get(ctx)
	if err != nil {
		t.Fatalf("should not get an error, but got: %v", err)
	}
	if dbConn == nil {
		t.Fatalf("db conn should not be nil")
	}
	if dbConn.pool != nil {
		t.Fatalf("db conn pool should be nil for appDebug")
	}
	dbConn.Recycle()
	if !dbConn.IsClosed() {
		t.Fatalf("db conn should be closed after recycle")
	}
}

func TestConnPoolPutWhilePoolIsClosed(t *testing.T) {
	connPool := newPool()
	defer func() {
		if recover() == nil {
			t.Fatalf("pool is closed, should get an error")
		}
	}()
	connPool.Put(nil)
}

func TestConnPoolSetCapacity(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	err := connPool.SetCapacity(-10)
	if err == nil {
		t.Fatalf("set capacity should return error for negative capacity")
	}
	err = connPool.SetCapacity(10)
	if err != nil {
		t.Fatalf("set capacity should succeed")
	}
	if connPool.Capacity() != 10 {
		t.Fatalf("capacity should be 10")
	}
}

func TestConnPoolStatJSON(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	connPool := newPool()
	if connPool.StatsJSON() != "{}" {
		t.Fatalf("pool is closed, stats json should be empty: {}")
	}
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	statsJSON := connPool.StatsJSON()
	if statsJSON == "" || statsJSON == "{}" {
		t.Fatalf("stats json should not be empty")
	}
}

func TestConnPoolStateWhilePoolIsClosed(t *testing.T) {
	connPool := newPool()
	if connPool.Capacity() != 0 {
		t.Fatalf("pool capacity should be 0 because it is still closed")
	}
	if connPool.Available() != 0 {
		t.Fatalf("pool available connections should be 0 because it is still closed")
	}
	if connPool.MaxCap() != 0 {
		t.Fatalf("pool max capacity should be 0 because it is still closed")
	}
	if connPool.WaitCount() != 0 {
		t.Fatalf("pool wait count should be 0 because it is still closed")
	}
	if connPool.WaitTime() != 0 {
		t.Fatalf("pool wait time should be 0 because it is still closed")
	}
	if connPool.IdleTimeout() != 0 {
		t.Fatalf("pool idle timeout should be 0 because it is still closed")
	}
}

func TestConnPoolStateWhilePoolIsOpen(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	idleTimeout := 10 * time.Second
	connPool := newPool()
	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
	defer connPool.Close()
	if connPool.Capacity() != 100 {
		t.Fatalf("pool capacity should be 100")
	}
	if connPool.MaxCap() != 100 {
		t.Fatalf("pool max capacity should be 100")
	}
	if connPool.WaitTime() != 0 {
		t.Fatalf("pool wait time should be 0")
	}
	if connPool.WaitCount() != 0 {
		t.Fatalf("pool wait count should be 0")
	}
	if connPool.IdleTimeout() != idleTimeout {
		t.Fatalf("pool idle timeout should be 0")
	}
	if connPool.Available() != 100 {
		t.Fatalf("pool available connections should be 100")
	}
	if connPool.Active() != 0 {
		t.Fatalf("pool active connections should be 0")
	}
	if connPool.InUse() != 0 {
		t.Fatalf("pool inUse connections should be 0")
	}
	dbConn, _ := connPool.Get(context.Background())
	if connPool.Available() != 99 {
		t.Fatalf("pool available connections should be 99")
	}
	if connPool.Active() != 1 {
		t.Fatalf("pool active connections should be 1")
	}
	if connPool.InUse() != 1 {
		t.Fatalf("pool inUse connections should be 1")
	}

	dbConn.Recycle()
	if connPool.Available() != 100 {
		t.Fatalf("pool available connections should be 100")
	}
	if connPool.Active() != 1 {
		t.Fatalf("pool active connections should be 1")
	}
	if connPool.InUse() != 0 {
		t.Fatalf("pool inUse connections should be 0")
	}
}

func newPool() *Pool {
	return New(tabletenv.NewTestEnv(nil, nil, "PoolTest"), "TestPool", tabletenv.ConnPoolConfig{
		Size:               100,
		IdleTimeoutSeconds: 10,
	})
}
<|endoftext|>


<commit_before>package llrb

import "fmt"

// Left and Right children must be present
func moveRedToLeft(nd *Node) *Node {
	fmt.Println("moveRedToLeft:", nd.Key)
	flipColor(nd)
	if isRed(nd.Right.Left) {
		nd.Right = rotateToRight(nd.Right)
		nd = rotateToLeft(nd)
		flipColor(nd)
	}
	return nd
}

// Left and Right children must be present
func moveRedToRight(nd *Node) *Node {
	fmt.Println("moveRedToRight:", nd.Key)
	flipColor(nd)
	if isRed(nd.Left.Left) {
		fmt.Println("moveRedToRight isRed(nd.Left.Left):", nd.Key)
		nd = rotateToRight(nd)
		flipColor(nd)
	}
	return nd
}

func fixUp(nd *Node) *Node {
	fmt.Println("fixUp", nd.Key)
	if isRed(nd.Right) {
		nd = rotateToLeft(nd)
	}
	if isRed(nd.Left) && isRed(nd.Left.Left) {
		nd = rotateToRight(nd)
	}
	if isRed(nd.Left) && isRed(nd.Right) {
		flipColor(nd)
	}
	return nd
}

// deleteMin code for LLRB 2-3 trees
func deleteMin(nd *Node) (*Node, Interface) {
	fmt.Println("deleteMin", nd.Key)
	if nd == nil {
		return nil, nil
	}
	if nd.Left == nil {
		return nil, nd.Key
	}
	if !isRed(nd.Left) && !isRed(nd.Left.Left) {
		nd = moveRedToLeft(nd)
	}
	var deleted Interface
	nd.Left, deleted = deleteMin(nd.Left)
	return fixUp(nd), deleted
}

// Delete deletes the node with a given key and returns the key.
// It returns nil if it does not exist in the tree.
//
//	Delete Algorithm:
//	1. Start 'delete' from tree Root.
//	//
//	2. Call 'delete' method recursively on each Node from binary search path.
//		- e.g. If the key to delete is greater than Root's key
//			, call 'delete' on Right Node.
//	//
//	// start
//	3. 'tree.delete(Node, Key)'
//
//	// end
//	//
//	4. If the tree's Root is not nil, set Root Black.
//	//
//	5. Return the Interface(nil if the key does not exist.)
//
func (tr *Tree) Delete(key Interface) Interface {
	var deleted Interface
	tr.Root, deleted = tr.delete(tr.Root, key)
	if tr.Root != nil {
		tr.Root.Black = true
	}
	return deleted
}

func (tr *Tree) delete(nd *Node, key Interface) (*Node, Interface) {
	fmt.Println("calling delete on", nd.Key, "for the key", key)
	if nd == nil {
		return nil, nil
	}

	var deleted Interface

	// if key is Less than nd.Key
	if key.Less(nd.Key) {

		// if key is Less than nd.Key
		// and nd.Left is empty
		if nd.Left == nil {

			// then nothing to delete
			// so return the nil
			return nd, nil
		}

		// if key is Less than nd.Key
		// and nd.Left is Black
		// and nd.Left.Left is Black
		if !isRed(nd.Left) && !isRed(nd.Left.Left) {

			// then moveRedToLeft(nd)
			nd = moveRedToLeft(nd)
		}

		// and recursively call tr.delete(nd.Left, key)
		nd.Left, deleted = tr.delete(nd.Left, key)

	} else {
		// if key is not Less than nd.Key
		//(or key is greater than or equal to nd.Key)
		//(or key >= nd.Key)

		// and nd.Left is Red
		if isRed(nd.Left) {

			// then rotateToRight(nd)
			nd = rotateToRight(nd)
		}

		// and nd.Key is not Less than key
		// (or nd.Key >= key)
		// (or key == nd.Key)
		// and nd.Right is empty
		if !nd.Key.Less(key) && nd.Right == nil {
			fmt.Println("!nd.Key.Less(key) && nd.Right == nil when", nd.Key)

			// then return nil to delete the key
			return nil, nd.Key
		}

		// and nd.Right is not empty
		// and nd.Right is Black
		// and nd.Right.Left is Black
		if nd.Right != nil && !isRed(nd.Right) && !isRed(nd.Right.Left) {
			fmt.Println("nd.Right != nil && !isRed(nd.Right) && !isRed(nd.Right.Left) when", nd.Key)

			// then moveRedToRight(nd)
			nd = moveRedToRight(nd)
		}

		// and key == nd.Key
		if !nd.Key.Less(key) {

			var subDeleted Interface

			// then deleteMin(nd.Right)
			nd.Right, subDeleted = deleteMin(nd.Right)
			if subDeleted == nil {
				panic("Unexpected nil value")
			}

			// and update nd.Key with deleteMin(nd.Right)
			deleted, nd.Key = nd.Key, subDeleted

		} else {
			// if updated nd.Key is Less than key (nd.Key < key)
			fmt.Println("nd.Right, deleted = tr.delete(nd.Right, key) at", nd.Key)
			nd.Right, deleted = tr.delete(nd.Right, key)
		}
	}

	return fixUp(nd), deleted
}
<commit_msg>try new line<commit_after>package llrb

import "fmt"

// Left and Right children must be present
func moveRedToLeft(nd *Node) *Node {
	fmt.Println("moveRedToLeft:", nd.Key)
	flipColor(nd)
	if isRed(nd.Right.Left) {
		nd.Right = rotateToRight(nd.Right)
		nd = rotateToLeft(nd)
		flipColor(nd)
	}
	return nd
}

// Left and Right children must be present
func moveRedToRight(nd *Node) *Node {
	fmt.Println("moveRedToRight:", nd.Key)
	flipColor(nd)
	if isRed(nd.Left.Left) {
		fmt.Println("moveRedToRight isRed(nd.Left.Left):", nd.Key)
		nd = rotateToRight(nd)
		flipColor(nd)
	}
	return nd
}

func fixUp(nd *Node) *Node {
	fmt.Println("fixUp", nd.Key)
	if isRed(nd.Right) {
		nd = rotateToLeft(nd)
	}
	if isRed(nd.Left) && isRed(nd.Left.Left) {
		nd = rotateToRight(nd)
	}
	if isRed(nd.Left) && isRed(nd.Right) {
		flipColor(nd)
	}
	return nd
}

// deleteMin code for LLRB 2-3 trees
func deleteMin(nd *Node) (*Node, Interface) {
	fmt.Println("deleteMin", nd.Key)
	if nd == nil {
		return nil, nil
	}
	if nd.Left == nil {
		return nil, nd.Key
	}
	if !isRed(nd.Left) && !isRed(nd.Left.Left) {
		nd = moveRedToLeft(nd)
	}
	var deleted Interface
	nd.Left, deleted = deleteMin(nd.Left)
	return fixUp(nd), deleted
}

// Delete deletes the node with a given key and returns the key.
// It returns nil if it does not exist in the tree.
//
//	Delete Algorithm:
//	1. Start 'delete' from tree Root.
//
//	2. Call 'delete' method recursively on each Node from binary search path.
//		- e.g. If the key to delete is greater than Root's key
//			, call 'delete' on Right Node.
//
//	# start
//	3. 'tree.delete(Node, Key)'
//
//	# end
//
//	4. If the tree's Root is not nil, set Root Black.
//
//	5. Return the Interface(nil if the key does not exist.)
//
func (tr *Tree) Delete(key Interface) Interface {
	var deleted Interface
	tr.Root, deleted = tr.delete(tr.Root, key)
	if tr.Root != nil {
		tr.Root.Black = true
	}
	return deleted
}

func (tr *Tree) delete(nd *Node, key Interface) (*Node, Interface) {
	fmt.Println("calling delete on", nd.Key, "for the key", key)
	if nd == nil {
		return nil, nil
	}

	var deleted Interface

	// if key is Less than nd.Key
	if key.Less(nd.Key) {

		// if key is Less than nd.Key
		// and nd.Left is empty
		if nd.Left == nil {

			// then nothing to delete
			// so return the nil
			return nd, nil
		}

		// if key is Less than nd.Key
		// and nd.Left is Black
		// and nd.Left.Left is Black
		if !isRed(nd.Left) && !isRed(nd.Left.Left) {

			// then moveRedToLeft(nd)
			nd = moveRedToLeft(nd)
		}

		// and recursively call tr.delete(nd.Left, key)
		nd.Left, deleted = tr.delete(nd.Left, key)

	} else {
		// if key is not Less than nd.Key
		//(or key is greater than or equal to nd.Key)
		//(or key >= nd.Key)

		// and nd.Left is Red
		if isRed(nd.Left) {

			// then rotateToRight(nd)
			nd = rotateToRight(nd)
		}

		// and nd.Key is not Less than key
		// (or nd.Key >= key)
		// (or key == nd.Key)
		// and nd.Right is empty
		if !nd.Key.Less(key) && nd.Right == nil {
			fmt.Println("!nd.Key.Less(key) && nd.Right == nil when", nd.Key)

			// then return nil to delete the key
			return nil, nd.Key
		}

		// and nd.Right is not empty
		// and nd.Right is Black
		// and nd.Right.Left is Black
		if nd.Right != nil && !isRed(nd.Right) && !isRed(nd.Right.Left) {
			fmt.Println("nd.Right != nil && !isRed(nd.Right) && !isRed(nd.Right.Left) when", nd.Key)

			// then moveRedToRight(nd)
			nd = moveRedToRight(nd)
		}

		// and key == nd.Key
		if !nd.Key.Less(key) {

			var subDeleted Interface

			// then deleteMin(nd.Right)
			nd.Right, subDeleted = deleteMin(nd.Right)
			if subDeleted == nil {
				panic("Unexpected nil value")
			}

			// and update nd.Key with deleteMin(nd.Right)
			deleted, nd.Key = nd.Key, subDeleted

		} else {
			// if updated nd.Key is Less than key (nd.Key < key)
			fmt.Println("nd.Right, deleted = tr.delete(nd.Right, key) at", nd.Key)
			nd.Right, deleted = tr.delete(nd.Right, key)
		}
	}

	return fixUp(nd), deleted
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Workiva, LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*
Package skiplist implements an n-dimensional rangetree based on a skip
list.  This should be faster than a straight slice implementation as
memcopy is avoided.

Time complexities revolve around the ability to quickly find items
in the n-dimensional skiplist.  That time can be defined by the
number of items in any dimension.  Let N1, N2,... Nn define the
number of dimensions.

Performance characteristics:
Space: O(n)
Search: O(log N1 + log N2 + ...log Nn) = O(log N1*N2*...Nn)
Insert: O(log N1 + log N2 + ...log Nn) = O(log N1*N2*...Nn)
Delete: O(log N1 + log N2 + ...log Nn) = O(log N1*N2*...Nn)
*/

package skiplist

import (
	"github.com/Workiva/go-datastructures/rangetree"
	"github.com/Workiva/go-datastructures/slice/skip"
)

// isLastDimension simply returns dimension == lastDimension-1.
// This panics if dimension >= lastDimension.
func isLastDimension(dimension, lastDimension uint64) bool {
	if dimension >= lastDimension { // useful in testing and denotes a serious problem
		panic(`Dimension is greater than possible dimensions.`)
	}

	return dimension == lastDimension-1
}

// needsDeletion returns a bool indicating if the provided value
// needs to be deleted based on the provided index and number.
func needsDeletion(value, index, number int64) bool {
	if number > 0 {
		return false
	}

	number = -number // get the magnitude
	offset := value - index

	return offset >= 0 && offset < number
}

// dimensionalBundle is an intermediate holder up to the last
// dimension and represents a wrapper around a skiplist.
type dimensionalBundle struct {
	key uint64
	sl  *skip.SkipList
}

// Key returns the key for this bundle.  Fulfills skip.Entry interface.
func (db *dimensionalBundle) Key() uint64 {
	return db.key
}

// lastBundle represents a bundle living at the last dimension
// of the tree.
type lastBundle struct {
	key   uint64
	entry rangetree.Entry
}

// Key returns the key for this bundle.  Fulfills skip.Entry interface.
func (lb *lastBundle) Key() uint64 {
	return lb.key
}

type skipListRT struct {
	top                *skip.SkipList
	dimensions, number uint64
}

func (rt *skipListRT) init(dimensions uint64) {
	rt.dimensions = dimensions
	rt.top = skip.New(uint64(0))
}

func (rt *skipListRT) add(entry rangetree.Entry) rangetree.Entry {
	var (
		value int64
		e     skip.Entry
		sl    = rt.top
		db    *dimensionalBundle
		lb    *lastBundle
	)

	for i := uint64(0); i < rt.dimensions; i++ {
		value = entry.ValueAtDimension(i)
		e = sl.Get(uint64(value))[0]
		if isLastDimension(i, rt.dimensions) {
			if e != nil { // this is an overwrite
				lb = e.(*lastBundle)
				oldEntry := lb.entry
				lb.entry = entry
				return oldEntry
			}

			// need to add new sl entry
			lb = &lastBundle{key: uint64(value), entry: entry}
			rt.number++
			sl.Insert(lb)
			return nil
		}

		if e == nil { // we need the intermediate dimension
			db = &dimensionalBundle{key: uint64(value), sl: skip.New(uint64(0))}
			sl.Insert(db)
		} else {
			db = e.(*dimensionalBundle)
		}

		sl = db.sl
	}

	panic(`Ran out of dimensions before for loop completed.`)
}

// Add will add the provided entries to the tree.  Any entries that
// were overwritten will be returned in the order in which they
// were overwritten.  If an entry's addition does not overwrite, a nil
// is returned for that cell for its index in the provided entries.
func (rt *skipListRT) Add(entries ...rangetree.Entry) rangetree.Entries {
	overwritten := make(rangetree.Entries, 0, len(entries))
	for _, e := range entries {
		overwritten = append(overwritten, rt.add(e))
	}

	return overwritten
}

func (rt *skipListRT) get(entry rangetree.Entry) rangetree.Entry {
	var (
		sl    = rt.top
		e     skip.Entry
		value uint64
	)
	for i := uint64(0); i < rt.dimensions; i++ {
		value = uint64(entry.ValueAtDimension(i))
		e = sl.Get(value)[0]
		if e == nil {
			return nil
		}

		if isLastDimension(i, rt.dimensions) {
			return e.(*lastBundle).entry
		}

		sl = e.(*dimensionalBundle).sl
	}

	panic(`Reached past for loop without finding last dimension.`)
}

// Get will return any rangetree.Entries matching the provided entries.
// Similar in functionality to a key lookup, this returns nil for any
// entry that could not be found.
func (rt *skipListRT) Get(entries ...rangetree.Entry) rangetree.Entries {
	results := make(rangetree.Entries, 0, len(entries))
	for _, e := range entries {
		results = append(results, rt.get(e))
	}

	return results
}

// Len returns the number of entries in the tree.
func (rt *skipListRT) Len() uint64 {
	return rt.number
}

// deleteRecursive is used by the delete logic.  The recursion depth
// only goes as far as the number of dimensions, so this shouldn't be an
// issue.
func (rt *skipListRT) deleteRecursive(sl *skip.SkipList, dimension uint64,
	entry rangetree.Entry) rangetree.Entry {

	value := entry.ValueAtDimension(dimension)
	if isLastDimension(dimension, rt.dimensions) {
		entries := sl.Delete(uint64(value))
		if entries[0] == nil {
			return nil
		}

		rt.number--
		return entries[0].(*lastBundle).entry
	}

	db, ok := sl.Get(uint64(value))[0].(*dimensionalBundle)
	if !ok { // value was not found
		return nil
	}

	result := rt.deleteRecursive(db.sl, dimension+1, entry)
	if result == nil {
		return nil
	}

	if db.sl.Len() == 0 {
		sl.Delete(db.key)
	}

	return result
}

func (rt *skipListRT) delete(entry rangetree.Entry) rangetree.Entry {
	return rt.deleteRecursive(rt.top, 0, entry)
}

// Delete will remove the provided entries from the tree.
func (rt *skipListRT) Delete(entries ...rangetree.Entry) {
	for _, e := range entries {
		rt.delete(e)
	}
}

func (rt *skipListRT) apply(sl *skip.SkipList, dimension uint64,
	interval rangetree.Interval, fn func(rangetree.Entry) bool) bool {

	lowValue, highValue := interval.LowAtDimension(dimension), interval.HighAtDimension(dimension)

	var e skip.Entry

	for iter := sl.Iter(uint64(lowValue)); iter.Next(); {
		e = iter.Value()
		if int64(e.Key()) >= highValue {
			break
		}

		if isLastDimension(dimension, rt.dimensions) {
			if !fn(e.(*lastBundle).entry) {
				return false
			}
		} else {

			if !rt.apply(e.(*dimensionalBundle).sl, dimension+1, interval, fn) {
				return false
			}
		}
	}

	return true
}

// Apply will call the provided function with each entry that exists
// within the provided range, in order.  Return false at any time to
// cancel iteration.  Altering the entry in such a way that its location
// changes will result in undefined behavior.
func (rt *skipListRT) Apply(interval rangetree.Interval, fn func(rangetree.Entry) bool) {
	rt.apply(rt.top, 0, interval, fn)
}

// Query will return a list of entries that fall within
// the provided interval.
func (rt *skipListRT) Query(interval rangetree.Interval) rangetree.Entries {
	entries := make(rangetree.Entries, 0, 100)
	rt.apply(rt.top, 0, interval, func(e rangetree.Entry) bool {
		entries = append(entries, e)
		return true
	})

	return entries
}

func (rt *skipListRT) flatten(sl *skip.SkipList, dimension uint64, entries *rangetree.Entries) {
	lastDimension := isLastDimension(dimension, rt.dimensions)
	for iter := sl.Iter(0); iter.Next(); {
		if lastDimension {
			*entries = append(*entries, iter.Value().(*lastBundle).entry)
		} else {
			rt.flatten(iter.Value().(*dimensionalBundle).sl, dimension+1, entries)
		}
	}
}

func (rt *skipListRT) insert(sl *skip.SkipList, dimension, insertDimension uint64,
	index, number int64, deleted, affected *rangetree.Entries) {

	var e skip.Entry
	lastDimension := isLastDimension(dimension, rt.dimensions)
	affectedDimension := dimension == insertDimension
	var iter skip.Iterator
	if dimension == insertDimension {
		iter = sl.Iter(uint64(index))
	} else {
		iter = sl.Iter(0)
	}

	var toDelete skip.Entries
	if number < 0 {
		toDelete = make(skip.Entries, 0, 100)
	}

	for iter.Next() {
		e = iter.Value()
		if !affectedDimension {
			rt.insert(e.(*dimensionalBundle).sl, dimension+1,
				insertDimension, index, number, deleted, affected,
			)
			continue
		}
		if needsDeletion(int64(e.Key()), index, number) {
			toDelete = append(toDelete, e)
			continue
		}

		if lastDimension {
			e.(*lastBundle).key += uint64(number)
			*affected = append(*affected, e.(*lastBundle).entry)
		} else {
			e.(*dimensionalBundle).key += uint64(number)
			rt.flatten(e.(*dimensionalBundle).sl, dimension+1, affected)
		}
	}

	if len(toDelete) > 0 {
		keys := make([]uint64, 0, len(toDelete))
		for _, e := range toDelete {
			if lastDimension {
				*deleted = append(*deleted, e.(*lastBundle).entry)
			} else {
				rt.flatten(e.(*dimensionalBundle).sl, dimension+1, deleted)
			}
			keys = append(keys, e.Key())
		}

		sl.Delete(keys...)
	}
}

// InsertAtDimension will increment items at and above the given index
// by the number provided.  Provide a negative number to to decrement.
// Returned are two lists.  The first list is a list of entries that
// were moved.  The second is a list entries that were deleted.  These
// lists are exclusive.
func (rt *skipListRT) InsertAtDimension(dimension uint64,
	index, number int64) (rangetree.Entries, rangetree.Entries) {

	if dimension >= rt.dimensions || number == 0 {
		return rangetree.Entries{}, rangetree.Entries{}
	}

	affected := make(rangetree.Entries, 0, 100)
	var deleted rangetree.Entries
	if number < 0 {
		deleted = make(rangetree.Entries, 0, 100)
	}

	rt.insert(rt.top, 0, dimension, index, number, &deleted, &affected)
	rt.number -= uint64(len(deleted))
	return affected, deleted
}

func new(dimensions uint64) *skipListRT {
	sl := &skipListRT{}
	sl.init(dimensions)
	return sl
}

// New will allocate, initialize, and return a new rangetree.RangeTree
// with the provided number of dimensions.
func New(dimensions uint64) rangetree.RangeTree {
	return new(dimensions)
}
<commit_msg>Fixed up the rangetree/skiplist package in light of the interface changes.<commit_after>/*
Copyright 2014 Workiva, LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*
Package skiplist implements an n-dimensional rangetree based on a skip
list.  This should be faster than a straight slice implementation as
memcopy is avoided.

Time complexities revolve around the ability to quickly find items
in the n-dimensional skiplist.  That time can be defined by the
number of items in any dimension.  Let N1, N2,... Nn define the
number of dimensions.

Performance characteristics:
Space: O(n)
Search: O(log N1 + log N2 + ...log Nn) = O(log N1*N2*...Nn)
Insert: O(log N1 + log N2 + ...log Nn) = O(log N1*N2*...Nn)
Delete: O(log N1 + log N2 + ...log Nn) = O(log N1*N2*...Nn)
*/

package skiplist

import (
	"github.com/Workiva/go-datastructures/rangetree"
	"github.com/Workiva/go-datastructures/slice/skip"
)

// keyed is required as in the rangetree code we often want to compare
// two different types of bundles and this allows us to do so without
// checking for each one.
type keyed interface {
	key() uint64
}

type skipEntry uint64

// Compare is required by the skip.Entry interface.
func (se skipEntry) Compare(other skip.Entry) int {
	otherSe := other.(skipEntry)
	if se == otherSe {
		return 0
	}

	if se > otherSe {
		return 1
	}

	return -1
}

func (se skipEntry) key() uint64 {
	return uint64(se)
}

// isLastDimension simply returns dimension == lastDimension-1.
// This panics if dimension >= lastDimension.
func isLastDimension(dimension, lastDimension uint64) bool {
	if dimension >= lastDimension { // useful in testing and denotes a serious problem
		panic(`Dimension is greater than possible dimensions.`)
	}

	return dimension == lastDimension-1
}

// needsDeletion returns a bool indicating if the provided value
// needs to be deleted based on the provided index and number.
func needsDeletion(value, index, number int64) bool {
	if number > 0 {
		return false
	}

	number = -number // get the magnitude
	offset := value - index

	return offset >= 0 && offset < number
}

// dimensionalBundle is an intermediate holder up to the last
// dimension and represents a wrapper around a skiplist.
type dimensionalBundle struct {
	id uint64
	sl *skip.SkipList
}

// Compare returns a value indicating the relative relationship and the
// provided bundle.
func (db *dimensionalBundle) Compare(e skip.Entry) int {
	keyed := e.(keyed)
	if db.id == keyed.key() {
		return 0
	}

	if db.id > keyed.key() {
		return 1
	}

	return -1
}

// key returns the key for this bundle.
func (db *dimensionalBundle) key() uint64 {
	return db.id
}

// lastBundle represents a bundle living at the last dimension
// of the tree.
type lastBundle struct {
	id    uint64
	entry rangetree.Entry
}

// Compare returns a value indicating the relative relationship and the
// provided bundle.
func (lb *lastBundle) Compare(e skip.Entry) int {
	keyed := e.(keyed)
	if lb.id == keyed.key() {
		return 0
	}

	if lb.id > keyed.key() {
		return 1
	}

	return -1
}

// Key returns the key for this bundle.
func (lb *lastBundle) key() uint64 {
	return lb.id
}

type skipListRT struct {
	top                *skip.SkipList
	dimensions, number uint64
}

func (rt *skipListRT) init(dimensions uint64) {
	rt.dimensions = dimensions
	rt.top = skip.New(uint64(0))
}

func (rt *skipListRT) add(entry rangetree.Entry) rangetree.Entry {
	var (
		value int64
		e     skip.Entry
		sl    = rt.top
		db    *dimensionalBundle
		lb    *lastBundle
	)

	for i := uint64(0); i < rt.dimensions; i++ {
		value = entry.ValueAtDimension(i)
		e = sl.Get(skipEntry(value))[0]
		if isLastDimension(i, rt.dimensions) {
			if e != nil { // this is an overwrite
				lb = e.(*lastBundle)
				oldEntry := lb.entry
				lb.entry = entry
				return oldEntry
			}

			// need to add new sl entry
			lb = &lastBundle{id: uint64(value), entry: entry}
			rt.number++
			sl.Insert(lb)
			return nil
		}

		if e == nil { // we need the intermediate dimension
			db = &dimensionalBundle{id: uint64(value), sl: skip.New(uint64(0))}
			sl.Insert(db)
		} else {
			db = e.(*dimensionalBundle)
		}

		sl = db.sl
	}

	panic(`Ran out of dimensions before for loop completed.`)
}

// Add will add the provided entries to the tree.  Any entries that
// were overwritten will be returned in the order in which they
// were overwritten.  If an entry's addition does not overwrite, a nil
// is returned for that cell for its index in the provided entries.
func (rt *skipListRT) Add(entries ...rangetree.Entry) rangetree.Entries {
	overwritten := make(rangetree.Entries, 0, len(entries))
	for _, e := range entries {
		overwritten = append(overwritten, rt.add(e))
	}

	return overwritten
}

func (rt *skipListRT) get(entry rangetree.Entry) rangetree.Entry {
	var (
		sl    = rt.top
		e     skip.Entry
		value uint64
	)
	for i := uint64(0); i < rt.dimensions; i++ {
		value = uint64(entry.ValueAtDimension(i))
		e = sl.Get(skipEntry(value))[0]
		if e == nil {
			return nil
		}

		if isLastDimension(i, rt.dimensions) {
			return e.(*lastBundle).entry
		}

		sl = e.(*dimensionalBundle).sl
	}

	panic(`Reached past for loop without finding last dimension.`)
}

// Get will return any rangetree.Entries matching the provided entries.
// Similar in functionality to a key lookup, this returns nil for any
// entry that could not be found.
func (rt *skipListRT) Get(entries ...rangetree.Entry) rangetree.Entries {
	results := make(rangetree.Entries, 0, len(entries))
	for _, e := range entries {
		results = append(results, rt.get(e))
	}

	return results
}

// Len returns the number of entries in the tree.
func (rt *skipListRT) Len() uint64 {
	return rt.number
}

// deleteRecursive is used by the delete logic.  The recursion depth
// only goes as far as the number of dimensions, so this shouldn't be an
// issue.
func (rt *skipListRT) deleteRecursive(sl *skip.SkipList, dimension uint64,
	entry rangetree.Entry) rangetree.Entry {

	value := entry.ValueAtDimension(dimension)
	if isLastDimension(dimension, rt.dimensions) {
		entries := sl.Delete(skipEntry(value))
		if entries[0] == nil {
			return nil
		}

		rt.number--
		return entries[0].(*lastBundle).entry
	}

	db, ok := sl.Get(skipEntry(value))[0].(*dimensionalBundle)
	if !ok { // value was not found
		return nil
	}

	result := rt.deleteRecursive(db.sl, dimension+1, entry)
	if result == nil {
		return nil
	}

	if db.sl.Len() == 0 {
		sl.Delete(db)
	}

	return result
}

func (rt *skipListRT) delete(entry rangetree.Entry) rangetree.Entry {
	return rt.deleteRecursive(rt.top, 0, entry)
}

// Delete will remove the provided entries from the tree.
func (rt *skipListRT) Delete(entries ...rangetree.Entry) {
	for _, e := range entries {
		rt.delete(e)
	}
}

func (rt *skipListRT) apply(sl *skip.SkipList, dimension uint64,
	interval rangetree.Interval, fn func(rangetree.Entry) bool) bool {

	lowValue, highValue := interval.LowAtDimension(dimension), interval.HighAtDimension(dimension)

	var e skip.Entry

	for iter := sl.Iter(skipEntry(lowValue)); iter.Next(); {
		e = iter.Value()
		if int64(e.(keyed).key()) >= highValue {
			break
		}

		if isLastDimension(dimension, rt.dimensions) {
			if !fn(e.(*lastBundle).entry) {
				return false
			}
		} else {

			if !rt.apply(e.(*dimensionalBundle).sl, dimension+1, interval, fn) {
				return false
			}
		}
	}

	return true
}

// Apply will call the provided function with each entry that exists
// within the provided range, in order.  Return false at any time to
// cancel iteration.  Altering the entry in such a way that its location
// changes will result in undefined behavior.
func (rt *skipListRT) Apply(interval rangetree.Interval, fn func(rangetree.Entry) bool) {
	rt.apply(rt.top, 0, interval, fn)
}

// Query will return a list of entries that fall within
// the provided interval.
func (rt *skipListRT) Query(interval rangetree.Interval) rangetree.Entries {
	entries := make(rangetree.Entries, 0, 100)
	rt.apply(rt.top, 0, interval, func(e rangetree.Entry) bool {
		entries = append(entries, e)
		return true
	})

	return entries
}

func (rt *skipListRT) flatten(sl *skip.SkipList, dimension uint64, entries *rangetree.Entries) {
	lastDimension := isLastDimension(dimension, rt.dimensions)
	for iter := sl.Iter(skipEntry(0)); iter.Next(); {
		if lastDimension {
			*entries = append(*entries, iter.Value().(*lastBundle).entry)
		} else {
			rt.flatten(iter.Value().(*dimensionalBundle).sl, dimension+1, entries)
		}
	}
}

func (rt *skipListRT) insert(sl *skip.SkipList, dimension, insertDimension uint64,
	index, number int64, deleted, affected *rangetree.Entries) {

	var e skip.Entry
	lastDimension := isLastDimension(dimension, rt.dimensions)
	affectedDimension := dimension == insertDimension
	var iter skip.Iterator
	if dimension == insertDimension {
		iter = sl.Iter(skipEntry(index))
	} else {
		iter = sl.Iter(skipEntry(0))
	}

	var toDelete skip.Entries
	if number < 0 {
		toDelete = make(skip.Entries, 0, 100)
	}

	for iter.Next() {
		e = iter.Value()
		if !affectedDimension {
			rt.insert(e.(*dimensionalBundle).sl, dimension+1,
				insertDimension, index, number, deleted, affected,
			)
			continue
		}
		if needsDeletion(int64(e.(keyed).key()), index, number) {
			toDelete = append(toDelete, e)
			continue
		}

		if lastDimension {
			e.(*lastBundle).id += uint64(number)
			*affected = append(*affected, e.(*lastBundle).entry)
		} else {
			e.(*dimensionalBundle).id += uint64(number)
			rt.flatten(e.(*dimensionalBundle).sl, dimension+1, affected)
		}
	}

	if len(toDelete) > 0 {
		for _, e := range toDelete {
			if lastDimension {
				*deleted = append(*deleted, e.(*lastBundle).entry)
			} else {
				rt.flatten(e.(*dimensionalBundle).sl, dimension+1, deleted)
			}
		}

		sl.Delete(toDelete...)
	}
}

// InsertAtDimension will increment items at and above the given index
// by the number provided.  Provide a negative number to to decrement.
// Returned are two lists.  The first list is a list of entries that
// were moved.  The second is a list entries that were deleted.  These
// lists are exclusive.
func (rt *skipListRT) InsertAtDimension(dimension uint64,
	index, number int64) (rangetree.Entries, rangetree.Entries) {

	if dimension >= rt.dimensions || number == 0 {
		return rangetree.Entries{}, rangetree.Entries{}
	}

	affected := make(rangetree.Entries, 0, 100)
	var deleted rangetree.Entries
	if number < 0 {
		deleted = make(rangetree.Entries, 0, 100)
	}

	rt.insert(rt.top, 0, dimension, index, number, &deleted, &affected)
	rt.number -= uint64(len(deleted))
	return affected, deleted
}

func new(dimensions uint64) *skipListRT {
	sl := &skipListRT{}
	sl.init(dimensions)
	return sl
}

// New will allocate, initialize, and return a new rangetree.RangeTree
// with the provided number of dimensions.
func New(dimensions uint64) rangetree.RangeTree {
	return new(dimensions)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"sort"
	"time"

	"github.com/gorilla/mux"
	"gopkg.in/mgo.v2/bson"
)

var usernameFilter = regexp.MustCompile("[^0-9a-zA-Z-_ ]")

var userListCache []byte
var userListCacheAge time.Time

func seenList() (map[string]time.Time, error) {
	var recent = map[string]time.Time{}
	if resp, err := http.Get("http://fofgaming.com:8890/seen.json"); err != nil {
		log.Println("error fetching seen.json:", err.Error())
		return recent, err
	} else {
		defer resp.Body.Close()
		var data = map[string]time.Time{}
		dec := json.NewDecoder(resp.Body)
		if err := dec.Decode(&data); err != nil {
			log.Println("error decoding seen.json:", err.Error())
			return recent, err
		}
		maxAge := time.Now().Add(0 - (30 * 24 * time.Hour))
		for id, t := range data {
			if t.After(maxAge) {
				recent[id] = t
			}
		}
	}
	return recent, nil
}

func getUserlist() []byte {
	if userListCacheAge.After(time.Now().Add(0 - time.Hour)) {
		return userListCache
	}
	resp, err := http.Get("http://127.0.0.1:8879/users.json")
	if err != nil {
		log.Printf("Error fetching user list: %s", err.Error())
		return userListCache
	}
	defer resp.Body.Close()

	recent, recentErr := seenList()

	var details struct {
		Members []struct {
			ID      string `json:"id"`
			Name    string `json:"name"`
			Bot     bool   `json:"is_bot"`
			Deleted bool   `json:"deleted"`
			Profile struct {
				FirstName string `json:"first_name"`
				Avatar    string `json:"image_192"`
			} `json:"profile"`
		} `json:"members"`
	}
	d := json.NewDecoder(resp.Body)
	if err := d.Decode(&details); err != nil {
		log.Printf("Error unmarshaling user list: %s", err.Error())
		return userListCache
	}

	var destinyLookup = []struct {
		UserID         string `bson:"userid"`
		DestinyAccount string `bson:"account"`
	}{}
	if err := mdb.DB("fof").C("userLookup").Find(nil).Select(bson.M{"_id": -1, "userid": 1, "account": 1}).All(&destinyLookup); err != nil {
		log.Printf("Error fetching destiny lookup docs for: %s", err.Error())
		return userListCache
	}

	rval := []map[string]string{}
	for _, m := range details.Members {
		if m.Bot {
			continue
		}
		if m.Deleted {
			continue
		}
		if recentErr != nil {
			if _, ok := recent[m.ID]; !ok {
				continue
			}
		}
		var destiny = ""
		for _, dUser := range destinyLookup {
			if dUser.UserID == m.ID {
				destiny = dUser.DestinyAccount
			}
		}
		rval = append(rval, map[string]string{
			"gamertag": m.Profile.FirstName,
			"username": m.Name,
			"avatar":   m.Profile.Avatar,
			"destiny":  destiny,
		})
	}
	if encoded, err := json.Marshal(rval); err != nil {
		return userListCache
	} else {
		userListCache = encoded
		userListCacheAge = time.Now()
		return encoded
	}
}

func userList(w http.ResponseWriter, r *http.Request) {
	setCORS(w, r)
	setJSON(w, r)
	w.Write(getUserlist())
}

func memberDoc(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	member := vars["member"]
	w.Header().Set("Content-Type", "application/json")
	var result []struct {
		Type string `bson:"type"`
		Raw  string `bson:"raw"`
	}
	if err := mdb.DB("fof").C("userData").Find(bson.M{"username": member}).Select(bson.M{"_id": -1, "type": 1, "raw": 1}).All(&result); err != nil {
		log.Printf("Error fetching member doc for %s: %s", member, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	var found = 0
	for _, v := range result {
		if len(v.Type) > 11 && v.Type == "character-" {
			continue
		}
		if found == 0 {
			fmt.Fprintf(w, "{\"%s\":%s", v.Type, v.Raw)
			found++
		} else {
			fmt.Fprintf(w, ",\"%s\":%s", v.Type, v.Raw)
		}
	}
	fmt.Fprint(w, "}")
}

func memberSubDoc(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	member := usernameFilter.ReplaceAllString(vars["member"], "")
	part := vars["key"]
	w.Header().Set("Content-Type", "application/json")
	var result struct {
		Raw []byte `bson:"raw"`
	}
	err := mdb.DB("fof").C("userData").Find(bson.M{
		"username": member,
		"type":     part,
	}).Select(bson.M{"_id": -1, "raw": 1}).One(&result)
	if err != nil {
		log.Printf("Error fetching member doc.%s for %s: %s", part, member, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	w.Write(result.Raw)
}

func memberSubDocKeys(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	member := usernameFilter.ReplaceAllString(vars["member"], "")
	w.Header().Set("Content-Type", "application/json")
	e := json.NewEncoder(w)
	var result []struct {
		Type string `bson:"type"`
	}
	err := mdb.DB("fof").C(
		"userData").Find(bson.M{"username": member}).Select(bson.M{"type": 1}).All(&result)
	if err != nil {
		log.Printf("Error fetching member doc.keys for %s: %s", member, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	var rval = []string{}
	for _, v := range result {
		rval = append(rval, v.Type)
	}
	sort.Strings(rval)
	e.Encode(rval)
}
<commit_msg>reverse error checking logic...<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"sort"
	"time"

	"github.com/gorilla/mux"
	"gopkg.in/mgo.v2/bson"
)

var usernameFilter = regexp.MustCompile("[^0-9a-zA-Z-_ ]")

var userListCache []byte
var userListCacheAge time.Time

func seenList() (map[string]time.Time, error) {
	var recent = map[string]time.Time{}
	if resp, err := http.Get("http://fofgaming.com:8890/seen.json"); err != nil {
		log.Println("error fetching seen.json:", err.Error())
		return recent, err
	} else {
		defer resp.Body.Close()
		var data = map[string]time.Time{}
		dec := json.NewDecoder(resp.Body)
		if err := dec.Decode(&data); err != nil {
			log.Println("error decoding seen.json:", err.Error())
			return recent, err
		}
		maxAge := time.Now().Add(0 - (30 * 24 * time.Hour))
		for id, t := range data {
			if t.After(maxAge) {
				recent[id] = t
			}
		}
	}
	return recent, nil
}

func getUserlist() []byte {
	if userListCacheAge.After(time.Now().Add(0 - time.Hour)) {
		return userListCache
	}
	resp, err := http.Get("http://127.0.0.1:8879/users.json")
	if err != nil {
		log.Printf("Error fetching user list: %s", err.Error())
		return userListCache
	}
	defer resp.Body.Close()

	recent, recentErr := seenList()

	var details struct {
		Members []struct {
			ID      string `json:"id"`
			Name    string `json:"name"`
			Bot     bool   `json:"is_bot"`
			Deleted bool   `json:"deleted"`
			Profile struct {
				FirstName string `json:"first_name"`
				Avatar    string `json:"image_192"`
			} `json:"profile"`
		} `json:"members"`
	}
	d := json.NewDecoder(resp.Body)
	if err := d.Decode(&details); err != nil {
		log.Printf("Error unmarshaling user list: %s", err.Error())
		return userListCache
	}

	var destinyLookup = []struct {
		UserID         string `bson:"userid"`
		DestinyAccount string `bson:"account"`
	}{}
	if err := mdb.DB("fof").C("userLookup").Find(nil).Select(bson.M{"_id": -1, "userid": 1, "account": 1}).All(&destinyLookup); err != nil {
		log.Printf("Error fetching destiny lookup docs for: %s", err.Error())
		return userListCache
	}

	rval := []map[string]string{}
	for _, m := range details.Members {
		if m.Bot {
			continue
		}
		if m.Deleted {
			continue
		}
		if recentErr == nil {
			if _, ok := recent[m.ID]; !ok {
				continue
			}
		}
		var destiny = ""
		for _, dUser := range destinyLookup {
			if dUser.UserID == m.ID {
				destiny = dUser.DestinyAccount
			}
		}
		rval = append(rval, map[string]string{
			"gamertag": m.Profile.FirstName,
			"username": m.Name,
			"avatar":   m.Profile.Avatar,
			"destiny":  destiny,
		})
	}
	if encoded, err := json.Marshal(rval); err != nil {
		return userListCache
	} else {
		userListCache = encoded
		userListCacheAge = time.Now()
		return encoded
	}
}

func userList(w http.ResponseWriter, r *http.Request) {
	setCORS(w, r)
	setJSON(w, r)
	w.Write(getUserlist())
}

func memberDoc(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	member := vars["member"]
	w.Header().Set("Content-Type", "application/json")
	var result []struct {
		Type string `bson:"type"`
		Raw  string `bson:"raw"`
	}
	if err := mdb.DB("fof").C("userData").Find(bson.M{"username": member}).Select(bson.M{"_id": -1, "type": 1, "raw": 1}).All(&result); err != nil {
		log.Printf("Error fetching member doc for %s: %s", member, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	var found = 0
	for _, v := range result {
		if len(v.Type) > 11 && v.Type == "character-" {
			continue
		}
		if found == 0 {
			fmt.Fprintf(w, "{\"%s\":%s", v.Type, v.Raw)
			found++
		} else {
			fmt.Fprintf(w, ",\"%s\":%s", v.Type, v.Raw)
		}
	}
	fmt.Fprint(w, "}")
}

func memberSubDoc(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	member := usernameFilter.ReplaceAllString(vars["member"], "")
	part := vars["key"]
	w.Header().Set("Content-Type", "application/json")
	var result struct {
		Raw []byte `bson:"raw"`
	}
	err := mdb.DB("fof").C("userData").Find(bson.M{
		"username": member,
		"type":     part,
	}).Select(bson.M{"_id": -1, "raw": 1}).One(&result)
	if err != nil {
		log.Printf("Error fetching member doc.%s for %s: %s", part, member, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	w.Write(result.Raw)
}

func memberSubDocKeys(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	member := usernameFilter.ReplaceAllString(vars["member"], "")
	w.Header().Set("Content-Type", "application/json")
	e := json.NewEncoder(w)
	var result []struct {
		Type string `bson:"type"`
	}
	err := mdb.DB("fof").C(
		"userData").Find(bson.M{"username": member}).Select(bson.M{"type": 1}).All(&result)
	if err != nil {
		log.Printf("Error fetching member doc.keys for %s: %s", member, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	var rval = []string{}
	for _, v := range result {
		rval = append(rval, v.Type)
	}
	sort.Strings(rval)
	e.Encode(rval)
}
<|endoftext|>


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

package buildbucket

import (
	"fmt"
	"net/http"
	"time"

	"golang.org/x/net/context"

	"go.chromium.org/gae/service/datastore"
	swarmbucketAPI "go.chromium.org/luci/common/api/buildbucket/swarmbucket/v1"
	swarmingAPI "go.chromium.org/luci/common/api/swarming/swarming/v1"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/milo/buildsource/swarming"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/milo/frontend/ui"
	"go.chromium.org/luci/server/auth"
)

func getPool(c context.Context, bid BuilderID) (*ui.MachinePool, error) {
	// Get PoolKey
	builderPool := model.BuilderPool{
		BuilderID: datastore.MakeKey(c, model.BuilderSummaryKind, bid.String()),
	}
	// These are eventually consistent, so just log an error and pass if not found.
	switch err := datastore.Get(c, &builderPool); {
	case datastore.IsErrNoSuchEntity(err):
		logging.Warningf(c, "builder pool not found")
		return nil, nil
	case err != nil:
		return nil, err
	}
	// Get BotPool
	botPool := model.BotPool{PoolID: builderPool.PoolKey.StringID()}
	switch err := datastore.Get(c, &botPool); {
	case datastore.IsErrNoSuchEntity(err):
		logging.Warningf(c, "bot pool not found")
		return nil, nil
	case err != nil:
		return nil, err
	}
	return ui.NewMachinePool(c, botPool.Bots), nil
}

// processBuilders parses out all of the builder pools from the Swarmbucket get_builders response,
// and saves the BuilderPool information into the datastore.
// It returns a list of PoolDescriptors that needs to be fetched and saved.
func processBuilders(c context.Context, r *swarmbucketAPI.SwarmingSwarmbucketApiGetBuildersResponseMessage) ([]model.PoolDescriptor, error) {
	var builderPools []model.BuilderPool
	var descriptors []model.PoolDescriptor
	seen := stringset.New(0)
	for _, bucket := range r.Buckets {
		for _, builder := range bucket.Builders {
			bid := NewBuilderID(bucket.Name, builder.Name)
			if bid.Project == "" {
				// This may happen if the bucket or builder does not fulfill the LUCI
				// naming convention.
				logging.Warningf(c, "invalid bucket/builder %q/%q, skipping", bucket.Name, builder.Name)
				continue
			}
			id := bid.String()
			descriptor := model.NewPoolDescriptor(bucket.SwarmingHostname, builder.SwarmingDimensions)
			dID := descriptor.PoolID()
			builderPools = append(builderPools, model.BuilderPool{
				BuilderID: datastore.MakeKey(c, model.BuilderSummaryKind, id),
				PoolKey:   datastore.MakeKey(c, model.BotPoolKind, dID),
			})
			if added := seen.Add(dID); added {
				descriptors = append(descriptors, descriptor)
			}
		}
	}
	return descriptors, datastore.Put(c, builderPools)
}

// parseBot parses a Swarming BotInfo response into the structure we will
// save into the datastore.  Since BotInfo doesn't have an explicit status
// field that matches Milo's abstraction of a Bot, the status is inferred:
// * A bot with TaskID is Busy
// * A bot that is dead or quarantined is Offline
// * Otherwise, it is implicitly connected and Idle.
func parseBot(c context.Context, swarmingHost string, botInfo *swarmingAPI.SwarmingRpcsBotInfo) (*model.Bot, error) {
	lastSeen, err := time.Parse(swarming.SwarmingTimeLayout, botInfo.LastSeenTs)
	if err != nil {
		return nil, err
	}
	result := &model.Bot{
		Name:     botInfo.BotId,
		URL:      fmt.Sprintf("https://%s/bot?id=%s", swarmingHost, botInfo.BotId),
		LastSeen: lastSeen,
	}

	switch {
	case botInfo.TaskId != "":
		result.Status = model.Busy
	case botInfo.IsDead || botInfo.Quarantined:
		result.Status = model.Offline
	default:
		// Defaults to idle.
	}
	return result, nil
}

// processBot retrieves the Bot pool details from Swarming for a given set of
// dimensions for its respective Swarming host, and saves the data into datastore.
func processBot(c context.Context, desc model.PoolDescriptor) error {
	t, err := auth.GetRPCTransport(c, auth.AsSelf)
	if err != nil {
		return err
	}
	sc, err := swarmingAPI.New(&http.Client{Transport: t})
	if err != nil {
		return err
	}
	sc.BasePath = fmt.Sprintf("https://%s/_ah/api/swarming/v1/", desc.Host())

	botList, err := sc.Bots.List().Dimensions(desc.Dimensions().Format()...).Do()
	if err != nil {
		return err
	}
	bots := make([]model.Bot, 0, len(botList.Items))
	for _, botInfo := range botList.Items {
		// Ignore deleted bots.
		if botInfo.Deleted {
			continue
		}
		bot, err := parseBot(c, desc.Host(), botInfo)
		if err != nil {
			return err
		}
		bots = append(bots, *bot)
	}
	// This is a large RPC, don't try to batch it.
	return datastore.Put(c, &model.BotPool{
		PoolID:     desc.PoolID(),
		Descriptor: desc,
		Bots:       bots,
		LastUpdate: clock.Now(c),
	})
}

// fetchBotPools resolves the descriptors into actual BotPool information.
// The input is a list of descriptors to fetch from swarming.
// Basically this just runs processBot() a bunch of times.
func processBots(c context.Context, descriptors []model.PoolDescriptor) error {
	return parallel.WorkPool(8, func(ch chan<- func() error) {
		for _, desc := range descriptors {
			desc := desc
			ch <- func() error {
				return processBot(c, desc)
			}
		}
	})
}

// UpdatePools is a cron job endpoint that:
// 1. Fetches all the builders from our associated Swarmbucket instance.
// 2. Consolidates all known descriptors (host+dimensions), saves BuilderPool.
// 3. Fetches and saves BotPool data from swarming for all known descriptors.
func UpdatePools(c context.Context) error {
	host, err := getHost(c)
	if err != nil {
		return err
	}
	sc, err := newSwarmbucketClient(c, host)
	if err != nil {
		return err
	}
	r, err := sc.GetBuilders().Do()
	if err != nil {
		return err
	}
	// Process builders and save them.  We get back the descriptors that we have
	// to fetch next.
	descriptors, err := processBuilders(c, r)
	if err != nil {
		return errors.Annotate(err, "processing builders").Err()
	}
	// And now also fetch and save the BotPools.
	return processBots(c, descriptors)
}
<commit_msg>[milo] Page through Swarming botlist call for bot pool info<commit_after>// Copyright 2018 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package buildbucket

import (
	"fmt"
	"net/http"
	"time"

	"golang.org/x/net/context"

	"go.chromium.org/gae/service/datastore"
	swarmbucketAPI "go.chromium.org/luci/common/api/buildbucket/swarmbucket/v1"
	swarmingAPI "go.chromium.org/luci/common/api/swarming/swarming/v1"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/milo/buildsource/swarming"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/milo/frontend/ui"
	"go.chromium.org/luci/server/auth"
)

func getPool(c context.Context, bid BuilderID) (*ui.MachinePool, error) {
	// Get PoolKey
	builderPool := model.BuilderPool{
		BuilderID: datastore.MakeKey(c, model.BuilderSummaryKind, bid.String()),
	}
	// These are eventually consistent, so just log an error and pass if not found.
	switch err := datastore.Get(c, &builderPool); {
	case datastore.IsErrNoSuchEntity(err):
		logging.Warningf(c, "builder pool not found")
		return nil, nil
	case err != nil:
		return nil, err
	}
	// Get BotPool
	botPool := model.BotPool{PoolID: builderPool.PoolKey.StringID()}
	switch err := datastore.Get(c, &botPool); {
	case datastore.IsErrNoSuchEntity(err):
		logging.Warningf(c, "bot pool not found")
		return nil, nil
	case err != nil:
		return nil, err
	}
	return ui.NewMachinePool(c, botPool.Bots), nil
}

// processBuilders parses out all of the builder pools from the Swarmbucket get_builders response,
// and saves the BuilderPool information into the datastore.
// It returns a list of PoolDescriptors that needs to be fetched and saved.
func processBuilders(c context.Context, r *swarmbucketAPI.SwarmingSwarmbucketApiGetBuildersResponseMessage) ([]model.PoolDescriptor, error) {
	var builderPools []model.BuilderPool
	var descriptors []model.PoolDescriptor
	seen := stringset.New(0)
	for _, bucket := range r.Buckets {
		for _, builder := range bucket.Builders {
			bid := NewBuilderID(bucket.Name, builder.Name)
			if bid.Project == "" {
				// This may happen if the bucket or builder does not fulfill the LUCI
				// naming convention.
				logging.Warningf(c, "invalid bucket/builder %q/%q, skipping", bucket.Name, builder.Name)
				continue
			}
			id := bid.String()
			descriptor := model.NewPoolDescriptor(bucket.SwarmingHostname, builder.SwarmingDimensions)
			dID := descriptor.PoolID()
			builderPools = append(builderPools, model.BuilderPool{
				BuilderID: datastore.MakeKey(c, model.BuilderSummaryKind, id),
				PoolKey:   datastore.MakeKey(c, model.BotPoolKind, dID),
			})
			if added := seen.Add(dID); added {
				descriptors = append(descriptors, descriptor)
			}
		}
	}
	return descriptors, datastore.Put(c, builderPools)
}

// parseBot parses a Swarming BotInfo response into the structure we will
// save into the datastore.  Since BotInfo doesn't have an explicit status
// field that matches Milo's abstraction of a Bot, the status is inferred:
// * A bot with TaskID is Busy
// * A bot that is dead or quarantined is Offline
// * Otherwise, it is implicitly connected and Idle.
func parseBot(c context.Context, swarmingHost string, botInfo *swarmingAPI.SwarmingRpcsBotInfo) (*model.Bot, error) {
	lastSeen, err := time.Parse(swarming.SwarmingTimeLayout, botInfo.LastSeenTs)
	if err != nil {
		return nil, err
	}
	result := &model.Bot{
		Name:     botInfo.BotId,
		URL:      fmt.Sprintf("https://%s/bot?id=%s", swarmingHost, botInfo.BotId),
		LastSeen: lastSeen,
	}

	switch {
	case botInfo.TaskId != "":
		result.Status = model.Busy
	case botInfo.IsDead || botInfo.Quarantined:
		result.Status = model.Offline
	default:
		// Defaults to idle.
	}
	return result, nil
}

// processBot retrieves the Bot pool details from Swarming for a given set of
// dimensions for its respective Swarming host, and saves the data into datastore.
func processBot(c context.Context, desc model.PoolDescriptor) error {
	t, err := auth.GetRPCTransport(c, auth.AsSelf)
	if err != nil {
		return err
	}
	sc, err := swarmingAPI.New(&http.Client{Transport: t})
	if err != nil {
		return err
	}
	sc.BasePath = fmt.Sprintf("https://%s/_ah/api/swarming/v1/", desc.Host())

	var bots []model.Bot
	bl := sc.Bots.List().Dimensions(desc.Dimensions().Format()...)
	// Keep fetching until the cursor is empty.
	for {
		botList, err := bl.Do()
		if err != nil {
			return err
		}
		for _, botInfo := range botList.Items {
			// Ignore deleted bots.
			if botInfo.Deleted {
				continue
			}
			bot, err := parseBot(c, desc.Host(), botInfo)
			if err != nil {
				return err
			}
			bots = append(bots, *bot)
		}

		if botList.Cursor == "" {
			break
		}
		bl = bl.Cursor(botList.Cursor)
	}
	// This is a large RPC, don't try to batch it.
	return datastore.Put(c, &model.BotPool{
		PoolID:     desc.PoolID(),
		Descriptor: desc,
		Bots:       bots,
		LastUpdate: clock.Now(c),
	})
}

// fetchBotPools resolves the descriptors into actual BotPool information.
// The input is a list of descriptors to fetch from swarming.
// Basically this just runs processBot() a bunch of times.
func processBots(c context.Context, descriptors []model.PoolDescriptor) error {
	return parallel.WorkPool(8, func(ch chan<- func() error) {
		for _, desc := range descriptors {
			desc := desc
			ch <- func() error {
				return processBot(c, desc)
			}
		}
	})
}

// UpdatePools is a cron job endpoint that:
// 1. Fetches all the builders from our associated Swarmbucket instance.
// 2. Consolidates all known descriptors (host+dimensions), saves BuilderPool.
// 3. Fetches and saves BotPool data from swarming for all known descriptors.
func UpdatePools(c context.Context) error {
	host, err := getHost(c)
	if err != nil {
		return err
	}
	sc, err := newSwarmbucketClient(c, host)
	if err != nil {
		return err
	}
	r, err := sc.GetBuilders().Do()
	if err != nil {
		return err
	}
	// Process builders and save them.  We get back the descriptors that we have
	// to fetch next.
	descriptors, err := processBuilders(c, r)
	if err != nil {
		return errors.Annotate(err, "processing builders").Err()
	}
	// And now also fetch and save the BotPools.
	return processBots(c, descriptors)
}
<|endoftext|>


<commit_before>/*
                       WARNING WARNING WARNING

                Attention all potential contributors

   This testfile is not in the best state. We've been slowly transitioning
   from the built in "testing" package to using Ginkgo. As you can see, we've
   changed the format, but a lot of the setup, test body, descriptions, etc
   are either hardcoded, completely lacking, or misleading.

   For example:

   Describe("Testing with ginkgo"...)      // This is not a great description
   It("TestDoesSoemthing"...)              // This is a horrible description

   Describe("create-user command"...       // Describe the actual object under test
   It("creates a user when provided ..."   // this is more descriptive

   For good examples of writing Ginkgo tests for the cli, refer to

   src/github.com/cloudfoundry/cli/cf/commands/application/delete_app_test.go
   src/github.com/cloudfoundry/cli/cf/terminal/ui_test.go
   src/github.com/cloudfoundry/loggregator_consumer/consumer_test.go
*/

package serviceauthtoken_test

import (
	. "github.com/cloudfoundry/cli/cf/commands/serviceauthtoken"
	"github.com/cloudfoundry/cli/cf/models"
	testapi "github.com/cloudfoundry/cli/testhelpers/api"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	. "github.com/cloudfoundry/cli/testhelpers/matchers"
)

func callListServiceAuthTokens(requirementsFactory *testreq.FakeReqFactory, authTokenRepo *testapi.FakeAuthTokenRepo) (ui *testterm.FakeUI) {
	ui = &testterm.FakeUI{}

	config := testconfig.NewRepositoryWithDefaults()

	cmd := NewListServiceAuthTokens(ui, config, authTokenRepo)
	testcmd.RunCommand(cmd, []string{}, requirementsFactory)

	return
}

var _ = Describe("Testing with ginkgo", func() {
	It("TestListServiceAuthTokensRequirements", func() {
		authTokenRepo := &testapi.FakeAuthTokenRepo{}
		requirementsFactory := &testreq.FakeReqFactory{}

		requirementsFactory.LoginSuccess = false
		callListServiceAuthTokens(requirementsFactory, authTokenRepo)
		Expect(testcmd.CommandDidPassRequirements).To(BeFalse())

		requirementsFactory.LoginSuccess = true
		callListServiceAuthTokens(requirementsFactory, authTokenRepo)
		Expect(testcmd.CommandDidPassRequirements).To(BeTrue())
	})
	It("TestListServiceAuthTokens", func() {

		requirementsFactory := &testreq.FakeReqFactory{LoginSuccess: true}
		authTokenRepo := &testapi.FakeAuthTokenRepo{}
		authToken := models.ServiceAuthTokenFields{}
		authToken.Label = "a label"
		authToken.Provider = "a provider"
		authToken2 := models.ServiceAuthTokenFields{}
		authToken2.Label = "a second label"
		authToken2.Provider = "a second provider"
		authTokenRepo.FindAllAuthTokens = []models.ServiceAuthTokenFields{authToken, authToken2}

		ui := callListServiceAuthTokens(requirementsFactory, authTokenRepo)
		Expect(ui.Outputs).To(ContainSubstrings(
			[]string{"Getting service auth tokens as", "my-user"},
			[]string{"OK"},
			[]string{"label", "provider"},
			[]string{"a label", "a provider"},
			[]string{"a second label", "a second provider"},
		))
	})
})
<commit_msg>Cleanup service-auth-tokens test<commit_after>package serviceauthtoken_test

import (
	"github.com/cloudfoundry/cli/cf/configuration"
	"github.com/cloudfoundry/cli/cf/models"
	testapi "github.com/cloudfoundry/cli/testhelpers/api"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"

	. "github.com/cloudfoundry/cli/cf/commands/serviceauthtoken"
	. "github.com/cloudfoundry/cli/testhelpers/matchers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

func callListServiceAuthTokens(requirementsFactory *testreq.FakeReqFactory, authTokenRepo *testapi.FakeAuthTokenRepo) (ui *testterm.FakeUI) {
	ui = &testterm.FakeUI{}

	config := testconfig.NewRepositoryWithDefaults()

	cmd := NewListServiceAuthTokens(ui, config, authTokenRepo)
	testcmd.RunCommand(cmd, []string{}, requirementsFactory)

	return
}

var _ = Describe("service-auth-tokens command", func() {
	var (
		ui                  *testterm.FakeUI
		configRepo          configuration.ReadWriter
		authTokenRepo       *testapi.FakeAuthTokenRepo
		requirementsFactory *testreq.FakeReqFactory
	)

	BeforeEach(func() {
		ui = &testterm.FakeUI{Inputs: []string{"y"}}
		authTokenRepo = &testapi.FakeAuthTokenRepo{}
		configRepo = testconfig.NewRepositoryWithDefaults()
		requirementsFactory = &testreq.FakeReqFactory{}
	})

	runCommand := func(args ...string) {
		testcmd.RunCommand(NewListServiceAuthTokens(ui, configRepo, authTokenRepo), args, requirementsFactory)
	}

	Describe("requirements", func() {
		It("fails when not logged in", func() {
			runCommand()
			Expect(testcmd.CommandDidPassRequirements).To(BeFalse())
		})
	})

	Context("when logged in and some service auth tokens exist", func() {
		BeforeEach(func() {
			requirementsFactory.LoginSuccess = true

			authTokenRepo.FindAllAuthTokens = []models.ServiceAuthTokenFields{
				models.ServiceAuthTokenFields{Label: "a label", Provider: "a provider"},
				models.ServiceAuthTokenFields{Label: "a second label", Provider: "a second provider"},
			}
		})

		It("shows you the service auth tokens", func() {
			runCommand()

			Expect(ui.Outputs).To(ContainSubstrings(
				[]string{"Getting service auth tokens as", "my-user"},
				[]string{"OK"},
				[]string{"label", "provider"},
				[]string{"a label", "a provider"},
				[]string{"a second label", "a second provider"},
			))
		})
	})
})
<|endoftext|>