<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package executor

import (
	"sort"

	"github.com/juju/errors"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/model"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"golang.org/x/net/context"
)

// DirtyDB stores uncommitted write operations for a transaction.
// It is stored and retrieved by context.Value and context.SetValue method.
type DirtyDB struct {
	// tables is a map whose key is tableID.
	tables map[int64]*DirtyTable
}

// AddRow adds a row to the DirtyDB.
func (udb *DirtyDB) AddRow(tid, handle int64, row []types.Datum) {
	dt := udb.GetDirtyTable(tid)
	for i := range row {
		if row[i].Kind() == types.KindString {
			row[i].SetBytes(row[i].GetBytes())
		}
	}
	dt.addedRows[handle] = row
}

// DeleteRow deletes a row from the DirtyDB.
func (udb *DirtyDB) DeleteRow(tid int64, handle int64) {
	dt := udb.GetDirtyTable(tid)
	delete(dt.addedRows, handle)
	dt.deletedRows[handle] = struct{}{}
}

// TruncateTable truncates a table.
func (udb *DirtyDB) TruncateTable(tid int64) {
	dt := udb.GetDirtyTable(tid)
	dt.addedRows = make(map[int64]Row)
	dt.truncated = true
}

// GetDirtyTable gets the DirtyTable by id from the DirtyDB.
func (udb *DirtyDB) GetDirtyTable(tid int64) *DirtyTable {
	dt, ok := udb.tables[tid]
	if !ok {
		dt = &DirtyTable{
			addedRows:   make(map[int64]Row),
			deletedRows: make(map[int64]struct{}),
		}
		udb.tables[tid] = dt
	}
	return dt
}

// DirtyTable stores uncommitted write operation for a transaction.
type DirtyTable struct {
	// addedRows ...
	// the key is handle.
	addedRows   map[int64]Row
	deletedRows map[int64]struct{}
	truncated   bool
}

// GetDirtyDB returns the DirtyDB bind to the context.
func GetDirtyDB(ctx sessionctx.Context) *DirtyDB {
	var udb *DirtyDB
	x := ctx.GetSessionVars().TxnCtx.DirtyDB
	if x == nil {
		udb = &DirtyDB{tables: make(map[int64]*DirtyTable)}
		ctx.GetSessionVars().TxnCtx.DirtyDB = udb
	} else {
		udb = x.(*DirtyDB)
	}
	return udb
}

// UnionScanExec merges the rows from dirty table and the rows from XAPI request.
type UnionScanExec struct {
	baseExecutor

	dirty *DirtyTable
	// usedIndex is the column offsets of the index which Src executor has used.
	usedIndex  []int
	desc       bool
	conditions []expression.Expression
	columns    []*model.ColumnInfo

	// belowHandleIndex is the handle's position of the below scan plan.
	belowHandleIndex int

	addedRows   []Row
	cursor      int
	sortErr     error
	snapshotRow Row
}

// Next implements Execution Next interface.
func (us *UnionScanExec) Next(ctx context.Context) (Row, error) {
	row, err := us.getOneRow(ctx)
	return row, errors.Trace(err)
}

// NextChunk implements the Executor NextChunk interface.
func (us *UnionScanExec) NextChunk(ctx context.Context, chk *chunk.Chunk) error {
	chk.Reset()
	mutableRow := chunk.MutRowFromTypes(us.retTypes())
	for i, batchSize := 0, us.ctx.GetSessionVars().MaxChunkSize; i < batchSize; i++ {
		row, err := us.getOneRow(ctx)
		if err != nil {
			return errors.Trace(err)
		}
		// no more data.
		if row == nil {
			return nil
		}
		mutableRow.SetDatums(row...)
		chk.AppendRow(mutableRow.ToRow())
	}
	return nil
}

// getOneRow gets one result row from dirty table or child.
func (us *UnionScanExec) getOneRow(ctx context.Context) (Row, error) {
	for {
		snapshotRow, err := us.getSnapshotRow(ctx)
		if err != nil {
			return nil, errors.Trace(err)
		}
		addedRow := us.getAddedRow()
		var row Row
		var isSnapshotRow bool
		if addedRow == nil {
			row = snapshotRow
			isSnapshotRow = true
		} else if snapshotRow == nil {
			row = addedRow
		} else {
			isSnapshotRow, err = us.shouldPickFirstRow(snapshotRow, addedRow)
			if err != nil {
				return nil, errors.Trace(err)
			}
			if isSnapshotRow {
				row = snapshotRow
			} else {
				row = addedRow
			}
		}
		if row == nil {
			return nil, nil
		}

		if isSnapshotRow {
			us.snapshotRow = nil
		} else {
			us.cursor++
		}
		return row, nil
	}
}

func (us *UnionScanExec) getSnapshotRow(ctx context.Context) (Row, error) {
	if us.dirty.truncated {
		return nil, nil
	}
	var err error
	if us.snapshotRow == nil {
		for {
			us.snapshotRow, err = us.children[0].Next(ctx)
			if err != nil {
				return nil, errors.Trace(err)
			}
			if us.snapshotRow == nil {
				break
			}
			snapshotHandle := us.snapshotRow[us.belowHandleIndex].GetInt64()
			if _, ok := us.dirty.deletedRows[snapshotHandle]; ok {
				continue
			}
			if _, ok := us.dirty.addedRows[snapshotHandle]; ok {
				// If src handle appears in added rows, it means there is conflict and the transaction will fail to
				// commit, but for simplicity, we don't handle it here.
				continue
			}
			break
		}
	}
	return us.snapshotRow, nil
}

func (us *UnionScanExec) getAddedRow() Row {
	var addedRow Row
	if us.cursor < len(us.addedRows) {
		addedRow = us.addedRows[us.cursor]
	}
	return addedRow
}

// shouldPickFirstRow picks the suitable row in order.
// The value returned is used to determine whether to pick the first input row.
func (us *UnionScanExec) shouldPickFirstRow(a, b Row) (bool, error) {
	var isFirstRow bool
	addedCmpSrc, err := us.compare(a, b)
	if err != nil {
		return isFirstRow, errors.Trace(err)
	}
	// Compare result will never be 0.
	if us.desc {
		if addedCmpSrc > 0 {
			isFirstRow = true
		}
	} else {
		if addedCmpSrc < 0 {
			isFirstRow = true
		}
	}
	return isFirstRow, nil
}

func (us *UnionScanExec) compare(a, b Row) (int, error) {
	sc := us.ctx.GetSessionVars().StmtCtx
	for _, colOff := range us.usedIndex {
		aColumn := a[colOff]
		bColumn := b[colOff]
		cmp, err := aColumn.CompareDatum(sc, &bColumn)
		if err != nil {
			return 0, errors.Trace(err)
		}
		if cmp != 0 {
			return cmp, nil
		}
	}
	aHandle := a[us.belowHandleIndex].GetInt64()
	bHandle := b[us.belowHandleIndex].GetInt64()
	var cmp int
	if aHandle == bHandle {
		cmp = 0
	} else if aHandle > bHandle {
		cmp = 1
	} else {
		cmp = -1
	}
	return cmp, nil
}

func (us *UnionScanExec) buildAndSortAddedRows() error {
	us.addedRows = make([]Row, 0, len(us.dirty.addedRows))
	for h, data := range us.dirty.addedRows {
		newData := make(types.DatumRow, 0, us.schema.Len())
		for _, col := range us.columns {
			if col.ID == model.ExtraHandleID {
				newData = append(newData, types.NewIntDatum(h))
			} else {
				newData = append(newData, data[col.Offset])
			}
		}
		matched, err := expression.EvalBool(us.ctx, us.conditions, newData)
		if err != nil {
			return errors.Trace(err)
		}
		if !matched {
			continue
		}
		us.addedRows = append(us.addedRows, newData)
	}
	if us.desc {
		sort.Sort(sort.Reverse(us))
	} else {
		sort.Sort(us)
	}
	if us.sortErr != nil {
		return errors.Trace(us.sortErr)
	}
	return nil
}

// Len implements sort.Interface interface.
func (us *UnionScanExec) Len() int {
	return len(us.addedRows)
}

// Less implements sort.Interface interface.
func (us *UnionScanExec) Less(i, j int) bool {
	cmp, err := us.compare(us.addedRows[i], us.addedRows[j])
	if err != nil {
		us.sortErr = errors.Trace(err)
		return true
	}
	return cmp < 0
}

// Swap implements sort.Interface interface.
func (us *UnionScanExec) Swap(i, j int) {
	us.addedRows[i], us.addedRows[j] = us.addedRows[j], us.addedRows[i]
}
<commit_msg>executor: remove Next function for UnionScanExec (#5996)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package executor

import (
	"sort"

	"github.com/juju/errors"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/model"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"golang.org/x/net/context"
)

// DirtyDB stores uncommitted write operations for a transaction.
// It is stored and retrieved by context.Value and context.SetValue method.
type DirtyDB struct {
	// tables is a map whose key is tableID.
	tables map[int64]*DirtyTable
}

// AddRow adds a row to the DirtyDB.
func (udb *DirtyDB) AddRow(tid, handle int64, row []types.Datum) {
	dt := udb.GetDirtyTable(tid)
	for i := range row {
		if row[i].Kind() == types.KindString {
			row[i].SetBytes(row[i].GetBytes())
		}
	}
	dt.addedRows[handle] = row
}

// DeleteRow deletes a row from the DirtyDB.
func (udb *DirtyDB) DeleteRow(tid int64, handle int64) {
	dt := udb.GetDirtyTable(tid)
	delete(dt.addedRows, handle)
	dt.deletedRows[handle] = struct{}{}
}

// TruncateTable truncates a table.
func (udb *DirtyDB) TruncateTable(tid int64) {
	dt := udb.GetDirtyTable(tid)
	dt.addedRows = make(map[int64]Row)
	dt.truncated = true
}

// GetDirtyTable gets the DirtyTable by id from the DirtyDB.
func (udb *DirtyDB) GetDirtyTable(tid int64) *DirtyTable {
	dt, ok := udb.tables[tid]
	if !ok {
		dt = &DirtyTable{
			addedRows:   make(map[int64]Row),
			deletedRows: make(map[int64]struct{}),
		}
		udb.tables[tid] = dt
	}
	return dt
}

// DirtyTable stores uncommitted write operation for a transaction.
type DirtyTable struct {
	// addedRows ...
	// the key is handle.
	addedRows   map[int64]Row
	deletedRows map[int64]struct{}
	truncated   bool
}

// GetDirtyDB returns the DirtyDB bind to the context.
func GetDirtyDB(ctx sessionctx.Context) *DirtyDB {
	var udb *DirtyDB
	x := ctx.GetSessionVars().TxnCtx.DirtyDB
	if x == nil {
		udb = &DirtyDB{tables: make(map[int64]*DirtyTable)}
		ctx.GetSessionVars().TxnCtx.DirtyDB = udb
	} else {
		udb = x.(*DirtyDB)
	}
	return udb
}

// UnionScanExec merges the rows from dirty table and the rows from XAPI request.
type UnionScanExec struct {
	baseExecutor

	dirty *DirtyTable
	// usedIndex is the column offsets of the index which Src executor has used.
	usedIndex  []int
	desc       bool
	conditions []expression.Expression
	columns    []*model.ColumnInfo

	// belowHandleIndex is the handle's position of the below scan plan.
	belowHandleIndex int

	addedRows   []Row
	cursor      int
	sortErr     error
	snapshotRow Row
}

// NextChunk implements the Executor NextChunk interface.
func (us *UnionScanExec) NextChunk(ctx context.Context, chk *chunk.Chunk) error {
	chk.Reset()
	mutableRow := chunk.MutRowFromTypes(us.retTypes())
	for i, batchSize := 0, us.ctx.GetSessionVars().MaxChunkSize; i < batchSize; i++ {
		row, err := us.getOneRow(ctx)
		if err != nil {
			return errors.Trace(err)
		}
		// no more data.
		if row == nil {
			return nil
		}
		mutableRow.SetDatums(row...)
		chk.AppendRow(mutableRow.ToRow())
	}
	return nil
}

// getOneRow gets one result row from dirty table or child.
func (us *UnionScanExec) getOneRow(ctx context.Context) (Row, error) {
	for {
		snapshotRow, err := us.getSnapshotRow(ctx)
		if err != nil {
			return nil, errors.Trace(err)
		}
		addedRow := us.getAddedRow()
		var row Row
		var isSnapshotRow bool
		if addedRow == nil {
			row = snapshotRow
			isSnapshotRow = true
		} else if snapshotRow == nil {
			row = addedRow
		} else {
			isSnapshotRow, err = us.shouldPickFirstRow(snapshotRow, addedRow)
			if err != nil {
				return nil, errors.Trace(err)
			}
			if isSnapshotRow {
				row = snapshotRow
			} else {
				row = addedRow
			}
		}
		if row == nil {
			return nil, nil
		}

		if isSnapshotRow {
			us.snapshotRow = nil
		} else {
			us.cursor++
		}
		return row, nil
	}
}

func (us *UnionScanExec) getSnapshotRow(ctx context.Context) (Row, error) {
	if us.dirty.truncated {
		return nil, nil
	}
	var err error
	if us.snapshotRow == nil {
		for {
			us.snapshotRow, err = us.children[0].Next(ctx)
			if err != nil {
				return nil, errors.Trace(err)
			}
			if us.snapshotRow == nil {
				break
			}
			snapshotHandle := us.snapshotRow[us.belowHandleIndex].GetInt64()
			if _, ok := us.dirty.deletedRows[snapshotHandle]; ok {
				continue
			}
			if _, ok := us.dirty.addedRows[snapshotHandle]; ok {
				// If src handle appears in added rows, it means there is conflict and the transaction will fail to
				// commit, but for simplicity, we don't handle it here.
				continue
			}
			break
		}
	}
	return us.snapshotRow, nil
}

func (us *UnionScanExec) getAddedRow() Row {
	var addedRow Row
	if us.cursor < len(us.addedRows) {
		addedRow = us.addedRows[us.cursor]
	}
	return addedRow
}

// shouldPickFirstRow picks the suitable row in order.
// The value returned is used to determine whether to pick the first input row.
func (us *UnionScanExec) shouldPickFirstRow(a, b Row) (bool, error) {
	var isFirstRow bool
	addedCmpSrc, err := us.compare(a, b)
	if err != nil {
		return isFirstRow, errors.Trace(err)
	}
	// Compare result will never be 0.
	if us.desc {
		if addedCmpSrc > 0 {
			isFirstRow = true
		}
	} else {
		if addedCmpSrc < 0 {
			isFirstRow = true
		}
	}
	return isFirstRow, nil
}

func (us *UnionScanExec) compare(a, b Row) (int, error) {
	sc := us.ctx.GetSessionVars().StmtCtx
	for _, colOff := range us.usedIndex {
		aColumn := a[colOff]
		bColumn := b[colOff]
		cmp, err := aColumn.CompareDatum(sc, &bColumn)
		if err != nil {
			return 0, errors.Trace(err)
		}
		if cmp != 0 {
			return cmp, nil
		}
	}
	aHandle := a[us.belowHandleIndex].GetInt64()
	bHandle := b[us.belowHandleIndex].GetInt64()
	var cmp int
	if aHandle == bHandle {
		cmp = 0
	} else if aHandle > bHandle {
		cmp = 1
	} else {
		cmp = -1
	}
	return cmp, nil
}

func (us *UnionScanExec) buildAndSortAddedRows() error {
	us.addedRows = make([]Row, 0, len(us.dirty.addedRows))
	for h, data := range us.dirty.addedRows {
		newData := make(types.DatumRow, 0, us.schema.Len())
		for _, col := range us.columns {
			if col.ID == model.ExtraHandleID {
				newData = append(newData, types.NewIntDatum(h))
			} else {
				newData = append(newData, data[col.Offset])
			}
		}
		matched, err := expression.EvalBool(us.ctx, us.conditions, newData)
		if err != nil {
			return errors.Trace(err)
		}
		if !matched {
			continue
		}
		us.addedRows = append(us.addedRows, newData)
	}
	if us.desc {
		sort.Sort(sort.Reverse(us))
	} else {
		sort.Sort(us)
	}
	if us.sortErr != nil {
		return errors.Trace(us.sortErr)
	}
	return nil
}

// Len implements sort.Interface interface.
func (us *UnionScanExec) Len() int {
	return len(us.addedRows)
}

// Less implements sort.Interface interface.
func (us *UnionScanExec) Less(i, j int) bool {
	cmp, err := us.compare(us.addedRows[i], us.addedRows[j])
	if err != nil {
		us.sortErr = errors.Trace(err)
		return true
	}
	return cmp < 0
}

// Swap implements sort.Interface interface.
func (us *UnionScanExec) Swap(i, j int) {
	us.addedRows[i], us.addedRows[j] = us.addedRows[j], us.addedRows[i]
}
<|endoftext|>


<commit_before>package command

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"strings"

	consul "github.com/hashicorp/consul/api"
)

const (
	localConsulURL  = "127.0.0.1:8500"
	defaultFileName = "data.json"
)

func getKVClient(ipaddress string) (*consul.KV, error) {
	config := consul.DefaultConfig()
	config.Address = ipaddress

	client, err := consul.NewClient(config)
	if err != nil {
		return nil, err
	}
	return client.KV(), nil
}

func readJSON(filename string, prefix string) (consul.KVPairs, error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	inpairs := map[string]string{}
	err = json.Unmarshal(data, &inpairs)
	if err != nil {
		return nil, err
	}

	var kvs consul.KVPairs
	for key, value := range inpairs {
		if prefix == "" || strings.HasPrefix(key, prefix) {
			kvs = append(kvs, &consul.KVPair{Key: key, Value: []byte(value)})
		}
	}
	return kvs, nil
}

func saveKVsToFile(fullKVs consul.KVPairs, fileName string) error {
	kvs := map[string]string{}
	for _, element := range fullKVs {
		kvs[element.Key] = string(element.Value)
	}

	data, err := json.MarshalIndent(kvs, "", "  ")
	if err != nil {
		return err
	}

	file, err := os.Create(fileName)
	if err != nil {
		return err
	}

	if _, err := file.Write([]byte(data)[:]); err != nil {
		return err
	}

	return nil
}
<commit_msg>trim http / https off consul addresses<commit_after>package command

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"strings"

	consul "github.com/hashicorp/consul/api"
)

const (
	localConsulURL  = "127.0.0.1:8500"
	defaultFileName = "data.json"
	http            = "http://"
	https           = "https://"
)

func getKVClient(ipaddress string) (*consul.KV, error) {
	ipaddress = strings.TrimPrefix(ipaddress, http)
	ipaddress = strings.TrimPrefix(ipaddress, https)

	config := consul.DefaultConfig()
	config.Address = ipaddress

	client, err := consul.NewClient(config)
	if err != nil {
		return nil, err
	}
	return client.KV(), nil
}

func readJSON(filename string, prefix string) (consul.KVPairs, error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	inpairs := map[string]string{}
	err = json.Unmarshal(data, &inpairs)
	if err != nil {
		return nil, err
	}

	var kvs consul.KVPairs
	for key, value := range inpairs {
		if prefix == "" || strings.HasPrefix(key, prefix) {
			kvs = append(kvs, &consul.KVPair{Key: key, Value: []byte(value)})
		}
	}
	return kvs, nil
}

func saveKVsToFile(fullKVs consul.KVPairs, fileName string) error {
	kvs := map[string]string{}
	for _, element := range fullKVs {
		kvs[element.Key] = string(element.Value)
	}

	data, err := json.MarshalIndent(kvs, "", "  ")
	if err != nil {
		return err
	}

	file, err := os.Create(fileName)
	if err != nil {
		return err
	}

	if _, err := file.Write([]byte(data)[:]); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package imap

import (
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

var registeredFetchParams []fetchParamDefinition
var peekRE *regexp.Regexp

// ErrUnrecognisedParameter indicates that the parameter requested in a FETCH
// command is unrecognised or not implemented in this IMAP server
var ErrUnrecognisedParameter = errors.New("Unrecognised Parameter")

type fetchParamDefinition struct {
	re      *regexp.Regexp
	handler func([]string, *Conn, Message, bool) string
}

func cmdFetch(args commandArgs, c *Conn) {
	start, _ := strconv.Atoi(args.Arg(1))

	searchByUID := args.Arg(0) == "UID "

	// Fetch the messages
	var msg Message
	if searchByUID {
		fmt.Printf("Searching by UID\n")
		msg = c.selectedMailbox.MessageByUid(uint32(start))
	} else {
		msg = c.selectedMailbox.MessageBySequenceNumber(uint32(start))
	}

	fetchParamString := args.Arg(3)
	if searchByUID && !strings.Contains(fetchParamString, "UID") {
		fetchParamString += " UID"
	}

	fetchParams, err := fetch(fetchParamString, c, msg)
	if err != nil {
		if err == ErrUnrecognisedParameter {
			c.writeResponse(args.Id(), "BAD Unrecognised Parameter")
			return
		} else {
			c.writeResponse(args.Id(), "BAD")
			return
		}
	}

	fullReply := fmt.Sprintf("%d FETCH (%s)",
		msg.SequenceNumber(),
		fetchParams)

	c.writeResponse("", fullReply)
	if searchByUID {
		c.writeResponse(args.Id(), "OK UID FETCH Completed")
	} else {
		c.writeResponse(args.Id(), "OK FETCH Completed")
	}
}

// Fetch requested params from a given message
// eg fetch("UID BODY[TEXT] RFC822.SIZE", c, message)
func fetch(params string, c *Conn, m Message) (string, error) {
	// Prepare the list of responses
	responseParams := make([]string, 0)

	paramList := splitParams(params)
	for _, param := range paramList {
		paramResponse, err := fetchParam(param, c, m)
		if err != nil {
			return "", err
		}
		responseParams = append(responseParams, paramResponse)
	}
	return strings.Join(responseParams, " "), nil
}

// Match a single fetch parameter and return the data
func fetchParam(param string, c *Conn, m Message) (string, error) {
	peek := false
	if peekRE.MatchString(param) {
		peek = true
	}
	// Search through the parameter list until a parameter handler is found
	for _, element := range registeredFetchParams {
		if element.re.MatchString(param) {
			return element.handler(element.re.FindStringSubmatch(param), c, m, peek), nil
		}
	}
	return "", ErrUnrecognisedParameter
}

// Register all supported fetch parameters
func init() {
	peekRE = regexp.MustCompile("\\.PEEK")
	registeredFetchParams = make([]fetchParamDefinition, 0)
	registerFetchParam("UID", fetchUID)
	registerFetchParam("FLAGS", fetchFlags)
	registerFetchParam("RFC822\\.SIZE", fetchRfcSize)
	registerFetchParam("INTERNALDATE", fetchInternalDate)
	registerFetchParam("BODY(?:\\.PEEK)?\\[HEADER\\]", fetchHeaders)
	registerFetchParam("BODY(?:\\.PEEK)?"+
		"\\[HEADER\\.FIELDS \\(([A-z\\s-]+)\\)\\]", fetchHeaderSpecificFields)
	registerFetchParam("BODY(?:\\.PEEK)?\\[TEXT\\]", fetchBody)
	registerFetchParam("BODY(?:\\.PEEK)?\\[\\]", fetchFullText)
}

func registerFetchParam(regex string, handler func([]string, *Conn, Message, bool) string) {
	newParam := fetchParamDefinition{
		re:      regexp.MustCompile(regex),
		handler: handler,
	}
	registeredFetchParams = append(registeredFetchParams, newParam)
}

// Fetch the UID of the mail message
func fetchUID(args []string, c *Conn, m Message, peekOnly bool) string {
	return fmt.Sprintf("UID %d", m.Uid())
}

func fetchFlags(args []string, c *Conn, m Message, peekOnly bool) string {
	flags := append(messageFlags(m), m.Keywords()...)
	flagList := strings.Join(flags, " ")
	return fmt.Sprintf("FLAGS (%s)", flagList)
}

func fetchRfcSize(args []string, c *Conn, m Message, peekOnly bool) string {
	return fmt.Sprintf("RFC822.SIZE %d", m.Size())
}

func fetchInternalDate(args []string, c *Conn, m Message, peekOnly bool) string {
	dateStr := m.InternalDate().Format(internalDate)
	return fmt.Sprintf("INTERNALDATE \"%s\"", dateStr)
}

func fetchHeaders(args []string, c *Conn, m Message, peekOnly bool) string {
	hdr := fmt.Sprintf("\r\n%s\r\n\r\n", m.Header())
	hdrLen := len(hdr)

	peekStr := ""
	if peekOnly {
		peekStr = ".PEEK"
	}

	return fmt.Sprintf("BODY%s[HEADER] {%d}%s", peekStr, hdrLen, hdr)
}

func fetchHeaderSpecificFields(args []string, c *Conn, m Message, peekOnly bool) string {
	if !peekOnly {
		fmt.Printf("TODO: Peek not requested, mark all as non-recent\n")
	}
	fields := strings.Split(args[1], " ")
	hdrs := m.Header()
	requestedHeaders := make(MIMEHeader)
	replyFieldList := make([]string, len(fields))
	for i, key := range fields {
		replyFieldList[i] = "\"" + key + "\""
		// If the key exists in the headers, copy it over
		if k, v, ok := hdrs.FindKey(key); ok {
			requestedHeaders[k] = v
		}
	}
	hdr := fmt.Sprintf("\r\n%s\r\n\r\n", requestedHeaders)
	hdrLen := len(hdr)

	return fmt.Sprintf("BODY[HEADER.FIELDS (%s)] {%d}%s",
		strings.Join(replyFieldList, " "),
		hdrLen,
		hdr)

}

func fetchBody(args []string, c *Conn, m Message, peekOnly bool) string {
	body := fmt.Sprintf("\r\n%s\r\n", m.Body())
	bodyLen := len(body)

	return fmt.Sprintf("BODY[TEXT] {%d}%s",
		bodyLen, body)
}

func fetchFullText(args []string, c *Conn, m Message, peekOnly bool) string {
	mail := fmt.Sprintf("\r\n%s\r\n\r\n%s\r\n", m.Header(), m.Body())
	mailLen := len(mail)

	return fmt.Sprintf("BODY[] {%d}%s",
		mailLen, mail)
}
<commit_msg>Outdent else statement. #8<commit_after>package imap

import (
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

var registeredFetchParams []fetchParamDefinition
var peekRE *regexp.Regexp

// ErrUnrecognisedParameter indicates that the parameter requested in a FETCH
// command is unrecognised or not implemented in this IMAP server
var ErrUnrecognisedParameter = errors.New("Unrecognised Parameter")

type fetchParamDefinition struct {
	re      *regexp.Regexp
	handler func([]string, *Conn, Message, bool) string
}

func cmdFetch(args commandArgs, c *Conn) {
	start, _ := strconv.Atoi(args.Arg(1))

	searchByUID := args.Arg(0) == "UID "

	// Fetch the messages
	var msg Message
	if searchByUID {
		fmt.Printf("Searching by UID\n")
		msg = c.selectedMailbox.MessageByUid(uint32(start))
	} else {
		msg = c.selectedMailbox.MessageBySequenceNumber(uint32(start))
	}

	fetchParamString := args.Arg(3)
	if searchByUID && !strings.Contains(fetchParamString, "UID") {
		fetchParamString += " UID"
	}

	fetchParams, err := fetch(fetchParamString, c, msg)
	if err != nil {
		if err == ErrUnrecognisedParameter {
			c.writeResponse(args.Id(), "BAD Unrecognised Parameter")
			return
		}

		c.writeResponse(args.Id(), "BAD")
		return
	}

	fullReply := fmt.Sprintf("%d FETCH (%s)",
		msg.SequenceNumber(),
		fetchParams)

	c.writeResponse("", fullReply)
	if searchByUID {
		c.writeResponse(args.Id(), "OK UID FETCH Completed")
	} else {
		c.writeResponse(args.Id(), "OK FETCH Completed")
	}
}

// Fetch requested params from a given message
// eg fetch("UID BODY[TEXT] RFC822.SIZE", c, message)
func fetch(params string, c *Conn, m Message) (string, error) {
	// Prepare the list of responses
	responseParams := make([]string, 0)

	paramList := splitParams(params)
	for _, param := range paramList {
		paramResponse, err := fetchParam(param, c, m)
		if err != nil {
			return "", err
		}
		responseParams = append(responseParams, paramResponse)
	}
	return strings.Join(responseParams, " "), nil
}

// Match a single fetch parameter and return the data
func fetchParam(param string, c *Conn, m Message) (string, error) {
	peek := false
	if peekRE.MatchString(param) {
		peek = true
	}
	// Search through the parameter list until a parameter handler is found
	for _, element := range registeredFetchParams {
		if element.re.MatchString(param) {
			return element.handler(element.re.FindStringSubmatch(param), c, m, peek), nil
		}
	}
	return "", ErrUnrecognisedParameter
}

// Register all supported fetch parameters
func init() {
	peekRE = regexp.MustCompile("\\.PEEK")
	registeredFetchParams = make([]fetchParamDefinition, 0)
	registerFetchParam("UID", fetchUID)
	registerFetchParam("FLAGS", fetchFlags)
	registerFetchParam("RFC822\\.SIZE", fetchRfcSize)
	registerFetchParam("INTERNALDATE", fetchInternalDate)
	registerFetchParam("BODY(?:\\.PEEK)?\\[HEADER\\]", fetchHeaders)
	registerFetchParam("BODY(?:\\.PEEK)?"+
		"\\[HEADER\\.FIELDS \\(([A-z\\s-]+)\\)\\]", fetchHeaderSpecificFields)
	registerFetchParam("BODY(?:\\.PEEK)?\\[TEXT\\]", fetchBody)
	registerFetchParam("BODY(?:\\.PEEK)?\\[\\]", fetchFullText)
}

func registerFetchParam(regex string, handler func([]string, *Conn, Message, bool) string) {
	newParam := fetchParamDefinition{
		re:      regexp.MustCompile(regex),
		handler: handler,
	}
	registeredFetchParams = append(registeredFetchParams, newParam)
}

// Fetch the UID of the mail message
func fetchUID(args []string, c *Conn, m Message, peekOnly bool) string {
	return fmt.Sprintf("UID %d", m.Uid())
}

func fetchFlags(args []string, c *Conn, m Message, peekOnly bool) string {
	flags := append(messageFlags(m), m.Keywords()...)
	flagList := strings.Join(flags, " ")
	return fmt.Sprintf("FLAGS (%s)", flagList)
}

func fetchRfcSize(args []string, c *Conn, m Message, peekOnly bool) string {
	return fmt.Sprintf("RFC822.SIZE %d", m.Size())
}

func fetchInternalDate(args []string, c *Conn, m Message, peekOnly bool) string {
	dateStr := m.InternalDate().Format(internalDate)
	return fmt.Sprintf("INTERNALDATE \"%s\"", dateStr)
}

func fetchHeaders(args []string, c *Conn, m Message, peekOnly bool) string {
	hdr := fmt.Sprintf("\r\n%s\r\n\r\n", m.Header())
	hdrLen := len(hdr)

	peekStr := ""
	if peekOnly {
		peekStr = ".PEEK"
	}

	return fmt.Sprintf("BODY%s[HEADER] {%d}%s", peekStr, hdrLen, hdr)
}

func fetchHeaderSpecificFields(args []string, c *Conn, m Message, peekOnly bool) string {
	if !peekOnly {
		fmt.Printf("TODO: Peek not requested, mark all as non-recent\n")
	}
	fields := strings.Split(args[1], " ")
	hdrs := m.Header()
	requestedHeaders := make(MIMEHeader)
	replyFieldList := make([]string, len(fields))
	for i, key := range fields {
		replyFieldList[i] = "\"" + key + "\""
		// If the key exists in the headers, copy it over
		if k, v, ok := hdrs.FindKey(key); ok {
			requestedHeaders[k] = v
		}
	}
	hdr := fmt.Sprintf("\r\n%s\r\n\r\n", requestedHeaders)
	hdrLen := len(hdr)

	return fmt.Sprintf("BODY[HEADER.FIELDS (%s)] {%d}%s",
		strings.Join(replyFieldList, " "),
		hdrLen,
		hdr)

}

func fetchBody(args []string, c *Conn, m Message, peekOnly bool) string {
	body := fmt.Sprintf("\r\n%s\r\n", m.Body())
	bodyLen := len(body)

	return fmt.Sprintf("BODY[TEXT] {%d}%s",
		bodyLen, body)
}

func fetchFullText(args []string, c *Conn, m Message, peekOnly bool) string {
	mail := fmt.Sprintf("\r\n%s\r\n\r\n%s\r\n", m.Header(), m.Body())
	mailLen := len(mail)

	return fmt.Sprintf("BODY[] {%d}%s",
		mailLen, mail)
}
<|endoftext|>


<commit_before>package commands

// Init sits here

import (
	"os"
	"path/filepath"

	"github.com/spf13/cobra"
)

func init() {

	// Define Deploy Flags
	deployCmd.Flags().StringArrayVarP(&run.tplSources, "template", "t", []string{}, "path to template file(s) Or stack::url")
	deployCmd.Flags().BoolVarP(&run.rollback, "disable-rollback", "", false, "Set Stack to rollback on deployment failures")
	deployCmd.Flags().BoolVarP(&run.all, "all", "A", false, "deploy all stacks with defined Sources in config")

	// Define Git Deploy Flags
	gitDeployCmd.Flags().BoolVarP(&run.rollback, "disable-rollback", "", false, "Set Stack to rollback on deployment failures")
	gitDeployCmd.Flags().BoolVarP(&run.all, "all", "A", false, "deploy all stacks with defined Sources in config")
	gitDeployCmd.Flags().StringVarP(&run.gituser, "user", "u", "", "git username")
	gitDeployCmd.Flags().StringVarP(&run.gitpass, "password", "", "", "git password")
	gitDeployCmd.Flags().StringVarP(&run.gitrsa, "ssh-rsa", "", filepath.Join(os.Getenv("HOME"), ".ssh/id_rsa"), "path to git SSH id_rsa")

	// Define Git Status Command
	gitStatusCmd.Flags().StringVarP(&run.gitrsa, "ssh-rsa", "", filepath.Join(os.Getenv("HOME"), ".ssh/id_rsa"), "path to git SSH id_rsa")
	gitStatusCmd.Flags().StringVarP(&run.gitpass, "password", "", "", "git password")
	gitStatusCmd.Flags().StringVarP(&run.gituser, "user", "u", "", "git username")

	// Define Terminate Flags
	terminateCmd.Flags().BoolVarP(&run.all, "all", "A", false, "terminate all stacks")

	// Define Output Flags
	outputsCmd.Flags().StringVarP(&run.profile, "profile", "p", "default", "configured aws profile")

	// Define Root Flags
	RootCmd.Flags().BoolVarP(&run.version, "version", "", false, "print current/running version")
	RootCmd.PersistentFlags().BoolVarP(&run.colors, "no-colors", "", false, "disable colors in outputs")
	RootCmd.PersistentFlags().StringVarP(&run.profile, "profile", "p", "default", "configured aws profile")
	RootCmd.PersistentFlags().StringVarP(&run.region, "region", "r", "", "configured aws region: if blank, the region is acquired via the profile")
	RootCmd.PersistentFlags().BoolVarP(&run.debug, "debug", "", false, "Run in debug mode...")

	// Define Lambda Invoke Flags
	invokeCmd.Flags().StringVarP(&run.funcEvent, "event", "e", "", "JSON Event data for AWS Lambda invoke")
	invokeCmd.Flags().BoolVarP(&run.lambdAsync, "async", "x", false, "invoke lambda function asynchronously ")

	// Define Changes Command
	changeCmd.AddCommand(create, rm, list, execute, desc)

	// Define Protect Command
	protectCmd.Flags().BoolVarP(&run.protectOff, "off", "", false, "set termination protection to off")
	protectCmd.Flags().BoolVarP(&run.all, "all", "A", false, "protect all stacks")

	// Add Config --config common flag
	for _, cmd := range []interface{}{
		checkCmd,
		updateCmd,
		outputsCmd,
		statusCmd,
		terminateCmd,
		generateCmd,
		deployCmd,
		gitDeployCmd,
		gitStatusCmd,
		policyCmd,
		valuesCmd,
		shellCmd,
		protectCmd,
	} {
		cmd.(*cobra.Command).Flags().StringVarP(&run.cfgSource, "config", "c", defaultConfig(), "path to config file")
	}

	// Add Template --template common flag
	for _, cmd := range []interface{}{
		generateCmd,
		updateCmd,
		checkCmd,
	} {
		cmd.(*cobra.Command).Flags().StringVarP(&run.tplSource, "template", "t", "", "path to template source Or stack::source")
	}

	for _, cmd := range []interface{}{
		create,
		list,
		rm,
		execute,
		desc,
	} {
		cmd.(*cobra.Command).Flags().StringVarP(&run.cfgSource, "config", "c", defaultConfig(), "path to config file [Required]")
		cmd.(*cobra.Command).Flags().StringVarP(&run.stackName, "stack", "s", "", "Qaz local project Stack Name [Required]")
	}

	create.Flags().StringVarP(&run.tplSource, "template", "t", "", "path to template file Or stack::url")
	changeCmd.Flags().StringVarP(&run.cfgSource, "config", "c", defaultConfig(), "path to config file")

	// add commands
	RootCmd.AddCommand(
		generateCmd,
		deployCmd,
		terminateCmd,
		statusCmd,
		outputsCmd,
		initCmd,
		updateCmd,
		checkCmd,
		exportsCmd,
		invokeCmd,
		changeCmd,
		policyCmd,
		gitDeployCmd,
		gitStatusCmd,
		valuesCmd,
		shellCmd,
		protectCmd,
		completionCmd,
	)

}
<commit_msg>Square brackets in description interferred with zsh completion syntax<commit_after>package commands

// Init sits here

import (
	"os"
	"path/filepath"

	"github.com/spf13/cobra"
)

func init() {

	// Define Deploy Flags
	deployCmd.Flags().StringArrayVarP(&run.tplSources, "template", "t", []string{}, "path to template file(s) Or stack::url")
	deployCmd.Flags().BoolVarP(&run.rollback, "disable-rollback", "", false, "Set Stack to rollback on deployment failures")
	deployCmd.Flags().BoolVarP(&run.all, "all", "A", false, "deploy all stacks with defined Sources in config")

	// Define Git Deploy Flags
	gitDeployCmd.Flags().BoolVarP(&run.rollback, "disable-rollback", "", false, "Set Stack to rollback on deployment failures")
	gitDeployCmd.Flags().BoolVarP(&run.all, "all", "A", false, "deploy all stacks with defined Sources in config")
	gitDeployCmd.Flags().StringVarP(&run.gituser, "user", "u", "", "git username")
	gitDeployCmd.Flags().StringVarP(&run.gitpass, "password", "", "", "git password")
	gitDeployCmd.Flags().StringVarP(&run.gitrsa, "ssh-rsa", "", filepath.Join(os.Getenv("HOME"), ".ssh/id_rsa"), "path to git SSH id_rsa")

	// Define Git Status Command
	gitStatusCmd.Flags().StringVarP(&run.gitrsa, "ssh-rsa", "", filepath.Join(os.Getenv("HOME"), ".ssh/id_rsa"), "path to git SSH id_rsa")
	gitStatusCmd.Flags().StringVarP(&run.gitpass, "password", "", "", "git password")
	gitStatusCmd.Flags().StringVarP(&run.gituser, "user", "u", "", "git username")

	// Define Terminate Flags
	terminateCmd.Flags().BoolVarP(&run.all, "all", "A", false, "terminate all stacks")

	// Define Output Flags
	outputsCmd.Flags().StringVarP(&run.profile, "profile", "p", "default", "configured aws profile")

	// Define Root Flags
	RootCmd.Flags().BoolVarP(&run.version, "version", "", false, "print current/running version")
	RootCmd.PersistentFlags().BoolVarP(&run.colors, "no-colors", "", false, "disable colors in outputs")
	RootCmd.PersistentFlags().StringVarP(&run.profile, "profile", "p", "default", "configured aws profile")
	RootCmd.PersistentFlags().StringVarP(&run.region, "region", "r", "", "configured aws region: if blank, the region is acquired via the profile")
	RootCmd.PersistentFlags().BoolVarP(&run.debug, "debug", "", false, "Run in debug mode...")

	// Define Lambda Invoke Flags
	invokeCmd.Flags().StringVarP(&run.funcEvent, "event", "e", "", "JSON Event data for AWS Lambda invoke")
	invokeCmd.Flags().BoolVarP(&run.lambdAsync, "async", "x", false, "invoke lambda function asynchronously ")

	// Define Changes Command
	changeCmd.AddCommand(create, rm, list, execute, desc)

	// Define Protect Command
	protectCmd.Flags().BoolVarP(&run.protectOff, "off", "", false, "set termination protection to off")
	protectCmd.Flags().BoolVarP(&run.all, "all", "A", false, "protect all stacks")

	// Add Config --config common flag
	for _, cmd := range []interface{}{
		checkCmd,
		updateCmd,
		outputsCmd,
		statusCmd,
		terminateCmd,
		generateCmd,
		deployCmd,
		gitDeployCmd,
		gitStatusCmd,
		policyCmd,
		valuesCmd,
		shellCmd,
		protectCmd,
	} {
		cmd.(*cobra.Command).Flags().StringVarP(&run.cfgSource, "config", "c", defaultConfig(), "path to config file")
	}

	// Add Template --template common flag
	for _, cmd := range []interface{}{
		generateCmd,
		updateCmd,
		checkCmd,
	} {
		cmd.(*cobra.Command).Flags().StringVarP(&run.tplSource, "template", "t", "", "path to template source Or stack::source")
	}

	for _, cmd := range []interface{}{
		create,
		list,
		rm,
		execute,
		desc,
	} {
		cmd.(*cobra.Command).Flags().StringVarP(&run.cfgSource, "config", "c", defaultConfig(), "path to config file (Required)")
		cmd.(*cobra.Command).Flags().StringVarP(&run.stackName, "stack", "s", "", "Qaz local project Stack Name (Required)")
	}

	create.Flags().StringVarP(&run.tplSource, "template", "t", "", "path to template file Or stack::url")
	changeCmd.Flags().StringVarP(&run.cfgSource, "config", "c", defaultConfig(), "path to config file")

	// add commands
	RootCmd.AddCommand(
		generateCmd,
		deployCmd,
		terminateCmd,
		statusCmd,
		outputsCmd,
		initCmd,
		updateCmd,
		checkCmd,
		exportsCmd,
		invokeCmd,
		changeCmd,
		policyCmd,
		gitDeployCmd,
		gitStatusCmd,
		valuesCmd,
		shellCmd,
		protectCmd,
		completionCmd,
	)

}
<|endoftext|>


<commit_before>package main

import (
	"testing"

	"github.com/AlexanderThaller/logger"
)

func Test_RunNoDataPath(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "NoDataPath")

	command := new(Command)
	err := command.Run()

	got := err.Error()
	expected := "the datapath can not be empty"

	testerr_output(t, l, err, got, expected)
}

func Test_RunDates(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Dates")

	action := ActionDates

	expected := `2014-10-30
2014-10-31
2014-11-01
`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunList(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List")

	action := ActionList
	expected := `TestNotes
TestTodos
TestTracks
`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunListProject(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project")

	action := ActionList
	command, buffer := testCommand(action)
	command.Project = "TestNotes"

	expected := "= TestNotes -- List\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== 2014-10-30T21:36:31.49146148+01:00
Test1

== 2014-10-30T21:36:33.49871531+01:00
Test2

== 2014-10-30T21:36:35.138412374+01:00
Test3

== 2014-10-30T21:36:36.810478305+01:00
Test4

== 2014-10-30T21:36:38.450479686+01:00
Test5

`

	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunListProjectNoDataDir(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project", "NoDataDir")

	action := ActionList
	command, buffer := testCommand(action)
	command.DataPath = TestDataPathFail

	expected := ""
	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunListProjectNoNotes(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project", "NoNotes")

	action := ActionList
	command, buffer := testCommand(action)
	command.Project = "TestTodos"

	expected := "= TestTodos -- List\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `* Test1
* Test2
* Test3
* Test4
* Test5
* Test7

`

	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunListProjectNoNotesNoTodos(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project", "NoNotesNoTodos")

	action := ActionList
	command, buffer := testCommand(action)
	command.Project = "TestTracks"

	expected := "= TestTracks -- List\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `* 2014-11-01T00:46:27.250010094+01:00 -- Test1
* 2014-11-01T00:46:31.186052306+01:00 -- Test2
* 2014-11-01T00:46:32.794131714+01:00 -- Test3
* 2014-11-01T00:46:34.322047221+01:00 -- Test4
* 2014-11-01T00:46:35.658221386+01:00 -- Test5
* 2014-11-01T00:46:57.953493565+01:00 -- Test7

`

	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunListProjectNoExists(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project", "NoExists")

	action := ActionList
	command, buffer := testCommand(action)
	command.Project = "DOES NOT EXIST"

	expected := ""
	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunNotes(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Notes")
	l.SetLevel(logger.Info)

	action := ActionNotes
	command, buffer := testCommand(action)
	command.Project = "TestNotes"

	expected := "= TestNotes -- Notes\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== 2014-10-30T21:36:31.49146148+01:00
Test1

== 2014-10-30T21:36:33.49871531+01:00
Test2

== 2014-10-30T21:36:35.138412374+01:00
Test3

== 2014-10-30T21:36:36.810478305+01:00
Test4

== 2014-10-30T21:36:38.450479686+01:00
Test5

`

	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunNotesProject(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Notes", "Project")
	l.SetLevel(logger.Info)

	action := ActionNotes
	expected := "= Lablog -- Notes\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestNotes

=== 2014-10-30T21:36:31.49146148+01:00
Test1

=== 2014-10-30T21:36:33.49871531+01:00
Test2

=== 2014-10-30T21:36:35.138412374+01:00
Test3

=== 2014-10-30T21:36:36.810478305+01:00
Test4

=== 2014-10-30T21:36:38.450479686+01:00
Test5

`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunProjects(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Projects")

	action := ActionProjects
	expected := `TestNotes
TestTodos
TestTracks
`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunTodos(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Todos")
	l.SetLevel(logger.Info)

	action := ActionTodos
	expected := "= Lablog -- Todos\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestTodos

* Test1
* Test2
* Test3
* Test4
* Test5
* Test7

`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunTracks(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Tracks")
	l.SetLevel(logger.Info)

	action := ActionTracks
	expected := "= Lablog -- Tracks\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestTracks

* 2014-11-01T00:46:27.250010094+01:00 -- Test1
* 2014-11-01T00:46:31.186052306+01:00 -- Test2
* 2014-11-01T00:46:32.794131714+01:00 -- Test3
* 2014-11-01T00:46:34.322047221+01:00 -- Test4
* 2014-11-01T00:46:35.658221386+01:00 -- Test5
* 2014-11-01T00:46:38.385921461+01:00
* 2014-11-01T00:46:51.833861322+01:00 -- Test6
* 2014-11-01T00:46:53.713763133+01:00 -- Test6
* 2014-11-01T00:46:55.609688812+01:00 -- Test7
* 2014-11-01T00:46:56.713668447+01:00 -- Test7
* 2014-11-01T00:46:57.953493565+01:00 -- Test7
* 2014-11-01T00:54:54.778921093+01:00

`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunTracksActive(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "TracksActive")
	l.SetLevel(logger.Info)

	action := ActionTracksActive
	expected := "= Lablog -- TracksActive\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestTracks

* 2014-11-01T00:46:27.250010094+01:00 -- Test1
* 2014-11-01T00:46:31.186052306+01:00 -- Test2
* 2014-11-01T00:46:32.794131714+01:00 -- Test3
* 2014-11-01T00:46:34.322047221+01:00 -- Test4
* 2014-11-01T00:46:35.658221386+01:00 -- Test5
* 2014-11-01T00:46:57.953493565+01:00 -- Test7

`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunTracksDurations(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "TracksDurations")
	l.SetLevel(logger.Info)

	action := ActionTracksDurations
	expected := "= Lablog -- Durations\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestTracks

* 8m16.392999632s
* Test6 -- 1.879901811s
* Test7 -- 1.103979635s

`

	testCommandRunOutput(t, l, action, expected)
}
<commit_msg>fixed some tests and added some benchmarks.<commit_after>package main

import (
	"testing"

	"github.com/AlexanderThaller/logger"
)

func Test_RunNoDataPath(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "NoDataPath")

	command := new(Command)
	err := command.Run()

	got := err.Error()
	expected := "the datapath can not be empty"

	testerr_output(t, l, err, got, expected)
}

func Test_RunDates(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Dates")

	action := ActionDates

	expected := `2014-10-30
2014-10-31
2014-11-01
`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunList(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List")

	action := ActionList
	expected := `TestNotes
TestTodos
TestTracks
`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunListProject(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project")

	action := ActionList
	command, buffer := testCommand(action)
	command.Project = "TestNotes"

	expected := "= TestNotes -- List\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== 2014-10-30T21:36:31.49146148+01:00
Test1

== 2014-10-30T21:36:33.49871531+01:00
Test2

== 2014-10-30T21:36:35.138412374+01:00
Test3

== 2014-10-30T21:36:36.810478305+01:00
Test4

== 2014-10-30T21:36:38.450479686+01:00
Test5

`

	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunListProjectNoDataDir(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project", "NoDataDir")

	action := ActionList
	command, buffer := testCommand(action)
	command.DataPath = TestDataPathFail

	expected := ""
	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunListProjectNoNotes(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project", "NoNotes")

	action := ActionList
	command, buffer := testCommand(action)
	command.Project = "TestTodos"

	expected := "= TestTodos -- List\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `* Test1
* Test2
* Test3
* Test4
* Test5
* Test7

`

	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunListProjectNoNotesNoTodos(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project", "NoNotesNoTodos")

	action := ActionList
	command, buffer := testCommand(action)
	command.Project = "TestTracks"

	expected := "= TestTracks -- List\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `* 2014-11-01T00:46:27.250010094+01:00 -- Test1
* 2014-11-01T00:46:31.186052306+01:00 -- Test2
* 2014-11-01T00:46:32.794131714+01:00 -- Test3
* 2014-11-01T00:46:34.322047221+01:00 -- Test4
* 2014-11-01T00:46:35.658221386+01:00 -- Test5
* 2014-11-01T00:46:57.953493565+01:00 -- Test7

`

	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunListProjectNoExists(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "List", "Project", "NoExists")

	action := ActionList
	command, buffer := testCommand(action)
	command.Project = "DOES NOT EXIST"

	expected := ""
	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunNotes(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Notes")
	l.SetLevel(logger.Info)

	action := ActionNotes
	expected := "= Lablog -- Notes\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestNotes

=== 2014-10-30T21:36:31.49146148+01:00
Test1

=== 2014-10-30T21:36:33.49871531+01:00
Test2

=== 2014-10-30T21:36:35.138412374+01:00
Test3

=== 2014-10-30T21:36:36.810478305+01:00
Test4

=== 2014-10-30T21:36:38.450479686+01:00
Test5

`

	testCommandRunOutput(t, l, action, expected)
}

func Benchmark_RunNotes(b *testing.B) {
	action := ActionNotes
	command, _ := testCommand(action)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		command.Run()
	}
}

func Test_RunNotesProject(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Notes", "Project")
	l.SetLevel(logger.Info)

	action := ActionNotes
	command, buffer := testCommand(action)
	command.Project = "TestNotes"

	expected := "= TestNotes -- Notes\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== 2014-10-30T21:36:31.49146148+01:00
Test1

== 2014-10-30T21:36:33.49871531+01:00
Test2

== 2014-10-30T21:36:35.138412374+01:00
Test3

== 2014-10-30T21:36:36.810478305+01:00
Test4

== 2014-10-30T21:36:38.450479686+01:00
Test5

`

	err := command.Run()
	got := buffer.String()

	testerr_output(t, l, err, got, expected)
}

func Test_RunProjects(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Projects")

	action := ActionProjects
	expected := `TestNotes
TestTodos
TestTracks
`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunTodos(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Todos")
	l.SetLevel(logger.Info)

	action := ActionTodos
	expected := "= Lablog -- Todos\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestTodos

* Test1
* Test2
* Test3
* Test4
* Test5
* Test7

`

	testCommandRunOutput(t, l, action, expected)
}

func Benchmark_RunTodos(b *testing.B) {
	action := ActionTodos
	command, _ := testCommand(action)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		command.Run()
	}
}

func Test_RunTracks(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "Tracks")
	l.SetLevel(logger.Info)

	action := ActionTracks
	expected := "= Lablog -- Tracks\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestTracks

* 2014-11-01T00:46:27.250010094+01:00 -- Test1
* 2014-11-01T00:46:31.186052306+01:00 -- Test2
* 2014-11-01T00:46:32.794131714+01:00 -- Test3
* 2014-11-01T00:46:34.322047221+01:00 -- Test4
* 2014-11-01T00:46:35.658221386+01:00 -- Test5
* 2014-11-01T00:46:38.385921461+01:00
* 2014-11-01T00:46:51.833861322+01:00 -- Test6
* 2014-11-01T00:46:53.713763133+01:00 -- Test6
* 2014-11-01T00:46:55.609688812+01:00 -- Test7
* 2014-11-01T00:46:56.713668447+01:00 -- Test7
* 2014-11-01T00:46:57.953493565+01:00 -- Test7
* 2014-11-01T00:54:54.778921093+01:00

`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunTracksActive(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "TracksActive")
	l.SetLevel(logger.Info)

	action := ActionTracksActive
	expected := "= Lablog -- TracksActive\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestTracks

* 2014-11-01T00:46:27.250010094+01:00 -- Test1
* 2014-11-01T00:46:31.186052306+01:00 -- Test2
* 2014-11-01T00:46:32.794131714+01:00 -- Test3
* 2014-11-01T00:46:34.322047221+01:00 -- Test4
* 2014-11-01T00:46:35.658221386+01:00 -- Test5
* 2014-11-01T00:46:57.953493565+01:00 -- Test7

`

	testCommandRunOutput(t, l, action, expected)
}

func Test_RunTracksDurations(t *testing.T) {
	l := logger.New(Name, "Test", "Command", "Run", "TracksDurations")
	l.SetLevel(logger.Info)

	action := ActionTracksDurations
	expected := "= Lablog -- Durations\n"
	expected += AsciiDocSettings + "\n\n"
	expected += `== TestTracks

* 8m16.392999632s
* Test6 -- 1.879901811s
* Test7 -- 1.103979635s

`

	testCommandRunOutput(t, l, action, expected)
}
<|endoftext|>


<commit_before>package windows_test

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/apcera/nats"
	"github.com/cloudfoundry/bosh-agent/agent/action"

	. "github.com/onsi/gomega"
)

const (
	prepareTemplate = `{
    "arguments": [
        {
            "deployment": "test",
            "job": {
                "name": "test-job",
                "template": "test-job",
                "templates": [
                    {
                        "name": "say-hello",
						"blobstore_id": "%s",
						"sha1": "eb9bebdb1f11494b27440ec6ccbefba00e713cd9"
                    }
                ]
            },
            "packages": {},
            "rendered_templates_archive": {
                "blobstore_id": "%s",
                "sha1": "6760d464064ee036db9898c736ff71c6d4457792"
            }
        }
    ],
    "method": "prepare",
    "reply_to": "%s"
}`
	errandTemplate = `
	{"protocol":2,"method":"run_errand","arguments":[],"reply_to":"%s"}
	`
	applyTemplate = `
{
    "arguments": [
        {
            "configuration_hash": "foo",
            "deployment": "hello-world-windows-deployment",
            "id": "62236318-6632-4318-94c7-c3dd6e8e5698",
            "index": 0,
            "job": {
                "blobstore_id": "%[1]s",
                "name": "say-hello",
                "sha1": "eb6e6c8bd1b1bc3dd91c741ec5c628b61a4d8f1d",
                "template": "say-hello",
                "templates": [
                    {
                        "blobstore_id": "%[1]s",
                        "name": "say-hello",
                        "sha1": "eb6e6c8bd1b1bc3dd91c741ec5c628b61a4d8f1d",
                        "version": "8fe0a4982b28ffe4e59d7c1e573c4f30a526770d"
                    }
                ],
                "version": "8fe0a4982b28ffe4e59d7c1e573c4f30a526770d"
            },
            "networks": {},
			"rendered_templates_archive": {
					"blobstore_id": "%[2]s",
					"sha1": "6760d464064ee036db9898c736ff71c6d4457792"
			}
        }
    ],
    "method": "apply",
    "protocol": 2,
    "reply_to": "%[3]s"
}
	`

	fetchLogsTemplate = `
{
    "arguments": [
        "job",
        null
    ],
    "method": "fetch_logs",
    "protocol": 2,
    "reply_to": "%s"
}
	`

	startTemplate = `
{
	"arguments":[],
	"method":"start",
	"protocol":2,
	"reply_to":"%s"
}
	`
	stopTemplate = `
{
	"protocol":2,
	"method":"stop",
	"arguments":[],
	"reply_to":"%s"
}
	`
)

type natsClient struct {
	nc  *nats.Conn
	sub *nats.Subscription
}

func NewNatsClient() *natsClient {
	return &natsClient{}
}

func (n *natsClient) Setup() error {
	var err error
	n.nc, err = nats.Connect(natsURI())
	if err != nil {
		return err
	}

	n.sub, err = n.nc.SubscribeSync(senderID)
	return err
}

func (n *natsClient) Cleanup() {
	_, err := n.RunStop()
	Expect(err).NotTo(HaveOccurred())

	n.nc.Close()
}

func (n *natsClient) PrepareJob(templateID, renderedTemplateArchiveID string) error {
	message := fmt.Sprintf(prepareTemplate, templateID, renderedTemplateArchiveID, senderID)
	prepareResponse, err := n.SendMessage(message)
	if err != nil {
		return err
	}

	check := n.checkStatus(prepareResponse["value"]["agent_task_id"])
	Eventually(check, 30*time.Second, 1*time.Second).Should(Equal("prepared"))

	message = fmt.Sprintf(applyTemplate, templateID, renderedTemplateArchiveID, senderID)
	applyResponse, err := n.SendMessage(message)
	if err != nil {
		return err
	}

	check = n.checkStatus(applyResponse["value"]["agent_task_id"])
	Eventually(check, 30*time.Second, 1*time.Second).Should(Equal("applied"))

	return nil
}

func (n *natsClient) RunStart() (map[string]string, error) {
	message := fmt.Sprintf(startTemplate, senderID)
	rawResponse, err := n.SendRawMessage(message)
	if err != nil {
		return map[string]string{}, err
	}

	response := map[string]string{}
	err = json.Unmarshal(rawResponse, &response)
	return response, err
}

func (n *natsClient) RunStop() (map[string]map[string]string, error) {
	message := fmt.Sprintf(stopTemplate, senderID)
	rawResponse, err := n.SendRawMessage(message)
	if err != nil {
		return map[string]map[string]string{}, err
	}

	response := map[string]map[string]string{}

	err = json.Unmarshal(rawResponse, &response)
	return response, err
}

func (n *natsClient) RunErrand() (map[string]map[string]string, error) {
	message := fmt.Sprintf(errandTemplate, senderID)
	return n.SendMessage(message)
}

func (n *natsClient) RunFetchLogs() (map[string]map[string]string, error) {
	message := fmt.Sprintf(fetchLogsTemplate, senderID)
	return n.SendMessage(message)
}

func (n *natsClient) CheckFetchLogsStatus(taskID string) func() (map[string]string, error) {
	return func() (map[string]string, error) {
		var result map[string]map[string]string
		valueResponse, err := n.getTask(taskID)
		if err != nil {
			return map[string]string{}, err
		}

		err = json.Unmarshal(valueResponse, &result)
		if err != nil {
			return map[string]string{}, err
		}

		return result["value"], nil
	}
}

func (n *natsClient) CheckErrandResultStatus(taskID string) func() (action.ErrandResult, error) {
	return func() (action.ErrandResult, error) {
		var result map[string]action.ErrandResult
		valueResponse, err := n.getTask(taskID)
		if err != nil {
			return action.ErrandResult{}, err
		}

		err = json.Unmarshal(valueResponse, &result)
		if err != nil {
			return action.ErrandResult{}, err
		}

		return result["value"], nil
	}
}

func (n *natsClient) SendRawMessage(message string) ([]byte, error) {
	err := n.nc.Publish(agentID, []byte(message))
	if err != nil {
		return nil, err
	}

	raw, err := n.sub.NextMsg(5 * time.Second)
	if err != nil {
		return nil, err
	}

	return raw.Data, nil
}

func (n *natsClient) SendMessage(message string) (map[string]map[string]string, error) {
	rawMessage, err := n.SendRawMessage(message)
	if err != nil {
		return nil, err
	}

	response := map[string]map[string]string{}
	err = json.Unmarshal(rawMessage, &response)
	return response, err
}

func (n *natsClient) getTask(taskID string) ([]byte, error) {
	message := fmt.Sprintf(`{"method": "get_task", "arguments": ["%s"], "reply_to": "%s"}`, taskID, senderID)
	return n.SendRawMessage(message)
}

func (n *natsClient) checkStatus(taskID string) func() (string, error) {
	return func() (string, error) {
		var result map[string]string
		valueResponse, err := n.getTask(taskID)
		if err != nil {
			return "", err
		}

		err = json.Unmarshal(valueResponse, &result)
		if err != nil {
			return "", err
		}

		return result["value"], nil
	}
}
<commit_msg>Increase the timeout for receiving a response to the 'start' task<commit_after>package windows_test

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/apcera/nats"
	"github.com/cloudfoundry/bosh-agent/agent/action"

	. "github.com/onsi/gomega"
)

const (
	prepareTemplate = `{
    "arguments": [
        {
            "deployment": "test",
            "job": {
                "name": "test-job",
                "template": "test-job",
                "templates": [
                    {
                        "name": "say-hello",
						"blobstore_id": "%s",
						"sha1": "eb9bebdb1f11494b27440ec6ccbefba00e713cd9"
                    }
                ]
            },
            "packages": {},
            "rendered_templates_archive": {
                "blobstore_id": "%s",
                "sha1": "6760d464064ee036db9898c736ff71c6d4457792"
            }
        }
    ],
    "method": "prepare",
    "reply_to": "%s"
}`
	errandTemplate = `
	{"protocol":2,"method":"run_errand","arguments":[],"reply_to":"%s"}
	`
	applyTemplate = `
{
    "arguments": [
        {
            "configuration_hash": "foo",
            "deployment": "hello-world-windows-deployment",
            "id": "62236318-6632-4318-94c7-c3dd6e8e5698",
            "index": 0,
            "job": {
                "blobstore_id": "%[1]s",
                "name": "say-hello",
                "sha1": "eb6e6c8bd1b1bc3dd91c741ec5c628b61a4d8f1d",
                "template": "say-hello",
                "templates": [
                    {
                        "blobstore_id": "%[1]s",
                        "name": "say-hello",
                        "sha1": "eb6e6c8bd1b1bc3dd91c741ec5c628b61a4d8f1d",
                        "version": "8fe0a4982b28ffe4e59d7c1e573c4f30a526770d"
                    }
                ],
                "version": "8fe0a4982b28ffe4e59d7c1e573c4f30a526770d"
            },
            "networks": {},
			"rendered_templates_archive": {
					"blobstore_id": "%[2]s",
					"sha1": "6760d464064ee036db9898c736ff71c6d4457792"
			}
        }
    ],
    "method": "apply",
    "protocol": 2,
    "reply_to": "%[3]s"
}
	`

	fetchLogsTemplate = `
{
    "arguments": [
        "job",
        null
    ],
    "method": "fetch_logs",
    "protocol": 2,
    "reply_to": "%s"
}
	`

	startTemplate = `
{
	"arguments":[],
	"method":"start",
	"protocol":2,
	"reply_to":"%s"
}
	`
	stopTemplate = `
{
	"protocol":2,
	"method":"stop",
	"arguments":[],
	"reply_to":"%s"
}
	`
)

type natsClient struct {
	nc  *nats.Conn
	sub *nats.Subscription
}

func NewNatsClient() *natsClient {
	return &natsClient{}
}

func (n *natsClient) Setup() error {
	var err error
	n.nc, err = nats.Connect(natsURI())
	if err != nil {
		return err
	}

	n.sub, err = n.nc.SubscribeSync(senderID)
	return err
}

func (n *natsClient) Cleanup() {
	err := n.RunStop()
	Expect(err).NotTo(HaveOccurred())

	n.nc.Close()
}

func (n *natsClient) PrepareJob(templateID, renderedTemplateArchiveID string) error {
	message := fmt.Sprintf(prepareTemplate, templateID, renderedTemplateArchiveID, senderID)
	prepareResponse, err := n.SendMessage(message)
	if err != nil {
		return err
	}

	check := n.checkStatus(prepareResponse["value"]["agent_task_id"])
	Eventually(check, 30*time.Second, 1*time.Second).Should(Equal("prepared"))

	message = fmt.Sprintf(applyTemplate, templateID, renderedTemplateArchiveID, senderID)
	applyResponse, err := n.SendMessage(message)
	if err != nil {
		return err
	}

	check = n.checkStatus(applyResponse["value"]["agent_task_id"])
	Eventually(check, 30*time.Second, 1*time.Second).Should(Equal("applied"))

	return nil
}

func (n *natsClient) RunStart() (map[string]string, error) {
	message := fmt.Sprintf(startTemplate, senderID)
	rawResponse, err := n.SendRawMessageWithTimeout(message, time.Minute)
	if err != nil {
		return map[string]string{}, err
	}

	response := map[string]string{}
	err = json.Unmarshal(rawResponse, &response)
	return response, err
}

func (n *natsClient) RunStop() error {
	message := fmt.Sprintf(stopTemplate, senderID)
	rawResponse, err := n.SendRawMessage(message)
	if err != nil {
		return err
	}

	response := map[string]map[string]string{}

	err = json.Unmarshal(rawResponse, &response)
	if err != nil {
		return err
	}

	check := n.checkStatus(response["value"]["agent_task_id"])
	Eventually(check, 30*time.Second, 1*time.Second).Should(Equal("stopped"))

	return nil
}

func (n *natsClient) RunErrand() (map[string]map[string]string, error) {
	message := fmt.Sprintf(errandTemplate, senderID)
	return n.SendMessage(message)
}

func (n *natsClient) RunFetchLogs() (map[string]map[string]string, error) {
	message := fmt.Sprintf(fetchLogsTemplate, senderID)
	return n.SendMessage(message)
}

func (n *natsClient) CheckFetchLogsStatus(taskID string) func() (map[string]string, error) {
	return func() (map[string]string, error) {
		var result map[string]map[string]string
		valueResponse, err := n.getTask(taskID)
		if err != nil {
			return map[string]string{}, err
		}

		err = json.Unmarshal(valueResponse, &result)
		if err != nil {
			return map[string]string{}, err
		}

		return result["value"], nil
	}
}

func (n *natsClient) CheckErrandResultStatus(taskID string) func() (action.ErrandResult, error) {
	return func() (action.ErrandResult, error) {
		var result map[string]action.ErrandResult
		valueResponse, err := n.getTask(taskID)
		if err != nil {
			return action.ErrandResult{}, err
		}

		err = json.Unmarshal(valueResponse, &result)
		if err != nil {
			return action.ErrandResult{}, err
		}

		return result["value"], nil
	}
}

func (n *natsClient) SendRawMessageWithTimeout(message string, timeout time.Duration) ([]byte, error) {
	err := n.nc.Publish(agentID, []byte(message))
	if err != nil {
		return nil, err
	}

	raw, err := n.sub.NextMsg(timeout)
	if err != nil {
		return nil, err
	}

	return raw.Data, nil
}

func (n *natsClient) SendRawMessage(message string) ([]byte, error) {
	return n.SendRawMessageWithTimeout(message, 5*time.Second)
}

func (n *natsClient) SendMessage(message string) (map[string]map[string]string, error) {
	rawMessage, err := n.SendRawMessage(message)
	if err != nil {
		return nil, err
	}

	response := map[string]map[string]string{}
	err = json.Unmarshal(rawMessage, &response)
	return response, err
}

func (n *natsClient) getTask(taskID string) ([]byte, error) {
	message := fmt.Sprintf(`{"method": "get_task", "arguments": ["%s"], "reply_to": "%s"}`, taskID, senderID)
	return n.SendRawMessage(message)
}

func (n *natsClient) checkStatus(taskID string) func() (string, error) {
	return func() (string, error) {
		var result map[string]string
		valueResponse, err := n.getTask(taskID)
		if err != nil {
			return "", err
		}

		err = json.Unmarshal(valueResponse, &result)
		if err != nil {
			return "", err
		}

		return result["value"], nil
	}
}
<|endoftext|>


<commit_before>package common

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/sirupsen/logrus"
)

const (
	RequestMethodGET  = "GET"
	RequestMethodPOST = "POST"

	SignatureMethodHMacSha256 = "HmacSHA256"
)

type Client struct {
	*http.Client

	credential CredentialInterface
	opts       Opts
}

type Opts struct {
	Method          string
	Region          string
	Host            string
	Path            string
	SignatureMethod string
	Schema          string

	Logger *logrus.Logger
}

type CredentialInterface interface {
	GetSecretId() (string, error)
	GetSecretKey() (string, error)

	Values() (CredentialValues, error)
}

type CredentialValues map[string]string

type Credential struct {
	SecretId  string
	SecretKey string
}

func (cred Credential) GetSecretId() (string, error) {
	return cred.SecretId, nil
}

func (cred Credential) GetSecretKey() (string, error) {
	return cred.SecretKey, nil
}

func (cred Credential) Values() (CredentialValues, error) {
	return CredentialValues{}, nil
}

func NewClient(credential CredentialInterface, opts Opts) (*Client, error) {
	if opts.Method == "" {
		opts.Method = RequestMethodGET
	}
	if opts.SignatureMethod == "" {
		opts.SignatureMethod = SignatureMethodHMacSha256
	}
	if opts.Schema == "" {
		opts.Schema = "https"
	}
	if opts.Logger == nil {
		opts.Logger = logrus.New()
	}
	return &Client{
		&http.Client{},
		credential,
		opts,
	}, nil
}

func (client *Client) Invoke(action string, args interface{}, response interface{}) error {
	switch client.opts.Method {
	case "GET":
		return client.InvokeWithGET(action, args, response)
	default:
		return client.InvokeWithPOST(action, args, response)
	}
}

func (client *Client) initCommonArgs(args *url.Values) {
	args.Set("Region", client.opts.Region)
	args.Set("Timestamp", fmt.Sprint(uint(time.Now().Unix())))
	args.Set("Nonce", fmt.Sprint(uint(rand.Int())))
	args.Set("SignatureMethod", client.opts.SignatureMethod)
}

func (client *Client) signGetRequest(secretId, secretKey string, values *url.Values) string {

	values.Set("SecretId", secretId)

	keys := make([]string, 0, len(*values))
	for k := range *values {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	kvs := make([]string, 0, len(keys))
	for _, k := range keys {
		kvs = append(kvs, fmt.Sprintf("%s=%s", k, values.Get(k)))
	}
	queryStr := strings.Join(kvs, "&")
	reqStr := fmt.Sprintf("GET%s%s?%s", client.opts.Host, client.opts.Path, queryStr)

	mac := hmac.New(sha256.New, []byte(secretKey))
	mac.Write([]byte(reqStr))
	signature := mac.Sum(nil)

	b64Encoded := base64.StdEncoding.EncodeToString(signature)

	return b64Encoded
}

func (client *Client) InvokeWithGET(action string, args interface{}, response interface{}) error {
	reqValues := url.Values{}

	credValues, err := client.credential.Values()
	if err != nil {
		return makeClientError(err)
	}

	for k, v := range credValues {
		reqValues.Set(k, v)
	}

	err = EncodeStruct(args, &reqValues)
	if err != nil {
		return makeClientError(err)
	}
	reqValues.Set("Action", action)
	client.initCommonArgs(&reqValues)

	secretId, err := client.credential.GetSecretId()
	if err != nil {
		return makeClientError(err)
	}
	secretKey, err := client.credential.GetSecretKey()
	if err != nil {
		return makeClientError(err)
	}

	signature := client.signGetRequest(secretId, secretKey, &reqValues)
	reqValues.Set("Signature", signature)

	reqQuery := reqValues.Encode()

	req, err := http.NewRequest(
		"GET",
		fmt.Sprintf("%s://%s%s?%s", client.opts.Schema, client.opts.Host, client.opts.Path, reqQuery),
		nil,
	)

	if err != nil {
		return makeClientError(err)
	}

	resp, err := client.Do(req)
	if err != nil {
		return makeClientError(err)
	}

	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return makeClientError(err)
	}

	client.opts.Logger.WithField("Action", action).Infof(
		"%s %s %d %s", "GET", req.URL, resp.StatusCode, body,
	)

	legacyErrorResponse := LegacyAPIError{}

	if err = json.Unmarshal(body, &legacyErrorResponse); err != nil {
		return makeClientError(err)
	}

	versionErrorResponse := VersionAPIError{}

	if err = json.Unmarshal(body, &versionErrorResponse); err != nil {
		return makeClientError(err)
	}

	if legacyErrorResponse.Code != NoErr || (legacyErrorResponse.CodeDesc != "" && legacyErrorResponse.CodeDesc != NoErrCodeDesc) {
		client.opts.Logger.WithField("Action", action).Errorf(
			"%s %s %d %s %v", "GET", req.URL, resp.StatusCode, body, legacyErrorResponse,
		)
		return legacyErrorResponse
	}

	if versionErrorResponse.Response.Error.Code != "" {
		client.opts.Logger.WithField("Action", action).Errorf(
			"%s %s %d %s %v", "GET", req.URL, resp.StatusCode, body, versionErrorResponse,
		)
		return versionErrorResponse
	}

	if err = json.Unmarshal(body, response); err != nil {
		return makeClientError(err)
	}

	return nil
}

func (client *Client) InvokeWithPOST(action string, args interface{}, response interface{}) error {
	return nil
}
<commit_msg>add client timeout<commit_after>package common

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/sirupsen/logrus"
)

const (
	RequestMethodGET  = "GET"
	RequestMethodPOST = "POST"

	SignatureMethodHMacSha256 = "HmacSHA256"
)

type Client struct {
	*http.Client

	credential CredentialInterface
	opts       Opts
}

type Opts struct {
	Method          string
	Region          string
	Host            string
	Path            string
	SignatureMethod string
	Schema          string

	Logger *logrus.Logger
}

type CredentialInterface interface {
	GetSecretId() (string, error)
	GetSecretKey() (string, error)

	Values() (CredentialValues, error)
}

type CredentialValues map[string]string

type Credential struct {
	SecretId  string
	SecretKey string
}

func (cred Credential) GetSecretId() (string, error) {
	return cred.SecretId, nil
}

func (cred Credential) GetSecretKey() (string, error) {
	return cred.SecretKey, nil
}

func (cred Credential) Values() (CredentialValues, error) {
	return CredentialValues{}, nil
}

func NewClient(credential CredentialInterface, opts Opts) (*Client, error) {
	if opts.Method == "" {
		opts.Method = RequestMethodGET
	}
	if opts.SignatureMethod == "" {
		opts.SignatureMethod = SignatureMethodHMacSha256
	}
	if opts.Schema == "" {
		opts.Schema = "https"
	}
	if opts.Logger == nil {
		opts.Logger = logrus.New()
	}
	return &Client{
		&http.Client{
			Timeout: time.Second * 60,
		},
		credential,
		opts,
	}, nil
}

func (client *Client) Invoke(action string, args interface{}, response interface{}) error {
	switch client.opts.Method {
	case "GET":
		return client.InvokeWithGET(action, args, response)
	default:
		return client.InvokeWithPOST(action, args, response)
	}
}

func (client *Client) initCommonArgs(args *url.Values) {
	args.Set("Region", client.opts.Region)
	args.Set("Timestamp", fmt.Sprint(uint(time.Now().Unix())))
	args.Set("Nonce", fmt.Sprint(uint(rand.Int())))
	args.Set("SignatureMethod", client.opts.SignatureMethod)
}

func (client *Client) signGetRequest(secretId, secretKey string, values *url.Values) string {

	values.Set("SecretId", secretId)

	keys := make([]string, 0, len(*values))
	for k := range *values {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	kvs := make([]string, 0, len(keys))
	for _, k := range keys {
		kvs = append(kvs, fmt.Sprintf("%s=%s", k, values.Get(k)))
	}
	queryStr := strings.Join(kvs, "&")
	reqStr := fmt.Sprintf("GET%s%s?%s", client.opts.Host, client.opts.Path, queryStr)

	mac := hmac.New(sha256.New, []byte(secretKey))
	mac.Write([]byte(reqStr))
	signature := mac.Sum(nil)

	b64Encoded := base64.StdEncoding.EncodeToString(signature)

	return b64Encoded
}

func (client *Client) InvokeWithGET(action string, args interface{}, response interface{}) error {
	reqValues := url.Values{}

	credValues, err := client.credential.Values()
	if err != nil {
		return makeClientError(err)
	}

	for k, v := range credValues {
		reqValues.Set(k, v)
	}

	err = EncodeStruct(args, &reqValues)
	if err != nil {
		return makeClientError(err)
	}
	reqValues.Set("Action", action)
	client.initCommonArgs(&reqValues)

	secretId, err := client.credential.GetSecretId()
	if err != nil {
		return makeClientError(err)
	}
	secretKey, err := client.credential.GetSecretKey()
	if err != nil {
		return makeClientError(err)
	}

	signature := client.signGetRequest(secretId, secretKey, &reqValues)
	reqValues.Set("Signature", signature)

	reqQuery := reqValues.Encode()

	req, err := http.NewRequest(
		"GET",
		fmt.Sprintf("%s://%s%s?%s", client.opts.Schema, client.opts.Host, client.opts.Path, reqQuery),
		nil,
	)

	if err != nil {
		return makeClientError(err)
	}

	resp, err := client.Do(req)
	if err != nil {
		return makeClientError(err)
	}

	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return makeClientError(err)
	}

	client.opts.Logger.WithField("Action", action).Infof(
		"%s %s %d %s", "GET", req.URL, resp.StatusCode, body,
	)

	legacyErrorResponse := LegacyAPIError{}

	if err = json.Unmarshal(body, &legacyErrorResponse); err != nil {
		return makeClientError(err)
	}

	versionErrorResponse := VersionAPIError{}

	if err = json.Unmarshal(body, &versionErrorResponse); err != nil {
		return makeClientError(err)
	}

	if legacyErrorResponse.Code != NoErr || (legacyErrorResponse.CodeDesc != "" && legacyErrorResponse.CodeDesc != NoErrCodeDesc) {
		client.opts.Logger.WithField("Action", action).Errorf(
			"%s %s %d %s %v", "GET", req.URL, resp.StatusCode, body, legacyErrorResponse,
		)
		return legacyErrorResponse
	}

	if versionErrorResponse.Response.Error.Code != "" {
		client.opts.Logger.WithField("Action", action).Errorf(
			"%s %s %d %s %v", "GET", req.URL, resp.StatusCode, body, versionErrorResponse,
		)
		return versionErrorResponse
	}

	if err = json.Unmarshal(body, response); err != nil {
		return makeClientError(err)
	}

	return nil
}

func (client *Client) InvokeWithPOST(action string, args interface{}, response interface{}) error {
	return nil
}
<|endoftext|>


<commit_before>// Package common contains common utilities that are shared among other packages.
// See each sub-package for detail.
package common

import (
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"v2ray.com/core/common/errors"
)

//go:generate errorgen

var (
	// ErrNoClue is for the situation that existing information is not enough to make a decision. For example, Router may return this error when there is no suitable route.
	ErrNoClue = errors.New("not enough information for making a decision")
)

// Must panics if err is not nil.
func Must(err error) {
	if err != nil {
		panic(err)
	}
}

// Must2 panics if the second parameter is not nil, otherwise returns the first parameter.
func Must2(v interface{}, err error) interface{} {
	Must(err)
	return v
}

// Error2 returns the err from the 2nd parameter.
func Error2(v interface{}, err error) error {
	return err
}

// envFile returns the name of the Go environment configuration file.
// Copy from https://github.com/golang/go/blob/c4f2a9788a7be04daf931ac54382fbe2cb754938/src/cmd/go/internal/cfg/cfg.go#L150-L166
func envFile() (string, error) {
	if file := os.Getenv("GOENV"); file != "" {
		if file == "off" {
			return "", fmt.Errorf("GOENV=off")
		}
		return file, nil
	}
	dir, err := os.UserConfigDir()
	if err != nil {
		return "", err
	}
	if dir == "" {
		return "", fmt.Errorf("missing user-config dir")
	}
	return filepath.Join(dir, "go", "env"), nil
}

// GetRuntimeEnv returns the value of runtime environment variable,
// that is set by running following command: `go env -w key=value`.
func GetRuntimeEnv(key string) (string, error) {
	file, err := envFile()
	if err != nil {
		return "", err
	}
	if file == "" {
		return "", fmt.Errorf("missing runtime env file")
	}
	var data []byte
	var runtimeEnv string
	data, readErr := ioutil.ReadFile(file)
	if readErr != nil {
		return "", readErr
	}
	envStrings := strings.Split(string(data), "\n")
	for _, envItem := range envStrings {
		envItem = strings.TrimSuffix(envItem, "\r")
		envKeyValue := strings.Split(envItem, "=")
		if strings.EqualFold(strings.TrimSpace(envKeyValue[0]), key) {
			runtimeEnv = strings.TrimSpace(envKeyValue[1])
		}
	}
	return runtimeEnv, nil
}

// GetGOBIN returns GOBIN environment variable as a string. It will NOT be empty.
func GetGOBIN() string {
	// The one set by user explicitly by `export GOBIN=/path` or `env GOBIN=/path command`
	GOBIN := os.Getenv("GOBIN")
	if GOBIN == "" {
		var err error
		// The one set by user by running `go env -w GOBIN=/path`
		GOBIN, err = GetRuntimeEnv("GOBIN")
		if err != nil {
			// The default one that Golang uses
			return filepath.Join(build.Default.GOPATH, "bin")
		}
		if GOBIN == "" {
			return filepath.Join(build.Default.GOPATH, "bin")
		}
		return GOBIN
	}
	return GOBIN
}

// GetGOPATH returns GOPATH environment variable as a string. It will NOT be empty.
func GetGOPATH() string {
	// The one set by user explicitly by `export GOPATH=/path` or `env GOPATH=/path command`
	GOPATH := os.Getenv("GOPATH")
	if GOPATH == "" {
		var err error
		// The one set by user by running `go env -w GOPATH=/path`
		GOPATH, err = GetRuntimeEnv("GOPATH")
		if err != nil {
			// The default one that Golang uses
			return build.Default.GOPATH
		}
		if GOPATH == "" {
			return build.Default.GOPATH
		}
		return GOPATH
	}
	return GOPATH
}

// GetModuleName returns the value of module in `go.mod` file.
func GetModuleName(pathToProjectRoot string) (string, error) {
	gomodPath := filepath.Join(pathToProjectRoot, "go.mod")
	gomodBytes, err := ioutil.ReadFile(gomodPath)
	if err != nil {
		return "", err
	}
	gomodContent := string(gomodBytes)
	moduleIdx := strings.Index(gomodContent, "module ")
	newLineIdx := strings.Index(gomodContent, "\n")

	var moduleName string
	if moduleIdx >= 0 {
		if newLineIdx >= 0 {
			moduleName = strings.TrimSpace(gomodContent[moduleIdx+6 : newLineIdx])
			moduleName = strings.TrimSuffix(moduleName, "\r")
		} else {
			moduleName = strings.TrimSpace(gomodContent[moduleIdx+6:])
		}
	} else {
		return "", fmt.Errorf("can not get module path in `%s`", gomodPath)
	}
	return moduleName, nil
}
<commit_msg>Support to read go.mod recursively to get module path<commit_after>// Package common contains common utilities that are shared among other packages.
// See each sub-package for detail.
package common

import (
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"v2ray.com/core/common/errors"
)

//go:generate errorgen

var (
	// ErrNoClue is for the situation that existing information is not enough to make a decision. For example, Router may return this error when there is no suitable route.
	ErrNoClue = errors.New("not enough information for making a decision")
)

// Must panics if err is not nil.
func Must(err error) {
	if err != nil {
		panic(err)
	}
}

// Must2 panics if the second parameter is not nil, otherwise returns the first parameter.
func Must2(v interface{}, err error) interface{} {
	Must(err)
	return v
}

// Error2 returns the err from the 2nd parameter.
func Error2(v interface{}, err error) error {
	return err
}

// envFile returns the name of the Go environment configuration file.
// Copy from https://github.com/golang/go/blob/c4f2a9788a7be04daf931ac54382fbe2cb754938/src/cmd/go/internal/cfg/cfg.go#L150-L166
func envFile() (string, error) {
	if file := os.Getenv("GOENV"); file != "" {
		if file == "off" {
			return "", fmt.Errorf("GOENV=off")
		}
		return file, nil
	}
	dir, err := os.UserConfigDir()
	if err != nil {
		return "", err
	}
	if dir == "" {
		return "", fmt.Errorf("missing user-config dir")
	}
	return filepath.Join(dir, "go", "env"), nil
}

// GetRuntimeEnv returns the value of runtime environment variable,
// that is set by running following command: `go env -w key=value`.
func GetRuntimeEnv(key string) (string, error) {
	file, err := envFile()
	if err != nil {
		return "", err
	}
	if file == "" {
		return "", fmt.Errorf("missing runtime env file")
	}
	var data []byte
	var runtimeEnv string
	data, readErr := ioutil.ReadFile(file)
	if readErr != nil {
		return "", readErr
	}
	envStrings := strings.Split(string(data), "\n")
	for _, envItem := range envStrings {
		envItem = strings.TrimSuffix(envItem, "\r")
		envKeyValue := strings.Split(envItem, "=")
		if strings.EqualFold(strings.TrimSpace(envKeyValue[0]), key) {
			runtimeEnv = strings.TrimSpace(envKeyValue[1])
		}
	}
	return runtimeEnv, nil
}

// GetGOBIN returns GOBIN environment variable as a string. It will NOT be empty.
func GetGOBIN() string {
	// The one set by user explicitly by `export GOBIN=/path` or `env GOBIN=/path command`
	GOBIN := os.Getenv("GOBIN")
	if GOBIN == "" {
		var err error
		// The one set by user by running `go env -w GOBIN=/path`
		GOBIN, err = GetRuntimeEnv("GOBIN")
		if err != nil {
			// The default one that Golang uses
			return filepath.Join(build.Default.GOPATH, "bin")
		}
		if GOBIN == "" {
			return filepath.Join(build.Default.GOPATH, "bin")
		}
		return GOBIN
	}
	return GOBIN
}

// GetGOPATH returns GOPATH environment variable as a string. It will NOT be empty.
func GetGOPATH() string {
	// The one set by user explicitly by `export GOPATH=/path` or `env GOPATH=/path command`
	GOPATH := os.Getenv("GOPATH")
	if GOPATH == "" {
		var err error
		// The one set by user by running `go env -w GOPATH=/path`
		GOPATH, err = GetRuntimeEnv("GOPATH")
		if err != nil {
			// The default one that Golang uses
			return build.Default.GOPATH
		}
		if GOPATH == "" {
			return build.Default.GOPATH
		}
		return GOPATH
	}
	return GOPATH
}

// GetModuleName returns the value of module in `go.mod` file.
func GetModuleName(pathToProjectRoot string) (string, error) {
	var moduleName string
	loopPath := pathToProjectRoot
	for {
		if idx := strings.LastIndex(loopPath, string(filepath.Separator)); idx >= 0 {
			gomodPath := filepath.Join(loopPath, "go.mod")
			gomodBytes, err := ioutil.ReadFile(gomodPath)
			if err != nil {
				loopPath = loopPath[:idx]
				continue
			}

			gomodContent := string(gomodBytes)
			moduleIdx := strings.Index(gomodContent, "module ")
			newLineIdx := strings.Index(gomodContent, "\n")

			if moduleIdx >= 0 {
				if newLineIdx >= 0 {
					moduleName = strings.TrimSpace(gomodContent[moduleIdx+6 : newLineIdx])
					moduleName = strings.TrimSuffix(moduleName, "\r")
				} else {
					moduleName = strings.TrimSpace(gomodContent[moduleIdx+6:])
				}
				return moduleName, nil
			}
			return "", fmt.Errorf("can not get module path in `%s`", gomodPath)
		}
		break
	}
	return moduleName, fmt.Errorf("no `go.mod` file in every parent directory of `%s`", pathToProjectRoot)
}
<|endoftext|>


<commit_before>// Package common contains common types, constants and functions used over different demoinfocs packages.
// Some constants prefixes:
// MVPReason - the reason why someone got the MVP award.
// HG - HitGroup - where a bullet hit the player.
// EE - EquipmentElement - basically the weapon identifiers.
// RER - RoundEndReason - why the round ended (bomb exploded, defused, time ran out. . .).
// EC - EquipmentClass - type of equipment (pistol, smg, heavy. . .).
package common

import (
	"strings"
)

var eqNameToWeapon map[string]EquipmentElement

var eqElementToName map[EquipmentElement]string

func init() {
	initEqNameToWeapon()
	initEqEementToName()
}

func initEqNameToWeapon() {
	eqNameToWeapon = make(map[string]EquipmentElement)

	eqNameToWeapon["ak47"] = EqAK47
	eqNameToWeapon["aug"] = EqAUG
	eqNameToWeapon["awp"] = EqAWP
	eqNameToWeapon["bizon"] = EqBizon
	eqNameToWeapon["c4"] = EqBomb
	eqNameToWeapon["deagle"] = EqDeagle
	eqNameToWeapon["decoy"] = EqDecoy
	eqNameToWeapon["decoygrenade"] = EqDecoy
	eqNameToWeapon["decoyprojectile"] = EqDecoy
	eqNameToWeapon["elite"] = EqDualBarettas
	eqNameToWeapon["famas"] = EqFamas
	eqNameToWeapon["fiveseven"] = EqFiveSeven
	eqNameToWeapon["flashbang"] = EqFlash
	eqNameToWeapon["g3sg1"] = EqG3SG1
	eqNameToWeapon["galil"] = EqGalil
	eqNameToWeapon["galilar"] = EqGalil
	eqNameToWeapon["glock"] = EqGlock
	eqNameToWeapon["hegrenade"] = EqHE
	eqNameToWeapon["hkp2000"] = EqP2000
	eqNameToWeapon["incgrenade"] = EqIncendiary
	eqNameToWeapon["incendiarygrenade"] = EqIncendiary
	eqNameToWeapon["m249"] = EqM249
	eqNameToWeapon["m4a1"] = EqM4A4
	eqNameToWeapon["mac10"] = EqMac10
	eqNameToWeapon["mag7"] = EqSwag7
	eqNameToWeapon["molotov"] = EqMolotov
	eqNameToWeapon["molotovgrenade"] = EqMolotov
	eqNameToWeapon["molotovprojectile"] = EqMolotov
	eqNameToWeapon["mp7"] = EqMP7
	eqNameToWeapon["mp9"] = EqMP9
	eqNameToWeapon["negev"] = EqNegev
	eqNameToWeapon["nova"] = EqNova
	eqNameToWeapon["p250"] = EqP250
	eqNameToWeapon["p90"] = EqP90
	eqNameToWeapon["sawedoff"] = EqSawedOff
	eqNameToWeapon["scar20"] = EqScar20
	eqNameToWeapon["sg556"] = EqSG556
	eqNameToWeapon["smokegrenade"] = EqSmoke
	eqNameToWeapon["smokegrenadeprojectile"] = EqSmoke
	eqNameToWeapon["taser"] = EqZeus
	eqNameToWeapon["tec9"] = EqTec9
	eqNameToWeapon["ump45"] = EqUMP
	eqNameToWeapon["xm1014"] = EqXM1014
	eqNameToWeapon["m4a1_silencer"] = EqM4A1
	eqNameToWeapon["m4a1_silencer_off"] = EqM4A1
	eqNameToWeapon["cz75a"] = EqCZ
	eqNameToWeapon["usp"] = EqUSP
	eqNameToWeapon["usp_silencer"] = EqUSP
	eqNameToWeapon["usp_silencer_off"] = EqUSP
	eqNameToWeapon["world"] = EqWorld
	eqNameToWeapon["inferno"] = EqIncendiary
	eqNameToWeapon["revolver"] = EqRevolver
	eqNameToWeapon["vest"] = EqKevlar
	eqNameToWeapon["vesthelm"] = EqHelmet
	eqNameToWeapon["defuser"] = EqDefuseKit

	// These don't exist and / or used to crash the game with the give command
	eqNameToWeapon["scar17"] = EqUnknown
	eqNameToWeapon["sensorgrenade"] = EqUnknown
	eqNameToWeapon["mp5navy"] = EqUnknown
	eqNameToWeapon["p228"] = EqUnknown
	eqNameToWeapon["scout"] = EqUnknown
	eqNameToWeapon["sg550"] = EqUnknown
	eqNameToWeapon["sg552"] = EqUnknown // This one still crashes the game :)
	eqNameToWeapon["tmp"] = EqUnknown
	eqNameToWeapon["worldspawn"] = EqUnknown
}

func initEqEementToName() {
	eqElementToName = make(map[EquipmentElement]string)
	eqElementToName[EqAK47] = "AK-47"
	eqElementToName[EqAUG] = "AUG"
	eqElementToName[EqAWP] = "AWP"
	eqElementToName[EqBizon] = "PP-Bizon"
	eqElementToName[EqBomb] = "C4"
	eqElementToName[EqDeagle] = "Desert Eagle"
	eqElementToName[EqDecoy] = "Decoy Grenade"
	eqElementToName[EqDualBarettas] = "Dual Barettas"
	eqElementToName[EqFamas] = "FAMAS"
	eqElementToName[EqFiveSeven] = "Five-SeveN"
	eqElementToName[EqFlash] = "Flashbang"
	eqElementToName[EqG3SG1] = "G3SG1"
	eqElementToName[EqGalil] = "Galil AR"
	eqElementToName[EqGlock] = "Glock-18"
	eqElementToName[EqHE] = "HE Grenade"
	eqElementToName[EqP2000] = "P2000"
	eqElementToName[EqIncendiary] = "Incendiary Grenade"
	eqElementToName[EqM249] = "M249"
	eqElementToName[EqM4A4] = "M4A1"
	eqElementToName[EqMac10] = "MAC-10"
	eqElementToName[EqSwag7] = "MAG-7"
	eqElementToName[EqMolotov] = "Molotov"
	eqElementToName[EqMP7] = "MP7"
	eqElementToName[EqMP9] = "MP9"
	eqElementToName[EqNegev] = "Negev"
	eqElementToName[EqNova] = "Nova"
	eqElementToName[EqP250] = "p250"
	eqElementToName[EqP90] = "P90"
	eqElementToName[EqSawedOff] = "Sawed-Off"
	eqElementToName[EqScar20] = "SCAR-20"
	eqElementToName[EqSG553] = "SG 553"
	eqElementToName[EqSmoke] = "Smoke Grenade"
	eqElementToName[EqScout] = "SSG 08"
	eqElementToName[EqZeus] = "Zeus x27"
	eqElementToName[EqTec9] = "Tec-9"
	eqElementToName[EqUMP] = "UMP-45"
	eqElementToName[EqXM1014] = "XM1014"
	eqElementToName[EqM4A1] = "M4A1"
	eqElementToName[EqCZ] = "CZ75 Auto"
	eqElementToName[EqUSP] = "USP-S"
	eqElementToName[EqWorld] = "World"
	eqElementToName[EqRevolver] = "R8 Revolver"
	eqElementToName[EqKevlar] = "Kevlar Vest"
	eqElementToName[EqHelmet] = "Kevlar + Helmet"
	eqElementToName[EqDefuseKit] = "Defuse Kit"
	eqElementToName[EqKnife] = "Knife"
}

// MapEquipment creates an EquipmentElement from the name of the weapon / equipment.
func MapEquipment(eqName string) EquipmentElement {
	eqName = strings.TrimPrefix(eqName, weaponPrefix)

	var wep EquipmentElement
	if strings.Contains(eqName, "knife") || strings.Contains(eqName, "bayonet") {
		wep = EqKnife
	} else {
		// If the eqName isn't known it will be EqUnknown as that is the default value for EquipmentElement
		var ok bool
		if wep, ok = eqNameToWeapon[eqName]; !ok {
			// TODO: Return error / warning for unmapped weapons
		}
	}

	return wep
}
<commit_msg>common: fix typo initEqEementToName -> initEqElementToName.<commit_after>// Package common contains common types, constants and functions used over different demoinfocs packages.
// Some constants prefixes:
// MVPReason - the reason why someone got the MVP award.
// HG - HitGroup - where a bullet hit the player.
// EE - EquipmentElement - basically the weapon identifiers.
// RER - RoundEndReason - why the round ended (bomb exploded, defused, time ran out. . .).
// EC - EquipmentClass - type of equipment (pistol, smg, heavy. . .).
package common

import (
	"strings"
)

var eqNameToWeapon map[string]EquipmentElement

var eqElementToName map[EquipmentElement]string

func init() {
	initEqNameToWeapon()
	initEqElementToName()
}

func initEqNameToWeapon() {
	eqNameToWeapon = make(map[string]EquipmentElement)

	eqNameToWeapon["ak47"] = EqAK47
	eqNameToWeapon["aug"] = EqAUG
	eqNameToWeapon["awp"] = EqAWP
	eqNameToWeapon["bizon"] = EqBizon
	eqNameToWeapon["c4"] = EqBomb
	eqNameToWeapon["deagle"] = EqDeagle
	eqNameToWeapon["decoy"] = EqDecoy
	eqNameToWeapon["decoygrenade"] = EqDecoy
	eqNameToWeapon["decoyprojectile"] = EqDecoy
	eqNameToWeapon["elite"] = EqDualBarettas
	eqNameToWeapon["famas"] = EqFamas
	eqNameToWeapon["fiveseven"] = EqFiveSeven
	eqNameToWeapon["flashbang"] = EqFlash
	eqNameToWeapon["g3sg1"] = EqG3SG1
	eqNameToWeapon["galil"] = EqGalil
	eqNameToWeapon["galilar"] = EqGalil
	eqNameToWeapon["glock"] = EqGlock
	eqNameToWeapon["hegrenade"] = EqHE
	eqNameToWeapon["hkp2000"] = EqP2000
	eqNameToWeapon["incgrenade"] = EqIncendiary
	eqNameToWeapon["incendiarygrenade"] = EqIncendiary
	eqNameToWeapon["m249"] = EqM249
	eqNameToWeapon["m4a1"] = EqM4A4
	eqNameToWeapon["mac10"] = EqMac10
	eqNameToWeapon["mag7"] = EqSwag7
	eqNameToWeapon["molotov"] = EqMolotov
	eqNameToWeapon["molotovgrenade"] = EqMolotov
	eqNameToWeapon["molotovprojectile"] = EqMolotov
	eqNameToWeapon["mp7"] = EqMP7
	eqNameToWeapon["mp9"] = EqMP9
	eqNameToWeapon["negev"] = EqNegev
	eqNameToWeapon["nova"] = EqNova
	eqNameToWeapon["p250"] = EqP250
	eqNameToWeapon["p90"] = EqP90
	eqNameToWeapon["sawedoff"] = EqSawedOff
	eqNameToWeapon["scar20"] = EqScar20
	eqNameToWeapon["sg556"] = EqSG556
	eqNameToWeapon["smokegrenade"] = EqSmoke
	eqNameToWeapon["smokegrenadeprojectile"] = EqSmoke
	eqNameToWeapon["taser"] = EqZeus
	eqNameToWeapon["tec9"] = EqTec9
	eqNameToWeapon["ump45"] = EqUMP
	eqNameToWeapon["xm1014"] = EqXM1014
	eqNameToWeapon["m4a1_silencer"] = EqM4A1
	eqNameToWeapon["m4a1_silencer_off"] = EqM4A1
	eqNameToWeapon["cz75a"] = EqCZ
	eqNameToWeapon["usp"] = EqUSP
	eqNameToWeapon["usp_silencer"] = EqUSP
	eqNameToWeapon["usp_silencer_off"] = EqUSP
	eqNameToWeapon["world"] = EqWorld
	eqNameToWeapon["inferno"] = EqIncendiary
	eqNameToWeapon["revolver"] = EqRevolver
	eqNameToWeapon["vest"] = EqKevlar
	eqNameToWeapon["vesthelm"] = EqHelmet
	eqNameToWeapon["defuser"] = EqDefuseKit

	// These don't exist and / or used to crash the game with the give command
	eqNameToWeapon["scar17"] = EqUnknown
	eqNameToWeapon["sensorgrenade"] = EqUnknown
	eqNameToWeapon["mp5navy"] = EqUnknown
	eqNameToWeapon["p228"] = EqUnknown
	eqNameToWeapon["scout"] = EqUnknown
	eqNameToWeapon["sg550"] = EqUnknown
	eqNameToWeapon["sg552"] = EqUnknown // This one still crashes the game :)
	eqNameToWeapon["tmp"] = EqUnknown
	eqNameToWeapon["worldspawn"] = EqUnknown
}

func initEqElementToName() {
	eqElementToName = make(map[EquipmentElement]string)
	eqElementToName[EqAK47] = "AK-47"
	eqElementToName[EqAUG] = "AUG"
	eqElementToName[EqAWP] = "AWP"
	eqElementToName[EqBizon] = "PP-Bizon"
	eqElementToName[EqBomb] = "C4"
	eqElementToName[EqDeagle] = "Desert Eagle"
	eqElementToName[EqDecoy] = "Decoy Grenade"
	eqElementToName[EqDualBarettas] = "Dual Barettas"
	eqElementToName[EqFamas] = "FAMAS"
	eqElementToName[EqFiveSeven] = "Five-SeveN"
	eqElementToName[EqFlash] = "Flashbang"
	eqElementToName[EqG3SG1] = "G3SG1"
	eqElementToName[EqGalil] = "Galil AR"
	eqElementToName[EqGlock] = "Glock-18"
	eqElementToName[EqHE] = "HE Grenade"
	eqElementToName[EqP2000] = "P2000"
	eqElementToName[EqIncendiary] = "Incendiary Grenade"
	eqElementToName[EqM249] = "M249"
	eqElementToName[EqM4A4] = "M4A1"
	eqElementToName[EqMac10] = "MAC-10"
	eqElementToName[EqSwag7] = "MAG-7"
	eqElementToName[EqMolotov] = "Molotov"
	eqElementToName[EqMP7] = "MP7"
	eqElementToName[EqMP9] = "MP9"
	eqElementToName[EqNegev] = "Negev"
	eqElementToName[EqNova] = "Nova"
	eqElementToName[EqP250] = "p250"
	eqElementToName[EqP90] = "P90"
	eqElementToName[EqSawedOff] = "Sawed-Off"
	eqElementToName[EqScar20] = "SCAR-20"
	eqElementToName[EqSG553] = "SG 553"
	eqElementToName[EqSmoke] = "Smoke Grenade"
	eqElementToName[EqScout] = "SSG 08"
	eqElementToName[EqZeus] = "Zeus x27"
	eqElementToName[EqTec9] = "Tec-9"
	eqElementToName[EqUMP] = "UMP-45"
	eqElementToName[EqXM1014] = "XM1014"
	eqElementToName[EqM4A1] = "M4A1"
	eqElementToName[EqCZ] = "CZ75 Auto"
	eqElementToName[EqUSP] = "USP-S"
	eqElementToName[EqWorld] = "World"
	eqElementToName[EqRevolver] = "R8 Revolver"
	eqElementToName[EqKevlar] = "Kevlar Vest"
	eqElementToName[EqHelmet] = "Kevlar + Helmet"
	eqElementToName[EqDefuseKit] = "Defuse Kit"
	eqElementToName[EqKnife] = "Knife"
}

// MapEquipment creates an EquipmentElement from the name of the weapon / equipment.
func MapEquipment(eqName string) EquipmentElement {
	eqName = strings.TrimPrefix(eqName, weaponPrefix)

	var wep EquipmentElement
	if strings.Contains(eqName, "knife") || strings.Contains(eqName, "bayonet") {
		wep = EqKnife
	} else {
		// If the eqName isn't known it will be EqUnknown as that is the default value for EquipmentElement
		var ok bool
		if wep, ok = eqNameToWeapon[eqName]; !ok {
			// TODO: Return error / warning for unmapped weapons
		}
	}

	return wep
}
<|endoftext|>