<commit_before>package handlers

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/dom-bot/itchy-guacamole/deck"
	"github.com/dom-bot/itchy-guacamole/score"
	"github.com/dom-bot/itchy-guacamole/veto"
	"github.com/julienschmidt/httprouter"
)

var (
	rnd           = rand.New(rand.NewSource(time.Now().UnixNano()))
	availableSets deck.Sets
)

type makeDeckRequest struct {
	Sets            []deck.Set        `json:"sets"`
	Weights         score.Weights     `json:"weights"`
	VetoProbability *veto.Probability `json:"veto_probability,omitempty"`
}

type deckHardware struct {
	CoinTokens         bool `json:"coin_tokens"`
	VictoryTokens      bool `json:"victory_tokens"`
	MinusOneCardTokens bool `json:"minus_one_card_tokens"`
	MinusOneCoinTokens bool `json:"minus_one_coin_tokens"`
	JourneyTokens      bool `json:"journey_tokens"`
	TavernMats         bool `json:"tavern_mats"`
	TradeRouteMats     bool `json:"trade_route_mats"`
	NativeVillageMats  bool `json:"native_village_mats"`
}

type deckResponse struct {
	ID                   string       `json:"id"`
	Cards                []deck.Card  `json:"cards"`
	Events               []deck.Card  `json:"events"`
	ColoniesAndPlatinums bool         `json:"colonies_and_platinums"`
	Shelters             bool         `json:"shelters"`
	Potions              bool         `json:"potions"`
	Spoils               bool         `json:"spoils"`
	Ruins                bool         `json:"ruins"`
	Hardware             deckHardware `json:"hardware"`
}

func init() {
	setString := os.Getenv("SETS")
	if setString == "" {
		availableSets.Add(deck.Dominion)
	} else {
		exps := strings.Split(setString, `,`)

		for _, exp := range exps {
			availableSets.Add(deck.Set(exp))
		}
	}

	fmt.Printf("Using Sets: %+v\n", availableSets)
}

func getSets(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	sets := make([]string, 0, 10)

	setString := os.Getenv("SETS")
	if setString == "" {
		sets = append(sets, string(deck.Dominion))

	} else {
		sets = strings.Split(setString, `,`)
	}

	enc := json.NewEncoder(w)
	_ = enc.Encode(sets)

	w.Header().Set("Content-Type", "application/json")
}

func getDeck(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	idBase64 := ps.ByName("id")
	if idBase64 == "" {
		http.Error(w, "missing 'id'", http.StatusBadRequest)
		return
	}

	id, err := base64.URLEncoding.DecodeString(idBase64)
	if idBase64 == "" {
		http.Error(w, fmt.Sprintf("unable to decode ID: %s", err), http.StatusBadRequest)
		return
	}

	d, err := deck.NewDeckFromID(id)
	if idBase64 == "" {
		http.Error(w, fmt.Sprintf("unable to deserialize ID: %s", err), http.StatusBadRequest)
		return
	}

	resp := deckResponse{
		ID:                   base64.URLEncoding.EncodeToString(d.ID()),
		Cards:                d.Cards,
		Events:               d.Events,
		ColoniesAndPlatinums: d.ColoniesAndPlatinums,
		Shelters:             d.Shelters,
		Potions:              d.Potions(),
		Spoils:               d.Spoils(),
		Ruins:                d.Ruins(),
		Hardware: deckHardware{
			CoinTokens:         d.CoinTokens(),
			VictoryTokens:      d.VictoryTokens(),
			MinusOneCardTokens: d.MinusOneCardTokens(),
			MinusOneCoinTokens: d.MinusOneCoinTokens(),
			JourneyTokens:      d.JourneyTokens(),
			TavernMats:         d.TavernMats(),
			TradeRouteMats:     d.TradeRouteMats(),
			NativeVillageMats:  d.NativeVillageMats(),
		},
	}
	enc := json.NewEncoder(w)
	_ = enc.Encode(resp)

	w.Header().Set("Content-Type", "application/json")
}

func makeDeck(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	var (
		sets            = availableSets
		requestedSets   deck.Sets
		vetoProbability veto.Probability
		req             makeDeckRequest
		maxScore        uint
		d               deck.Deck
	)

	decoder := json.NewDecoder(r.Body)
	err := decoder.Decode(&req)
	if err != nil {
		http.Error(w, fmt.Sprintf("Error decoding JSON body: %s", err), http.StatusBadRequest)
		return
	}

	for _, set := range req.Sets {
		requestedSets.Add(set)
	}
	if !requestedSets.Empty() {
		sets.Intersect(requestedSets)
	}
	if sets.Empty() {
		http.Error(w, "Can't generate a deck from no sets", http.StatusBadRequest)
		return
	}

	if req.VetoProbability == nil {
		fmt.Println("Using default probs")
		// Defaults
		vetoProbability = veto.Probability{
			WhenTooExpensive:     0.9,
			WhenNoTrashing:       0.9,
			WhenNoChaining:       0.3,
			WhenTooManySets:      0.7,
			WhenTooManyMechanics: 0.8,
			WhenTooManyAttacks:   0.7,
		}
	} else {
		vetoProbability = *req.VetoProbability
	}

	count := 0
GenerateDeck:
	candidateDeck := deck.NewRandomDeck(sets)
	if veto.TooExpensive(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.NoTrashing(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.NoChaining(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.TooManySets(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.TooManyMechanics(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.TooManyAttacks(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	count++

	candidateScore := score.Evaluate(req.Weights, candidateDeck)
	if candidateScore > maxScore {
		d = candidateDeck
		maxScore = candidateScore
	}
	if count < 100 {
		goto GenerateDeck
	}

	resp := deckResponse{
		ID:                   base64.URLEncoding.EncodeToString(d.ID()),
		Cards:                d.Cards,
		Events:               d.Events,
		ColoniesAndPlatinums: d.ColoniesAndPlatinums,
		Shelters:             d.Shelters,
		Potions:              d.Potions(),
		Spoils:               d.Spoils(),
		Ruins:                d.Ruins(),
		Hardware: deckHardware{
			CoinTokens:         d.CoinTokens(),
			VictoryTokens:      d.VictoryTokens(),
			MinusOneCardTokens: d.MinusOneCardTokens(),
			MinusOneCoinTokens: d.MinusOneCoinTokens(),
			JourneyTokens:      d.JourneyTokens(),
			TavernMats:         d.TavernMats(),
			TradeRouteMats:     d.TradeRouteMats(),
			NativeVillageMats:  d.NativeVillageMats(),
		},
	}
	enc := json.NewEncoder(w)
	_ = enc.Encode(resp)

	w.Header().Set("Content-Type", "application/json")
}

func indexRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	contents, err := ioutil.ReadFile("app/public/index.html")
	if err != nil {
		log.Println(err)
	}

	fmt.Fprint(w, string(contents))
}
<commit_msg>Use defaults when zero-value too<commit_after>package handlers

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/dom-bot/itchy-guacamole/deck"
	"github.com/dom-bot/itchy-guacamole/score"
	"github.com/dom-bot/itchy-guacamole/veto"
	"github.com/julienschmidt/httprouter"
)

var (
	rnd           = rand.New(rand.NewSource(time.Now().UnixNano()))
	availableSets deck.Sets
)

type makeDeckRequest struct {
	Sets            []deck.Set        `json:"sets"`
	Weights         score.Weights     `json:"weights"`
	VetoProbability *veto.Probability `json:"veto_probability,omitempty"`
}

type deckHardware struct {
	CoinTokens         bool `json:"coin_tokens"`
	VictoryTokens      bool `json:"victory_tokens"`
	MinusOneCardTokens bool `json:"minus_one_card_tokens"`
	MinusOneCoinTokens bool `json:"minus_one_coin_tokens"`
	JourneyTokens      bool `json:"journey_tokens"`
	TavernMats         bool `json:"tavern_mats"`
	TradeRouteMats     bool `json:"trade_route_mats"`
	NativeVillageMats  bool `json:"native_village_mats"`
}

type deckResponse struct {
	ID                   string       `json:"id"`
	Cards                []deck.Card  `json:"cards"`
	Events               []deck.Card  `json:"events"`
	ColoniesAndPlatinums bool         `json:"colonies_and_platinums"`
	Shelters             bool         `json:"shelters"`
	Potions              bool         `json:"potions"`
	Spoils               bool         `json:"spoils"`
	Ruins                bool         `json:"ruins"`
	Hardware             deckHardware `json:"hardware"`
}

func init() {
	setString := os.Getenv("SETS")
	if setString == "" {
		availableSets.Add(deck.Dominion)
	} else {
		exps := strings.Split(setString, `,`)

		for _, exp := range exps {
			availableSets.Add(deck.Set(exp))
		}
	}

	fmt.Printf("Using Sets: %+v\n", availableSets)
}

func getSets(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	sets := make([]string, 0, 10)

	setString := os.Getenv("SETS")
	if setString == "" {
		sets = append(sets, string(deck.Dominion))

	} else {
		sets = strings.Split(setString, `,`)
	}

	enc := json.NewEncoder(w)
	_ = enc.Encode(sets)

	w.Header().Set("Content-Type", "application/json")
}

func getDeck(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	idBase64 := ps.ByName("id")
	if idBase64 == "" {
		http.Error(w, "missing 'id'", http.StatusBadRequest)
		return
	}

	id, err := base64.URLEncoding.DecodeString(idBase64)
	if idBase64 == "" {
		http.Error(w, fmt.Sprintf("unable to decode ID: %s", err), http.StatusBadRequest)
		return
	}

	d, err := deck.NewDeckFromID(id)
	if idBase64 == "" {
		http.Error(w, fmt.Sprintf("unable to deserialize ID: %s", err), http.StatusBadRequest)
		return
	}

	resp := deckResponse{
		ID:                   base64.URLEncoding.EncodeToString(d.ID()),
		Cards:                d.Cards,
		Events:               d.Events,
		ColoniesAndPlatinums: d.ColoniesAndPlatinums,
		Shelters:             d.Shelters,
		Potions:              d.Potions(),
		Spoils:               d.Spoils(),
		Ruins:                d.Ruins(),
		Hardware: deckHardware{
			CoinTokens:         d.CoinTokens(),
			VictoryTokens:      d.VictoryTokens(),
			MinusOneCardTokens: d.MinusOneCardTokens(),
			MinusOneCoinTokens: d.MinusOneCoinTokens(),
			JourneyTokens:      d.JourneyTokens(),
			TavernMats:         d.TavernMats(),
			TradeRouteMats:     d.TradeRouteMats(),
			NativeVillageMats:  d.NativeVillageMats(),
		},
	}
	enc := json.NewEncoder(w)
	_ = enc.Encode(resp)

	w.Header().Set("Content-Type", "application/json")
}

func makeDeck(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	var (
		sets            = availableSets
		requestedSets   deck.Sets
		vetoProbability veto.Probability
		req             makeDeckRequest
		maxScore        uint
		d               deck.Deck
	)

	decoder := json.NewDecoder(r.Body)
	err := decoder.Decode(&req)
	if err != nil {
		http.Error(w, fmt.Sprintf("Error decoding JSON body: %s", err), http.StatusBadRequest)
		return
	}

	for _, set := range req.Sets {
		requestedSets.Add(set)
	}
	if !requestedSets.Empty() {
		sets.Intersect(requestedSets)
	}
	if sets.Empty() {
		http.Error(w, "Can't generate a deck from no sets", http.StatusBadRequest)
		return
	}

	if req.VetoProbability == nil || (req.VetoProbability == veto.Probability{}) {
		fmt.Println("Using default probs")
		// Defaults
		vetoProbability = veto.Probability{
			WhenTooExpensive:     0.9,
			WhenNoTrashing:       0.9,
			WhenNoChaining:       0.3,
			WhenTooManySets:      0.7,
			WhenTooManyMechanics: 0.8,
			WhenTooManyAttacks:   0.7,
		}
	} else {
		vetoProbability = *req.VetoProbability
	}

	count := 0
GenerateDeck:
	candidateDeck := deck.NewRandomDeck(sets)
	if veto.TooExpensive(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.NoTrashing(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.NoChaining(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.TooManySets(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.TooManyMechanics(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	if veto.TooManyAttacks(vetoProbability, candidateDeck) {
		goto GenerateDeck
	}
	count++

	candidateScore := score.Evaluate(req.Weights, candidateDeck)
	if candidateScore > maxScore {
		d = candidateDeck
		maxScore = candidateScore
	}
	if count < 100 {
		goto GenerateDeck
	}

	resp := deckResponse{
		ID:                   base64.URLEncoding.EncodeToString(d.ID()),
		Cards:                d.Cards,
		Events:               d.Events,
		ColoniesAndPlatinums: d.ColoniesAndPlatinums,
		Shelters:             d.Shelters,
		Potions:              d.Potions(),
		Spoils:               d.Spoils(),
		Ruins:                d.Ruins(),
		Hardware: deckHardware{
			CoinTokens:         d.CoinTokens(),
			VictoryTokens:      d.VictoryTokens(),
			MinusOneCardTokens: d.MinusOneCardTokens(),
			MinusOneCoinTokens: d.MinusOneCoinTokens(),
			JourneyTokens:      d.JourneyTokens(),
			TavernMats:         d.TavernMats(),
			TradeRouteMats:     d.TradeRouteMats(),
			NativeVillageMats:  d.NativeVillageMats(),
		},
	}
	enc := json.NewEncoder(w)
	_ = enc.Encode(resp)

	w.Header().Set("Content-Type", "application/json")
}

func indexRoute(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	contents, err := ioutil.ReadFile("app/public/index.html")
	if err != nil {
		log.Println(err)
	}

	fmt.Fprint(w, string(contents))
}
<|endoftext|>


<commit_before>package lfs

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/rubyist/tracerx"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
)

const (
	mediaType = "application/vnd.git-lfs+json; charset-utf-8"
)

// The apiEvent* statuses (and the apiEvent channel) are used by
// UploadQueue to know when it is OK to process uploads concurrently.
const (
	apiEventFail = iota
	apiEventSuccess
)

var (
	lfsMediaTypeRE             = regexp.MustCompile(`\Aapplication/vnd\.git\-lfs\+json(;|\z)`)
	jsonMediaTypeRE            = regexp.MustCompile(`\Aapplication/json(;|\z)`)
	objectRelationDoesNotExist = errors.New("relation does not exist")
	hiddenHeaders              = map[string]bool{
		"Authorization": true,
	}

	defaultErrors = map[int]string{
		400: "Client error: %s",
		401: "Authorization error: %s\nCheck that you have proper access to the repository",
		403: "Authorization error: %s\nCheck that you have proper access to the repository",
		404: "Repository or object not found: %s\nCheck that it exists and that you have proper access to it",
		500: "Server error: %s",
	}

	apiEvent = make(chan int)
)

type objectResource struct {
	Oid   string                   `json:"oid,omitempty"`
	Size  int64                    `json:"size,omitempty"`
	Links map[string]*linkRelation `json:"_links,omitempty"`
}

func (o *objectResource) NewRequest(relation, method string) (*http.Request, Creds, error) {
	rel, ok := o.Rel(relation)
	if !ok {
		return nil, nil, objectRelationDoesNotExist
	}

	req, creds, err := newClientRequest(method, rel.Href)
	if err != nil {
		return nil, nil, err
	}

	for h, v := range rel.Header {
		req.Header.Set(h, v)
	}

	return req, creds, nil
}

func (o *objectResource) Rel(name string) (*linkRelation, bool) {
	if o.Links == nil {
		return nil, false
	}

	rel, ok := o.Links[name]
	return rel, ok
}

type linkRelation struct {
	Href   string            `json:"href"`
	Header map[string]string `json:"header,omitempty"`
}

type ClientError struct {
	Message          string `json:"message"`
	DocumentationUrl string `json:"documentation_url,omitempty"`
	RequestId        string `json:"request_id,omitempty"`
}

func (e *ClientError) Error() string {
	msg := e.Message
	if len(e.DocumentationUrl) > 0 {
		msg += "\nDocs: " + e.DocumentationUrl
	}
	if len(e.RequestId) > 0 {
		msg += "\nRequest ID: " + e.RequestId
	}
	return msg
}

func Download(oid string) (io.ReadCloser, int64, *WrappedError) {
	req, creds, err := newApiRequest("GET", oid)
	if err != nil {
		return nil, 0, Error(err)
	}

	res, obj, wErr := doApiRequest(req, creds)
	if wErr != nil {
		return nil, 0, wErr
	}

	req, creds, err = obj.NewRequest("download", "GET")
	if err != nil {
		return nil, 0, Error(err)
	}

	res, wErr = doHttpRequest(req, creds)
	if wErr != nil {
		return nil, 0, wErr
	}

	return res.Body, res.ContentLength, nil
}

type byteCloser struct {
	*bytes.Reader
}

func (b *byteCloser) Close() error {
	return nil
}

func Upload(oidPath, filename string, cb CopyCallback) *WrappedError {
	oid := filepath.Base(oidPath)
	file, err := os.Open(oidPath)
	if err != nil {
		return Error(err)
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return Error(err)
	}

	reqObj := &objectResource{
		Oid:  oid,
		Size: stat.Size(),
	}

	by, err := json.Marshal(reqObj)
	if err != nil {
		return Error(err)
	}

	req, creds, err := newApiRequest("POST", oid)
	if err != nil {
		return Error(err)
	}

	req.Header.Set("Content-Type", mediaType)
	req.Header.Set("Content-Length", strconv.Itoa(len(by)))
	req.ContentLength = int64(len(by))
	req.Body = &byteCloser{bytes.NewReader(by)}

	tracerx.Printf("api: uploading %s (%s)", filename, oid)
	res, obj, wErr := doApiRequest(req, creds)
	if wErr != nil {
		sendApiEvent(apiEventFail)
		return wErr
	}

	sendApiEvent(apiEventSuccess)

	reader := &CallbackReader{
		C:         cb,
		TotalSize: reqObj.Size,
		Reader:    file,
	}

	if res.StatusCode == 200 {
		// Drain the reader to update any progress bars
		io.Copy(ioutil.Discard, reader)
		return nil
	}

	req, creds, err = obj.NewRequest("upload", "PUT")
	if err != nil {
		return Error(err)
	}

	if len(req.Header.Get("Content-Type")) == 0 {
		req.Header.Set("Content-Type", "application/octet-stream")
	}
	req.Header.Set("Content-Length", strconv.FormatInt(reqObj.Size, 10))
	req.ContentLength = reqObj.Size

	req.Body = ioutil.NopCloser(reader)

	res, wErr = doHttpRequest(req, creds)
	if wErr != nil {
		return wErr
	}

	if res.StatusCode > 299 {
		return Errorf(nil, "Invalid status for %s %s: %d", req.Method, req.URL, res.StatusCode)
	}

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	req, creds, err = obj.NewRequest("verify", "POST")
	if err == objectRelationDoesNotExist {
		return nil
	} else if err != nil {
		return Error(err)
	}

	req.Header.Set("Content-Type", mediaType)
	req.Header.Set("Content-Length", strconv.Itoa(len(by)))
	req.ContentLength = int64(len(by))
	req.Body = ioutil.NopCloser(bytes.NewReader(by))
	res, wErr = doHttpRequest(req, creds)

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	return wErr
}

func doHttpRequest(req *http.Request, creds Creds) (*http.Response, *WrappedError) {
	res, err := DoHTTP(Config, req)

	var wErr *WrappedError

	if err != nil {
		wErr = Errorf(err, "Error for %s %s", res.Request.Method, res.Request.URL)
	} else {
		if creds != nil {
			saveCredentials(creds, res)
		}

		wErr = handleResponse(res)
	}

	if wErr != nil {
		if res != nil {
			setErrorResponseContext(wErr, res)
		} else {
			setErrorRequestContext(wErr, req)
		}
	}

	return res, wErr
}

func doApiRequestWithRedirects(req *http.Request, creds Creds, via []*http.Request) (*http.Response, *WrappedError) {
	res, wErr := doHttpRequest(req, creds)
	if wErr != nil {
		return nil, wErr
	}

	if res.StatusCode == 307 {
		redirectedReq, redirectedCreds, err := newClientRequest(req.Method, res.Header.Get("Location"))
		if err != nil {
			return nil, Errorf(err, err.Error())
		}

		via = append(via, req)
		if seeker, ok := req.Body.(io.Seeker); ok {
			if _, err := seeker.Seek(0, 0); err != nil {
				return nil, Error(err)
			}
			redirectedReq.Body = req.Body
			redirectedReq.ContentLength = req.ContentLength
		} else {
			return nil, Errorf(nil, "Request body needs to be an io.Seeker to handle redirects.")
		}

		if err = checkRedirect(redirectedReq, via); err != nil {
			return nil, Errorf(err, err.Error())
		}

		return doApiRequestWithRedirects(redirectedReq, redirectedCreds, via)
	}

	return res, nil
}

func doApiRequest(req *http.Request, creds Creds) (*http.Response, *objectResource, *WrappedError) {
	via := make([]*http.Request, 0, 4)
	res, wErr := doApiRequestWithRedirects(req, creds, via)
	if wErr != nil {
		return nil, nil, wErr
	}

	obj := &objectResource{}
	wErr = decodeApiResponse(res, obj)

	if wErr != nil {
		setErrorResponseContext(wErr, res)
		return nil, nil, wErr
	}

	return res, obj, nil
}

func handleResponse(res *http.Response) *WrappedError {
	if res.StatusCode < 400 {
		return nil
	}

	defer func() {
		io.Copy(ioutil.Discard, res.Body)
		res.Body.Close()
	}()

	cliErr := &ClientError{}
	wErr := decodeApiResponse(res, cliErr)
	if wErr == nil {
		if len(cliErr.Message) == 0 {
			wErr = defaultError(res)
		} else {
			wErr = Error(cliErr)
		}
	}

	wErr.Panic = res.StatusCode > 499 && res.StatusCode != 501 && res.StatusCode != 509
	return wErr
}

func decodeApiResponse(res *http.Response, obj interface{}) *WrappedError {
	ctype := res.Header.Get("Content-Type")
	if !(lfsMediaTypeRE.MatchString(ctype) || jsonMediaTypeRE.MatchString(ctype)) {
		return nil
	}

	err := json.NewDecoder(res.Body).Decode(obj)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	if err != nil {
		return Errorf(err, "Unable to parse HTTP response for %s %s", res.Request.Method, res.Request.URL)
	}

	return nil
}

func defaultError(res *http.Response) *WrappedError {
	var msgFmt string

	if f, ok := defaultErrors[res.StatusCode]; ok {
		msgFmt = f
	} else if res.StatusCode < 500 {
		msgFmt = defaultErrors[400] + fmt.Sprintf(" from HTTP %d", res.StatusCode)
	} else {
		msgFmt = defaultErrors[500] + fmt.Sprintf(" from HTTP %d", res.StatusCode)
	}

	return Error(fmt.Errorf(msgFmt, res.Request.URL))
}

func saveCredentials(creds Creds, res *http.Response) {
	if creds == nil {
		return
	}

	if res.StatusCode < 300 {
		execCreds(creds, "approve")
	} else if res.StatusCode == 401 {
		execCreds(creds, "reject")
	}
}

func newApiRequest(method, oid string) (*http.Request, Creds, error) {
	endpoint := Config.Endpoint()
	objectOid := oid
	operation := "download"
	if method == "POST" {
		objectOid = ""
		operation = "upload"
	}

	res, err := sshAuthenticate(endpoint, operation, oid)
	if err != nil {
		tracerx.Printf("ssh: attempted with %s.  Error: %s",
			endpoint.SshUserAndHost, err.Error(),
		)
	}

	if len(res.Href) > 0 {
		endpoint.Url = res.Href
	}

	u, err := ObjectUrl(endpoint, objectOid)
	if err != nil {
		return nil, nil, err
	}

	req, creds, err := newClientRequest(method, u.String())
	if err == nil {
		req.Header.Set("Accept", mediaType)
		if res.Header != nil {
			for key, value := range res.Header {
				req.Header.Set(key, value)
			}
		}
	}
	return req, creds, err
}

func newClientRequest(method, rawurl string) (*http.Request, Creds, error) {
	req, err := http.NewRequest(method, rawurl, nil)
	if err != nil {
		return req, nil, err
	}

	req.Header.Set("User-Agent", UserAgent)
	creds, err := getCreds(req)
	return req, creds, err
}

func getCreds(req *http.Request) (Creds, error) {
	if len(req.Header.Get("Authorization")) > 0 {
		return nil, nil
	}

	apiUrl, err := Config.ObjectUrl("")
	if err != nil {
		return nil, err
	}

	if req.URL.Scheme == apiUrl.Scheme &&
		req.URL.Host == apiUrl.Host {
		creds, err := credentials(req.URL)
		if err != nil {
			return nil, err
		}

		token := fmt.Sprintf("%s:%s", creds["username"], creds["password"])
		auth := "Basic " + base64.URLEncoding.EncodeToString([]byte(token))
		req.Header.Set("Authorization", auth)
		return creds, nil
	}

	return nil, nil
}

func setErrorRequestContext(err *WrappedError, req *http.Request) {
	err.Set("Endpoint", Config.Endpoint().Url)
	err.Set("URL", fmt.Sprintf("%s %s", req.Method, req.URL.String()))
	setErrorHeaderContext(err, "Response", req.Header)
}

func setErrorResponseContext(err *WrappedError, res *http.Response) {
	err.Set("Status", res.Status)
	setErrorHeaderContext(err, "Request", res.Header)
	setErrorRequestContext(err, res.Request)
}

func setErrorHeaderContext(err *WrappedError, prefix string, head http.Header) {
	for key, _ := range head {
		contextKey := fmt.Sprintf("%s:%s", prefix, key)
		if _, skip := hiddenHeaders[key]; skip {
			err.Set(contextKey, "--")
		} else {
			err.Set(contextKey, head.Get(key))
		}
	}
}

func sendApiEvent(event int) {
	select {
	case apiEvent <- event:
	default:
	}
}
<commit_msg>Decreased nesting<commit_after>package lfs

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/rubyist/tracerx"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
)

const (
	mediaType = "application/vnd.git-lfs+json; charset-utf-8"
)

// The apiEvent* statuses (and the apiEvent channel) are used by
// UploadQueue to know when it is OK to process uploads concurrently.
const (
	apiEventFail = iota
	apiEventSuccess
)

var (
	lfsMediaTypeRE             = regexp.MustCompile(`\Aapplication/vnd\.git\-lfs\+json(;|\z)`)
	jsonMediaTypeRE            = regexp.MustCompile(`\Aapplication/json(;|\z)`)
	objectRelationDoesNotExist = errors.New("relation does not exist")
	hiddenHeaders              = map[string]bool{
		"Authorization": true,
	}

	defaultErrors = map[int]string{
		400: "Client error: %s",
		401: "Authorization error: %s\nCheck that you have proper access to the repository",
		403: "Authorization error: %s\nCheck that you have proper access to the repository",
		404: "Repository or object not found: %s\nCheck that it exists and that you have proper access to it",
		500: "Server error: %s",
	}

	apiEvent = make(chan int)
)

type objectResource struct {
	Oid   string                   `json:"oid,omitempty"`
	Size  int64                    `json:"size,omitempty"`
	Links map[string]*linkRelation `json:"_links,omitempty"`
}

func (o *objectResource) NewRequest(relation, method string) (*http.Request, Creds, error) {
	rel, ok := o.Rel(relation)
	if !ok {
		return nil, nil, objectRelationDoesNotExist
	}

	req, creds, err := newClientRequest(method, rel.Href)
	if err != nil {
		return nil, nil, err
	}

	for h, v := range rel.Header {
		req.Header.Set(h, v)
	}

	return req, creds, nil
}

func (o *objectResource) Rel(name string) (*linkRelation, bool) {
	if o.Links == nil {
		return nil, false
	}

	rel, ok := o.Links[name]
	return rel, ok
}

type linkRelation struct {
	Href   string            `json:"href"`
	Header map[string]string `json:"header,omitempty"`
}

type ClientError struct {
	Message          string `json:"message"`
	DocumentationUrl string `json:"documentation_url,omitempty"`
	RequestId        string `json:"request_id,omitempty"`
}

func (e *ClientError) Error() string {
	msg := e.Message
	if len(e.DocumentationUrl) > 0 {
		msg += "\nDocs: " + e.DocumentationUrl
	}
	if len(e.RequestId) > 0 {
		msg += "\nRequest ID: " + e.RequestId
	}
	return msg
}

func Download(oid string) (io.ReadCloser, int64, *WrappedError) {
	req, creds, err := newApiRequest("GET", oid)
	if err != nil {
		return nil, 0, Error(err)
	}

	res, obj, wErr := doApiRequest(req, creds)
	if wErr != nil {
		return nil, 0, wErr
	}

	req, creds, err = obj.NewRequest("download", "GET")
	if err != nil {
		return nil, 0, Error(err)
	}

	res, wErr = doHttpRequest(req, creds)
	if wErr != nil {
		return nil, 0, wErr
	}

	return res.Body, res.ContentLength, nil
}

type byteCloser struct {
	*bytes.Reader
}

func (b *byteCloser) Close() error {
	return nil
}

func Upload(oidPath, filename string, cb CopyCallback) *WrappedError {
	oid := filepath.Base(oidPath)
	file, err := os.Open(oidPath)
	if err != nil {
		return Error(err)
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return Error(err)
	}

	reqObj := &objectResource{
		Oid:  oid,
		Size: stat.Size(),
	}

	by, err := json.Marshal(reqObj)
	if err != nil {
		return Error(err)
	}

	req, creds, err := newApiRequest("POST", oid)
	if err != nil {
		return Error(err)
	}

	req.Header.Set("Content-Type", mediaType)
	req.Header.Set("Content-Length", strconv.Itoa(len(by)))
	req.ContentLength = int64(len(by))
	req.Body = &byteCloser{bytes.NewReader(by)}

	tracerx.Printf("api: uploading %s (%s)", filename, oid)
	res, obj, wErr := doApiRequest(req, creds)
	if wErr != nil {
		sendApiEvent(apiEventFail)
		return wErr
	}

	sendApiEvent(apiEventSuccess)

	reader := &CallbackReader{
		C:         cb,
		TotalSize: reqObj.Size,
		Reader:    file,
	}

	if res.StatusCode == 200 {
		// Drain the reader to update any progress bars
		io.Copy(ioutil.Discard, reader)
		return nil
	}

	req, creds, err = obj.NewRequest("upload", "PUT")
	if err != nil {
		return Error(err)
	}

	if len(req.Header.Get("Content-Type")) == 0 {
		req.Header.Set("Content-Type", "application/octet-stream")
	}
	req.Header.Set("Content-Length", strconv.FormatInt(reqObj.Size, 10))
	req.ContentLength = reqObj.Size

	req.Body = ioutil.NopCloser(reader)

	res, wErr = doHttpRequest(req, creds)
	if wErr != nil {
		return wErr
	}

	if res.StatusCode > 299 {
		return Errorf(nil, "Invalid status for %s %s: %d", req.Method, req.URL, res.StatusCode)
	}

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	req, creds, err = obj.NewRequest("verify", "POST")
	if err == objectRelationDoesNotExist {
		return nil
	} else if err != nil {
		return Error(err)
	}

	req.Header.Set("Content-Type", mediaType)
	req.Header.Set("Content-Length", strconv.Itoa(len(by)))
	req.ContentLength = int64(len(by))
	req.Body = ioutil.NopCloser(bytes.NewReader(by))
	res, wErr = doHttpRequest(req, creds)

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	return wErr
}

func doHttpRequest(req *http.Request, creds Creds) (*http.Response, *WrappedError) {
	res, err := DoHTTP(Config, req)

	var wErr *WrappedError

	if err != nil {
		wErr = Errorf(err, "Error for %s %s", res.Request.Method, res.Request.URL)
	} else {
		if creds != nil {
			saveCredentials(creds, res)
		}

		wErr = handleResponse(res)
	}

	if wErr != nil {
		if res != nil {
			setErrorResponseContext(wErr, res)
		} else {
			setErrorRequestContext(wErr, req)
		}
	}

	return res, wErr
}

func doApiRequestWithRedirects(req *http.Request, creds Creds, via []*http.Request) (*http.Response, *WrappedError) {
	res, wErr := doHttpRequest(req, creds)
	if wErr != nil {
		return nil, wErr
	}

	if res.StatusCode == 307 {
		redirectedReq, redirectedCreds, err := newClientRequest(req.Method, res.Header.Get("Location"))
		if err != nil {
			return nil, Errorf(err, err.Error())
		}

		via = append(via, req)
		seeker, ok := req.Body.(io.Seeker)
		if !ok {
			return nil, Errorf(nil, "Request body needs to be an io.Seeker to handle redirects.")
		}

		if _, err := seeker.Seek(0, 0); err != nil {
			return nil, Error(err)
		}
		redirectedReq.Body = req.Body
		redirectedReq.ContentLength = req.ContentLength

		if err = checkRedirect(redirectedReq, via); err != nil {
			return nil, Errorf(err, err.Error())
		}

		return doApiRequestWithRedirects(redirectedReq, redirectedCreds, via)
	}

	return res, nil
}

func doApiRequest(req *http.Request, creds Creds) (*http.Response, *objectResource, *WrappedError) {
	via := make([]*http.Request, 0, 4)
	res, wErr := doApiRequestWithRedirects(req, creds, via)
	if wErr != nil {
		return nil, nil, wErr
	}

	obj := &objectResource{}
	wErr = decodeApiResponse(res, obj)

	if wErr != nil {
		setErrorResponseContext(wErr, res)
		return nil, nil, wErr
	}

	return res, obj, nil
}

func handleResponse(res *http.Response) *WrappedError {
	if res.StatusCode < 400 {
		return nil
	}

	defer func() {
		io.Copy(ioutil.Discard, res.Body)
		res.Body.Close()
	}()

	cliErr := &ClientError{}
	wErr := decodeApiResponse(res, cliErr)
	if wErr == nil {
		if len(cliErr.Message) == 0 {
			wErr = defaultError(res)
		} else {
			wErr = Error(cliErr)
		}
	}

	wErr.Panic = res.StatusCode > 499 && res.StatusCode != 501 && res.StatusCode != 509
	return wErr
}

func decodeApiResponse(res *http.Response, obj interface{}) *WrappedError {
	ctype := res.Header.Get("Content-Type")
	if !(lfsMediaTypeRE.MatchString(ctype) || jsonMediaTypeRE.MatchString(ctype)) {
		return nil
	}

	err := json.NewDecoder(res.Body).Decode(obj)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	if err != nil {
		return Errorf(err, "Unable to parse HTTP response for %s %s", res.Request.Method, res.Request.URL)
	}

	return nil
}

func defaultError(res *http.Response) *WrappedError {
	var msgFmt string

	if f, ok := defaultErrors[res.StatusCode]; ok {
		msgFmt = f
	} else if res.StatusCode < 500 {
		msgFmt = defaultErrors[400] + fmt.Sprintf(" from HTTP %d", res.StatusCode)
	} else {
		msgFmt = defaultErrors[500] + fmt.Sprintf(" from HTTP %d", res.StatusCode)
	}

	return Error(fmt.Errorf(msgFmt, res.Request.URL))
}

func saveCredentials(creds Creds, res *http.Response) {
	if creds == nil {
		return
	}

	if res.StatusCode < 300 {
		execCreds(creds, "approve")
	} else if res.StatusCode == 401 {
		execCreds(creds, "reject")
	}
}

func newApiRequest(method, oid string) (*http.Request, Creds, error) {
	endpoint := Config.Endpoint()
	objectOid := oid
	operation := "download"
	if method == "POST" {
		objectOid = ""
		operation = "upload"
	}

	res, err := sshAuthenticate(endpoint, operation, oid)
	if err != nil {
		tracerx.Printf("ssh: attempted with %s.  Error: %s",
			endpoint.SshUserAndHost, err.Error(),
		)
	}

	if len(res.Href) > 0 {
		endpoint.Url = res.Href
	}

	u, err := ObjectUrl(endpoint, objectOid)
	if err != nil {
		return nil, nil, err
	}

	req, creds, err := newClientRequest(method, u.String())
	if err == nil {
		req.Header.Set("Accept", mediaType)
		if res.Header != nil {
			for key, value := range res.Header {
				req.Header.Set(key, value)
			}
		}
	}
	return req, creds, err
}

func newClientRequest(method, rawurl string) (*http.Request, Creds, error) {
	req, err := http.NewRequest(method, rawurl, nil)
	if err != nil {
		return req, nil, err
	}

	req.Header.Set("User-Agent", UserAgent)
	creds, err := getCreds(req)
	return req, creds, err
}

func getCreds(req *http.Request) (Creds, error) {
	if len(req.Header.Get("Authorization")) > 0 {
		return nil, nil
	}

	apiUrl, err := Config.ObjectUrl("")
	if err != nil {
		return nil, err
	}

	if req.URL.Scheme == apiUrl.Scheme &&
		req.URL.Host == apiUrl.Host {
		creds, err := credentials(req.URL)
		if err != nil {
			return nil, err
		}

		token := fmt.Sprintf("%s:%s", creds["username"], creds["password"])
		auth := "Basic " + base64.URLEncoding.EncodeToString([]byte(token))
		req.Header.Set("Authorization", auth)
		return creds, nil
	}

	return nil, nil
}

func setErrorRequestContext(err *WrappedError, req *http.Request) {
	err.Set("Endpoint", Config.Endpoint().Url)
	err.Set("URL", fmt.Sprintf("%s %s", req.Method, req.URL.String()))
	setErrorHeaderContext(err, "Response", req.Header)
}

func setErrorResponseContext(err *WrappedError, res *http.Response) {
	err.Set("Status", res.Status)
	setErrorHeaderContext(err, "Request", res.Header)
	setErrorRequestContext(err, res.Request)
}

func setErrorHeaderContext(err *WrappedError, prefix string, head http.Header) {
	for key, _ := range head {
		contextKey := fmt.Sprintf("%s:%s", prefix, key)
		if _, skip := hiddenHeaders[key]; skip {
			err.Set(contextKey, "--")
		} else {
			err.Set(contextKey, head.Get(key))
		}
	}
}

func sendApiEvent(event int) {
	select {
	case apiEvent <- event:
	default:
	}
}
<|endoftext|>


<commit_before>package uploads

import (
	"time"

	"math"

	"fmt"

	r "github.com/dancannon/gorethink"
	"github.com/materials-commons/gohandy/file"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/pkg/domain"
	"github.com/willf/bitset"
)

var _ = fmt.Println

// A IDRequest requests a new upload id be created for
// the given parameters.
type IDRequest struct {
	User        string
	DirectoryID string
	ProjectID   string
	FileName    string
	FileSize    int64
	Checksum    string
	ChunkSize   int32
	FileMTime   time.Time
	Host        string
	Birthtime   time.Time
}

// IDService creates new upload requests
type IDService interface {
	ID(req IDRequest) (*schema.Upload, error)
	Delete(requestID, user string) error
	UploadsForProject(projectID, user string) ([]schema.Upload, error)
}

// idService implements the IDService interface using
// the dai services.
type idService struct {
	dirs        dai.Dirs
	projects    dai.Projects
	uploads     dai.Uploads
	files       dai.Files
	access      domain.Access
	fops        file.Operations
	tracker     *blockTracker
	requestPath requestPath
}

// NewIDServiceUsingSession creates a new idService that connects to the database using
// the given session.
func NewIDService(session *r.Session) *idService {
	access := domain.NewAccess(dai.NewRProjects(session), dai.NewRFiles(session), dai.NewRUsers(session))
	return &idService{
		dirs:        dai.NewRDirs(session),
		projects:    dai.NewRProjects(session),
		uploads:     dai.NewRUploads(session),
		files:       dai.NewRFiles(session),
		access:      access,
		fops:        file.OS,
		tracker:     requestBlockTracker,
		requestPath: &mcdirRequestPath{},
	}
}

// ID will create a new Upload request or return an existing one.
func (s *idService) ID(req IDRequest) (*schema.Upload, error) {
	var (
		err  error
		proj *schema.Project
		dir  *schema.Directory
	)

	// Check that project exists and user has access
	if proj, err = s.getProjectValidatingAccess(req.ProjectID, req.User); err != nil {
		return nil, err
	}

	// Check that directory exists and user has access
	if dir, err = s.getDirectoryValidatingAccess(req.DirectoryID, req.ProjectID, req.User); err != nil {
		return nil, err
	}

	return s.createUploadRequest(req, proj, dir)
}

// createUploadRequest will create an upload. It checks if there is a matching file or an existing upload.
// If it finds a matching file it returns a finished upload. If it finds an existing upload it returns it,
// otherwise it returns a new upload.
func (s *idService) createUploadRequest(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	upload, err := s.findExisting(req, proj, dir)
	switch {
	case err == app.ErrNotFound:
		return s.createNewUploadRequest(req, proj, dir)
	case err != nil:
		return nil, err
	default:
		return upload, nil
	}
}

// findExisting checks if there is an outstanding upload request, or a file that already matches
// the upload request. If it finds a match it will return the corresponding upload request. In
// the case of an already uploaded file it will return an upload request with all blocks marked
// as uploaded.
func (s *idService) findExisting(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	if _, err := s.files.ByChecksum(req.Checksum); err == nil {
		return s.createFinishedUpload(req, proj, dir)
	}
	return s.findMatchingUploadRequest(req)
}

// createFinishedUpload will create an upload entry with all blocks marked as uploaded.
func (s *idService) createFinishedUpload(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	// Create a new upload request and then set all blocks as already uploaded.
	if upload, err := s.createNewUploadRequest(req, proj, dir); err != nil {
		return nil, err
	} else {
		s.tracker.markAllBlocks(upload.ID)
		s.tracker.setIsExistingFile(upload.ID, true)
		upload.SetFBlocks(s.tracker.getBlocks(upload.ID))
		return upload, nil
	}
}

// findMatchingUploadRequest will search the set of upload requests and see if there
// is already an outstanding upload request that matches this one.
func (s *idService) findMatchingUploadRequest(req IDRequest) (*schema.Upload, error) {
	searchParams := dai.UploadSearch{
		ProjectID:   req.ProjectID,
		DirectoryID: req.DirectoryID,
		FileName:    req.FileName,
		Checksum:    req.Checksum,
	}

	if existingUpload, err := s.uploads.Search(searchParams); err == nil {
		// Found existing
		if bset := s.tracker.getBlocks(existingUpload.ID); bset != nil {
			existingUpload.File.Blocks = bset
		}
		return existingUpload, nil
	}
	return nil, app.ErrNotFound
}

// createNewUploadRequest creates a new upload request and inserts it into the database.
func (s *idService) createNewUploadRequest(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	n := uint(numBlocks(req.FileSize, req.ChunkSize))
	upload := schema.CUpload().
		Owner(req.User).
		Project(req.ProjectID, proj.Name).
		ProjectOwner(proj.Owner).
		Directory(req.DirectoryID, dir.Name).
		Host(req.Host).
		FName(req.FileName).
		FSize(req.FileSize).
		FChecksum(req.Checksum).
		FRemoteMTime(req.FileMTime).
		FBlocks(bitset.New(n)).
		Create()
	u, err := s.uploads.Insert(&upload)
	if err != nil {
		fmt.Println("insert err", err)
		return nil, err
	}

	if err := s.initUpload(u.ID, req.FileSize, req.ChunkSize); err != nil {
		fmt.Println("initUpload err", err)
		s.uploads.Delete(u.ID)
		return nil, err
	}
	return u, nil
}

// getProjectValidatingAccess retrieves the project with the given projectID. It checks that the
// given user has access to that project.
func (s *idService) getProjectValidatingAccess(projectID, user string) (*schema.Project, error) {
	project, err := s.projects.ByID(projectID)
	switch {
	case err != nil:
		return nil, err
	case !s.access.AllowedByOwner(projectID, user):
		return nil, app.ErrNoAccess
	default:
		return project, nil
	}
}

// getDirectoryValidatingAccess retrieves the directory with the given directoryID. It checks access to the
// directory and validates that the directory exists in the given project.
func (s *idService) getDirectoryValidatingAccess(directoryID, projectID, user string) (*schema.Directory, error) {
	dir, err := s.dirs.ByID(directoryID)
	switch {
	case err != nil:
		return nil, err
	case !s.projects.HasDirectory(projectID, directoryID):
		return nil, app.ErrInvalid
	case !s.access.AllowedByOwner(projectID, user):
		return nil, app.ErrNoAccess
	default:
		return dir, nil
	}
}

// initUpload
func (s *idService) initUpload(id string, fileSize int64, chunkSize int32) error {
	if err := s.requestPath.mkdirFromID(id); err != nil {
		return err
	}

	s.tracker.load(id, numBlocks(fileSize, chunkSize))
	return nil
}

// numBlocks
func numBlocks(fileSize int64, chunkSize int32) int {
	// round up to nearest number of blocks
	d := float64(fileSize) / float64(chunkSize)
	n := int(math.Ceil(d))
	return n
}

// Delete will delete the given requestID if the user has access
// to delete that request. Owners of requests can delete their
// own requests. Project owners can delete any request, even if
// they don't own the request.
func (s *idService) Delete(requestID, user string) error {
	upload, err := s.uploads.ByID(requestID)
	switch {
	case err != nil:
		return err
	case !s.access.AllowedByOwner(upload.ProjectID, user):
		return app.ErrNoAccess
	default:
		if err := s.uploads.Delete(requestID); err != nil {
			return err
		}
		// Delete the directory where chunks were being written
		s.fops.RemoveAll(app.MCDir.UploadDir(requestID))
		return nil
	}
}

// ListForProject will return all the uploads associated with a project.
func (s *idService) UploadsForProject(projectID, user string) ([]schema.Upload, error) {
	_, err := s.getProjectValidatingAccess(projectID, user)
	switch {
	case err == app.ErrNotFound:
		// Invalid project
		return nil, app.ErrInvalid
	case err != nil:
		return nil, err
	default:
		return s.uploads.ForProject(projectID)
	}
}
<commit_msg>Update method documentation to match method name.<commit_after>package uploads

import (
	"time"

	"math"

	"fmt"

	r "github.com/dancannon/gorethink"
	"github.com/materials-commons/gohandy/file"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/pkg/domain"
	"github.com/willf/bitset"
)

var _ = fmt.Println

// A IDRequest requests a new upload id be created for
// the given parameters.
type IDRequest struct {
	User        string
	DirectoryID string
	ProjectID   string
	FileName    string
	FileSize    int64
	Checksum    string
	ChunkSize   int32
	FileMTime   time.Time
	Host        string
	Birthtime   time.Time
}

// IDService creates new upload requests
type IDService interface {
	ID(req IDRequest) (*schema.Upload, error)
	Delete(requestID, user string) error
	UploadsForProject(projectID, user string) ([]schema.Upload, error)
}

// idService implements the IDService interface using
// the dai services.
type idService struct {
	dirs        dai.Dirs
	projects    dai.Projects
	uploads     dai.Uploads
	files       dai.Files
	access      domain.Access
	fops        file.Operations
	tracker     *blockTracker
	requestPath requestPath
}

// NewIDService creates a new idService that connects to the database using
// the given session.
func NewIDService(session *r.Session) *idService {
	access := domain.NewAccess(dai.NewRProjects(session), dai.NewRFiles(session), dai.NewRUsers(session))
	return &idService{
		dirs:        dai.NewRDirs(session),
		projects:    dai.NewRProjects(session),
		uploads:     dai.NewRUploads(session),
		files:       dai.NewRFiles(session),
		access:      access,
		fops:        file.OS,
		tracker:     requestBlockTracker,
		requestPath: &mcdirRequestPath{},
	}
}

// ID will create a new Upload request or return an existing one.
func (s *idService) ID(req IDRequest) (*schema.Upload, error) {
	var (
		err  error
		proj *schema.Project
		dir  *schema.Directory
	)

	// Check that project exists and user has access
	if proj, err = s.getProjectValidatingAccess(req.ProjectID, req.User); err != nil {
		return nil, err
	}

	// Check that directory exists and user has access
	if dir, err = s.getDirectoryValidatingAccess(req.DirectoryID, req.ProjectID, req.User); err != nil {
		return nil, err
	}

	return s.createUploadRequest(req, proj, dir)
}

// createUploadRequest will create an upload. It checks if there is a matching file or an existing upload.
// If it finds a matching file it returns a finished upload. If it finds an existing upload it returns it,
// otherwise it returns a new upload.
func (s *idService) createUploadRequest(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	upload, err := s.findExisting(req, proj, dir)
	switch {
	case err == app.ErrNotFound:
		return s.createNewUploadRequest(req, proj, dir)
	case err != nil:
		return nil, err
	default:
		return upload, nil
	}
}

// findExisting checks if there is an outstanding upload request, or a file that already matches
// the upload request. If it finds a match it will return the corresponding upload request. In
// the case of an already uploaded file it will return an upload request with all blocks marked
// as uploaded.
func (s *idService) findExisting(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	if _, err := s.files.ByChecksum(req.Checksum); err == nil {
		return s.createFinishedUpload(req, proj, dir)
	}
	return s.findMatchingUploadRequest(req)
}

// createFinishedUpload will create an upload entry with all blocks marked as uploaded.
func (s *idService) createFinishedUpload(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	// Create a new upload request and then set all blocks as already uploaded.
	if upload, err := s.createNewUploadRequest(req, proj, dir); err != nil {
		return nil, err
	} else {
		s.tracker.markAllBlocks(upload.ID)
		s.tracker.setIsExistingFile(upload.ID, true)
		upload.SetFBlocks(s.tracker.getBlocks(upload.ID))
		return upload, nil
	}
}

// findMatchingUploadRequest will search the set of upload requests and see if there
// is already an outstanding upload request that matches this one.
func (s *idService) findMatchingUploadRequest(req IDRequest) (*schema.Upload, error) {
	searchParams := dai.UploadSearch{
		ProjectID:   req.ProjectID,
		DirectoryID: req.DirectoryID,
		FileName:    req.FileName,
		Checksum:    req.Checksum,
	}

	if existingUpload, err := s.uploads.Search(searchParams); err == nil {
		// Found existing
		if bset := s.tracker.getBlocks(existingUpload.ID); bset != nil {
			existingUpload.File.Blocks = bset
		}
		return existingUpload, nil
	}
	return nil, app.ErrNotFound
}

// createNewUploadRequest creates a new upload request and inserts it into the database.
func (s *idService) createNewUploadRequest(req IDRequest, proj *schema.Project, dir *schema.Directory) (*schema.Upload, error) {
	n := uint(numBlocks(req.FileSize, req.ChunkSize))
	upload := schema.CUpload().
		Owner(req.User).
		Project(req.ProjectID, proj.Name).
		ProjectOwner(proj.Owner).
		Directory(req.DirectoryID, dir.Name).
		Host(req.Host).
		FName(req.FileName).
		FSize(req.FileSize).
		FChecksum(req.Checksum).
		FRemoteMTime(req.FileMTime).
		FBlocks(bitset.New(n)).
		Create()
	u, err := s.uploads.Insert(&upload)
	if err != nil {
		fmt.Println("insert err", err)
		return nil, err
	}

	if err := s.initUpload(u.ID, req.FileSize, req.ChunkSize); err != nil {
		fmt.Println("initUpload err", err)
		s.uploads.Delete(u.ID)
		return nil, err
	}
	return u, nil
}

// getProjectValidatingAccess retrieves the project with the given projectID. It checks that the
// given user has access to that project.
func (s *idService) getProjectValidatingAccess(projectID, user string) (*schema.Project, error) {
	project, err := s.projects.ByID(projectID)
	switch {
	case err != nil:
		return nil, err
	case !s.access.AllowedByOwner(projectID, user):
		return nil, app.ErrNoAccess
	default:
		return project, nil
	}
}

// getDirectoryValidatingAccess retrieves the directory with the given directoryID. It checks access to the
// directory and validates that the directory exists in the given project.
func (s *idService) getDirectoryValidatingAccess(directoryID, projectID, user string) (*schema.Directory, error) {
	dir, err := s.dirs.ByID(directoryID)
	switch {
	case err != nil:
		return nil, err
	case !s.projects.HasDirectory(projectID, directoryID):
		return nil, app.ErrInvalid
	case !s.access.AllowedByOwner(projectID, user):
		return nil, app.ErrNoAccess
	default:
		return dir, nil
	}
}

// initUpload
func (s *idService) initUpload(id string, fileSize int64, chunkSize int32) error {
	if err := s.requestPath.mkdirFromID(id); err != nil {
		return err
	}

	s.tracker.load(id, numBlocks(fileSize, chunkSize))
	return nil
}

// numBlocks
func numBlocks(fileSize int64, chunkSize int32) int {
	// round up to nearest number of blocks
	d := float64(fileSize) / float64(chunkSize)
	n := int(math.Ceil(d))
	return n
}

// Delete will delete the given requestID if the user has access
// to delete that request. Owners of requests can delete their
// own requests. Project owners can delete any request, even if
// they don't own the request.
func (s *idService) Delete(requestID, user string) error {
	upload, err := s.uploads.ByID(requestID)
	switch {
	case err != nil:
		return err
	case !s.access.AllowedByOwner(upload.ProjectID, user):
		return app.ErrNoAccess
	default:
		if err := s.uploads.Delete(requestID); err != nil {
			return err
		}
		// Delete the directory where chunks were being written
		s.fops.RemoveAll(app.MCDir.UploadDir(requestID))
		return nil
	}
}

// ListForProject will return all the uploads associated with a project.
func (s *idService) UploadsForProject(projectID, user string) ([]schema.Upload, error) {
	_, err := s.getProjectValidatingAccess(projectID, user)
	switch {
	case err == app.ErrNotFound:
		// Invalid project
		return nil, app.ErrInvalid
	case err != nil:
		return nil, err
	default:
		return s.uploads.ForProject(projectID)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"time"
)

type Game struct {
	clients        []*Client
	pendingClients chan *Client
	world          *World
}

func NewGame(w *World) *Game {
	g := new(Game)
	g.clients = make([]*Client, 0)
	g.pendingClients = make(chan *Client, 10)
// 	g.world = NewWorld(0)
	g.world = w
	return g
}

func (g *Game) Connect(c *Client) {
	g.pendingClients <- c
}

func (g *Game) Disconnect(c *Client) {
	i := g.findClient(c)
	if i == -1 {
		log.Println("[WARN] Attempt to disconnect client who is not connected.")
		return
	}
	g.clients[i] = g.clients[len(g.clients) - 1]
	g.clients = g.clients[:len(g.clients) - 1]
	g.Announce(c.name + " has left the game :(")
}

func (g *Game) Announce(message string) {
	g.Chat("SERVER", message)
}

func (g *Game) Chat(user string, message string) {
	log.Println("[CHAT]", user + ":", message)
	g.Broadcast(&MsgChat{
		DisplayName: user,
		Message: message,
	})
}

func (g *Game) Broadcast(m Message) {
	for _, c := range g.clients {
		c.Send(m)
	}
}

func (g *Game) BroadcastLossy(m Message) {
	for _, c := range g.clients {
		c.SendLossy(m)
	}
}

func (g *Game) findClientByName(name string) *Client {
	for _, c := range g.clients {
		if c.name == name {
			return c
		}
	}
	return nil
}

func (g *Game) findClient(c *Client) int {
	for i, other := range g.clients {
		if other == c {
			 return i
		}
	}
	return -1
}

func (g *Game) Run() {
	updateTicker := time.Tick(time.Second / 60)
	for {
		<-updateTicker
		g.Tick()
	}
}

func (g *Game) Tick() {
	for {
		select {
		case c := <-g.pendingClients:
			g.clients = append(g.clients, c)
			g.Announce(c.name + " has joined the game!")
		default:
			break;
		}
	}
	for _, c := range g.clients {
		c.Tick(g)
		select {
		case <-c.Errors:
			g.Disconnect(c)
		default:
		}
	}
	g.world.Tick(g)
}
<commit_msg>Infinite loop repaired<commit_after>package main

import (
	"log"
	"time"
)

type Game struct {
	clients        []*Client
	pendingClients chan *Client
	world          *World
}

func NewGame(w *World) *Game {
	g := new(Game)
	g.clients = make([]*Client, 0)
	g.pendingClients = make(chan *Client, 10)
// 	g.world = NewWorld(0)
	g.world = w
	return g
}

// Thread safe
func (g *Game) Connect(c *Client) {
	g.pendingClients <- c
}

// Not thread safe (This better be called on the Games main thread)
func (g *Game) handlePendingClients() {
	for {
		select {
		case c := <-g.pendingClients:
			g.clients = append(g.clients, c)
			g.Announce(c.name + " has joined the game!")
		default:
			return
		}
	}
}

func (g *Game) Disconnect(c *Client) {
	i := g.findClient(c)
	if i == -1 {
		log.Println("[WARN] Attempt to disconnect client who is not connected.")
		return
	}
	g.clients[i] = g.clients[len(g.clients) - 1]
	g.clients = g.clients[:len(g.clients) - 1]
	g.Announce(c.name + " has left the game :(")
}

func (g *Game) Announce(message string) {
	g.Chat("SERVER", message)
}

func (g *Game) Chat(user string, message string) {
	log.Println("[CHAT]", user + ":", message)
	g.Broadcast(&MsgChat{
		DisplayName: user,
		Message: message,
	})
}

func (g *Game) Broadcast(m Message) {
	for _, c := range g.clients {
		c.Send(m)
	}
}

func (g *Game) BroadcastLossy(m Message) {
	for _, c := range g.clients {
		c.SendLossy(m)
	}
}

func (g *Game) findClientByName(name string) *Client {
	for _, c := range g.clients {
		if c.name == name {
			return c
		}
	}
	return nil
}

func (g *Game) findClient(c *Client) int {
	for i, other := range g.clients {
		if other == c {
			 return i
		}
	}
	return -1
}

func (g *Game) Run() {
	updateTicker := time.Tick(time.Second / 60)
	for {
		<-updateTicker
		g.Tick()
	}
}

func (g *Game) Tick() {
	g.handlePendingClients()
	for _, c := range g.clients {
		c.Tick(g)
		select {
		case <-c.Errors:
			g.Disconnect(c)
		default:
		}
	}
	g.world.Tick(g)
}
<|endoftext|>


<commit_before>package river

import (
	"flag"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/siddontang/go-mysql-elasticsearch/elastic"
	"github.com/siddontang/go-mysql/client"
	. "gopkg.in/check.v1"
)

var my_addr = flag.String("my_addr", "127.0.0.1:3306", "MySQL addr")
var es_addr = flag.String("es_addr", "127.0.0.1:9200", "Elasticsearch addr")

func Test(t *testing.T) {
	TestingT(t)
}

type riverTestSuite struct {
	c *client.Conn
	r *River
}

var _ = Suite(&riverTestSuite{})

func (s *riverTestSuite) SetUpSuite(c *C) {
	var err error
	s.c, err = client.Connect(*my_addr, "root", "", "test")
	c.Assert(err, IsNil)

	s.testExecute(c, "SET SESSION binlog_format = 'ROW'")

	schema := `
        CREATE TABLE IF NOT EXISTS %s (
            id INT,
            title VARCHAR(256),
            content VARCHAR(256),
            mylist VARCHAR(256),
            tenum ENUM("e1", "e2", "e3"),
            tset SET("a", "b", "c"),
            tbit BIT(1) default 1,
            PRIMARY KEY(id)) ENGINE=INNODB;
    `

	s.testExecute(c, "DROP TABLE IF EXISTS test_river")
	s.testExecute(c, fmt.Sprintf(schema, "test_river"))
	s.testExecute(c, fmt.Sprintf(schema, "test_for_id"))

	for i := 0; i < 10; i++ {
		table := fmt.Sprintf("test_river_%04d", i)
		s.testExecute(c, fmt.Sprintf("DROP TABLE IF EXISTS %s", table))
		s.testExecute(c, fmt.Sprintf(schema, table))
	}

	cfg := new(Config)
	cfg.MyAddr = *my_addr
	cfg.MyUser = "root"
	cfg.MyPassword = ""
	cfg.MyCharset = "utf8"
	cfg.ESAddr = *es_addr

	cfg.ServerID = 1001
	cfg.Flavor = "mysql"

	cfg.DataDir = "/tmp/test_river"
	cfg.DumpExec = "mysqldump"

	cfg.StatAddr = "127.0.0.1:12800"
	cfg.BulkSize = 1
	cfg.FlushBulkTime = TomlDuration{3 * time.Millisecond}

	os.RemoveAll(cfg.DataDir)

	cfg.Sources = []SourceConfig{SourceConfig{Schema: "test", Tables: []string{"test_river", "test_river_[0-9]{4}", "test_for_id"}}}

	cfg.Rules = []*Rule{
		&Rule{Schema: "test",
			Table:        "test_river",
			Index:        "river",
			Type:         "river",
			FieldMapping: map[string]string{"title": "es_title", "mylist": "es_mylist,list"},
		},

		&Rule{Schema: "test",
			Table:        "test_for_id",
			Index:        "river",
			Type:         "river",
			ID:           []string{"id", "title"},
			FieldMapping: map[string]string{"title": "es_title", "mylist": "es_mylist,list"},
		},

		&Rule{Schema: "test",
			Table:        "test_river_[0-9]{4}",
			Index:        "river",
			Type:         "river",
			FieldMapping: map[string]string{"title": "es_title", "mylist": "es_mylist,list"},
		},
	}

	s.r, err = NewRiver(cfg)
	c.Assert(err, IsNil)

	err = s.r.es.DeleteIndex("river")
	c.Assert(err, IsNil)
}

func (s *riverTestSuite) TearDownSuite(c *C) {
	if s.c != nil {
		s.c.Close()
	}

	if s.r != nil {
		s.r.Close()
	}
}

func (s *riverTestSuite) TestConfig(c *C) {
	str := `
my_addr = "127.0.0.1:3306"
my_user = "root"
my_pass = ""
my_charset = "utf8"
es_addr = "127.0.0.1:9200"
es_user = ""
es_pass = ""
data_dir = "./var"

[[source]]
schema = "test"

tables = ["test_river", "test_river_[0-9]{4}", "test_for_id"]

[[rule]]
schema = "test"
table = "test_river"
index = "river"
type = "river"
parent = "pid"

    [rule.field]
    title = "es_title"
    mylist = "es_mylist,list"

[[rule]]
schema = "test"
table = "test_for_id"
index = "river"
type = "river"
parent = "pid"
id = ["id", "title"]
    [rule.field]
    title = "es_title"
    mylist = "es_mylist,list"

[[rule]]
schema = "test"
table = "test_river_[0-9]{4}"
index = "river"
type = "river"

    [rule.field]
    title = "es_title"
    mylist = "es_mylist,list"

`

	cfg, err := NewConfig(str)
	c.Assert(err, IsNil)
	c.Assert(cfg.Sources, HasLen, 1)
	c.Assert(cfg.Sources[0].Tables, HasLen, 3)
	c.Assert(cfg.Rules, HasLen, 3)
}

func (s *riverTestSuite) testExecute(c *C, query string, args ...interface{}) {
	_, err := s.c.Execute(query, args...)
	c.Assert(err, IsNil)
}

func (s *riverTestSuite) testPrepareData(c *C) {
	s.testExecute(c, "INSERT INTO test_river (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", 1, "first", "hello go 1", "e1", "a,b")
	s.testExecute(c, "INSERT INTO test_river (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", 2, "second", "hello mysql 2", "e2", "b,c")
	s.testExecute(c, "INSERT INTO test_river (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", 3, "third", "hello elaticsearch 3", "e3", "c")
	s.testExecute(c, "INSERT INTO test_river (id, title, content, tenum, tset, tbit) VALUES (?, ?, ?, ?, ?, ?)", 4, "fouth", "hello go-mysql-elasticserach 4", "e1", "a,b,c", 0)
	s.testExecute(c, "INSERT INTO test_for_id (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", 1, "first", "hello go 1", "e1", "a,b")

	for i := 0; i < 10; i++ {
		table := fmt.Sprintf("test_river_%04d", i)
		s.testExecute(c, fmt.Sprintf("INSERT INTO %s (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", table), 5+i, "abc", "hello", "e1", "a,b,c")
	}
}

func (s *riverTestSuite) testElasticGet(c *C, id string) *elastic.Response {
	index := "river"
	docType := "river"

	r, err := s.r.es.Get(index, docType, id)
	if err != nil {
		r := new(elastic.Response)
		r.Found = false
		return r
	}

	return r
}

func testWaitSyncDone(c *C, r *River) {
	<-r.canal.WaitDumpDone()

	err := r.canal.CatchMasterPos(10 * time.Second)
	c.Assert(err, IsNil)

	for i := 0; i < 1000; i++ {
		if len(r.syncCh) == 0 {
			return
		}

		time.Sleep(10 * time.Millisecond)
	}

	c.Fatalf("wait 1s but still have %d items to be synced", len(r.syncCh))
}

func (s *riverTestSuite) TestRiver(c *C) {
	s.testPrepareData(c)

	s.r.Start()

	testWaitSyncDone(c, s.r)

	var r *elastic.Response
	r = s.testElasticGet(c, "1")
	c.Assert(r.Found, Equals, true)
	c.Assert(r.Source["tenum"], Equals, "e1")
	c.Assert(r.Source["tset"], Equals, "a,b")

	r = s.testElasticGet(c, "1:first")
	c.Assert(r.Found, Equals, true)

	r = s.testElasticGet(c, "100")
	c.Assert(r.Found, Equals, false)

	for i := 0; i < 10; i++ {
		r = s.testElasticGet(c, fmt.Sprintf("%d", 5+i))
		c.Assert(r.Found, Equals, true)
		c.Assert(r.Source["es_title"], Equals, "abc")
	}

	s.testExecute(c, "UPDATE test_river SET title = ?, tenum = ?, tset = ?, mylist = ? WHERE id = ?", "second 2", "e3", "a,b,c", "a,b,c", 2)
	s.testExecute(c, "DELETE FROM test_river WHERE id = ?", 1)
	s.testExecute(c, "UPDATE test_river SET title = ?, id = ? WHERE id = ?", "second 30", 30, 3)

	// so we can insert invalid data
	s.testExecute(c, `SET SESSION sql_mode="NO_ENGINE_SUBSTITUTION";`)

	// bad insert
	s.testExecute(c, "UPDATE test_river SET title = ?, tenum = ?, tset = ? WHERE id = ?", "second 2", "e5", "a,b,c,d", 4)

	for i := 0; i < 10; i++ {
		table := fmt.Sprintf("test_river_%04d", i)
		s.testExecute(c, fmt.Sprintf("UPDATE %s SET title = ? WHERE id = ?", table), "hello", 5+i)
	}

	testWaitSyncDone(c, s.r)

	r = s.testElasticGet(c, "1")
	c.Assert(r.Found, Equals, false)

	r = s.testElasticGet(c, "2")
	c.Assert(r.Found, Equals, true)
	c.Assert(r.Source["es_title"], Equals, "second 2")
	c.Assert(r.Source["tenum"], Equals, "e3")
	c.Assert(r.Source["tset"], Equals, "a,b,c")
	c.Assert(r.Source["es_mylist"], DeepEquals, []interface{}{"a", "b", "c"})
	c.Assert(r.Source["tbit"], Equals, float64(1))

	r = s.testElasticGet(c, "4")
	c.Assert(r.Found, Equals, true)
	c.Assert(r.Source["tenum"], Equals, "")
	c.Assert(r.Source["tset"], Equals, "a,b,c")
	c.Assert(r.Source["tbit"], Equals, float64(0))

	r = s.testElasticGet(c, "3")
	c.Assert(r.Found, Equals, false)

	r = s.testElasticGet(c, "30")
	c.Assert(r.Found, Equals, true)
	c.Assert(r.Source["es_title"], Equals, "second 30")

	for i := 0; i < 10; i++ {
		r = s.testElasticGet(c, fmt.Sprintf("%d", 5+i))
		c.Assert(r.Found, Equals, true)
		c.Assert(r.Source["es_title"], Equals, "hello")
	}
}
<commit_msg>Update river_test.go<commit_after>package river

import (
	"flag"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/siddontang/go-mysql-elasticsearch/elastic"
	"github.com/siddontang/go-mysql/client"
	. "gopkg.in/check.v1"
)

var my_addr = flag.String("my_addr", "127.0.0.1:3306", "MySQL addr")
var es_addr = flag.String("es_addr", "127.0.0.1:9200", "Elasticsearch addr")

func Test(t *testing.T) {
	TestingT(t)
}

type riverTestSuite struct {
	c *client.Conn
	r *River
}

var _ = Suite(&riverTestSuite{})

func (s *riverTestSuite) SetUpSuite(c *C) {
	var err error
	s.c, err = client.Connect(*my_addr, "root", "", "test")
	c.Assert(err, IsNil)

	s.testExecute(c, "SET SESSION binlog_format = 'ROW'")

	schema := `
        CREATE TABLE IF NOT EXISTS %s (
            id INT,
            title VARCHAR(256),
            content VARCHAR(256),
            mylist VARCHAR(256),
            tenum ENUM("e1", "e2", "e3"),
            tset SET("a", "b", "c"),
            tbit BIT(1) default 1,
            PRIMARY KEY(id)) ENGINE=INNODB;
    `

	s.testExecute(c, "DROP TABLE IF EXISTS test_river")
	s.testExecute(c, fmt.Sprintf(schema, "test_river"))
	s.testExecute(c, fmt.Sprintf(schema, "test_for_id"))

	for i := 0; i < 10; i++ {
		table := fmt.Sprintf("test_river_%04d", i)
		s.testExecute(c, fmt.Sprintf("DROP TABLE IF EXISTS %s", table))
		s.testExecute(c, fmt.Sprintf(schema, table))
	}

	cfg := new(Config)
	cfg.MyAddr = *my_addr
	cfg.MyUser = "root"
	cfg.MyPassword = ""
	cfg.MyCharset = "utf8"
	cfg.ESAddr = *es_addr

	cfg.ServerID = 1001
	cfg.Flavor = "mysql"

	cfg.DataDir = "/tmp/test_river"
	cfg.DumpExec = "mysqldump"

	cfg.StatAddr = "127.0.0.1:12800"
	cfg.BulkSize = 1
	cfg.FlushBulkTime = TomlDuration{3 * time.Millisecond}

	os.RemoveAll(cfg.DataDir)

	cfg.Sources = []SourceConfig{SourceConfig{Schema: "test", Tables: []string{"test_river", "test_river_[0-9]{4}", "test_for_id"}}}

	cfg.Rules = []*Rule{
		&Rule{Schema: "test",
			Table:        "test_river",
			Index:        "river",
			Type:         "river",
			FieldMapping: map[string]string{"title": "es_title", "mylist": "es_mylist,list"},
		},

		&Rule{Schema: "test",
			Table:        "test_for_id",
			Index:        "river",
			Type:         "river",
			ID:           []string{"id", "title"},
			FieldMapping: map[string]string{"title": "es_title", "mylist": "es_mylist,list"},
		},

		&Rule{Schema: "test",
			Table:        "test_river_[0-9]{4}",
			Index:        "river",
			Type:         "river",
			FieldMapping: map[string]string{"title": "es_title", "mylist": "es_mylist,list"},
		},
	}

	s.r, err = NewRiver(cfg)
	c.Assert(err, IsNil)

	err = s.r.es.DeleteIndex("river")
	c.Assert(err, IsNil)
}

func (s *riverTestSuite) TearDownSuite(c *C) {
	if s.c != nil {
		s.c.Close()
	}

	if s.r != nil {
		s.r.Close()
	}
}

func (s *riverTestSuite) TestConfig(c *C) {
	str := `
my_addr = "127.0.0.1:3306"
my_user = "root"
my_pass = ""
my_charset = "utf8"
es_addr = "127.0.0.1:9200"
es_user = ""
es_pass = ""
data_dir = "./var"

[[source]]
schema = "test"

tables = ["test_river", "test_river_[0-9]{4}", "test_for_id"]

[[rule]]
schema = "test"
table = "test_river"
index = "river"
type = "river"
parent = "pid"

    [rule.field]
    title = "es_title"
    mylist = "es_mylist,list"

[[rule]]
schema = "test"
table = "test_for_id"
index = "river"
type = "river"
parent = "pid"
id = ["id", "title"]
    [rule.field]
    title = "es_title"
    mylist = "es_mylist,list"

[[rule]]
schema = "test"
table = "test_river_[0-9]{4}"
index = "river"
type = "river"

    [rule.field]
    title = "es_title"
    mylist = "es_mylist,list"

`

	cfg, err := NewConfig(str)
	c.Assert(err, IsNil)
	c.Assert(cfg.Sources, HasLen, 1)
	c.Assert(cfg.Sources[0].Tables, HasLen, 3)
	c.Assert(cfg.Rules, HasLen, 3)
}

func (s *riverTestSuite) testExecute(c *C, query string, args ...interface{}) {
	_, err := s.c.Execute(query, args...)
	c.Assert(err, IsNil)
}

func (s *riverTestSuite) testPrepareData(c *C) {
	s.testExecute(c, "INSERT INTO test_river (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", 1, "first", "hello go 1", "e1", "a,b")
	s.testExecute(c, "INSERT INTO test_river (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", 2, "second", "hello mysql 2", "e2", "b,c")
	s.testExecute(c, "INSERT INTO test_river (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", 3, "third", "hello elaticsearch 3", "e3", "c")
	s.testExecute(c, "INSERT INTO test_river (id, title, content, tenum, tset, tbit) VALUES (?, ?, ?, ?, ?, ?)", 4, "fouth", "hello go-mysql-elasticserach 4", "e1", "a,b,c", 0)
	s.testExecute(c, "INSERT INTO test_for_id (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", 1, "first", "hello go 1", "e1", "a,b")

	for i := 0; i < 10; i++ {
		table := fmt.Sprintf("test_river_%04d", i)
		s.testExecute(c, fmt.Sprintf("INSERT INTO %s (id, title, content, tenum, tset) VALUES (?, ?, ?, ?, ?)", table), 5+i, "abc", "hello", "e1", "a,b,c")
	}
}

func (s *riverTestSuite) testElasticGet(c *C, id string) *elastic.Response {
	index := "river"
	docType := "river"

	r, err := s.r.es.Get(index, docType, id)
	c.Assert(err, IsNil)

	return r
}

func testWaitSyncDone(c *C, r *River) {
	<-r.canal.WaitDumpDone()

	err := r.canal.CatchMasterPos(10 * time.Second)
	c.Assert(err, IsNil)

	for i := 0; i < 1000; i++ {
		if len(r.syncCh) == 0 {
			return
		}

		time.Sleep(10 * time.Millisecond)
	}

	c.Fatalf("wait 1s but still have %d items to be synced", len(r.syncCh))
}

func (s *riverTestSuite) TestRiver(c *C) {
	s.testPrepareData(c)

	s.r.Start()

	testWaitSyncDone(c, s.r)

	var r *elastic.Response
	r = s.testElasticGet(c, "1")
	c.Assert(r.Found, Equals, true)
	c.Assert(r.Source["tenum"], Equals, "e1")
	c.Assert(r.Source["tset"], Equals, "a,b")

	r = s.testElasticGet(c, "1:first")
	c.Assert(r.Found, Equals, true)

	r = s.testElasticGet(c, "100")
	c.Assert(r.Found, Equals, false)

	for i := 0; i < 10; i++ {
		r = s.testElasticGet(c, fmt.Sprintf("%d", 5+i))
		c.Assert(r.Found, Equals, true)
		c.Assert(r.Source["es_title"], Equals, "abc")
	}

	s.testExecute(c, "UPDATE test_river SET title = ?, tenum = ?, tset = ?, mylist = ? WHERE id = ?", "second 2", "e3", "a,b,c", "a,b,c", 2)
	s.testExecute(c, "DELETE FROM test_river WHERE id = ?", 1)
	s.testExecute(c, "UPDATE test_river SET title = ?, id = ? WHERE id = ?", "second 30", 30, 3)

	// so we can insert invalid data
	s.testExecute(c, `SET SESSION sql_mode="NO_ENGINE_SUBSTITUTION";`)

	// bad insert
	s.testExecute(c, "UPDATE test_river SET title = ?, tenum = ?, tset = ? WHERE id = ?", "second 2", "e5", "a,b,c,d", 4)

	for i := 0; i < 10; i++ {
		table := fmt.Sprintf("test_river_%04d", i)
		s.testExecute(c, fmt.Sprintf("UPDATE %s SET title = ? WHERE id = ?", table), "hello", 5+i)
	}

	testWaitSyncDone(c, s.r)

	r = s.testElasticGet(c, "1")
	c.Assert(r.Found, Equals, false)

	r = s.testElasticGet(c, "2")
	c.Assert(r.Found, Equals, true)
	c.Assert(r.Source["es_title"], Equals, "second 2")
	c.Assert(r.Source["tenum"], Equals, "e3")
	c.Assert(r.Source["tset"], Equals, "a,b,c")
	c.Assert(r.Source["es_mylist"], DeepEquals, []interface{}{"a", "b", "c"})
	c.Assert(r.Source["tbit"], Equals, float64(1))

	r = s.testElasticGet(c, "4")
	c.Assert(r.Found, Equals, true)
	c.Assert(r.Source["tenum"], Equals, "")
	c.Assert(r.Source["tset"], Equals, "a,b,c")
	c.Assert(r.Source["tbit"], Equals, float64(0))

	r = s.testElasticGet(c, "3")
	c.Assert(r.Found, Equals, false)

	r = s.testElasticGet(c, "30")
	c.Assert(r.Found, Equals, true)
	c.Assert(r.Source["es_title"], Equals, "second 30")

	for i := 0; i < 10; i++ {
		r = s.testElasticGet(c, fmt.Sprintf("%d", 5+i))
		c.Assert(r.Found, Equals, true)
		c.Assert(r.Source["es_title"], Equals, "hello")
	}
}
<|endoftext|>


<commit_before>package command

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/hashicorp/otto/appfile"
	"github.com/hashicorp/otto/directory"
	"github.com/hashicorp/otto/otto"
	"github.com/hashicorp/otto/ui"
	"github.com/mitchellh/cli"
)

const (
	// DefaultAppfile is the default filename for the Appfile
	DefaultAppfile = "Appfile"

	// DefaultOutputDir is the default filename for the output directory
	DefaultOutputDir = ".otto"

	// DefaultDataDir is the default directory for the directory
	// data if a directory in the Appfile isn't specified.
	DefaultDataDir = "otto-data"
)

// FlagSetFlags is an enum to define what flags are present in the
// default FlagSet returned by Meta.FlagSet
type FlagSetFlags uint

const (
	FlagSetNone    FlagSetFlags = 0
	FlagSetAppfile FlagSetFlags = iota
)

// Meta are the meta-options that are available on all or most commands.
type Meta struct {
	CoreConfig *otto.CoreConfig
	Ui         cli.Ui

	// These are fields set by flags
	flagAppfile string
}

// Appfile loads the Appfile according to the path given by the
// -appfile flag.
func (m *Meta) Appfile() (*appfile.File, error) {
	// Get the path to where the Appfile lives
	path := m.flagAppfile
	if path == "" {
		path = "."
	}

	// Verify the path is valid
	fi, err := os.Stat(path)
	if err != nil {
		return nil, fmt.Errorf(
			"Error checking Appfile path: %s", err)
	}
	if fi.IsDir() {
		path = filepath.Join(path, DefaultAppfile)
	}

	// Load the appfile
	app, err := appfile.ParseFile(path)
	if err != nil {
		return nil, fmt.Errorf(
			"Error parsing Appfile: %s", err)
	}

	return app, nil
}

// Core returns the core for the given Appfile. The file where the
// Appfile was loaded from should be set in appfile.File.Path. This
// root appfile path will be used as the default output directory
// for Otto.
func (m *Meta) Core(f *appfile.File) (*otto.Core, error) {
	outputDir := DefaultOutputDir
	dir, err := m.Directory(f)
	if err != nil {
		return nil, err
	}

	config := *m.CoreConfig
	config.Appfile = f
	config.Directory = dir
	config.OutputDir = filepath.Join(filepath.Dir(f.Path), outputDir)
	config.Ui = m.OttoUi()

	return otto.NewCore(&config)
}

// Directory returns the Otto directory backend for the given
// Appfile. If no directory backend is specified, a local folder
// will be used.
func (m *Meta) Directory(f *appfile.File) (directory.Backend, error) {
	// TODO: Appfile can't specify directory configuration

	return &directory.FolderBackend{Dir: DefaultDataDir}, nil
}

// FlagSet returns a FlagSet with the common flags that every
// command implements. The exact behavior of FlagSet can be configured
// using the flags as the second parameter.
func (m *Meta) FlagSet(n string, fs FlagSetFlags) *flag.FlagSet {
	f := flag.NewFlagSet(n, flag.ContinueOnError)

	if fs&FlagSetAppfile != 0 {
		f.StringVar(&m.flagAppfile, "appfile", "", "")
	}

	// Create an io.Writer that writes to our Ui properly for errors.
	// This is kind of a hack, but it does the job. Basically: create
	// a pipe, use a scanner to break it into lines, and output each line
	// to the UI. Do this forever.
	errR, errW := io.Pipe()
	errScanner := bufio.NewScanner(errR)
	go func() {
		for errScanner.Scan() {
			m.Ui.Error(errScanner.Text())
		}
	}()
	f.SetOutput(errW)

	return f
}

// OttoUi returns the ui.Ui object.
func (m *Meta) OttoUi() ui.Ui {
	return NewUi(m.Ui)
}
<commit_msg>command: support OTTO_APPFILE env var<commit_after>package command

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/hashicorp/otto/appfile"
	"github.com/hashicorp/otto/directory"
	"github.com/hashicorp/otto/otto"
	"github.com/hashicorp/otto/ui"
	"github.com/mitchellh/cli"
)

const (
	// DefaultAppfile is the default filename for the Appfile
	DefaultAppfile = "Appfile"

	// DefaultOutputDir is the default filename for the output directory
	DefaultOutputDir = ".otto"

	// DefaultDataDir is the default directory for the directory
	// data if a directory in the Appfile isn't specified.
	DefaultDataDir = "otto-data"

	// EnvAppFile is the environment variable to point to an appfile.
	EnvAppFile = "OTTO_APPFILE"
)

// FlagSetFlags is an enum to define what flags are present in the
// default FlagSet returned by Meta.FlagSet
type FlagSetFlags uint

const (
	FlagSetNone    FlagSetFlags = 0
	FlagSetAppfile FlagSetFlags = iota
)

// Meta are the meta-options that are available on all or most commands.
type Meta struct {
	CoreConfig *otto.CoreConfig
	Ui         cli.Ui

	// These are fields set by flags
	flagAppfile string
}

// Appfile loads the Appfile according to the path given by the
// -appfile flag.
func (m *Meta) Appfile() (*appfile.File, error) {
	// Get the path to where the Appfile lives
	path := m.flagAppfile
	if path == "" {
		path = os.Getenv(EnvAppFile)
	}
	if path == "" {
		path = "."
	}

	// Verify the path is valid
	fi, err := os.Stat(path)
	if err != nil {
		return nil, fmt.Errorf(
			"Error checking Appfile path: %s", err)
	}
	if fi.IsDir() {
		path = filepath.Join(path, DefaultAppfile)
	}

	// Load the appfile
	app, err := appfile.ParseFile(path)
	if err != nil {
		return nil, fmt.Errorf(
			"Error parsing Appfile: %s", err)
	}

	return app, nil
}

// Core returns the core for the given Appfile. The file where the
// Appfile was loaded from should be set in appfile.File.Path. This
// root appfile path will be used as the default output directory
// for Otto.
func (m *Meta) Core(f *appfile.File) (*otto.Core, error) {
	outputDir := DefaultOutputDir
	dir, err := m.Directory(f)
	if err != nil {
		return nil, err
	}

	config := *m.CoreConfig
	config.Appfile = f
	config.Directory = dir
	config.OutputDir = filepath.Join(filepath.Dir(f.Path), outputDir)
	config.Ui = m.OttoUi()

	return otto.NewCore(&config)
}

// Directory returns the Otto directory backend for the given
// Appfile. If no directory backend is specified, a local folder
// will be used.
func (m *Meta) Directory(f *appfile.File) (directory.Backend, error) {
	// TODO: Appfile can't specify directory configuration

	return &directory.FolderBackend{Dir: DefaultDataDir}, nil
}

// FlagSet returns a FlagSet with the common flags that every
// command implements. The exact behavior of FlagSet can be configured
// using the flags as the second parameter.
func (m *Meta) FlagSet(n string, fs FlagSetFlags) *flag.FlagSet {
	f := flag.NewFlagSet(n, flag.ContinueOnError)

	if fs&FlagSetAppfile != 0 {
		f.StringVar(&m.flagAppfile, "appfile", "", "")
	}

	// Create an io.Writer that writes to our Ui properly for errors.
	// This is kind of a hack, but it does the job. Basically: create
	// a pipe, use a scanner to break it into lines, and output each line
	// to the UI. Do this forever.
	errR, errW := io.Pipe()
	errScanner := bufio.NewScanner(errR)
	go func() {
		for errScanner.Scan() {
			m.Ui.Error(errScanner.Text())
		}
	}()
	f.SetOutput(errW)

	return f
}

// OttoUi returns the ui.Ui object.
func (m *Meta) OttoUi() ui.Ui {
	return NewUi(m.Ui)
}
<|endoftext|>


<commit_before>// Copyright (c) Christoph Berger. All rights reserved.
// Use of this source code is governed by the BSD (3-Clause)
// License that can be found in the LICENSE.txt file.
//
// This source code may import third-party source code whose
// licenses are provided in the respective license files.

package start

import (
	"fmt"
	"os"
	"testing"

	flag "github.com/ogier/pflag"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAdd(t *testing.T) {
	Convey("When adding a command to the command map, then...", t, func() {
		Commands.Add(&Command{
			Name: "cmd1",
		})

		Convey("the command map should contain it", func() {
			So(Commands["cmd1"], ShouldNotBeNil)
			So(Commands["cmd1"].Name, ShouldEqual, "cmd1")

		})

		Convey("adding a command with the same name should fail", func() {
			err := Commands.Add(&Command{
				Name: "cmd1",
			})
			So(err, ShouldNotBeNil)
		})

		Convey("adding a subcommand to the command should work", func() {
			err := Commands.Add(&Command{
				Parent: "cmd1",
				Name:   "subcmd1",
			})
			So(err, ShouldBeNil)
			So(len(Commands["cmd1"].children), ShouldEqual, 1)
			So(Commands["cmd1"].children["subcmd1"], ShouldNotBeNil)
			So(Commands["cmd1"].children["subcmd1"].Name, ShouldEqual, "subcmd1")

			Convey("adding the same subcommand twice should fail", func() {
				err := Commands.Add(&Command{
					Parent: "cmd1",
					Name:   "subcmd1",
				})
				So(err, ShouldNotBeNil)
			})

		})

		Reset(func() {
			Commands = make(CommandMap)
		})
	})

}

func TestCommands(t *testing.T) {

	var yes bool
	var size int
	var global string

	// ContinueOnError is required when running goconvey as server; otherwise, unrecognized
	// flags that are passed to the test executable will cause an error:
	// "unknown shorthand flag: 't' in -test.v=true"
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
	// To suppress warnings resulting from standard flags -test and -json,
	// read -t and -j into dummy flags.
	// The flags used for actual testing must not use -t or -j shorthands.
	var testflag string
	var jsonflag string
	flag.StringVarP(&testflag, "t", "t", "t", "t")
	flag.StringVarP(&jsonflag, "j", "j", "j", "j")

	flag.BoolVarP(&yes, "yes", "y", false, "A boolean flag")
	flag.IntVarP(&size, "size", "s", 23, "An int flag")
	flag.StringVarP(&global, "global", "g", "global flag", "A global string flag")

	Commands = make(CommandMap) // Ensure the Commands map starts empty for the test

	os.Args = []string{os.Args[0]}

	Convey("Ensure that the test flags exist", t, func() {
		if err := Parse(); err != nil {
			fmt.Println(err)
		}
		So(flag.Lookup("yes").Name, ShouldEqual, "yes")
		So(flag.Lookup("size").Name, ShouldEqual, "size")
	})

	Convey("When setting up some commands, then...", t, func() {
		SetDescription("This is the test application for the start package.")

		Add(&Command{
			Name:  "test",
			Short: "A test command",
			Long:  "Command test helps testing the start package. It accepts all flags.",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the test command.")
				return nil
			},
		})

		Add(&Command{
			Name:  "flags",
			Flags: []string{"yes", "size"},
			Short: "A test command with flags",
			Long:  "Command flags helps testing flags.",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the testflags command.")
				fmt.Printf("--yes is %v", yes)
				fmt.Printf("--size is %v", size)
				return nil
			},
		})

		Add(&Command{
			Name:  "do",
			Short: "A command with subcommands",
			Long: `Command do helps testing subcommands.
Usage:
	do something
	do nothing`,
		})

		Add(&Command{
			Parent: "do",
			Name:   "something",
			Short:  "A subcommand that does something.",
			Long:   "do something does something",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the do something command.")
				return nil
			},
		})

		Add(&Command{
			Parent: "do",
			Name:   "nothing",
			Short:  "A subcommand that does nothing.",
			Long:   "do nothing does nothing",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the do nothing command.")
				return nil
			},
		})

		if err := Parse(); err != nil {
			fmt.Println(err)
		}

		Convey("readCommand should identify all of them correctly", func() {
			cmd, err := readCommand([]string{"test", "arg1", "arg2"})
			So(cmd, ShouldNotBeNil)
			So(cmd.Name, ShouldEqual, "test")
			So(cmd.Args, ShouldResemble, []string{"arg1", "arg2"})
			So(err, ShouldBeNil)

			cmd, err = readCommand([]string{"do", "something", "arg1"})
			So(cmd, ShouldNotBeNil)
			So(cmd.Name, ShouldEqual, "something")
			So(cmd.Args, ShouldResemble, []string{"arg1"})
			So(err, ShouldBeNil)

			cmd, err = readCommand([]string{"do", "nothing"})
			So(cmd, ShouldNotBeNil)
			So(cmd.Name, ShouldEqual, "nothing")
			So(cmd.Args, ShouldResemble, []string{})
			So(err, ShouldBeNil)
		})

		Convey("readCommand should return the Usage command if no valid command was passed in", func() {
			cmd, err := readCommand([]string{"invalid", "arg1"})
			So(cmd, ShouldNotBeNil)
			So(cmd.Cmd, ShouldEqual, Usage)
			So(err, ShouldBeNil)
		})

		Convey("Usage() should print the usage", func() {
			Usage(nil)
			Usage(Commands["test"])
			Usage(Commands["flags"])
			Usage(Commands["do"])
			Usage(Commands["do"].children["something"])
			Usage(Commands["do"].children["nothing"])

		})
		Reset(func() {
			Commands = make(CommandMap)
		})
	})

}

func TestCheckFlags(t *testing.T) {
	var first int
	var second int
	var third int
	var global int
	var rejectedFlags map[string]bool

	// ContinueOnError is required when running goconvey as server; otherwise, unrecognized
	// flags that are passed to the test executable will cause an error:
	// "unknown shorthand flag: 't' in -test.v=true"
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
	// To suppress warnings resulting from standard flags -test and -json,
	// read -t and -j into dummy flags.
	// The flags used for actual testing must not use -t or -j shorthands.
	var testflag string
	var jsonflag string
	flag.StringVarP(&testflag, "t", "t", "t", "t")
	flag.StringVarP(&jsonflag, "j", "j", "j", "j")

	flag.IntVarP(&first, "first", "f", 1, "The first flag")
	flag.IntVarP(&second, "second", "s", 2, "The second flag")
	flag.IntVarP(&third, "third", "h", 3, "The third flag")
	flag.IntVarP(&global, "global", "g", 4, "The global flag")

	os.Args = []string{os.Args[0], "--first=10", "--second=20", "--third=30", "--global=40", "anargument", "anotherarg"}

	Commands = make(CommandMap) // clear the commands map for this test
	Add(&Command{
		Name:  "cmd12",
		Flags: []string{"first", "second"},
	})
	Add(&Command{
		Name:  "cmd23",
		Flags: []string{"second", "third"},
	})
	Add(&Command{
		Name:  "cmd123",
		Flags: []string{"first", "second", "third"},
	})

	if err := Parse(); err != nil {
		fmt.Println(err)
	}

	Convey("A command should accept its own flags and all global flags", t, func() {
		rejectedFlags = checkFlags(Commands["cmd123"])
		So(len(rejectedFlags), ShouldEqual, 0)
	})
	Convey("A command should reject the flags that belong to the other command only", t, func() {
		rejectedFlags = checkFlags(Commands["cmd23"])
		So(len(rejectedFlags), ShouldEqual, 1)
		So(rejectedFlags["first"], ShouldEqual, true)

		rejectedFlags = checkFlags(Commands["cmd12"])
		So(len(rejectedFlags), ShouldEqual, 1)
		So(rejectedFlags["third"], ShouldEqual, true)
	})
}
<commit_msg>Removed: Test for returning the Usage command if no command is passed in. Reason: the Usage command is returned within an anonymous closure.<commit_after>// Copyright (c) Christoph Berger. All rights reserved.
// Use of this source code is governed by the BSD (3-Clause)
// License that can be found in the LICENSE.txt file.
//
// This source code may import third-party source code whose
// licenses are provided in the respective license files.

package start

import (
	"fmt"
	"os"
	"testing"

	flag "github.com/ogier/pflag"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAdd(t *testing.T) {
	Convey("When adding a command to the command map, then...", t, func() {
		Commands.Add(&Command{
			Name: "cmd1",
		})

		Convey("the command map should contain it", func() {
			So(Commands["cmd1"], ShouldNotBeNil)
			So(Commands["cmd1"].Name, ShouldEqual, "cmd1")

		})

		Convey("adding a command with the same name should fail", func() {
			err := Commands.Add(&Command{
				Name: "cmd1",
			})
			So(err, ShouldNotBeNil)
		})

		Convey("adding a subcommand to the command should work", func() {
			err := Commands.Add(&Command{
				Parent: "cmd1",
				Name:   "subcmd1",
			})
			So(err, ShouldBeNil)
			So(len(Commands["cmd1"].children), ShouldEqual, 1)
			So(Commands["cmd1"].children["subcmd1"], ShouldNotBeNil)
			So(Commands["cmd1"].children["subcmd1"].Name, ShouldEqual, "subcmd1")

			Convey("adding the same subcommand twice should fail", func() {
				err := Commands.Add(&Command{
					Parent: "cmd1",
					Name:   "subcmd1",
				})
				So(err, ShouldNotBeNil)
			})

		})

		Reset(func() {
			Commands = make(CommandMap)
		})
	})

}

func TestCommands(t *testing.T) {

	var yes bool
	var size int
	var global string

	// ContinueOnError is required when running goconvey as server; otherwise, unrecognized
	// flags that are passed to the test executable will cause an error:
	// "unknown shorthand flag: 't' in -test.v=true"
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
	// To suppress warnings resulting from standard flags -test and -json,
	// read -t and -j into dummy flags.
	// The flags used for actual testing must not use -t or -j shorthands.
	var testflag string
	var jsonflag string
	flag.StringVarP(&testflag, "t", "t", "t", "t")
	flag.StringVarP(&jsonflag, "j", "j", "j", "j")

	flag.BoolVarP(&yes, "yes", "y", false, "A boolean flag")
	flag.IntVarP(&size, "size", "s", 23, "An int flag")
	flag.StringVarP(&global, "global", "g", "global flag", "A global string flag")

	Commands = make(CommandMap) // Ensure the Commands map starts empty for the test

	os.Args = []string{os.Args[0]}

	Convey("Ensure that the test flags exist", t, func() {
		if err := Parse(); err != nil {
			fmt.Println(err)
		}
		So(flag.Lookup("yes").Name, ShouldEqual, "yes")
		So(flag.Lookup("size").Name, ShouldEqual, "size")
	})

	Convey("When setting up some commands, then...", t, func() {
		SetDescription("This is the test application for the start package.")

		Add(&Command{
			Name:  "test",
			Short: "A test command",
			Long:  "Command test helps testing the start package. It accepts all flags.",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the test command.")
				return nil
			},
		})

		Add(&Command{
			Name:  "flags",
			Flags: []string{"yes", "size"},
			Short: "A test command with flags",
			Long:  "Command flags helps testing flags.",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the testflags command.")
				fmt.Printf("--yes is %v", yes)
				fmt.Printf("--size is %v", size)
				return nil
			},
		})

		Add(&Command{
			Name:  "do",
			Short: "A command with subcommands",
			Long: `Command do helps testing subcommands.
Usage:
	do something
	do nothing`,
		})

		Add(&Command{
			Parent: "do",
			Name:   "something",
			Short:  "A subcommand that does something.",
			Long:   "do something does something",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the do something command.")
				return nil
			},
		})

		Add(&Command{
			Parent: "do",
			Name:   "nothing",
			Short:  "A subcommand that does nothing.",
			Long:   "do nothing does nothing",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the do nothing command.")
				return nil
			},
		})

		if err := Parse(); err != nil {
			fmt.Println(err)
		}

		Convey("readCommand should identify all of them correctly", func() {
			cmd, err := readCommand([]string{"test", "arg1", "arg2"})
			So(cmd, ShouldNotBeNil)
			So(cmd.Name, ShouldEqual, "test")
			So(cmd.Args, ShouldResemble, []string{"arg1", "arg2"})
			So(err, ShouldBeNil)

			cmd, err = readCommand([]string{"do", "something", "arg1"})
			So(cmd, ShouldNotBeNil)
			So(cmd.Name, ShouldEqual, "something")
			So(cmd.Args, ShouldResemble, []string{"arg1"})
			So(err, ShouldBeNil)

			cmd, err = readCommand([]string{"do", "nothing"})
			So(cmd, ShouldNotBeNil)
			So(cmd.Name, ShouldEqual, "nothing")
			So(cmd.Args, ShouldResemble, []string{})
			So(err, ShouldBeNil)
		})

		Convey("Usage() should print the usage", func() {
			Usage(nil)
			Usage(Commands["test"])
			Usage(Commands["flags"])
			Usage(Commands["do"])
			Usage(Commands["do"].children["something"])
			Usage(Commands["do"].children["nothing"])

		})
		Reset(func() {
			Commands = make(CommandMap)
		})
	})

}

func TestCheckFlags(t *testing.T) {
	var first int
	var second int
	var third int
	var global int
	var rejectedFlags map[string]bool

	// ContinueOnError is required when running goconvey as server; otherwise, unrecognized
	// flags that are passed to the test executable will cause an error:
	// "unknown shorthand flag: 't' in -test.v=true"
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
	// To suppress warnings resulting from standard flags -test and -json,
	// read -t and -j into dummy flags.
	// The flags used for actual testing must not use -t or -j shorthands.
	var testflag string
	var jsonflag string
	flag.StringVarP(&testflag, "t", "t", "t", "t")
	flag.StringVarP(&jsonflag, "j", "j", "j", "j")

	flag.IntVarP(&first, "first", "f", 1, "The first flag")
	flag.IntVarP(&second, "second", "s", 2, "The second flag")
	flag.IntVarP(&third, "third", "h", 3, "The third flag")
	flag.IntVarP(&global, "global", "g", 4, "The global flag")

	os.Args = []string{os.Args[0], "--first=10", "--second=20", "--third=30", "--global=40", "anargument", "anotherarg"}

	Commands = make(CommandMap) // clear the commands map for this test
	Add(&Command{
		Name:  "cmd12",
		Flags: []string{"first", "second"},
	})
	Add(&Command{
		Name:  "cmd23",
		Flags: []string{"second", "third"},
	})
	Add(&Command{
		Name:  "cmd123",
		Flags: []string{"first", "second", "third"},
	})

	if err := Parse(); err != nil {
		fmt.Println(err)
	}

	Convey("A command should accept its own flags and all global flags", t, func() {
		rejectedFlags = checkFlags(Commands["cmd123"])
		So(len(rejectedFlags), ShouldEqual, 0)
	})
	Convey("A command should reject the flags that belong to the other command only", t, func() {
		rejectedFlags = checkFlags(Commands["cmd23"])
		So(len(rejectedFlags), ShouldEqual, 1)
		So(rejectedFlags["first"], ShouldEqual, true)

		rejectedFlags = checkFlags(Commands["cmd12"])
		So(len(rejectedFlags), ShouldEqual, 1)
		So(rejectedFlags["third"], ShouldEqual, true)
	})
}
<|endoftext|>