<commit_before>// Copyright 2015-present Oursky Ltd.
//
// 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 skydbtest

import (
	"fmt"
	"reflect"
	"time"

	"github.com/skygeario/skygear-server/pkg/server/skydb"
)

// MapConn is a naive memory implementation of skydb.Conn
type MapConn struct {
	UserMap                map[string]skydb.AuthInfo
	AssetMap               map[string]skydb.Asset
	InternalPublicDB       skydb.Database
	usernameMap            map[string]skydb.AuthInfo
	emailMap               map[string]skydb.AuthInfo
	recordAccessMap        map[string]skydb.RecordACL
	recordDefaultAccessMap map[string]skydb.RecordACL
	fieldAccess            skydb.FieldACL
	skydb.Conn
}

// NewMapConn returns a new MapConn.
func NewMapConn() *MapConn {
	return &MapConn{
		UserMap:                map[string]skydb.AuthInfo{},
		recordAccessMap:        map[string]skydb.RecordACL{},
		recordDefaultAccessMap: map[string]skydb.RecordACL{},
		fieldAccess:            skydb.FieldACL{},
		AssetMap:               map[string]skydb.Asset{},
	}
}

// CreateAuth creates a AuthInfo in UserMap.
func (conn *MapConn) CreateAuth(authinfo *skydb.AuthInfo) error {
	if _, existed := conn.UserMap[authinfo.ID]; existed {
		return skydb.ErrUserDuplicated
	}

	conn.UserMap[authinfo.ID] = *authinfo
	return nil
}

// GetAuth returns a AuthInfo in UserMap.
func (conn *MapConn) GetAuth(id string, authinfo *skydb.AuthInfo) error {
	u, ok := conn.UserMap[id]
	if !ok {
		return skydb.ErrUserNotFound
	}

	*authinfo = u
	return nil
}

// GetAuthByPrincipalID returns a AuthInfo by its principalID.
func (conn *MapConn) GetAuthByPrincipalID(principalID string, authinfo *skydb.AuthInfo) error {
	for _, u := range conn.UserMap {
		if _, ok := u.ProviderInfo[principalID]; ok {
			*authinfo = u
			return nil
		}
	}

	return skydb.ErrUserNotFound
}

// UpdateAuth updates an existing AuthInfo in UserMap.
func (conn *MapConn) UpdateAuth(authinfo *skydb.AuthInfo) error {
	if _, ok := conn.UserMap[authinfo.ID]; !ok {
		return skydb.ErrUserNotFound
	}

	conn.UserMap[authinfo.ID] = *authinfo
	return nil
}

// DeleteAuth remove an existing in UserMap.
func (conn *MapConn) DeleteAuth(id string) error {
	if _, ok := conn.UserMap[id]; !ok {
		return skydb.ErrUserNotFound
	}

	delete(conn.UserMap, id)
	return nil
}

// GetAdminRoles is not implemented.
func (conn *MapConn) GetAdminRoles() ([]string, error) {
	return []string{
		"admin",
	}, nil
}

// SetAdminRoles is not implemented.
func (conn *MapConn) SetAdminRoles(roles []string) error {
	panic("not implemented")
}

// GetDefaultRoles always return user for testing
func (conn *MapConn) GetDefaultRoles() ([]string, error) {
	return []string{
		"user",
	}, nil
}

// SetDefaultRoles is not implemented.
func (conn *MapConn) SetDefaultRoles(roles []string) error {
	panic("not implemented")
}

// SetRecordAccess sets record creation access
func (conn *MapConn) SetRecordAccess(recordType string, acl skydb.RecordACL) error {
	conn.recordAccessMap[recordType] = acl
	return nil
}

// SetRecordDefaultAccess sets record creation access
func (conn *MapConn) SetRecordDefaultAccess(recordType string, acl skydb.RecordACL) error {
	conn.recordDefaultAccessMap[recordType] = acl
	return nil
}

// GetRecordAccess returns record creation access of a specific type
func (conn *MapConn) GetRecordAccess(recordType string) (skydb.RecordACL, error) {
	acl, gotIt := conn.recordAccessMap[recordType]
	if !gotIt {
		acl = skydb.NewRecordACL([]skydb.RecordACLEntry{})
	}

	return acl, nil
}

// GetRecordDefaultAccess returns record default access of a specific type
func (conn *MapConn) GetRecordDefaultAccess(recordType string) (skydb.RecordACL, error) {
	acl, gotIt := conn.recordDefaultAccessMap[recordType]
	if !gotIt {
		return nil, nil
	}
	return acl, nil
}

// SetRecordFieldAccess sets record field access for all types
func (conn *MapConn) SetRecordFieldAccess(acl skydb.FieldACL) error {
	conn.fieldAccess = acl
	return nil
}

// GetRecordFieldAccess returns record field access for all types
func (conn *MapConn) GetRecordFieldAccess() (skydb.FieldACL, error) {
	return conn.fieldAccess, nil
}

// GetAsset is not implemented.
func (conn *MapConn) GetAsset(name string, asset *skydb.Asset) error {
	panic("not implemented")
}

// SaveAsset is not implemented.
func (conn *MapConn) SaveAsset(asset *skydb.Asset) error {
	panic("not implemented")
}

// GetAssets always returns empty array.
func (conn *MapConn) GetAssets(names []string) ([]skydb.Asset, error) {
	assets := []skydb.Asset{}
	for _, v := range names {
		asset, ok := conn.AssetMap[v]
		if ok {
			assets = append(assets, asset)
		}
	}
	return assets, nil
}

// QueryRelation is not implemented.
func (conn *MapConn) QueryRelation(user string, name string, direction string, config skydb.QueryConfig) []skydb.AuthInfo {
	panic("not implemented")
}

// QueryRelationCount is not implemented.
func (conn *MapConn) QueryRelationCount(user string, name string, direction string) (uint64, error) {
	panic("not implemented")
}

// AddRelation is not implemented.
func (conn *MapConn) AddRelation(user string, name string, targetUser string) error {
	panic("not implemented")
}

// RemoveRelation is not implemented.
func (conn *MapConn) RemoveRelation(user string, name string, targetUser string) error {
	panic("not implemented")
}

// GetDevice is not implemented.
func (conn *MapConn) GetDevice(id string, device *skydb.Device) error {
	panic("not implemented")
}

// QueryDevicesByUser is not implemented.
func (conn *MapConn) QueryDevicesByUser(user string) ([]skydb.Device, error) {
	panic("not implemented")
}

// SaveDevice is not implemented.
func (conn *MapConn) SaveDevice(device *skydb.Device) error {
	panic("not implemented")
}

// DeleteDevice is not implemented.
func (conn *MapConn) DeleteDevice(id string) error {
	panic("not implemented")
}

// DeleteDevicesByToken is not implemented.
func (conn *MapConn) DeleteDevicesByToken(token string, t time.Time) error {
	panic("not implemented")
}

// DeleteEmptyDevicesByTime is not implemented.
func (conn *MapConn) DeleteEmptyDevicesByTime(t time.Time) error {
	panic("not implemented")
}

// PublicDB is not implemented.
func (conn *MapConn) PublicDB() skydb.Database {
	return conn.InternalPublicDB
}

// PrivateDB is not implemented.
func (conn *MapConn) PrivateDB(userKey string) skydb.Database {
	panic("not implemented")
}

// Subscribe is not implemented.
func (conn *MapConn) Subscribe(recordEventChan chan skydb.RecordEvent) error {
	panic("not implemented")
}

// Close does nothing.
func (conn *MapConn) Close() error {
	// do nothing
	return nil
}

// RecordMap is a string=>Record map
type RecordMap map[string]skydb.Record

// SubscriptionMap is a string=>Subscription map
type SubscriptionMap map[string]skydb.Subscription

// RecordSchemaMap is a string=>RecordSchema map
type RecordSchemaMap map[string]skydb.RecordSchema

//recordType string, acl RecordACL

// MapDB is a naive memory implementation of skydb.Database.
type MapDB struct {
	RecordMap       RecordMap
	SubscriptionMap SubscriptionMap
	RecordSchemaMap RecordSchemaMap
	DBConn          skydb.Conn
	skydb.Database
}

// NewMapDB returns a new MapDB ready for use.
func NewMapDB() *MapDB {
	return &MapDB{
		RecordMap:       RecordMap{},
		SubscriptionMap: SubscriptionMap{},
		RecordSchemaMap: RecordSchemaMap{},
		DBConn:          &MapConn{},
	}
}

func (db *MapDB) IsReadOnly() bool { return false }

func (db *MapDB) DatabaseType() skydb.DatabaseType { return skydb.PublicDatabase }

// ID returns a mock Database ID.
func (db *MapDB) ID() string {
	return ""
}

func (db *MapDB) UserRecordType() string {
	return "user"
}

// Get returns a Record from RecordMap.
func (db *MapDB) Get(id skydb.RecordID, record *skydb.Record) error {
	r, ok := db.RecordMap[id.String()]
	if !ok {
		return skydb.ErrRecordNotFound
	}
	*record = r
	return nil

}

// Save assigns Record to RecordMap.
func (db *MapDB) Save(record *skydb.Record) error {
	recordID := record.ID.String()

	if origRecord, ok := db.RecordMap[recordID]; ok {
		*record = *origRecord.MergedCopy(record)
	}

	db.RecordMap[recordID] = *record
	return nil
}

// Delete remove the specified key from RecordMap.
func (db *MapDB) Delete(id skydb.RecordID) error {
	_, ok := db.RecordMap[id.String()]
	if !ok {
		return skydb.ErrRecordNotFound
	}
	delete(db.RecordMap, id.String())
	return nil
}

// Query is not implemented.
func (db *MapDB) Query(query *skydb.Query) (*skydb.Rows, error) {
	panic("skydbtest: MapDB.Query not supported")
}

// Extend store the type of the field.
func (db *MapDB) Extend(recordType string, schema skydb.RecordSchema) (bool, error) {
	if _, ok := db.RecordSchemaMap[recordType]; ok {
		for fieldName, fieldType := range schema {
			if _, ok := db.RecordSchemaMap[recordType][fieldName]; ok {
				ft := db.RecordSchemaMap[recordType][fieldName]
				if !reflect.DeepEqual(ft, fieldType) {
					return false, fmt.Errorf("Wrong type")
				}
			}
			db.RecordSchemaMap[recordType][fieldName] = fieldType
		}
	} else {
		db.RecordSchemaMap[recordType] = schema
	}
	return true, nil
}

func (db *MapDB) RenameSchema(recordType, oldColumnName, newColumnName string) error {
	if _, ok := db.RecordSchemaMap[recordType]; !ok {
		return fmt.Errorf("record type %s does not exist", recordType)
	}
	if _, ok := db.RecordSchemaMap[recordType][oldColumnName]; !ok {
		return fmt.Errorf("column %s does not exist", oldColumnName)
	}
	if _, ok := db.RecordSchemaMap[recordType][newColumnName]; ok {
		if !reflect.DeepEqual(
			db.RecordSchemaMap[recordType][oldColumnName],
			db.RecordSchemaMap[recordType][newColumnName],
		) {
			return fmt.Errorf("column type conflict")
		}
	}
	db.RecordSchemaMap[recordType][newColumnName] = db.RecordSchemaMap[recordType][oldColumnName]
	delete(db.RecordSchemaMap[recordType], oldColumnName)

	return nil
}

func (db *MapDB) DeleteSchema(recordType, columnName string) error {
	if _, ok := db.RecordSchemaMap[recordType]; !ok {
		return fmt.Errorf("record type %s does not exist", recordType)
	}
	if _, ok := db.RecordSchemaMap[recordType][columnName]; !ok {
		return fmt.Errorf("column %s does not exist", columnName)
	}
	delete(db.RecordSchemaMap[recordType], columnName)
	return nil
}

// GetSchema returns the record schema of a record type
func (db *MapDB) GetSchema(recordType string) (skydb.RecordSchema, error) {
	if _, ok := db.RecordSchemaMap[recordType]; !ok {
		return nil, fmt.Errorf("record type %s does not exist", recordType)
	}
	return db.RecordSchemaMap[recordType], nil
}

// GetRecordSchemas returns a list of all existing record type
func (db *MapDB) GetRecordSchemas() (map[string]skydb.RecordSchema, error) {
	return db.RecordSchemaMap, nil
}

// GetSubscription return a Subscription from SubscriptionMap.
func (db *MapDB) GetSubscription(name string, deviceID string, subscription *skydb.Subscription) error {
	s, ok := db.SubscriptionMap[deviceID+"/"+name]
	if !ok {
		return skydb.ErrSubscriptionNotFound
	}
	*subscription = s
	return nil
}

// SaveSubscription assigns to SubscriptionMap.
func (db *MapDB) SaveSubscription(subscription *skydb.Subscription) error {
	db.SubscriptionMap[subscription.DeviceID+"/"+subscription.ID] = *subscription
	return nil
}

// DeleteSubscription deletes the specified key from SubscriptionMap.
func (db *MapDB) DeleteSubscription(name string, deviceID string) error {
	key := deviceID + "/" + name
	_, ok := db.SubscriptionMap[key]
	if !ok {
		return skydb.ErrSubscriptionNotFound
	}
	delete(db.SubscriptionMap, key)
	return nil
}

// MockTxDatabase implements and records TxDatabase's methods and delegates other
// calls to underlying Database
type MockTxDatabase struct {
	DidBegin, DidCommit, DidRollback bool
	skydb.Database
}

func NewMockTxDatabase(backingDB skydb.Database) *MockTxDatabase {
	return &MockTxDatabase{Database: backingDB}
}

func (db *MockTxDatabase) Begin() error {
	db.DidBegin = true
	return nil
}

func (db *MockTxDatabase) Commit() error {
	db.DidCommit = true
	return nil
}

func (db *MockTxDatabase) Rollback() error {
	db.DidRollback = true
	return nil
}

var _ skydb.TxDatabase = &MockTxDatabase{}

var (
	_ skydb.Conn       = NewMapConn()
	_ skydb.Database   = NewMapDB()
	_ skydb.TxDatabase = &MockTxDatabase{}
)
<commit_msg>Remove unused vars<commit_after>// Copyright 2015-present Oursky Ltd.
//
// 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 skydbtest

import (
	"fmt"
	"reflect"
	"time"

	"github.com/skygeario/skygear-server/pkg/server/skydb"
)

// MapConn is a naive memory implementation of skydb.Conn
type MapConn struct {
	UserMap                map[string]skydb.AuthInfo
	AssetMap               map[string]skydb.Asset
	InternalPublicDB       skydb.Database
	recordAccessMap        map[string]skydb.RecordACL
	recordDefaultAccessMap map[string]skydb.RecordACL
	fieldAccess            skydb.FieldACL
	skydb.Conn
}

// NewMapConn returns a new MapConn.
func NewMapConn() *MapConn {
	return &MapConn{
		UserMap:                map[string]skydb.AuthInfo{},
		recordAccessMap:        map[string]skydb.RecordACL{},
		recordDefaultAccessMap: map[string]skydb.RecordACL{},
		fieldAccess:            skydb.FieldACL{},
		AssetMap:               map[string]skydb.Asset{},
	}
}

// CreateAuth creates a AuthInfo in UserMap.
func (conn *MapConn) CreateAuth(authinfo *skydb.AuthInfo) error {
	if _, existed := conn.UserMap[authinfo.ID]; existed {
		return skydb.ErrUserDuplicated
	}

	conn.UserMap[authinfo.ID] = *authinfo
	return nil
}

// GetAuth returns a AuthInfo in UserMap.
func (conn *MapConn) GetAuth(id string, authinfo *skydb.AuthInfo) error {
	u, ok := conn.UserMap[id]
	if !ok {
		return skydb.ErrUserNotFound
	}

	*authinfo = u
	return nil
}

// GetAuthByPrincipalID returns a AuthInfo by its principalID.
func (conn *MapConn) GetAuthByPrincipalID(principalID string, authinfo *skydb.AuthInfo) error {
	for _, u := range conn.UserMap {
		if _, ok := u.ProviderInfo[principalID]; ok {
			*authinfo = u
			return nil
		}
	}

	return skydb.ErrUserNotFound
}

// UpdateAuth updates an existing AuthInfo in UserMap.
func (conn *MapConn) UpdateAuth(authinfo *skydb.AuthInfo) error {
	if _, ok := conn.UserMap[authinfo.ID]; !ok {
		return skydb.ErrUserNotFound
	}

	conn.UserMap[authinfo.ID] = *authinfo
	return nil
}

// DeleteAuth remove an existing in UserMap.
func (conn *MapConn) DeleteAuth(id string) error {
	if _, ok := conn.UserMap[id]; !ok {
		return skydb.ErrUserNotFound
	}

	delete(conn.UserMap, id)
	return nil
}

// GetAdminRoles is not implemented.
func (conn *MapConn) GetAdminRoles() ([]string, error) {
	return []string{
		"admin",
	}, nil
}

// SetAdminRoles is not implemented.
func (conn *MapConn) SetAdminRoles(roles []string) error {
	panic("not implemented")
}

// GetDefaultRoles always return user for testing
func (conn *MapConn) GetDefaultRoles() ([]string, error) {
	return []string{
		"user",
	}, nil
}

// SetDefaultRoles is not implemented.
func (conn *MapConn) SetDefaultRoles(roles []string) error {
	panic("not implemented")
}

// SetRecordAccess sets record creation access
func (conn *MapConn) SetRecordAccess(recordType string, acl skydb.RecordACL) error {
	conn.recordAccessMap[recordType] = acl
	return nil
}

// SetRecordDefaultAccess sets record creation access
func (conn *MapConn) SetRecordDefaultAccess(recordType string, acl skydb.RecordACL) error {
	conn.recordDefaultAccessMap[recordType] = acl
	return nil
}

// GetRecordAccess returns record creation access of a specific type
func (conn *MapConn) GetRecordAccess(recordType string) (skydb.RecordACL, error) {
	acl, gotIt := conn.recordAccessMap[recordType]
	if !gotIt {
		acl = skydb.NewRecordACL([]skydb.RecordACLEntry{})
	}

	return acl, nil
}

// GetRecordDefaultAccess returns record default access of a specific type
func (conn *MapConn) GetRecordDefaultAccess(recordType string) (skydb.RecordACL, error) {
	acl, gotIt := conn.recordDefaultAccessMap[recordType]
	if !gotIt {
		return nil, nil
	}
	return acl, nil
}

// SetRecordFieldAccess sets record field access for all types
func (conn *MapConn) SetRecordFieldAccess(acl skydb.FieldACL) error {
	conn.fieldAccess = acl
	return nil
}

// GetRecordFieldAccess returns record field access for all types
func (conn *MapConn) GetRecordFieldAccess() (skydb.FieldACL, error) {
	return conn.fieldAccess, nil
}

// GetAsset is not implemented.
func (conn *MapConn) GetAsset(name string, asset *skydb.Asset) error {
	panic("not implemented")
}

// SaveAsset is not implemented.
func (conn *MapConn) SaveAsset(asset *skydb.Asset) error {
	panic("not implemented")
}

// GetAssets always returns empty array.
func (conn *MapConn) GetAssets(names []string) ([]skydb.Asset, error) {
	assets := []skydb.Asset{}
	for _, v := range names {
		asset, ok := conn.AssetMap[v]
		if ok {
			assets = append(assets, asset)
		}
	}
	return assets, nil
}

// QueryRelation is not implemented.
func (conn *MapConn) QueryRelation(user string, name string, direction string, config skydb.QueryConfig) []skydb.AuthInfo {
	panic("not implemented")
}

// QueryRelationCount is not implemented.
func (conn *MapConn) QueryRelationCount(user string, name string, direction string) (uint64, error) {
	panic("not implemented")
}

// AddRelation is not implemented.
func (conn *MapConn) AddRelation(user string, name string, targetUser string) error {
	panic("not implemented")
}

// RemoveRelation is not implemented.
func (conn *MapConn) RemoveRelation(user string, name string, targetUser string) error {
	panic("not implemented")
}

// GetDevice is not implemented.
func (conn *MapConn) GetDevice(id string, device *skydb.Device) error {
	panic("not implemented")
}

// QueryDevicesByUser is not implemented.
func (conn *MapConn) QueryDevicesByUser(user string) ([]skydb.Device, error) {
	panic("not implemented")
}

// SaveDevice is not implemented.
func (conn *MapConn) SaveDevice(device *skydb.Device) error {
	panic("not implemented")
}

// DeleteDevice is not implemented.
func (conn *MapConn) DeleteDevice(id string) error {
	panic("not implemented")
}

// DeleteDevicesByToken is not implemented.
func (conn *MapConn) DeleteDevicesByToken(token string, t time.Time) error {
	panic("not implemented")
}

// DeleteEmptyDevicesByTime is not implemented.
func (conn *MapConn) DeleteEmptyDevicesByTime(t time.Time) error {
	panic("not implemented")
}

// PublicDB is not implemented.
func (conn *MapConn) PublicDB() skydb.Database {
	return conn.InternalPublicDB
}

// PrivateDB is not implemented.
func (conn *MapConn) PrivateDB(userKey string) skydb.Database {
	panic("not implemented")
}

// Subscribe is not implemented.
func (conn *MapConn) Subscribe(recordEventChan chan skydb.RecordEvent) error {
	panic("not implemented")
}

// Close does nothing.
func (conn *MapConn) Close() error {
	// do nothing
	return nil
}

// RecordMap is a string=>Record map
type RecordMap map[string]skydb.Record

// SubscriptionMap is a string=>Subscription map
type SubscriptionMap map[string]skydb.Subscription

// RecordSchemaMap is a string=>RecordSchema map
type RecordSchemaMap map[string]skydb.RecordSchema

//recordType string, acl RecordACL

// MapDB is a naive memory implementation of skydb.Database.
type MapDB struct {
	RecordMap       RecordMap
	SubscriptionMap SubscriptionMap
	RecordSchemaMap RecordSchemaMap
	DBConn          skydb.Conn
	skydb.Database
}

// NewMapDB returns a new MapDB ready for use.
func NewMapDB() *MapDB {
	return &MapDB{
		RecordMap:       RecordMap{},
		SubscriptionMap: SubscriptionMap{},
		RecordSchemaMap: RecordSchemaMap{},
		DBConn:          &MapConn{},
	}
}

func (db *MapDB) IsReadOnly() bool { return false }

func (db *MapDB) DatabaseType() skydb.DatabaseType { return skydb.PublicDatabase }

// ID returns a mock Database ID.
func (db *MapDB) ID() string {
	return ""
}

func (db *MapDB) UserRecordType() string {
	return "user"
}

// Get returns a Record from RecordMap.
func (db *MapDB) Get(id skydb.RecordID, record *skydb.Record) error {
	r, ok := db.RecordMap[id.String()]
	if !ok {
		return skydb.ErrRecordNotFound
	}
	*record = r
	return nil

}

// Save assigns Record to RecordMap.
func (db *MapDB) Save(record *skydb.Record) error {
	recordID := record.ID.String()

	if origRecord, ok := db.RecordMap[recordID]; ok {
		*record = *origRecord.MergedCopy(record)
	}

	db.RecordMap[recordID] = *record
	return nil
}

// Delete remove the specified key from RecordMap.
func (db *MapDB) Delete(id skydb.RecordID) error {
	_, ok := db.RecordMap[id.String()]
	if !ok {
		return skydb.ErrRecordNotFound
	}
	delete(db.RecordMap, id.String())
	return nil
}

// Query is not implemented.
func (db *MapDB) Query(query *skydb.Query) (*skydb.Rows, error) {
	panic("skydbtest: MapDB.Query not supported")
}

// Extend store the type of the field.
func (db *MapDB) Extend(recordType string, schema skydb.RecordSchema) (bool, error) {
	if _, ok := db.RecordSchemaMap[recordType]; ok {
		for fieldName, fieldType := range schema {
			if _, ok := db.RecordSchemaMap[recordType][fieldName]; ok {
				ft := db.RecordSchemaMap[recordType][fieldName]
				if !reflect.DeepEqual(ft, fieldType) {
					return false, fmt.Errorf("Wrong type")
				}
			}
			db.RecordSchemaMap[recordType][fieldName] = fieldType
		}
	} else {
		db.RecordSchemaMap[recordType] = schema
	}
	return true, nil
}

func (db *MapDB) RenameSchema(recordType, oldColumnName, newColumnName string) error {
	if _, ok := db.RecordSchemaMap[recordType]; !ok {
		return fmt.Errorf("record type %s does not exist", recordType)
	}
	if _, ok := db.RecordSchemaMap[recordType][oldColumnName]; !ok {
		return fmt.Errorf("column %s does not exist", oldColumnName)
	}
	if _, ok := db.RecordSchemaMap[recordType][newColumnName]; ok {
		if !reflect.DeepEqual(
			db.RecordSchemaMap[recordType][oldColumnName],
			db.RecordSchemaMap[recordType][newColumnName],
		) {
			return fmt.Errorf("column type conflict")
		}
	}
	db.RecordSchemaMap[recordType][newColumnName] = db.RecordSchemaMap[recordType][oldColumnName]
	delete(db.RecordSchemaMap[recordType], oldColumnName)

	return nil
}

func (db *MapDB) DeleteSchema(recordType, columnName string) error {
	if _, ok := db.RecordSchemaMap[recordType]; !ok {
		return fmt.Errorf("record type %s does not exist", recordType)
	}
	if _, ok := db.RecordSchemaMap[recordType][columnName]; !ok {
		return fmt.Errorf("column %s does not exist", columnName)
	}
	delete(db.RecordSchemaMap[recordType], columnName)
	return nil
}

// GetSchema returns the record schema of a record type
func (db *MapDB) GetSchema(recordType string) (skydb.RecordSchema, error) {
	if _, ok := db.RecordSchemaMap[recordType]; !ok {
		return nil, fmt.Errorf("record type %s does not exist", recordType)
	}
	return db.RecordSchemaMap[recordType], nil
}

// GetRecordSchemas returns a list of all existing record type
func (db *MapDB) GetRecordSchemas() (map[string]skydb.RecordSchema, error) {
	return db.RecordSchemaMap, nil
}

// GetSubscription return a Subscription from SubscriptionMap.
func (db *MapDB) GetSubscription(name string, deviceID string, subscription *skydb.Subscription) error {
	s, ok := db.SubscriptionMap[deviceID+"/"+name]
	if !ok {
		return skydb.ErrSubscriptionNotFound
	}
	*subscription = s
	return nil
}

// SaveSubscription assigns to SubscriptionMap.
func (db *MapDB) SaveSubscription(subscription *skydb.Subscription) error {
	db.SubscriptionMap[subscription.DeviceID+"/"+subscription.ID] = *subscription
	return nil
}

// DeleteSubscription deletes the specified key from SubscriptionMap.
func (db *MapDB) DeleteSubscription(name string, deviceID string) error {
	key := deviceID + "/" + name
	_, ok := db.SubscriptionMap[key]
	if !ok {
		return skydb.ErrSubscriptionNotFound
	}
	delete(db.SubscriptionMap, key)
	return nil
}

// MockTxDatabase implements and records TxDatabase's methods and delegates other
// calls to underlying Database
type MockTxDatabase struct {
	DidBegin, DidCommit, DidRollback bool
	skydb.Database
}

func NewMockTxDatabase(backingDB skydb.Database) *MockTxDatabase {
	return &MockTxDatabase{Database: backingDB}
}

func (db *MockTxDatabase) Begin() error {
	db.DidBegin = true
	return nil
}

func (db *MockTxDatabase) Commit() error {
	db.DidCommit = true
	return nil
}

func (db *MockTxDatabase) Rollback() error {
	db.DidRollback = true
	return nil
}

var _ skydb.TxDatabase = &MockTxDatabase{}

var (
	_ skydb.Conn       = NewMapConn()
	_ skydb.Database   = NewMapDB()
	_ skydb.TxDatabase = &MockTxDatabase{}
)
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"time"

	"encoding/json"

	"github.com/cheggaaa/pb"
	"github.com/codegangsta/cli"
	"github.com/mattn/go-colorable"
	"github.com/mattn/go-isatty"
	"github.com/rylio/ytdl"
	log "github.com/sirupsen/logrus"
)

type options struct {
	noProgress  bool
	outputFile  string
	infoOnly    bool
	silent      bool
	debug       bool
	append      bool
	filters     []string
	downloadURL bool
	byteRange   string
	json        bool
	startOffset string
}

func main() {
	app := cli.NewApp()
	app.Name = "ytdl"
	app.HelpName = "ytdl"
	// Set our own custom args usage
	app.ArgsUsage = "[youtube url or video id]"
	app.Usage = "Download youtube videos"
	app.HideHelp = true
	app.Version = "0.5.0"

	app.Flags = []cli.Flag{
		cli.HelpFlag,
		cli.StringFlag{
			Name:  "output, o",
			Usage: "Write output to a file, passing - outputs to stdout",
			Value: "{{.Title}}.{{.Ext}}",
		},
		cli.BoolFlag{
			Name:  "info, i",
			Usage: "Only output video info",
		},
		cli.BoolFlag{
			Name:  "no-progress",
			Usage: "Disable the progress bar",
		},
		cli.BoolFlag{
			Name:  "silent, s",
			Usage: "Only output errors, also disables progress bar",
		},
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "Output debug log",
		},
		cli.BoolFlag{
			Name:  "append, a",
			Usage: "Append to output file instead of overwriting",
		},
		cli.StringSliceFlag{
			Name:  "filter, f",
			Usage: "Filter available formats, syntax: [format_key]:val1,val2",
		},
		cli.StringFlag{
			Name:  "range, r",
			Usage: "Download a specific range of bytes of the video, [start]-[end]",
		},
		cli.BoolFlag{
			Name:  "download-url, u",
			Usage: "Prints download url to stdout",
		},
		cli.BoolFlag{
			Name:  "json, j",
			Usage: "Print info json to stdout",
		},
		cli.StringFlag{
			Name:  "start-offset",
			Usage: "Offset the start of the video by time",
		},
	}

	app.Action = func(c *cli.Context) error {
		identifier := c.Args().First()
		if identifier == "" || c.Bool("help") {
			cli.ShowAppHelp(c)
		} else {
			options := options{
				noProgress:  c.Bool("no-progress"),
				outputFile:  c.String("output"),
				infoOnly:    c.Bool("info"),
				silent:      c.Bool("silent"),
				debug:       c.Bool("debug"),
				append:      c.Bool("append"),
				filters:     c.StringSlice("filter"),
				downloadURL: c.Bool("download-url"),
				byteRange:   c.String("range"),
				json:        c.Bool("json"),
				startOffset: c.String("start-offset"),
			}
			if len(options.filters) == 0 {
				options.filters = cli.StringSlice{
					fmt.Sprintf("%s:mp4", ytdl.FormatExtensionKey),
					fmt.Sprintf("!%s:", ytdl.FormatVideoEncodingKey),
					fmt.Sprintf("!%s:", ytdl.FormatAudioEncodingKey),
					fmt.Sprint("best"),
				}
			}
			handler(identifier, options)
		}
		return nil
	}
	app.Run(os.Args)
}

func handler(identifier string, options options) {
	var err error
	defer func() {
		if err != nil {
			log.SetOutput(os.Stderr)
			log.Fatal(err.Error())
		}
	}()

	var out io.Writer
	var logOut io.Writer = os.Stdout
	if runtime.GOOS == "windows" && isatty.IsTerminal(os.Stdout.Fd()) {
		logOut = colorable.NewColorableStdout()
		log.SetFormatter(&log.TextFormatter{ForceColors: true})
	}
	// if downloading to stdout, set log output to stderr, not sure if this is correct
	if options.outputFile == "-" {
		out = os.Stdout
		logOut = os.Stderr
	}
	log.SetOutput(logOut)

	// ouput only errors or not
	silent := options.outputFile == "" ||
		options.silent || options.infoOnly || options.downloadURL || options.json
	if silent {
		log.SetLevel(log.FatalLevel)
	} else if options.debug {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.InfoLevel)
	}

	// TODO: Show activity indicator
	log.Info("Fetching video info...")
	//fmt.Print("\u001b[0G")
	//fmt.Print("\u001b[2K")
	info, err := ytdl.GetVideoInfo(identifier)
	if err != nil {
		err = fmt.Errorf("Unable to fetch video info: %s", err.Error())
		return
	}

	if options.infoOnly {
		fmt.Println("Title:", info.Title)
		fmt.Println("Author:", info.Author)
		fmt.Println("Date Published:", info.DatePublished.Format("Jan 2 2006"))
		fmt.Println("Duration:", info.Duration)
		return
	} else if options.json {
		var data []byte
		data, err = json.MarshalIndent(info, "", "\t")
		if err != nil {
			err = fmt.Errorf("Unable to marshal json: %s", err.Error())
			return
		}
		fmt.Println(string(data))
		return
	}

	formats := info.Formats
	// parse filter arguments, and filter through formats
	for _, filter := range options.filters {
		filter, err := parseFilter(filter)
		if err == nil {
			formats = filter(formats)
		}
	}
	if len(formats) == 0 {
		err = fmt.Errorf("No formats available that match criteria")
		return
	}
	format := formats[0]
	downloadURL, err := info.GetDownloadURL(format)
	if err != nil {
		err = fmt.Errorf("Unable to get download url: %s", err.Error())
		return
	}
	if options.startOffset != "" {
		var offset time.Duration
		offset, err = time.ParseDuration(options.startOffset)
		query := downloadURL.Query()
		query.Set("begin", fmt.Sprint(int64(offset/time.Millisecond)))
		downloadURL.RawQuery = query.Encode()
	}
	if options.downloadURL {
		fmt.Print(downloadURL.String())
		// print new line character if outputing to terminal
		if isatty.IsTerminal(os.Stdout.Fd()) {
			fmt.Println()
		}
		return
	}
	if out == nil {
		var fileName string
		fileName, err = createFileName(options.outputFile, outputFileName{
			Title:         sanitizeFileNamePart(info.Title),
			Ext:           sanitizeFileNamePart(format.Extension),
			DatePublished: sanitizeFileNamePart(info.DatePublished.Format("2006-01-02")),
			Resolution:    sanitizeFileNamePart(format.Resolution),
			Author:        sanitizeFileNamePart(info.Author),
			Duration:      sanitizeFileNamePart(info.Duration.String()),
		})
		if err != nil {
			err = fmt.Errorf("Unable to parse output file file name: %s", err.Error())
			return
		}
		// Create file truncate if append flag is not set
		flags := os.O_CREATE | os.O_WRONLY
		if options.append {
			flags |= os.O_APPEND
		} else {
			flags |= os.O_TRUNC
		}
		var f *os.File
		// open as write only
		f, err = os.OpenFile(fileName, flags, 0666)
		if err != nil {
			err = fmt.Errorf("Unable to open output file: %s", err.Error())
			return
		}
		defer f.Close()
		out = f
	}

	log.Info("Downloading to ", out.(*os.File).Name())
	var req *http.Request
	req, err = http.NewRequest("GET", downloadURL.String(), nil)
	// if byte range flag is set, use http range header option
	if options.byteRange != "" || options.append {
		if options.byteRange == "" && out != os.Stdout {
			if stat, err := out.(*os.File).Stat(); err != nil {
				options.byteRange = strconv.FormatInt(stat.Size(), 10) + "-"
			}
		}
		if options.byteRange != "" {
			req.Header.Set("Range", "bytes="+options.byteRange)
		}
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil || resp.StatusCode < 200 || resp.StatusCode >= 300 {
		if err == nil {
			err = fmt.Errorf("Received status code %d from download url", resp.StatusCode)
		}
		err = fmt.Errorf("Unable to start download: %s", err.Error())
		return
	}
	defer resp.Body.Close()
	// if we aren't in silent mode or the no progress flag wasn't set,
	// initialize progress bar
	if !silent && !options.noProgress {
		progressBar := pb.New64(resp.ContentLength)
		progressBar.SetUnits(pb.U_BYTES)
		progressBar.ShowTimeLeft = true
		progressBar.ShowSpeed = true
		//	progressBar.RefreshRate = time.Millisecond * 1
		progressBar.Output = logOut
		progressBar.Start()
		defer progressBar.Finish()
		out = io.MultiWriter(out, progressBar)
	}
	_, err = io.Copy(out, resp.Body)
}
<commit_msg>add runtime download option feature<commit_after>package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"time"

	"encoding/json"

	"github.com/cheggaaa/pb"
	"github.com/codegangsta/cli"
	"github.com/mattn/go-colorable"
	"github.com/mattn/go-isatty"
	"github.com/rylio/ytdl"
	log "github.com/sirupsen/logrus"
	"github.com/olekukonko/tablewriter"
	"sort"
)

type byResolution [][]string

type options struct {
	noProgress     bool
	outputFile     string
	infoOnly       bool
	silent         bool
	debug          bool
	append         bool
	filters        []string
	downloadURL    bool
	byteRange      string
	json           bool
	startOffset    string
	downloadOption bool
}

func main() {
	app := cli.NewApp()
	app.Name = "ytdl"
	app.HelpName = "ytdl"
	// Set our own custom args usage
	app.ArgsUsage = "[youtube url or video id]"
	app.Usage = "Download youtube videos"
	app.HideHelp = true
	app.Version = "0.5.0"

	app.Flags = []cli.Flag{
		cli.HelpFlag,
		cli.StringFlag{
			Name:  "output, o",
			Usage: "Write output to a file, passing - outputs to stdout",
			Value: "{{.Title}}.{{.Ext}}",
		},
		cli.BoolFlag{
			Name:  "info, i",
			Usage: "Only output video info",
		},
		cli.BoolFlag{
			Name:  "no-progress",
			Usage: "Disable the progress bar",
		},
		cli.BoolFlag{
			Name:  "silent, s",
			Usage: "Only output errors, also disables progress bar",
		},
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "Output debug log",
		},
		cli.BoolFlag{
			Name:  "append, a",
			Usage: "Append to output file instead of overwriting",
		},
		cli.StringSliceFlag{
			Name:  "filter, f",
			Usage: "Filter available formats, syntax: [format_key]:val1,val2",
		},
		cli.StringFlag{
			Name:  "range, r",
			Usage: "Download a specific range of bytes of the video, [start]-[end]",
		},
		cli.BoolFlag{
			Name:  "download-url, u",
			Usage: "Prints download url to stdout",
		},
		cli.BoolFlag{
			Name:  "json, j",
			Usage: "Print info json to stdout",
		},
		cli.StringFlag{
			Name:  "start-offset",
			Usage: "Offset the start of the video by time",
		},
		cli.BoolFlag{
			Name:  "download-option, p",
			Usage: "Print video and audio download options",
		},
	}

	app.Action = func(c *cli.Context) error {
		identifier := c.Args().First()
		if identifier == "" || c.Bool("help") {
			cli.ShowAppHelp(c)
		} else {
			options := options{
				noProgress:  c.Bool("no-progress"),
				outputFile:  c.String("output"),
				infoOnly:    c.Bool("info"),
				silent:      c.Bool("silent"),
				debug:       c.Bool("debug"),
				append:      c.Bool("append"),
				filters:     c.StringSlice("filter"),
				downloadURL: c.Bool("download-url"),
				byteRange:   c.String("range"),
				json:        c.Bool("json"),
				startOffset: c.String("start-offset"),
				downloadOption:      c.Bool("download-option"),
			}
			if len(options.filters) == 0 {
				options.filters = cli.StringSlice{
					fmt.Sprintf("%s:mp4", ytdl.FormatExtensionKey),
					fmt.Sprintf("!%s:", ytdl.FormatVideoEncodingKey),
					fmt.Sprintf("!%s:", ytdl.FormatAudioEncodingKey),
					fmt.Sprint("best"),
				}
			}
			handler(identifier, options)
		}
		return nil
	}
	app.Run(os.Args)
}

func handler(identifier string, options options) {
	var itag int
	var err error
	defer func() {
		if err != nil {
			log.SetOutput(os.Stderr)
			log.Fatal(err.Error())
		}
	}()

	var out io.Writer
	var logOut io.Writer = os.Stdout
	if runtime.GOOS == "windows" && isatty.IsTerminal(os.Stdout.Fd()) {
		logOut = colorable.NewColorableStdout()
		log.SetFormatter(&log.TextFormatter{ForceColors: true})
	}
	// if downloading to stdout, set log output to stderr, not sure if this is correct
	if options.outputFile == "-" {
		out = os.Stdout
		logOut = os.Stderr
	}
	log.SetOutput(logOut)

	// ouput only errors or not
	silent := options.outputFile == "" ||
		options.silent || options.infoOnly || options.downloadURL || options.json
	if silent {
		log.SetLevel(log.FatalLevel)
	} else if options.debug {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.InfoLevel)
	}

	// TODO: Show activity indicator
	log.Info("Fetching video info...")
	//fmt.Print("\u001b[0G")
	//fmt.Print("\u001b[2K")
	info, err := ytdl.GetVideoInfo(identifier)
	if err != nil {
		err = fmt.Errorf("Unable to fetch video info: %s", err.Error())
		return
	}

	if options.infoOnly {
		fmt.Println("Title:", info.Title)
		fmt.Println("Author:", info.Author)
		fmt.Println("Date Published:", info.DatePublished.Format("Jan 2 2006"))
		fmt.Println("Duration:", info.Duration)
		return
	} else if options.json {
		var data []byte
		data, err = json.MarshalIndent(info, "", "\t")
		if err != nil {
			err = fmt.Errorf("Unable to marshal json: %s", err.Error())
			return
		}
		fmt.Println(string(data))
		return
	} else if options.downloadOption {
		var data [][]string

		for _, format := range info.Formats {
			data = append(data, []string{strconv.Itoa(format.Itag), format.Extension, format.Resolution, format.VideoEncoding, format.AudioEncoding, strconv.Itoa(format.AudioBitrate)})
		}
		sort.Sort(byResolution(data))

		table := tablewriter.NewWriter(os.Stdout)
		table.SetHeader([]string{"itag", "ext", "res", "vEncoding", "aEncoding", "aBitrate"})

		for _, v := range data {
			table.Append(v)
		}
		table.Render() // Send output

		fmt.Print("Enter the itag of the file you would like to download or enter 0 to abort: ")
		_, err := fmt.Scanf("%d", &itag)
		if err != nil {
			return
		} else if itag == 0 {
			return
		}
	}

	formats := info.Formats

	if itag != 0 {
		filter, err := parseFilter(fmt.Sprintf("itag:%d", itag))
		if err == nil {
			formats = filter(formats)
		}
	} else {
		// parse filter arguments, and filter through formats
		for _, filter := range options.filters {
			filter, err := parseFilter(filter)
			if err == nil {
				formats = filter(formats)
			}
		}
	}

	if len(formats) == 0 {
		err = fmt.Errorf("No formats available that match criteria")
		return
	}
	format := formats[0]
	downloadURL, err := info.GetDownloadURL(format)
	if err != nil {
		err = fmt.Errorf("Unable to get download url: %s", err.Error())
		return
	}
	if options.startOffset != "" {
		var offset time.Duration
		offset, err = time.ParseDuration(options.startOffset)
		query := downloadURL.Query()
		query.Set("begin", fmt.Sprint(int64(offset/time.Millisecond)))
		downloadURL.RawQuery = query.Encode()
	}
	if options.downloadURL {
		fmt.Print(downloadURL.String())
		// print new line character if outputing to terminal
		if isatty.IsTerminal(os.Stdout.Fd()) {
			fmt.Println()
		}
		return
	}
	if out == nil {
		var fileName string
		fileName, err = createFileName(options.outputFile, outputFileName{
			Title:         sanitizeFileNamePart(info.Title),
			Ext:           sanitizeFileNamePart(format.Extension),
			DatePublished: sanitizeFileNamePart(info.DatePublished.Format("2006-01-02")),
			Resolution:    sanitizeFileNamePart(format.Resolution),
			Author:        sanitizeFileNamePart(info.Author),
			Duration:      sanitizeFileNamePart(info.Duration.String()),
		})
		if err != nil {
			err = fmt.Errorf("Unable to parse output file file name: %s", err.Error())
			return
		}
		// Create file truncate if append flag is not set
		flags := os.O_CREATE | os.O_WRONLY
		if options.append {
			flags |= os.O_APPEND
		} else {
			flags |= os.O_TRUNC
		}
		var f *os.File
		// open as write only
		f, err = os.OpenFile(fileName, flags, 0666)
		if err != nil {
			err = fmt.Errorf("Unable to open output file: %s", err.Error())
			return
		}
		defer f.Close()
		out = f
	}

	log.Info("Downloading to ", out.(*os.File).Name())
	var req *http.Request
	req, err = http.NewRequest("GET", downloadURL.String(), nil)
	// if byte range flag is set, use http range header option
	if options.byteRange != "" || options.append {
		if options.byteRange == "" && out != os.Stdout {
			if stat, err := out.(*os.File).Stat(); err != nil {
				options.byteRange = strconv.FormatInt(stat.Size(), 10) + "-"
			}
		}
		if options.byteRange != "" {
			req.Header.Set("Range", "bytes="+options.byteRange)
		}
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil || resp.StatusCode < 200 || resp.StatusCode >= 300 {
		if err == nil {
			err = fmt.Errorf("Received status code %d from download url", resp.StatusCode)
		}
		err = fmt.Errorf("Unable to start download: %s", err.Error())
		return
	}
	defer resp.Body.Close()
	// if we aren't in silent mode or the no progress flag wasn't set,
	// initialize progress bar
	if !silent && !options.noProgress {
		progressBar := pb.New64(resp.ContentLength)
		progressBar.SetUnits(pb.U_BYTES)
		progressBar.ShowTimeLeft = true
		progressBar.ShowSpeed = true
		//	progressBar.RefreshRate = time.Millisecond * 1
		progressBar.Output = logOut
		progressBar.Start()
		defer progressBar.Finish()
		out = io.MultiWriter(out, progressBar)
	}
	_, err = io.Copy(out, resp.Body)
}

func (s byResolution) Len() int {
	return len(s)
}
func (s byResolution) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s byResolution) Less(i, j int) bool {
	if s[i][2] == "" {
		return false
	} else if s[j][2] == "" {
		return true
	}

	if len(s[i][2]) < len(s[j][2]) {
		return false
	} else if len(s[i][2]) > len(s[j][2]) {
		return true
	}

	return s[i][2] > s[j][2]
}<|endoftext|>


<commit_before>// Package report reports statistics about TSM files.
package report

import (
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/influxdata/influxdb/models"
	"github.com/influxdata/influxdb/tsdb/engine/tsm1"
	"github.com/retailnext/hllpp"
)

// Command represents the program execution for "influxd report".
type Command struct {
	Stderr io.Writer
	Stdout io.Writer

	dir      string
	pattern  string
	detailed bool
}

// NewCommand returns a new instance of Command.
func NewCommand() *Command {
	return &Command{
		Stderr: os.Stderr,
		Stdout: os.Stdout,
	}
}

// Run executes the command.
func (cmd *Command) Run(args ...string) error {
	fs := flag.NewFlagSet("report", flag.ExitOnError)
	fs.StringVar(&cmd.pattern, "pattern", "", "Include only files matching a pattern")

	fs.BoolVar(&cmd.detailed, "detailed", false, "Report detailed cardinality estimates")

	fs.SetOutput(cmd.Stdout)
	fs.Usage = cmd.printUsage

	if err := fs.Parse(args); err != nil {
		return err
	}
	cmd.dir = fs.Arg(0)

	err := cmd.isShardDir(cmd.dir)
	if cmd.detailed && err != nil {
		return fmt.Errorf("-detailed only supported for shard dirs.")
	}

	totalSeries := hllpp.New()
	tagCardinalities := map[string]*hllpp.HLLPP{}
	measCardinalities := map[string]*hllpp.HLLPP{}
	fieldCardinalities := map[string]*hllpp.HLLPP{}

	dbCardinalities := map[string]*hllpp.HLLPP{}

	start := time.Now()

	tw := tabwriter.NewWriter(cmd.Stdout, 8, 2, 1, ' ', 0)
	fmt.Fprintln(tw, strings.Join([]string{"DB", "RP", "Shard", "File", "Series", "New (Est)", "Load Time"}, "\t"))

	if err := cmd.WalkShardDirs(cmd.dir, func(db, rp, id, path string) error {
		if cmd.pattern != "" && strings.Contains(path, cmd.pattern) {
			return nil
		}

		file, err := os.OpenFile(path, os.O_RDONLY, 0600)
		if err != nil {
			fmt.Fprintf(cmd.Stderr, "error: %s: %v. Skipping.\n", path, err)
			return nil
		}

		loadStart := time.Now()
		reader, err := tsm1.NewTSMReader(file)
		if err != nil {
			fmt.Fprintf(cmd.Stderr, "error: %s: %v. Skipping.\n", file.Name(), err)
			return nil
		}
		loadTime := time.Since(loadStart)

		dbCount := dbCardinalities[db]
		if dbCount == nil {
			dbCount = hllpp.New()
			dbCardinalities[db] = dbCount
		}

		oldCount := dbCount.Count()

		seriesCount := reader.KeyCount()
		for i := 0; i < seriesCount; i++ {
			key, _ := reader.KeyAt(i)
			totalSeries.Add([]byte(key))
			dbCount.Add([]byte(key))

			if cmd.detailed {
				sep := strings.Index(string(key), "#!~#")
				seriesKey, field := key[:sep], key[sep+4:]
				measurement, tags := models.ParseKey(seriesKey)

				measCount := measCardinalities[measurement]
				if measCount == nil {
					measCount = hllpp.New()
					measCardinalities[measurement] = measCount
				}
				measCount.Add([]byte(key))

				fieldCount := fieldCardinalities[measurement]
				if fieldCount == nil {
					fieldCount = hllpp.New()
					fieldCardinalities[measurement] = fieldCount
				}
				fieldCount.Add([]byte(field))

				for _, t := range tags {
					tagCount := tagCardinalities[string(t.Key)]
					if tagCount == nil {
						tagCount = hllpp.New()
						tagCardinalities[string(t.Key)] = tagCount
					}
					tagCount.Add(t.Value)
				}
			}
		}
		reader.Close()

		fmt.Fprintln(tw, strings.Join([]string{
			db, rp, id,
			filepath.Base(file.Name()),
			strconv.FormatInt(int64(seriesCount), 10),
			strconv.FormatInt(int64(dbCount.Count()-oldCount), 10),

			loadTime.String(),
		}, "\t"))
		if cmd.detailed {
			tw.Flush()
		}
		return nil
	}); err != nil {
		return err
	}

	tw.Flush()
	println()
	fmt.Printf("Statistics\n")
	fmt.Printf("  Series:\n")
	for db, counts := range dbCardinalities {
		fmt.Printf("     - %s (est): %d (%d%%)\n", db, counts.Count(), int(float64(counts.Count())/float64(totalSeries.Count())*100))
	}
	fmt.Printf("  Total (est): %d\n", totalSeries.Count())

	if cmd.detailed {
		fmt.Printf("\n  Measurements (est):\n")
		for _, t := range sortKeys(measCardinalities) {
			fmt.Printf("    - %v: %d (%d%%)\n", t, measCardinalities[t].Count(), int((float64(measCardinalities[t].Count())/float64(totalSeries.Count()))*100))
		}

		fmt.Printf("\n  Fields (est):\n")
		for _, t := range sortKeys(fieldCardinalities) {
			fmt.Printf("    - %v: %d\n", t, fieldCardinalities[t].Count())
		}

		fmt.Printf("\n  Tags (est):\n")
		for _, t := range sortKeys(tagCardinalities) {
			fmt.Printf("    - %v: %d\n", t, tagCardinalities[t].Count())
		}
	}

	fmt.Printf("Completed in %s\n", time.Since(start))
	return nil
}

// sortKeys is a quick helper to return the sorted set of a map's keys
func sortKeys(vals map[string]*hllpp.HLLPP) (keys []string) {
	for k := range vals {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	return keys
}

func (cmd *Command) isShardDir(dir string) error {
	name := filepath.Base(dir)
	if id, err := strconv.Atoi(name); err != nil || id < 1 {
		return fmt.Errorf("not a valid shard dir: %v", dir)
	}

	return nil
}

func (cmd *Command) WalkShardDirs(root string, fn func(db, rp, id, path string) error) error {
	type location struct {
		db, rp, id, path string
	}

	var dirs []location
	if err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		if filepath.Ext(info.Name()) == "."+tsm1.TSMFileExtension {
			shardDir := filepath.Dir(path)

			if err := cmd.isShardDir(shardDir); err != nil {
				return err
			}
			absPath, err := filepath.Abs(path)
			if err != nil {
				return err
			}
			parts := strings.Split(absPath, string(filepath.Separator))
			db, rp, id := parts[len(parts)-4], parts[len(parts)-3], parts[len(parts)-2]
			dirs = append(dirs, location{db: db, rp: rp, id: id, path: path})
			return nil
		}
		return nil
	}); err != nil {
		return err
	}

	sort.Slice(dirs, func(i, j int) bool {
		a, _ := strconv.Atoi(dirs[i].id)
		b, _ := strconv.Atoi(dirs[j].id)
		return a < b
	})

	for _, shard := range dirs {
		if err := fn(shard.db, shard.rp, shard.id, shard.path); err != nil {
			return err
		}
	}
	return nil
}

// printUsage prints the usage message to STDERR.
func (cmd *Command) printUsage() {
	usage := `Displays shard level report.

Usage: influx_inspect report [flags]

    -pattern <pattern>
            Include only files matching a pattern.
    -detailed
            Report detailed cardinality estimates.
            Defaults to "false".
`

	fmt.Fprintf(cmd.Stdout, usage)
}
<commit_msg>Add exact counting<commit_after>// Package report reports statistics about TSM files.
package report

import (
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/influxdata/influxdb/models"
	"github.com/influxdata/influxdb/tsdb/engine/tsm1"
	"github.com/retailnext/hllpp"
)

// Command represents the program execution for "influxd report".
type Command struct {
	Stderr io.Writer
	Stdout io.Writer

	dir             string
	pattern         string
	detailed, exact bool
}

// NewCommand returns a new instance of Command.
func NewCommand() *Command {
	return &Command{
		Stderr: os.Stderr,
		Stdout: os.Stdout,
	}
}

// Run executes the command.
func (cmd *Command) Run(args ...string) error {
	fs := flag.NewFlagSet("report", flag.ExitOnError)
	fs.StringVar(&cmd.pattern, "pattern", "", "Include only files matching a pattern")
	fs.BoolVar(&cmd.detailed, "detailed", false, "Report detailed cardinality estimates")
	fs.BoolVar(&cmd.exact, "exact", false, "Report exact counts")

	fs.SetOutput(cmd.Stdout)
	fs.Usage = cmd.printUsage

	if err := fs.Parse(args); err != nil {
		return err
	}

	newCounterFn := newHLLCounter
	estTitle := " (est)"
	if cmd.exact {
		estTitle = ""
		newCounterFn = newExactCounter
	}

	cmd.dir = fs.Arg(0)

	err := cmd.isShardDir(cmd.dir)
	if cmd.detailed && err != nil {
		return fmt.Errorf("-detailed only supported for shard dirs.")
	}

	totalSeries := newCounterFn()
	tagCardinalities := map[string]counter{}
	measCardinalities := map[string]counter{}
	fieldCardinalities := map[string]counter{}

	dbCardinalities := map[string]counter{}

	start := time.Now()

	tw := tabwriter.NewWriter(cmd.Stdout, 8, 2, 1, ' ', 0)
	fmt.Fprintln(tw, strings.Join([]string{"DB", "RP", "Shard", "File", "Series", "New" + estTitle, "Load Time"}, "\t"))

	if err := cmd.WalkShardDirs(cmd.dir, func(db, rp, id, path string) error {
		if cmd.pattern != "" && strings.Contains(path, cmd.pattern) {
			return nil
		}

		file, err := os.OpenFile(path, os.O_RDONLY, 0600)
		if err != nil {
			fmt.Fprintf(cmd.Stderr, "error: %s: %v. Skipping.\n", path, err)
			return nil
		}

		loadStart := time.Now()
		reader, err := tsm1.NewTSMReader(file)
		if err != nil {
			fmt.Fprintf(cmd.Stderr, "error: %s: %v. Skipping.\n", file.Name(), err)
			return nil
		}
		loadTime := time.Since(loadStart)

		dbCount := dbCardinalities[db]
		if dbCount == nil {
			dbCount = newCounterFn()
			dbCardinalities[db] = dbCount
		}

		oldCount := dbCount.Count()

		seriesCount := reader.KeyCount()
		for i := 0; i < seriesCount; i++ {
			key, _ := reader.KeyAt(i)
			totalSeries.Add([]byte(key))
			dbCount.Add([]byte(key))

			if cmd.detailed {
				sep := strings.Index(string(key), "#!~#")
				seriesKey, field := key[:sep], key[sep+4:]
				measurement, tags := models.ParseKey(seriesKey)

				measCount := measCardinalities[measurement]
				if measCount == nil {
					measCount = newCounterFn()
					measCardinalities[measurement] = measCount
				}
				measCount.Add([]byte(key))

				fieldCount := fieldCardinalities[measurement]
				if fieldCount == nil {
					fieldCount = newCounterFn()
					fieldCardinalities[measurement] = fieldCount
				}
				fieldCount.Add([]byte(field))

				for _, t := range tags {
					tagCount := tagCardinalities[string(t.Key)]
					if tagCount == nil {
						tagCount = newCounterFn()
						tagCardinalities[string(t.Key)] = tagCount
					}
					tagCount.Add(t.Value)
				}
			}
		}
		reader.Close()

		fmt.Fprintln(tw, strings.Join([]string{
			db, rp, id,
			filepath.Base(file.Name()),
			strconv.FormatInt(int64(seriesCount), 10),
			strconv.FormatInt(int64(dbCount.Count()-oldCount), 10),

			loadTime.String(),
		}, "\t"))
		if cmd.detailed {
			tw.Flush()
		}
		return nil
	}); err != nil {
		return err
	}

	tw.Flush()
	println()
	fmt.Printf("Statistics\n")
	fmt.Printf("  Series:\n")
	for db, counts := range dbCardinalities {
		fmt.Printf("     - %s%s: %d (%d%%)\n", db, estTitle, counts.Count(), int(float64(counts.Count())/float64(totalSeries.Count())*100))
	}
	fmt.Printf("  Total%s: %d\n", estTitle, totalSeries.Count())

	if cmd.detailed {
		fmt.Printf("\n  Measurements (est):\n")
		for _, t := range sortKeys(measCardinalities) {
			fmt.Printf("    - %v: %d (%d%%)\n", t, measCardinalities[t].Count(), int((float64(measCardinalities[t].Count())/float64(totalSeries.Count()))*100))
		}

		fmt.Printf("\n  Fields (est):\n")
		for _, t := range sortKeys(fieldCardinalities) {
			fmt.Printf("    - %v: %d\n", t, fieldCardinalities[t].Count())
		}

		fmt.Printf("\n  Tags (est):\n")
		for _, t := range sortKeys(tagCardinalities) {
			fmt.Printf("    - %v: %d\n", t, tagCardinalities[t].Count())
		}
	}

	fmt.Printf("Completed in %s\n", time.Since(start))
	return nil
}

// sortKeys is a quick helper to return the sorted set of a map's keys
func sortKeys(vals map[string]counter) (keys []string) {
	for k := range vals {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	return keys
}

func (cmd *Command) isShardDir(dir string) error {
	name := filepath.Base(dir)
	if id, err := strconv.Atoi(name); err != nil || id < 1 {
		return fmt.Errorf("not a valid shard dir: %v", dir)
	}

	return nil
}

func (cmd *Command) WalkShardDirs(root string, fn func(db, rp, id, path string) error) error {
	type location struct {
		db, rp, id, path string
	}

	var dirs []location
	if err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		if filepath.Ext(info.Name()) == "."+tsm1.TSMFileExtension {
			shardDir := filepath.Dir(path)

			if err := cmd.isShardDir(shardDir); err != nil {
				return err
			}
			absPath, err := filepath.Abs(path)
			if err != nil {
				return err
			}
			parts := strings.Split(absPath, string(filepath.Separator))
			db, rp, id := parts[len(parts)-4], parts[len(parts)-3], parts[len(parts)-2]
			dirs = append(dirs, location{db: db, rp: rp, id: id, path: path})
			return nil
		}
		return nil
	}); err != nil {
		return err
	}

	sort.Slice(dirs, func(i, j int) bool {
		a, _ := strconv.Atoi(dirs[i].id)
		b, _ := strconv.Atoi(dirs[j].id)
		return a < b
	})

	for _, shard := range dirs {
		if err := fn(shard.db, shard.rp, shard.id, shard.path); err != nil {
			return err
		}
	}
	return nil
}

// printUsage prints the usage message to STDERR.
func (cmd *Command) printUsage() {
	usage := `Displays shard level report.

Usage: influx_inspect report [flags]

    -pattern <pattern>
            Include only files matching a pattern.
    -exact
            Report exact cardinality counts instead of estimates.  Note: this can use a lot of memory.
            Defaults to "false".
    -detailed
            Report detailed cardinality estimates.
            Defaults to "false".
`

	fmt.Fprintf(cmd.Stdout, usage)
}

// counter abstracts a a method of counting keys.
type counter interface {
	Add(key []byte)
	Count() uint64
}

// newHLLCounter returns an approximate counter using HyperLogLogs for cardinality estimation.
func newHLLCounter() counter {
	return hllpp.New()
}

// exactCounter returns an exact count for keys using counting all distinct items in a set.
type exactCounter struct {
	m map[string]struct{}
}

func (c *exactCounter) Add(key []byte) {
	c.m[string(key)] = struct{}{}
}

func (c *exactCounter) Count() uint64 {
	return uint64(len(c.m))
}

func newExactCounter() counter {
	return &exactCounter{
		m: make(map[string]struct{}),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/hashicorp/consul/command"
	"github.com/hashicorp/consul/command/agent"
	"github.com/mitchellh/cli"
	"os"
	"os/signal"
)

// Commands is the mapping of all the available Serf commands.
var Commands map[string]cli.CommandFactory

func init() {
	ui := &cli.BasicUi{Writer: os.Stdout}

	Commands = map[string]cli.CommandFactory{
		"agent": func() (cli.Command, error) {
			return &agent.Command{
				Revision:          GitCommit,
				Version:           Version,
				VersionPrerelease: VersionPrerelease,
				Ui:                ui,
				ShutdownCh:        make(chan struct{}),
			}, nil
		},

		"event": func() (cli.Command, error) {
			return &command.EventCommand{
				Ui: ui,
			}, nil
		},

		"exec": func() (cli.Command, error) {
			return &command.ExecCommand{
				ShutdownCh: makeShutdownCh(),
				Ui:         ui,
			}, nil
		},

		"force-leave": func() (cli.Command, error) {
			return &command.ForceLeaveCommand{
				Ui: ui,
			}, nil
		},

		"join": func() (cli.Command, error) {
			return &command.JoinCommand{
				Ui: ui,
			}, nil
		},

		"keygen": func() (cli.Command, error) {
			return &command.KeygenCommand{
				Ui: ui,
			}, nil
		},

		"leave": func() (cli.Command, error) {
			return &command.LeaveCommand{
				Ui: ui,
			}, nil
		},

		"members": func() (cli.Command, error) {
			return &command.MembersCommand{
				Ui: ui,
			}, nil
		},

		"monitor": func() (cli.Command, error) {
			return &command.MonitorCommand{
				ShutdownCh: makeShutdownCh(),
				Ui:         ui,
			}, nil
		},

		"info": func() (cli.Command, error) {
			return &command.InfoCommand{
				Ui: ui,
			}, nil
		},

		"reload": func() (cli.Command, error) {
			return &command.ReloadCommand{
				Ui: ui,
			}, nil
		},

		"version": func() (cli.Command, error) {
			return &command.VersionCommand{
				Revision:          GitCommit,
				Version:           GitDescribe,
				VersionPrerelease: VersionPrerelease,
				Ui:                ui,
			}, nil
		},

		"watch": func() (cli.Command, error) {
			return &command.WatchCommand{
				ShutdownCh: makeShutdownCh(),
				Ui:         ui,
			}, nil
		},
	}
}

// makeShutdownCh returns a channel that can be used for shutdown
// notifications for commands. This channel will send a message for every
// interrupt received.
func makeShutdownCh() <-chan struct{} {
	resultCh := make(chan struct{})

	signalCh := make(chan os.Signal, 4)
	signal.Notify(signalCh, os.Interrupt)
	go func() {
		for {
			<-signalCh
			resultCh <- struct{}{}
		}
	}()

	return resultCh
}
<commit_msg>Update to 'consul version' behaviour<commit_after>package main

import (
	"github.com/hashicorp/consul/command"
	"github.com/hashicorp/consul/command/agent"
	"github.com/mitchellh/cli"
	"os"
	"os/signal"
)

// Commands is the mapping of all the available Serf commands.
var Commands map[string]cli.CommandFactory

func init() {
	ui := &cli.BasicUi{Writer: os.Stdout}

	Commands = map[string]cli.CommandFactory{
		"agent": func() (cli.Command, error) {
			return &agent.Command{
				Revision:          GitCommit,
				Version:           Version,
				VersionPrerelease: VersionPrerelease,
				Ui:                ui,
				ShutdownCh:        make(chan struct{}),
			}, nil
		},

		"event": func() (cli.Command, error) {
			return &command.EventCommand{
				Ui: ui,
			}, nil
		},

		"exec": func() (cli.Command, error) {
			return &command.ExecCommand{
				ShutdownCh: makeShutdownCh(),
				Ui:         ui,
			}, nil
		},

		"force-leave": func() (cli.Command, error) {
			return &command.ForceLeaveCommand{
				Ui: ui,
			}, nil
		},

		"join": func() (cli.Command, error) {
			return &command.JoinCommand{
				Ui: ui,
			}, nil
		},

		"keygen": func() (cli.Command, error) {
			return &command.KeygenCommand{
				Ui: ui,
			}, nil
		},

		"leave": func() (cli.Command, error) {
			return &command.LeaveCommand{
				Ui: ui,
			}, nil
		},

		"members": func() (cli.Command, error) {
			return &command.MembersCommand{
				Ui: ui,
			}, nil
		},

		"monitor": func() (cli.Command, error) {
			return &command.MonitorCommand{
				ShutdownCh: makeShutdownCh(),
				Ui:         ui,
			}, nil
		},

		"info": func() (cli.Command, error) {
			return &command.InfoCommand{
				Ui: ui,
			}, nil
		},

		"reload": func() (cli.Command, error) {
			return &command.ReloadCommand{
				Ui: ui,
			}, nil
		},

		"version": func() (cli.Command, error) {
			ver := ""
			rel := ""
			if len(GitDescribe) == 0 {
				ver = Version
				rel = "dev"
			} else {
				ver = GitDescribe
				rel = VersionPrerelease
			}

			return &command.VersionCommand{
				Revision:          GitCommit,
				Version:           ver,
				VersionPrerelease: rel,
				Ui:                ui,
			}, nil
		},

		"watch": func() (cli.Command, error) {
			return &command.WatchCommand{
				ShutdownCh: makeShutdownCh(),
				Ui:         ui,
			}, nil
		},
	}
}

// makeShutdownCh returns a channel that can be used for shutdown
// notifications for commands. This channel will send a message for every
// interrupt received.
func makeShutdownCh() <-chan struct{} {
	resultCh := make(chan struct{})

	signalCh := make(chan os.Signal, 4)
	signal.Notify(signalCh, os.Interrupt)
	go func() {
		for {
			<-signalCh
			resultCh <- struct{}{}
		}
	}()

	return resultCh
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/codegangsta/cli"

	"github.com/yuuki1/grabeni/aws"
	. "github.com/yuuki1/grabeni/log"
)

var Commands = []cli.Command{
	commandStatus,
	commandList,
	commandGrab,
	commandAttach,
	commandDetach,
}

var commandStatus = cli.Command{
	Name:  "status",
	Usage: "Show ENI status",
	Description: `
    Show the information of the ENI identified with <eni-id>.
`,
	Action: doStatus,
}

var commandList = cli.Command{
	Name:  "list",
	Usage: "Show all ENIs",
	Description: `
    List the ENIs owned by your account.
`,
	Action: doList,
}

var commandGrab = cli.Command{
	Name:  "grab",
	Usage: "Detach and attach ENI",
	Description: `
    Detach the ENI identified with <eni-id> whether the eni has already attached or not. And then attach the ENI to EC2 instance identified with --instanceid option.
`,
	Action: doGrab,
	Flags: []cli.Flag{
		cli.IntFlag{Name: "d, deviceindex", Value: 1, Usage: "device index number"},
		cli.StringFlag{Name: "I, instanceid", Usage: "attach-targeted instance id"},
		cli.IntFlag{Name: "t, timeout", Value: 10, Usage: "each attach and detach API request timeout seconds"},
		cli.IntFlag{Name: "i, interval", Value: 2, Usage: "each attach and detach API request polling interval seconds"},
	},
}

var commandAttach = cli.Command{
	Name:  "attach",
	Usage: "Attach ENI",
	Description: `
    Just attach the ENI identified with <eni-id> to EC2 instance identified with --instanceid option.
`,
	Action: doAttach,
	Flags: []cli.Flag{
		cli.IntFlag{Name: "d, deviceindex", Value: 1, Usage: "device index number"},
		cli.StringFlag{Name: "I, instanceid", Usage: "attach-targeted instance id"},
		cli.IntFlag{Name: "t, timeout", Value: 10, Usage: "each attach and detach API request timeout seconds"},
		cli.IntFlag{Name: "i, interval", Value: 2, Usage: "each attach and detach API request polling interval seconds"},
	},
}

var commandDetach = cli.Command{
	Name:  "detach",
	Usage: "Detach ENI",
	Description: `
    Just detach the ENI identified with <eni-id>.
`,
	Action: doDetach,
	Flags: []cli.Flag{
		cli.IntFlag{Name: "t, timeout", Value: 10, Usage: "each attach and detach API request timeout seconds"},
		cli.IntFlag{Name: "i, interval", Value: 2, Usage: "each attach and detach API request polling interval seconds"},
	},
}

type commandDoc struct {
	Parent    string
	Arguments string
}

var commandDocs = map[string]commandDoc{
	"status": {"", "<eni-id>"},
	"list":  {"", ""},
	"grab": {"", "[--deviceindex | -d <device index> (default: 1)] [--instanceid | -i <instance id>] <eni-id>"},
	"attach": {"", "[--deviceindex | -d <device index> (default: 1)] [--instanceid | -i <instance id>] <eni-id>"},
	"detach":  {"", "<eni-id>"},
}

// Makes template conditionals to generate per-command documents.
func mkCommandsTemplate(genTemplate func(commandDoc) string) string {
	template := "{{if false}}"
	for _, command := range append(Commands) {
		template = template + fmt.Sprintf("{{else if (eq .Name %q)}}%s", command.Name, genTemplate(commandDocs[command.Name]))
	}
	return template + "{{end}}"
}

func init() {
	argsTemplate := mkCommandsTemplate(func(doc commandDoc) string { return doc.Arguments })
	parentTemplate := mkCommandsTemplate(func(doc commandDoc) string { return string(strings.TrimLeft(doc.Parent+" ", " ")) })

	cli.CommandHelpTemplate = `NAME:
    {{.Name}} - {{.Usage}}

USAGE:
    grabeni ` + parentTemplate + `{{.Name}} ` + argsTemplate + `
{{if (len .Description)}}
DESCRIPTION: {{.Description}}
{{end}}{{if (len .Flags)}}
OPTIONS:
    {{range .Flags}}{{.}}
    {{end}}
{{end}}`
}

func awsCli(c *cli.Context) *aws.Client {
	client, err := aws.NewClient(
		c.GlobalString("region"),
		c.GlobalString("accesskey"),
		c.GlobalString("secretkey"),
	)
	DieIf(err)
	return client
}

func fetchInstanceIDIfEmpty(c *cli.Context) string {
	if instanceID := c.String("instanceid"); instanceID != "" {
		return instanceID
	}

	instanceID, err := aws.GetInstanceID()
	if err != nil {
		cli.ShowCommandHelp(c, c.Command.Name)
		DieIf(err)
	}

	return instanceID
}

func doStatus(c *cli.Context) {
	if len(c.Args()) < 1 {
		cli.ShowCommandHelp(c, "status")
		os.Exit(1)
	}

	eniID := c.Args()[0]
	if eniID == "" {
		cli.ShowCommandHelp(c, "status")
		os.Exit(1)
	}

	eni, err := awsCli(c).DescribeENIByID(eniID)
	DieIf(err)
	if eni == nil {
		os.Exit(0)
	}

	fmt.Println(StatusHeader())
	fmt.Println(StatusRow(eni))
}

func doList(c *cli.Context) {
	enis, err := awsCli(c).DescribeENIs()
	DieIf(err)
	if enis == nil {
		os.Exit(0)
	}

	fmt.Println(StatusHeader())
	for _, eni := range enis {
		fmt.Println(StatusRow(eni))
	}
}

func doGrab(c *cli.Context) {
	if len(c.Args()) < 1 {
		cli.ShowCommandHelp(c, "grab")
		os.Exit(1)
	}

	eniID := c.Args()[0]
	if eniID == "" {
		cli.ShowCommandHelp(c, "grab")
		os.Exit(1)
	}

	instanceID := fetchInstanceIDIfEmpty(c)
	deviceIndex := c.Int("deviceindex")

	err, ok := awsCli(c).GrabENI(&aws.GrabENIParam{
		InterfaceID: eniID,
		InstanceID: instanceID,
		DeviceIndex: deviceIndex,
	}, &aws.RetryParam{
		TimeoutSec: int64(c.Int("timeout")),
		IntervalSec: int64(c.Int("interval")),
	})
	DieIf(err)
	if !ok {
		Logf("attached", "eni %s already attached to %s", eniID, instanceID)
		os.Exit(0)
	}

	Logf("grabbed", "eni %s attached to instance %s", eniID, instanceID)
}

func doAttach(c *cli.Context) {
	if len(c.Args()) < 1 {
		cli.ShowCommandHelp(c, "attach")
		os.Exit(1)
	}

	eniID := c.Args()[0]
	if eniID == "" {
		cli.ShowCommandHelp(c, "attach")
		os.Exit(1)
	}

	instanceID := fetchInstanceIDIfEmpty(c)
	deviceIndex := c.Int("deviceindex")

	err := awsCli(c).AttachENIWithRetry(&aws.AttachENIParam{
		InterfaceID: eniID,
		InstanceID: instanceID,
		DeviceIndex: deviceIndex,
	}, &aws.RetryParam{
		TimeoutSec: int64(c.Int("timeout")),
		IntervalSec: int64(c.Int("interval")),
	})
	DieIf(err)

	Logf("attached", "eni %s attached to instance %s", eniID, instanceID)
}

func doDetach(c *cli.Context) {
	if len(c.Args()) < 1 {
		cli.ShowCommandHelp(c, "detach")
		os.Exit(1)
	}

	eniID := c.Args()[0]
	if eniID == "" {
		cli.ShowCommandHelp(c, "detach")
		os.Exit(1)
	}

	err := awsCli(c).DetachENIWithRetry(&aws.DetachENIParam{
		InterfaceID: eniID,
	}, &aws.RetryParam{
		TimeoutSec: int64(c.Int("timeout")),
		IntervalSec: int64(c.Int("interval")),
	})
	DieIf(err)

	Logf("detached", "eni %s detached", eniID)
}
<commit_msg>Shorten code<commit_after>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/codegangsta/cli"

	"github.com/yuuki1/grabeni/aws"
	. "github.com/yuuki1/grabeni/log"
)

var Commands = []cli.Command{
	commandStatus,
	commandList,
	commandGrab,
	commandAttach,
	commandDetach,
}

var commandStatus = cli.Command{
	Name:  "status",
	Usage: "Show ENI status",
	Description: `
    Show the information of the ENI identified with <eni-id>.
`,
	Action: doStatus,
}

var commandList = cli.Command{
	Name:  "list",
	Usage: "Show all ENIs",
	Description: `
    List the ENIs owned by your account.
`,
	Action: doList,
}

var commandGrab = cli.Command{
	Name:  "grab",
	Usage: "Detach and attach ENI",
	Description: `
    Detach the ENI identified with <eni-id> whether the eni has already attached or not. And then attach the ENI to EC2 instance identified with --instanceid option.
`,
	Action: doGrab,
	Flags: []cli.Flag{
		cli.IntFlag{Name: "d, deviceindex", Value: 1, Usage: "device index number"},
		cli.StringFlag{Name: "I, instanceid", Usage: "attach-targeted instance id"},
		cli.IntFlag{Name: "t, timeout", Value: 10, Usage: "each attach and detach API request timeout seconds"},
		cli.IntFlag{Name: "i, interval", Value: 2, Usage: "each attach and detach API request polling interval seconds"},
	},
}

var commandAttach = cli.Command{
	Name:  "attach",
	Usage: "Attach ENI",
	Description: `
    Just attach the ENI identified with <eni-id> to EC2 instance identified with --instanceid option.
`,
	Action: doAttach,
	Flags: []cli.Flag{
		cli.IntFlag{Name: "d, deviceindex", Value: 1, Usage: "device index number"},
		cli.StringFlag{Name: "I, instanceid", Usage: "attach-targeted instance id"},
		cli.IntFlag{Name: "t, timeout", Value: 10, Usage: "each attach and detach API request timeout seconds"},
		cli.IntFlag{Name: "i, interval", Value: 2, Usage: "each attach and detach API request polling interval seconds"},
	},
}

var commandDetach = cli.Command{
	Name:  "detach",
	Usage: "Detach ENI",
	Description: `
    Just detach the ENI identified with <eni-id>.
`,
	Action: doDetach,
	Flags: []cli.Flag{
		cli.IntFlag{Name: "t, timeout", Value: 10, Usage: "each attach and detach API request timeout seconds"},
		cli.IntFlag{Name: "i, interval", Value: 2, Usage: "each attach and detach API request polling interval seconds"},
	},
}

type commandDoc struct {
	Parent    string
	Arguments string
}

var commandDocs = map[string]commandDoc{
	"status": {"", "<eni-id>"},
	"list":  {"", ""},
	"grab": {"", "[--deviceindex | -d <device index> (default: 1)] [--instanceid | -i <instance id>] <eni-id>"},
	"attach": {"", "[--deviceindex | -d <device index> (default: 1)] [--instanceid | -i <instance id>] <eni-id>"},
	"detach":  {"", "<eni-id>"},
}

// Makes template conditionals to generate per-command documents.
func mkCommandsTemplate(genTemplate func(commandDoc) string) string {
	template := "{{if false}}"
	for _, command := range append(Commands) {
		template = template + fmt.Sprintf("{{else if (eq .Name %q)}}%s", command.Name, genTemplate(commandDocs[command.Name]))
	}
	return template + "{{end}}"
}

func init() {
	argsTemplate := mkCommandsTemplate(func(doc commandDoc) string { return doc.Arguments })
	parentTemplate := mkCommandsTemplate(func(doc commandDoc) string { return string(strings.TrimLeft(doc.Parent+" ", " ")) })

	cli.CommandHelpTemplate = `NAME:
    {{.Name}} - {{.Usage}}

USAGE:
    grabeni ` + parentTemplate + `{{.Name}} ` + argsTemplate + `
{{if (len .Description)}}
DESCRIPTION: {{.Description}}
{{end}}{{if (len .Flags)}}
OPTIONS:
    {{range .Flags}}{{.}}
    {{end}}
{{end}}`
}

func awsCli(c *cli.Context) *aws.Client {
	client, err := aws.NewClient(
		c.GlobalString("region"),
		c.GlobalString("accesskey"),
		c.GlobalString("secretkey"),
	)
	DieIf(err)
	return client
}

func fetchInstanceIDIfEmpty(c *cli.Context) string {
	if instanceID := c.String("instanceid"); instanceID != "" {
		return instanceID
	}

	instanceID, err := aws.GetInstanceID()
	if err != nil {
		cli.ShowCommandHelp(c, c.Command.Name)
		DieIf(err)
	}

	return instanceID
}

func doStatus(c *cli.Context) {
	if len(c.Args()) < 1 {
		cli.ShowCommandHelp(c, "status")
		os.Exit(1)
	}

	eniID := c.Args()[0]
	if eniID == "" {
		cli.ShowCommandHelp(c, "status")
		os.Exit(1)
	}

	eni, err := awsCli(c).DescribeENIByID(eniID)
	DieIf(err)
	if eni == nil {
		os.Exit(0)
	}

	fmt.Println(StatusHeader())
	fmt.Println(StatusRow(eni))
}

func doList(c *cli.Context) {
	enis, err := awsCli(c).DescribeENIs()
	DieIf(err)
	if enis == nil {
		os.Exit(0)
	}

	fmt.Println(StatusHeader())
	for _, eni := range enis {
		fmt.Println(StatusRow(eni))
	}
}

func doGrab(c *cli.Context) {
	if len(c.Args()) < 1 {
		cli.ShowCommandHelp(c, "grab")
		os.Exit(1)
	}

	eniID := c.Args()[0]
	if eniID == "" {
		cli.ShowCommandHelp(c, "grab")
		os.Exit(1)
	}

	err, ok := awsCli(c).GrabENI(&aws.GrabENIParam{
		InterfaceID: eniID,
		InstanceID: fetchInstanceIDIfEmpty(c),
		DeviceIndex: c.Int("deviceindex"),
	}, &aws.RetryParam{
		TimeoutSec: int64(c.Int("timeout")),
		IntervalSec: int64(c.Int("interval")),
	})
	DieIf(err)
	if !ok {
		Logf("attached", "eni %s already attached to %s", eniID, instanceID)
		os.Exit(0)
	}

	Logf("grabbed", "eni %s attached to instance %s", eniID, instanceID)
}

func doAttach(c *cli.Context) {
	if len(c.Args()) < 1 {
		cli.ShowCommandHelp(c, "attach")
		os.Exit(1)
	}

	eniID := c.Args()[0]
	if eniID == "" {
		cli.ShowCommandHelp(c, "attach")
		os.Exit(1)
	}

	err := awsCli(c).AttachENIWithRetry(&aws.AttachENIParam{
		InterfaceID: eniID,
		InstanceID: fetchInstanceIDIfEmpty(c),
		DeviceIndex: c.Int("deviceindex"),
	}, &aws.RetryParam{
		TimeoutSec: int64(c.Int("timeout")),
		IntervalSec: int64(c.Int("interval")),
	})
	DieIf(err)

	Logf("attached", "eni %s attached to instance %s", eniID, instanceID)
}

func doDetach(c *cli.Context) {
	if len(c.Args()) < 1 {
		cli.ShowCommandHelp(c, "detach")
		os.Exit(1)
	}

	eniID := c.Args()[0]
	if eniID == "" {
		cli.ShowCommandHelp(c, "detach")
		os.Exit(1)
	}

	err := awsCli(c).DetachENIWithRetry(&aws.DetachENIParam{
		InterfaceID: eniID,
	}, &aws.RetryParam{
		TimeoutSec: int64(c.Int("timeout")),
		IntervalSec: int64(c.Int("interval")),
	})
	DieIf(err)

	Logf("detached", "eni %s detached", eniID)
}
<|endoftext|>


<commit_before>package cpu

import (
	"bufio"
	"os"
	"regexp"
	"strconv"
)

type Cpu struct{}

const name = "cpu"

func (self *Cpu) Name() string {
	return name
}

func (self *Cpu) Collect() (result interface{}, err error) {
	result, err = getCpuInfo()
	return
}

var cpuMap = map[string]string{
	"vendor_id":  "vendor_id",
	"model name": "model_name",
	"cpu cores":  "cpu_cores",
	"cpu MHz\t":  "mhz",
	"cache size": "cache_size",
	// "processor":   "processor",
	// "cpu family":  "family",
	// "model":       "model",
	// "stepping":    "stepping",
	// "physical id": "physical_id",
	// "flags":       "flags",
}

// fixme tri This collector saves only the last processor's info

func getCpuInfo() (cpuInfo map[string]string, err error) {
	file, err := os.Open("/proc/cpuinfo")

	if err != nil {
		return
	}

	var lines []string
	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	if scanner.Err() != nil {
		err = scanner.Err()
		return
	}

	cpuInfo = make(map[string]string)
	count := 0

	for _, line := range lines {
		pair := regexp.MustCompile("\t: ").Split(line, 2)

		key, ok := cpuMap[pair[0]]
		if pair[0] == "processor" {
			count += 1
		}
		if ok {
			cpuInfo[key] = pair[1]
		}
	}

	cpuInfo["total"] = strconv.Itoa(count)
	return
}
<commit_msg>Collect all cpus<commit_after>package cpu

import (
	"bufio"
	"os"
	"regexp"
	"strconv"
)

type Cpu struct{}

const name = "cpu"

func (self *Cpu) Name() string {
	return name
}

func (self *Cpu) Collect() (result interface{}, err error) {
	result, err = getCpuInfo()
	return
}

var cpuMap = map[string]string{
	"vendor_id":  "vendor_id",
	"model name": "model_name",
	"cpu cores":  "cpu_cores",
	"cpu MHz\t":  "mhz",
	"cache size": "cache_size",
	"processor":   "processor",
	"cpu family":  "family",
	"model\t":       "model",
	"stepping":    "stepping",
	"physical id": "physical_id",
}

// fixme tri This collector saves only the last processor's info

func getCpuInfo() (cpuInfo map[string]interface{}, err error) {
	file, err := os.Open("/proc/cpuinfo")

	if err != nil {
		return
	}

	var lines []string
	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	if scanner.Err() != nil {
		err = scanner.Err()
		return
	}

	cpuInfo = make(map[string]interface{})
	count := 1
	currentCPU := make(map[string]string)

	for _, line := range lines[1:] {
		pair := regexp.MustCompile("\t: ").Split(line, 2)

		key, ok := cpuMap[pair[0]]
		if pair[0] == "processor" { // next processor's information
			cpuInfo[strconv.Itoa(count)] = currentCPU
			count += 1
		}
		if ok {
			currentCPU[key] = pair[1]
		}
	}
	cpuInfo[strconv.Itoa(count)] = currentCPU

	return
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google Inc. All Rights Reserved.

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 ctc_lib

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"testing"

	"github.com/spf13/cobra"
)

type TestInterface struct {
	Greeting string
	Name     string
}

type TestSubcommandOutput struct {
	Breed string
	Size  string
}

var Greeting string
var Name string

func RunCommand(command *cobra.Command, args []string) (interface{}, error) {
	fmt.Println("Running Hello World Command")
	if Name == "" {
		return (*TestInterface)(nil), errors.New("Please supply Name Argument")
	}
	testOutput := TestInterface{
		Greeting: Greeting,
		Name:     Name,
	}
	return testOutput, nil
}

func TestContainerToolCommandTemplate(t *testing.T) {
	Version = "1.0.1"
	testCommand := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	var OutputBuffer bytes.Buffer
	testCommand.Flags().StringVarP(&Greeting, "greeting", "g", "Hello", "Greeting")
	testCommand.Flags().StringVarP(&Name, "name", "n", "", "Name")
	testCommand.Command.SetOutput(&OutputBuffer)
	testCommand.SetArgs([]string{"version", "--template", "{{.Version}}"})
	testCommand.Execute()
	// check template applies to the output
	if OutputBuffer.String() != "1.0.1\n" {
		t.Errorf("Expected to contain: \n %v\nGot:\n %v\n", "1.0.1", OutputBuffer.String())
	}
}

func TestContainerToolCommandOutput(t *testing.T) {
	testCommand := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	testCommand.Flags().StringVarP(&Greeting, "greeting", "g", "Hello", "Greeting")
	testCommand.Flags().StringVarP(&Name, "name", "n", "", "Name")
	var OutputBuffer bytes.Buffer
	testCommand.Command.SetOutput(&OutputBuffer)
	testCommand.SetArgs([]string{"--name=Sparks"})
	testCommand.Execute()
	var expectedOutput = TestInterface{
		Greeting: "Hello",
		Name:     "Sparks",
	}

	if expectedOutput != testCommand.Output {
		t.Errorf("Expected to contain: \n %v\nGot:\n %v\n", expectedOutput, testCommand.Output)
	}
}

func TestContainerToolCommandSubCommandOutput(t *testing.T) {
	testCommand := &ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	testCommand.SetArgs([]string{"details", "--template", "{{.Breed}}"})
	var OutputBuffer bytes.Buffer
	testCommand.Command.SetOutput(&OutputBuffer)
	testSubCommand := &ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use:   "details",
				Short: "More Info",
			},
		},
		Output: &TestSubcommandOutput{},
		RunO: func(command *cobra.Command, args []string) (interface{}, error) {
			return TestSubcommandOutput{
				Breed: "Chihuhua Mix",
				Size:  "Small",
			}, nil
		},
	}
	testCommand.AddCommand(testSubCommand)
	testCommand.Execute()
	var expectedOutput = TestSubcommandOutput{
		Breed: "Chihuhua Mix",
		Size:  "Small",
	}

	if testSubCommand.Output != expectedOutput {
		t.Errorf("Expected to contain: \n %v\nGot:\n %v\n", expectedOutput, testSubCommand.Output)
	}

	// check template applies to the output
	if OutputBuffer.String() != "Chihuhua Mix\n" {
		t.Errorf("Expected to contain: \n Chihuhua Mix \nGot:\n %v\n", OutputBuffer.String())
	}

}

func TestContainerToolCommandPanic(t *testing.T) {
	testCommand := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	testCommand.Flags().String("foo", "", "")
	testCommand.MarkFlagRequired("foo")
	if os.Getenv("TEST_EXIT_CODE") == "1" {
		testCommand.Execute()
		return
	}
	// Run the go test again with environment variable set to run the command.
	cmd := exec.Command(os.Args[0], "-test.run=TestContainerToolCommandPanic")
	cmd.Env = append(os.Environ(), "TEST_EXIT_CODE=1")
	err := cmd.Run()
	if e, ok := err.(*exec.ExitError); ok && !e.Success() {
		return
	}
	t.Fatalf("process ran with err %v, want non zero exit status", err)
}

func TestContainerToolCommandPanicWithNoExit(t *testing.T) {
	defer SetNoExitOnError(false)
	testCommand := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	SetNoExitOnError(true)

	testCommand.Flags().String("foo", "", "")
	testCommand.MarkFlagRequired("foo")
	err := testCommand.Execute()

	expected := fmt.Sprintf("required flag(s) %q not set", "foo")
	if err.Error() != expected {
		t.Errorf("Expected Error: \n %q \nGot:\n %q\n", expected, err.Error())
	}
}

func TestContainerToolCommandRunDefined(t *testing.T) {
	defer SetNoExitOnError(false)
	runDefined := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use:   "run",
				Short: "Invalid command Description",
				Run: func(command *cobra.Command, args []string) {
					fmt.Println("Run is defined")
				},
			},
		},
		Output: "",
		RunO: func(command *cobra.Command, args []string) (interface{}, error) {
			return nil, nil
		},
	}
	SetNoExitOnError(true)
	err := runDefined.ValidateCommand()
	expectedError := ("Cannot provide both Command.Run and RunO implementation." +
		"\nEither implement Command.Run implementation or RunO implemetation")
	if err.Error() != expectedError {
		t.Errorf("Expected Error: \n %q \nGot:\n %q\n", expectedError, err.Error())
	}
}
<commit_msg>Fix test after exitOnError Refactor<commit_after>/*
Copyright 2018 Google Inc. All Rights Reserved.

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 ctc_lib

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"testing"

	"github.com/spf13/cobra"
)

type TestInterface struct {
	Greeting string
	Name     string
}

type TestSubcommandOutput struct {
	Breed string
	Size  string
}

var Greeting string
var Name string

func RunCommand(command *cobra.Command, args []string) (interface{}, error) {
	fmt.Println("Running Hello World Command")
	if Name == "" {
		return (*TestInterface)(nil), errors.New("Please supply Name Argument")
	}
	testOutput := TestInterface{
		Greeting: Greeting,
		Name:     Name,
	}
	return testOutput, nil
}

func TestContainerToolCommandTemplate(t *testing.T) {
	Version = "1.0.1"
	testCommand := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	var OutputBuffer bytes.Buffer
	testCommand.Flags().StringVarP(&Greeting, "greeting", "g", "Hello", "Greeting")
	testCommand.Flags().StringVarP(&Name, "name", "n", "", "Name")
	testCommand.Command.SetOutput(&OutputBuffer)
	testCommand.SetArgs([]string{"version", "--template", "{{.Version}}"})
	testCommand.Execute()
	// check template applies to the output
	if OutputBuffer.String() != "1.0.1\n" {
		t.Errorf("Expected to contain: \n %v\nGot:\n %v\n", "1.0.1", OutputBuffer.String())
	}
}

func TestContainerToolCommandOutput(t *testing.T) {
	testCommand := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	testCommand.Flags().StringVarP(&Greeting, "greeting", "g", "Hello", "Greeting")
	testCommand.Flags().StringVarP(&Name, "name", "n", "", "Name")
	var OutputBuffer bytes.Buffer
	testCommand.Command.SetOutput(&OutputBuffer)
	testCommand.SetArgs([]string{"--name=Sparks"})
	testCommand.Execute()
	var expectedOutput = TestInterface{
		Greeting: "Hello",
		Name:     "Sparks",
	}

	if expectedOutput != testCommand.Output {
		t.Errorf("Expected to contain: \n %v\nGot:\n %v\n", expectedOutput, testCommand.Output)
	}
}

func TestContainerToolCommandSubCommandOutput(t *testing.T) {
	testCommand := &ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	testCommand.SetArgs([]string{"details", "--template", "{{.Breed}}"})
	var OutputBuffer bytes.Buffer
	testCommand.Command.SetOutput(&OutputBuffer)
	testSubCommand := &ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use:   "details",
				Short: "More Info",
			},
		},
		Output: &TestSubcommandOutput{},
		RunO: func(command *cobra.Command, args []string) (interface{}, error) {
			return TestSubcommandOutput{
				Breed: "Chihuhua Mix",
				Size:  "Small",
			}, nil
		},
	}
	testCommand.AddCommand(testSubCommand)
	testCommand.Execute()
	var expectedOutput = TestSubcommandOutput{
		Breed: "Chihuhua Mix",
		Size:  "Small",
	}

	if testSubCommand.Output != expectedOutput {
		t.Errorf("Expected to contain: \n %v\nGot:\n %v\n", expectedOutput, testSubCommand.Output)
	}

	// check template applies to the output
	if OutputBuffer.String() != "Chihuhua Mix\n" {
		t.Errorf("Expected to contain: \n Chihuhua Mix \nGot:\n %v\n", OutputBuffer.String())
	}

}

func TestContainerToolCommandPanic(t *testing.T) {
	testCommand := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	testCommand.Flags().String("foo", "", "")
	testCommand.MarkFlagRequired("foo")
	if os.Getenv("TEST_EXIT_CODE") == "1" {
		testCommand.Execute()
		return
	}
	// Run the go test again with environment variable set to run the command.
	cmd := exec.Command(os.Args[0], "-test.run=TestContainerToolCommandPanic")
	cmd.Env = append(os.Environ(), "TEST_EXIT_CODE=1")
	err := cmd.Run()
	if e, ok := err.(*exec.ExitError); ok && !e.Success() {
		return
	}
	t.Fatalf("process ran with err %v, want non zero exit status", err)
}

func TestContainerToolCommandPanicWithNoExit(t *testing.T) {
	defer SetExitOnError(true)
	testCommand := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use: "Hello Command",
			},
			Phase: "test",
		},
		Output: &TestInterface{},
		RunO:   RunCommand,
	}
	SetExitOnError(false)

	testCommand.Flags().String("foo", "", "")
	testCommand.MarkFlagRequired("foo")
	err := testCommand.Execute()

	expected := fmt.Sprintf("required flag(s) %q not set", "foo")
	if err.Error() != expected {
		t.Errorf("Expected Error: \n %q \nGot:\n %q\n", expected, err.Error())
	}
}

func TestContainerToolCommandRunDefined(t *testing.T) {
	defer SetExitOnError(true)
	runDefined := ContainerToolCommand{
		ContainerToolCommandBase: &ContainerToolCommandBase{
			Command: &cobra.Command{
				Use:   "run",
				Short: "Invalid command Description",
				Run: func(command *cobra.Command, args []string) {
					fmt.Println("Run is defined")
				},
			},
		},
		Output: "",
		RunO: func(command *cobra.Command, args []string) (interface{}, error) {
			return nil, nil
		},
	}
	SetExitOnError(false)
	err := runDefined.ValidateCommand()
	expectedError := ("Cannot provide both Command.Run and RunO implementation." +
		"\nEither implement Command.Run implementation or RunO implemetation")
	if err.Error() != expectedError {
		t.Errorf("Expected Error: \n %q \nGot:\n %q\n", expectedError, err.Error())
	}
}
<|endoftext|>


<commit_before>package httpbackend

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"time"

	"code.uber.internal/infra/kraken/core"
	"code.uber.internal/infra/kraken/lib/backend/backenderrors"
	"code.uber.internal/infra/kraken/utils/httputil"
)

// Config defines http post/get upload/download urls
// and http connnection parameters. The URLs come with string format
// specifiers and define how to pass sha256 parameters
type Config struct {
	UploadURL   string `yaml:"upload_url"`   // http upload post url
	DownloadURL string `yaml:"download_url"` // http download get url
	Timeout     time.Duration
}

// Client implements downloading/uploading object from/to S3
type Client struct {
	config Config
}

func (c Config) applyDefaults() Config {
	if c.Timeout == 0 {
		c.Timeout = 180 * time.Second
	}
	return c
}

// NewClient creates http client from input parameters
func NewClient(config Config) (*Client, error) {
	return &Client{config: config.applyDefaults()}, nil
}

// Stat always succeeds.
// TODO(codyg): Support stat URL.
func (c *Client) Stat(name string) (*core.BlobInfo, error) {
	return core.NewBlobInfo(0), nil
}

// Download downloads the content from a configured url and writes the data
// to dst.
func (c *Client) Download(name string, dst io.Writer) error {
	b := new(bytes.Buffer)

	// using Fprintf instead of Sprintf to handle formatting errors
	_, err := fmt.Fprintf(b, c.config.DownloadURL, name)
	if err != nil {
		return fmt.Errorf("could not create a url: %s", err)
	}

	resp, err := httputil.Get(b.String(), httputil.SendTimeout(c.config.Timeout))
	if err != nil {
		if httputil.IsNotFound(err) {
			return backenderrors.ErrBlobNotFound
		}
		return fmt.Errorf("could not get a content from http backend: %s", err)
	}

	defer resp.Body.Close()

	_, err = io.Copy(dst, resp.Body)
	if err != nil {
		return fmt.Errorf("could not copy response buffer: %s", err)
	}

	return err
}

// Upload is not supported.
func (c *Client) Upload(name string, src io.Reader) error {
	return errors.New("not supported")
}

// List is not supported.
func (c *Client) List(prefix string) ([]string, error) {
	return nil, errors.New("not supported")
}
<commit_msg>Add retry to httpbackend / clean up code<commit_after>package httpbackend

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"time"

	"code.uber.internal/infra/kraken/core"
	"code.uber.internal/infra/kraken/lib/backend/backenderrors"
	"code.uber.internal/infra/kraken/utils/httputil"
)

// Config defines http post/get upload/download urls
// and http connnection parameters. The URLs come with string format
// specifiers and define how to pass sha256 parameters
type Config struct {
	UploadURL   string `yaml:"upload_url"`   // http upload post url
	DownloadURL string `yaml:"download_url"` // http download get url
	Timeout     time.Duration
}

// Client implements downloading/uploading object from/to S3
type Client struct {
	config Config
}

func (c Config) applyDefaults() Config {
	if c.Timeout == 0 {
		c.Timeout = 180 * time.Second
	}
	return c
}

// NewClient creates http client from input parameters
func NewClient(config Config) (*Client, error) {
	return &Client{config: config.applyDefaults()}, nil
}

// Stat always succeeds.
// TODO(codyg): Support stat URL.
func (c *Client) Stat(name string) (*core.BlobInfo, error) {
	return core.NewBlobInfo(0), nil
}

// Download downloads the content from a configured url and writes the data
// to dst.
func (c *Client) Download(name string, dst io.Writer) error {
	// Use Fprintf instead of Sprintf to handle formatting errors.
	var b bytes.Buffer
	if _, err := fmt.Fprintf(&b, c.config.DownloadURL, name); err != nil {
		return fmt.Errorf("format url: %s", err)
	}
	resp, err := httputil.Get(
		b.String(),
		httputil.SendTimeout(c.config.Timeout),
		httputil.SendRetry())
	if err != nil {
		if httputil.IsNotFound(err) {
			return backenderrors.ErrBlobNotFound
		}
		return err
	}
	defer resp.Body.Close()
	if _, err := io.Copy(dst, resp.Body); err != nil {
		return fmt.Errorf("copy: %s", err)
	}
	return nil
}

// Upload is not supported.
func (c *Client) Upload(name string, src io.Reader) error {
	return errors.New("not supported")
}

// List is not supported.
func (c *Client) List(prefix string) ([]string, error) {
	return nil, errors.New("not supported")
}
<|endoftext|>