<commit_before>/*
 * Copyright (c) 2016 Yanko Bolanos
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 */
package main

import (
	"context"
	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/firehose"
	"math"
	"sync"
	"time"
)

var (
	REQUEST_SIZE_LIMIT int = 2e+6 // 2MB
	BATCH_LIMIT            = 500
)

type FirehoseStream struct {
	svc          *firehose.Firehose
	stream       string
	dataChan     chan []byte
	interval     time.Duration
	recordFormat []Attribute
	ctx          context.Context
	wg           sync.WaitGroup
}

func NewFirehoseStream(ctx context.Context, recordFormat []Attribute, accessKey, secretAccessKey, awsRegion, awsSTSRole, streamName string) *FirehoseStream {

	if awsRegion == "" {
		log.Fatal("Please Specify the region your firehose is.")
	}

	s := &FirehoseStream{}
	sess := &session.Session{}
	awsConfig := &aws.Config{Region: aws.String(awsRegion)}

	if awsSTSRole != "" {
		sess = session.Must(session.NewSession())
		creds := stscreds.NewCredentials(sess, awsSTSRole)
		awsConfig.Credentials = creds
	} else if accessKey == "" || secretAccessKey == "" {
		// try IAM
		sess = session.New(nil)
	} else {
		creds := credentials.NewStaticCredentials(accessKey, secretAccessKey, "")
		config := &aws.Config{
			Region:      aws.String(awsRegion),
			Credentials: creds,
		}
		sess = session.New(config)
	}

	s.svc = firehose.New(sess, awsConfig)
	s.stream = streamName
	s.dataChan = make(chan []byte, BATCH_LIMIT*5)
	s.interval = 5 * time.Second
	s.recordFormat = recordFormat
	s.ctx = ctx

	go s.intervalStreamer()

	return s
}

func (s *FirehoseStream) Stream(r *Record) error {
	s.dataChan <- r.RecordToCSV()
	return nil
}

func (s *FirehoseStream) Close() {
	s.wg.Wait()
}

func (s *FirehoseStream) intervalStreamer() {

	accum := []*firehose.Record{}
	sizeAccumulator := 0
	timer := time.NewTicker(s.interval)
	exit := false
LOOP:
	for {

		data := []byte{}
		flush := false

		select {
		case incomingData := <-s.dataChan:
			data = incomingData
		case <-timer.C:
			flush = true
			data = nil
		case <-s.ctx.Done():
			flush = true
			data = nil
			log.Printf("context done. Force Flush")
			exit = true
		}

		if data != nil {

			if len(data) > 1000000 {
				log.Printf("Data recorded exceeded 1MB. skipping: \n--\n%s\n--", data)
				continue
			}

			sizeAccumulator += len(data)
			record := &firehose.Record{
				Data: data,
			}
			accum = append(accum, record)
		}

		if (len(accum) == BATCH_LIMIT || sizeAccumulator > REQUEST_SIZE_LIMIT || flush) && len(accum) > 0 {

			dataCopy := make([]*firehose.Record, len(accum))
			copy(dataCopy, accum)

			s.uploadRecords(dataCopy, 0)

			accum = []*firehose.Record{}
			sizeAccumulator = 0

		}

		if exit {
			s.wg.Done()
			s.wg.Wait()
			break LOOP
		}

	}
}

func (s *FirehoseStream) uploadRecords(data []*firehose.Record, failCount int) {
	s.wg.Add(1)
	go s._uploadRecords(data, failCount)
}

func (s *FirehoseStream) _uploadRecords(data []*firehose.Record, failCount int) {

	defer s.wg.Done()

	var sleepTime = time.Duration(math.Min(60.0, float64(5*failCount))) * time.Second
	if sleepTime > time.Duration(0) {
		log.Warnf("Retrying %v records in %v seconds", len(data), sleepTime)
	}
	time.Sleep(sleepTime)

	params := &firehose.PutRecordBatchInput{
		DeliveryStreamName: aws.String(s.stream),
		Records:            data,
	}

	r, err := s.svc.PutRecordBatch(params)
	if err != nil {
		log.Error(err.Error())
		s.uploadRecords(data, failCount+1)
		return
	}

	if *r.FailedPutCount > 0 {

		// Build a new array with the records that failed.
		newData := []*firehose.Record{}
		for i, req := range r.RequestResponses {
			if req.ErrorCode != nil {
				newData = append(newData, data[i])
			}
		}

		s.uploadRecords(data, failCount+1)
	}

	if failCount > 0 {
		log.Warnf("%v records succeeded after %v retries", len(data), failCount)
	}

}

func (s *FirehoseStream) RecordFormat() []Attribute {
	return s.recordFormat
}
<commit_msg>fixes #4<commit_after>/*
 * Copyright (c) 2016 Yanko Bolanos
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 */
package main

import (
	"context"
	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/firehose"
	"math"
	"sync"
	"time"
)

var (
	REQUEST_SIZE_LIMIT int = 2e+6 // 2MB
	BATCH_LIMIT            = 500
)

type FirehoseStream struct {
	svc          *firehose.Firehose
	stream       string
	dataChan     chan []byte
	interval     time.Duration
	recordFormat []Attribute
	ctx          context.Context
	wg           sync.WaitGroup
}

func NewFirehoseStream(ctx context.Context, recordFormat []Attribute, accessKey, secretAccessKey, awsRegion, awsSTSRole, streamName string) *FirehoseStream {

	if awsRegion == "" {
		log.Fatal("Please Specify the region your firehose is.")
	}

	s := &FirehoseStream{}
	sess := &session.Session{}
	awsConfig := &aws.Config{Region: aws.String(awsRegion)}

	if awsSTSRole != "" {
		sess = session.Must(session.NewSession())
		creds := stscreds.NewCredentials(sess, awsSTSRole)
		awsConfig.Credentials = creds
	} else if accessKey == "" || secretAccessKey == "" {
		// try IAM
		sess = session.New(nil)
	} else {
		creds := credentials.NewStaticCredentials(accessKey, secretAccessKey, "")
		config := &aws.Config{
			Region:      aws.String(awsRegion),
			Credentials: creds,
		}
		sess = session.New(config)
	}

	s.svc = firehose.New(sess, awsConfig)
	s.stream = streamName
	s.dataChan = make(chan []byte, BATCH_LIMIT*5)
	s.interval = 5 * time.Second
	s.recordFormat = recordFormat
	s.ctx = ctx

	s.wg.Add(1)
	go s.intervalStreamer()

	return s
}

func (s *FirehoseStream) Stream(r *Record) error {
	s.dataChan <- r.RecordToCSV()
	return nil
}

func (s *FirehoseStream) Close() {
	s.wg.Wait()
}

func (s *FirehoseStream) intervalStreamer() {

	accum := []*firehose.Record{}
	sizeAccumulator := 0
	timer := time.NewTicker(s.interval)
	exit := false
LOOP:
	for {

		data := []byte{}
		flush := false

		select {
		case incomingData := <-s.dataChan:
			data = incomingData
		case <-timer.C:
			flush = true
			data = nil
		case <-s.ctx.Done():
			flush = true
			data = nil
			log.Printf("context done. Force Flush")
			exit = true
		}

		if data != nil {

			if len(data) > 1000000 {
				log.Printf("Data recorded exceeded 1MB. skipping: \n--\n%s\n--", data)
				continue
			}

			sizeAccumulator += len(data)
			record := &firehose.Record{
				Data: data,
			}
			accum = append(accum, record)
		}

		if (len(accum) == BATCH_LIMIT || sizeAccumulator > REQUEST_SIZE_LIMIT || flush) && len(accum) > 0 {

			dataCopy := make([]*firehose.Record, len(accum))
			copy(dataCopy, accum)

			s.uploadRecords(dataCopy, 0)

			accum = []*firehose.Record{}
			sizeAccumulator = 0

		}

		if exit {
			s.wg.Done()
			s.wg.Wait()
			break LOOP
		}

	}
}

func (s *FirehoseStream) uploadRecords(data []*firehose.Record, failCount int) {
	s.wg.Add(1)
	go s._uploadRecords(data, failCount)
}

func (s *FirehoseStream) _uploadRecords(data []*firehose.Record, failCount int) {

	defer s.wg.Done()

	var sleepTime = time.Duration(math.Min(60.0, float64(5*failCount))) * time.Second
	if sleepTime > time.Duration(0) {
		log.Warnf("Retrying %v records in %v seconds", len(data), sleepTime)
	}
	time.Sleep(sleepTime)

	params := &firehose.PutRecordBatchInput{
		DeliveryStreamName: aws.String(s.stream),
		Records:            data,
	}

	r, err := s.svc.PutRecordBatch(params)
	if err != nil {
		log.Error(err.Error())
		s.uploadRecords(data, failCount+1)
		return
	}

	if *r.FailedPutCount > 0 {

		// Build a new array with the records that failed.
		newData := []*firehose.Record{}
		for i, req := range r.RequestResponses {
			if req.ErrorCode != nil {
				newData = append(newData, data[i])
			}
		}

		s.uploadRecords(data, failCount+1)
	}

	if failCount > 0 {
		log.Warnf("%v records succeeded after %v retries", len(data), failCount)
	}

}

func (s *FirehoseStream) RecordFormat() []Attribute {
	return s.recordFormat
}
<|endoftext|>


<commit_before>package handler

import (
	"github.com/gorilla/mux"
	"github.com/justinas/alice"
	"github.com/rs/zerolog"
)

// NewMuxRouter sets up the mux.Router and registers routes to URL paths
// using the available handlers
func NewMuxRouter(logger zerolog.Logger, handlers Handlers) *mux.Router {
	// create a new gorilla/mux router
	rtr := mux.NewRouter()

	// Start a new alice handler chain
	c := alice.New()

	// add LoggerHandlerChain handler chain and zerolog logger to Context
	c = LoggerHandlerChain(logger, c)

	// send Router through PathPrefix method to validate any standard
	// subroutes you may want for your APIs. e.g. I always want to be
	// sure that every request has "/api" as part of it's path prefix
	// without having to put it into every handle path in my various
	// routing functions
	rtr = rtr.PathPrefix("/api").Subrouter()

	// Match only POST requests at /api/v1/movies
	// with Content-Type header = application/json
	rtr.Handle("/v1/movies",
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.CreateMovieHandler)).
		Methods("POST").
		Headers("Content-Type", "application/json")

	// Match only PUT requests having an ID at /api/v1/movies/{id}
	// with the Content-Type header = application/json
	rtr.Handle("/v1/movies/{id}",
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.UpdateMovieHandler)).
		Methods("PUT").
		Headers("Content-Type", "application/json")

	// Match only DELETE requests having an ID at /api/v1/movies/{id}
	rtr.Handle("/v1/movies/{id}",
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.DeleteMovieHandler)).
		Methods("DELETE")

	// Match only GET requests having an ID at /api/v1/movies/{id}
	rtr.Handle("/v1/movies/{id}",
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.FindMovieByIDHandler)).
		Methods("GET")

	// Match only GET requests /api/v1/movies
	rtr.Handle("/v1/movies",
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.FindAllMoviesHandler)).
		Methods("GET")

	// Match only GET requests at /api/v1/ping
	rtr.Handle("/v1/ping",
		c.Append(JSONContentTypeHandler).
			Then(handlers.PingHandler)).
		Methods("GET")

	return rtr
}
<commit_msg>use constants<commit_after>package handler

import (
	"net/http"

	"github.com/gorilla/mux"
	"github.com/justinas/alice"
	"github.com/rs/zerolog"
)

const (
	pathPrefix       string = "/api"
	moviesV1PathRoot string = "/v1/movies"
)

// NewMuxRouter sets up the mux.Router and registers routes to URL paths
// using the available handlers
func NewMuxRouter(logger zerolog.Logger, handlers Handlers) *mux.Router {
	// create a new gorilla/mux router
	rtr := mux.NewRouter()

	// Start a new alice handler chain
	c := alice.New()

	// add LoggerHandlerChain handler chain and zerolog logger to Context
	c = LoggerHandlerChain(logger, c)

	// send Router through PathPrefix method to validate any standard
	// subroutes you may want for your APIs. e.g. I always want to be
	// sure that every request has "/api" as part of it's path prefix
	// without having to put it into every handle path in my various
	// routing functions
	rtr = rtr.PathPrefix(pathPrefix).Subrouter()

	// Match only POST requests at /api/v1/movies
	// with Content-Type header = application/json
	rtr.Handle(moviesV1PathRoot,
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.CreateMovieHandler)).
		Methods(http.MethodPost).
		Headers("Content-Type", "application/json")

	// Match only PUT requests having an ID at /api/v1/movies/{id}
	// with the Content-Type header = application/json
	rtr.Handle(moviesV1PathRoot+"/{id}",
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.UpdateMovieHandler)).
		Methods(http.MethodPut).
		Headers("Content-Type", "application/json")

	// Match only DELETE requests having an ID at /api/v1/movies/{id}
	rtr.Handle(moviesV1PathRoot+"/{id}",
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.DeleteMovieHandler)).
		Methods(http.MethodDelete)

	// Match only GET requests having an ID at /api/v1/movies/{id}
	rtr.Handle(moviesV1PathRoot+"/{id}",
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.FindMovieByIDHandler)).
		Methods(http.MethodGet)

	// Match only GET requests /api/v1/movies
	rtr.Handle(moviesV1PathRoot,
		c.Append(AccessTokenHandler).
			Append(JSONContentTypeHandler).
			Then(handlers.FindAllMoviesHandler)).
		Methods(http.MethodGet)

	// Match only GET requests at /api/v1/ping
	rtr.Handle("/v1/ping",
		c.Append(JSONContentTypeHandler).
			Then(handlers.PingHandler)).
		Methods(http.MethodGet)

	return rtr
}
<|endoftext|>


<commit_before>package handlers

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"github.com/arschles/gbs/log"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/gorilla/mux"
)

const (
	defaultBuildImg = "quay.io/arschles/gbs-env:0.0.1"
)

func BuildURL() string {
	return fmt.Sprintf("/{%s}/{%s}/{%s}", site, org, repo)
}

type startBuildReq struct {
	BuildEnv string `json:"build_env"`
}

type startBuildResp struct {
	StatusURL string `json:"status_url"`
}

func Build(workdir string, dockerCl *docker.Client) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		flusher, ok := w.(http.Flusher)
		if !ok {
			httpErrf(w, http.StatusInternalServerError, "server doesn't support flushing output")
			return
		}

		site, ok := mux.Vars(r)[site]
		if !ok {
			httpErrf(w, http.StatusBadRequest, "missing site in path")
			return
		}

		org, ok := mux.Vars(r)[org]
		if !ok {
			httpErrf(w, http.StatusBadRequest, "missing org in path")
			return
		}

		repo, ok := mux.Vars(r)[repo]
		if !ok {
			httpErrf(w, http.StatusBadRequest, "missing repo in path")
			return
		}

		buildImg := defaultBuildImg
		req := new(startBuildReq)
		if err := json.NewDecoder(r.Body).Decode(req); err == nil {
			buildImg = req.BuildEnv
		}
		defer r.Body.Close()

		containerOpts := createContainerOpts(buildImg, workdir, site, org, repo)
		container, err := dockerCl.CreateContainer(containerOpts)
		if err != nil {
			log.Errf("creating container [%s]", err)
			httpErrf(w, http.StatusInternalServerError, "error creating container [%s]", err)
			return
		}

		hostConfig := &docker.HostConfig{Binds: []string{fmt.Sprintf("%s:%s", workdir, absPwd)}}
		if err := dockerCl.StartContainer(container.ID, hostConfig); err != nil {
			log.Errf("starting container [%s]", err)
			httpErrf(w, http.StatusInternalServerError, "error starting container [%s]", err)
			return
		}

		attachOpts, outputReader := attachToContainerOpts(container.ID)
		errCh := make(chan error)
		go func() {
			if err := dockerCl.AttachToContainer(attachOpts); err != nil {
				errCh <- err
			}
		}()

		go func(reader io.Reader) {
			scanner := bufio.NewScanner(reader)
			for scanner.Scan() {
				fmt.Fprintf(w, "%s\n", scanner.Text())
				flusher.Flush()
			}
			if err := scanner.Err(); err != nil {
				fmt.Fprintf(w, "error with scanner in attached container [%s]\n", err)
			}
		}(outputReader)

		code, err := dockerCl.WaitContainer(container.ID)
		if err != nil {
			log.Errf("waiting for container %s [%s]", container.ID, err)
			return
		}
		w.WriteHeader(http.StatusCreated)
		w.Write([]byte(fmt.Sprintf("exited with error code %d\n", code)))

		removeOpts := docker.RemoveContainerOptions{
			ID:            container.ID,
			RemoveVolumes: true,
			Force:         true,
		}
		if err := dockerCl.RemoveContainer(removeOpts); err != nil {
			log.Errf("removing container %s [%s]", container.ID, err)
		}
	})
}
<commit_msg>write created before flushing<commit_after>package handlers

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"github.com/arschles/gbs/log"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/gorilla/mux"
)

const (
	defaultBuildImg = "quay.io/arschles/gbs-env:0.0.1"
)

func BuildURL() string {
	return fmt.Sprintf("/{%s}/{%s}/{%s}", site, org, repo)
}

type startBuildReq struct {
	BuildEnv string `json:"build_env"`
}

type startBuildResp struct {
	StatusURL string `json:"status_url"`
}

func Build(workdir string, dockerCl *docker.Client) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		flusher, ok := w.(http.Flusher)
		if !ok {
			httpErrf(w, http.StatusInternalServerError, "server doesn't support flushing output")
			return
		}

		site, ok := mux.Vars(r)[site]
		if !ok {
			httpErrf(w, http.StatusBadRequest, "missing site in path")
			return
		}

		org, ok := mux.Vars(r)[org]
		if !ok {
			httpErrf(w, http.StatusBadRequest, "missing org in path")
			return
		}

		repo, ok := mux.Vars(r)[repo]
		if !ok {
			httpErrf(w, http.StatusBadRequest, "missing repo in path")
			return
		}

		buildImg := defaultBuildImg
		req := new(startBuildReq)
		if err := json.NewDecoder(r.Body).Decode(req); err == nil {
			buildImg = req.BuildEnv
		}
		defer r.Body.Close()

		containerOpts := createContainerOpts(buildImg, workdir, site, org, repo)
		container, err := dockerCl.CreateContainer(containerOpts)
		if err != nil {
			log.Errf("creating container [%s]", err)
			httpErrf(w, http.StatusInternalServerError, "error creating container [%s]", err)
			return
		}

		hostConfig := &docker.HostConfig{Binds: []string{fmt.Sprintf("%s:%s", workdir, absPwd)}}
		if err := dockerCl.StartContainer(container.ID, hostConfig); err != nil {
			log.Errf("starting container [%s]", err)
			httpErrf(w, http.StatusInternalServerError, "error starting container [%s]", err)
			return
		}

		attachOpts, outputReader := attachToContainerOpts(container.ID)
		errCh := make(chan error)
		go func() {
			if err := dockerCl.AttachToContainer(attachOpts); err != nil {
				errCh <- err
			}
		}()

		w.WriteHeader(http.StatusCreated)
		go func(reader io.Reader) {
			scanner := bufio.NewScanner(reader)
			for scanner.Scan() {
				fmt.Fprintf(w, "%s\n", scanner.Text())
				flusher.Flush()
			}
			if err := scanner.Err(); err != nil {
				fmt.Fprintf(w, "error with scanner in attached container [%s]\n", err)
			}
		}(outputReader)

		code, err := dockerCl.WaitContainer(container.ID)
		if err != nil {
			log.Errf("waiting for container %s [%s]", container.ID, err)
			return
		}
		w.Write([]byte(fmt.Sprintf("exited with error code %d\n", code)))

		removeOpts := docker.RemoveContainerOptions{
			ID:            container.ID,
			RemoveVolumes: true,
			Force:         true,
		}
		if err := dockerCl.RemoveContainer(removeOpts); err != nil {
			log.Errf("removing container %s [%s]", container.ID, err)
		}
	})
}
<|endoftext|>


<commit_before>package braintree

import (
	"testing"
	"time"

	"github.com/alexd765/braintree/btdate"
	"github.com/shopspring/decimal"
)

func TestCancelSubscription(t *testing.T) {
	t.Parallel()

	t.Run("existing", func(t *testing.T) {
		t.Parallel()
		customer, err := bt.Customer().Create(CustomerInput{
			FirstName: "first",
			CreditCard: &CreditCardInput{
				PaymentMethodNonce: "fake-valid-visa-nonce",
			},
		})
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}

		subscription, err := bt.Subscription().Create(
			SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
			},
		)
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}

		subscription, err = bt.Subscription().Cancel(subscription.ID)
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}
		if subscription.Status != SubscriptionStatusCanceled {
			t.Errorf("subscription.Status: got: %s, want: %s", subscription.Status, SubscriptionStatusCanceled)
		}
	})

	t.Run("nonExisting", func(t *testing.T) {
		t.Parallel()
		if _, err := bt.Subscription().Cancel("nonExisting"); err == nil || err.Error() != "404 Not Found" {
			t.Errorf("got: %v, want: 404 Not Found", err)
		}
	})
}

func TestCreateSubscription(t *testing.T) {
	t.Parallel()

	customer, err := bt.Customer().Create(CustomerInput{
		FirstName: "first",
		CreditCard: &CreditCardInput{
			PaymentMethodNonce: "fake-valid-visa-nonce",
		},
	})
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}

	tests := []struct {
		name                string
		input               SubscriptionInput
		wantStartDate       btdate.Date
		wantNextBillingDate btdate.Date
	}{
		{
			name: "normal",
			input: SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
			},
			wantStartDate:       btdate.Today(),
			wantNextBillingDate: btdate.FromTime(time.Now().UTC().AddDate(0, 1, 0)),
		},
		{
			name: "trial1day",
			input: SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
				TrialDuration:      1,
				TrialDurationUnit:  "day",
				TrialPeriod:        true,
			},
			wantStartDate:       btdate.Date{},
			wantNextBillingDate: btdate.FromTime(time.Now().UTC().AddDate(0, 0, 1)),
		},
		{
			name: "trial2weeks",
			input: SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
				TrialDuration:      2,
				TrialDurationUnit:  "month",
				TrialPeriod:        true,
			},
			wantStartDate:       btdate.Date{},
			wantNextBillingDate: btdate.FromTime(time.Now().UTC().AddDate(0, 2, 0)),
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()

			subscription, err := bt.Subscription().Create(test.input)
			if err != nil {
				t.Fatalf("unexpected err: %s", err)
			}
			if subscription.PlanID != "plan1" {
				t.Errorf("subscription.PlanID: got: %s, want: plan1", subscription.PlanID)
			}
			if subscription.BillingPeriodStartDate != test.wantStartDate {
				t.Errorf("subscription.BillingPeriodStartDate: got %s, want %s", subscription.BillingPeriodStartDate, test.wantStartDate)
			}
			if subscription.NextBillingDate != test.wantNextBillingDate {
				t.Errorf("subscription.NextBillingDate: got %s, want %s", subscription.NextBillingDate, test.wantNextBillingDate)
			}
		})
	}

	t.Run("withoutToken", func(t *testing.T) {
		_, err := bt.Subscription().Create(SubscriptionInput{PlanID: "plan1"})
		apiErr, ok := err.(*APIError)
		if !ok {
			t.Fatalf("expected APIError")
		}
		if apiErr == nil || apiErr.Code != 91903 {
			t.Errorf("got %v, want error code 91903", apiErr)
		}
	})
}

func TestFindSubscription(t *testing.T) {
	t.Parallel()

	t.Run("existing", func(t *testing.T) {
		t.Parallel()

		subscription, err := bt.Subscription().Find("sub1")
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}
		if !subscription.Price.Equals(decimal.NewFromFloat(5)) {
			t.Errorf("subscription price: got %s, want 5", subscription.Price)
		}
		if len(subscription.Transactions) == 0 {
			t.Fatalf("subscription.Transactions: got 0, want more")
		}
		if planID := subscription.Transactions[0].PlanID; planID != "plan1" {
			t.Errorf("subscription.Transactions[0].PlanID: got %s, want plan1", planID)
		}
	})

	t.Run("nonExisting", func(t *testing.T) {
		t.Parallel()

		subscription, err := bt.Subscription().Find("sub2")
		if err == nil || err.Error() != "404 Not Found" {
			t.Errorf("got: %v, want: 404 Not Found", err)
		}
		if subscription != nil {
			t.Errorf("got: %+v, want: <nil>", subscription)
		}
	})
}

func TestUpdateSubscription(t *testing.T) {
	t.Parallel()

	t.Run("existing", func(t *testing.T) {
		t.Parallel()

		customer, err := bt.Customer().Create(CustomerInput{
			FirstName: "first",
			CreditCard: &CreditCardInput{
				PaymentMethodNonce: "fake-valid-visa-nonce",
			},
		})
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}

		subscription, err := bt.Subscription().Create(
			SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
			},
		)
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}

		wantPrice := decimal.NewFromFloat(6)
		subscription, err = bt.Subscription().Update(
			SubscriptionInput{
				ID:    subscription.ID,
				Price: &wantPrice,
			},
		)
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}
		if !subscription.Price.Equals(decimal.NewFromFloat(6)) {
			t.Errorf("subscription.Price: got %s, want 6", subscription.Price)
		}
	})

	t.Run("nonExisting", func(t *testing.T) {
		t.Parallel()

		subscription, err := bt.Subscription().Update(SubscriptionInput{ID: "cus2"})
		if err == nil || err.Error() != "404 Not Found" {
			t.Errorf("got: %v, want: 404 Not Found", err)
		}
		if subscription != nil {
			t.Errorf("got: %+v, want: <nil>", subscription)
		}
	})
}
<commit_msg>generate past due subscriptions for testing<commit_after>package braintree

import (
	"fmt"
	"testing"
	"time"

	"github.com/alexd765/braintree/btdate"
	"github.com/shopspring/decimal"
)

func TestCancelSubscription(t *testing.T) {
	t.Parallel()

	t.Run("existing", func(t *testing.T) {
		t.Parallel()
		customer, err := bt.Customer().Create(CustomerInput{
			FirstName: "first",
			CreditCard: &CreditCardInput{
				PaymentMethodNonce: "fake-valid-visa-nonce",
			},
		})
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}

		subscription, err := bt.Subscription().Create(
			SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
			},
		)
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}

		subscription, err = bt.Subscription().Cancel(subscription.ID)
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}
		if subscription.Status != SubscriptionStatusCanceled {
			t.Errorf("subscription.Status: got: %s, want: %s", subscription.Status, SubscriptionStatusCanceled)
		}
	})

	t.Run("nonExisting", func(t *testing.T) {
		t.Parallel()
		if _, err := bt.Subscription().Cancel("nonExisting"); err == nil || err.Error() != "404 Not Found" {
			t.Errorf("got: %v, want: 404 Not Found", err)
		}
	})
}

func TestCreateSubscription(t *testing.T) {
	t.Parallel()

	customer, err := bt.Customer().Create(CustomerInput{
		FirstName: "first",
		CreditCard: &CreditCardInput{
			PaymentMethodNonce: "fake-valid-visa-nonce",
		},
	})
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}

	tests := []struct {
		name                string
		input               SubscriptionInput
		wantStartDate       btdate.Date
		wantNextBillingDate btdate.Date
	}{
		{
			name: "normal",
			input: SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
			},
			wantStartDate:       btdate.Today(),
			wantNextBillingDate: btdate.FromTime(time.Now().UTC().AddDate(0, 1, 0)),
		},
		{
			name: "trial1day",
			input: SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
				TrialDuration:      1,
				TrialDurationUnit:  "day",
				TrialPeriod:        true,
			},
			wantStartDate:       btdate.Date{},
			wantNextBillingDate: btdate.FromTime(time.Now().UTC().AddDate(0, 0, 1)),
		},
		{
			name: "trial2weeks",
			input: SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
				TrialDuration:      2,
				TrialDurationUnit:  "month",
				TrialPeriod:        true,
			},
			wantStartDate:       btdate.Date{},
			wantNextBillingDate: btdate.FromTime(time.Now().UTC().AddDate(0, 2, 0)),
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()

			subscription, err := bt.Subscription().Create(test.input)
			if err != nil {
				t.Fatalf("unexpected err: %s", err)
			}
			if subscription.PlanID != "plan1" {
				t.Errorf("subscription.PlanID: got: %s, want: plan1", subscription.PlanID)
			}
			if subscription.BillingPeriodStartDate != test.wantStartDate {
				t.Errorf("subscription.BillingPeriodStartDate: got %s, want %s", subscription.BillingPeriodStartDate, test.wantStartDate)
			}
			if subscription.NextBillingDate != test.wantNextBillingDate {
				t.Errorf("subscription.NextBillingDate: got %s, want %s", subscription.NextBillingDate, test.wantNextBillingDate)
			}
		})
	}

	t.Run("withoutToken", func(t *testing.T) {
		_, err := bt.Subscription().Create(SubscriptionInput{PlanID: "plan1"})
		apiErr, ok := err.(*APIError)
		if !ok {
			t.Fatalf("expected APIError")
		}
		if apiErr == nil || apiErr.Code != 91903 {
			t.Errorf("got %v, want error code 91903", apiErr)
		}
	})
}

func TestFindSubscription(t *testing.T) {
	t.Parallel()

	t.Run("existing", func(t *testing.T) {
		t.Parallel()

		subscription, err := bt.Subscription().Find("sub1")
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}
		if !subscription.Price.Equals(decimal.NewFromFloat(5)) {
			t.Errorf("subscription price: got %s, want 5", subscription.Price)
		}
		if len(subscription.Transactions) == 0 {
			t.Fatalf("subscription.Transactions: got 0, want more")
		}
		if planID := subscription.Transactions[0].PlanID; planID != "plan1" {
			t.Errorf("subscription.Transactions[0].PlanID: got %s, want plan1", planID)
		}
	})

	t.Run("nonExisting", func(t *testing.T) {
		t.Parallel()

		subscription, err := bt.Subscription().Find("sub2")
		if err == nil || err.Error() != "404 Not Found" {
			t.Errorf("got: %v, want: 404 Not Found", err)
		}
		if subscription != nil {
			t.Errorf("got: %+v, want: <nil>", subscription)
		}
	})
}

func TestUpdateSubscription(t *testing.T) {
	t.Parallel()

	t.Run("existing", func(t *testing.T) {
		t.Parallel()

		customer, err := bt.Customer().Create(CustomerInput{
			FirstName: "first",
			CreditCard: &CreditCardInput{
				PaymentMethodNonce: "fake-valid-visa-nonce",
			},
		})
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}

		subscription, err := bt.Subscription().Create(
			SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
			},
		)
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}

		wantPrice := decimal.NewFromFloat(6)
		subscription, err = bt.Subscription().Update(
			SubscriptionInput{
				ID:    subscription.ID,
				Price: &wantPrice,
			},
		)
		if err != nil {
			t.Fatalf("unexpected err: %s", err)
		}
		if !subscription.Price.Equals(decimal.NewFromFloat(6)) {
			t.Errorf("subscription.Price: got %s, want 6", subscription.Price)
		}
	})

	t.Run("nonExisting", func(t *testing.T) {
		t.Parallel()

		subscription, err := bt.Subscription().Update(SubscriptionInput{ID: "cus2"})
		if err == nil || err.Error() != "404 Not Found" {
			t.Errorf("got: %v, want: 404 Not Found", err)
		}
		if subscription != nil {
			t.Errorf("got: %+v, want: <nil>", subscription)
		}
	})
}

// TestGeneratePastDueSubscriptions will generate subscriptions for a user
//"pastDue" with variying trial durations. When the trial is over the charge
// attempt will fail and the subscription state will be Past Due.
func TestGeneratePastDueSubscriptions(t *testing.T) {
	t.Skip("Not really a test.")

	customer, err := bt.Customer().Find("pastDue")
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}

	twoThousand := decimal.NewFromFloat(2000)
	for day := 1; day <= 14; day++ {
		fmt.Printf("day %d\n", day)
		for count := 0; count < 10; count++ {
			_, err := bt.Subscription().Create(SubscriptionInput{
				PlanID:             "plan1",
				PaymentMethodToken: customer.CreditCards[0].Token,
				Price:              &twoThousand,
				TrialDuration:      day,
				TrialDurationUnit:  "day",
				TrialPeriod:        true,
			})
			if err != nil {
				t.Fatalf("unexpected err: %s", err)
			}
		}
	}
}
<|endoftext|>


<commit_before>package ecs

import (
	"encoding/json"
	"strconv"
	"time"

	"github.com/denverdino/aliyungo/common"
	"github.com/denverdino/aliyungo/util"
)

// InstanceStatus represents instance status
type InstanceStatus string

// Constants of InstanceStatus
const (
	Creating = InstanceStatus("Creating")
	Running  = InstanceStatus("Running")
	Starting = InstanceStatus("Starting")

	Stopped  = InstanceStatus("Stopped")
	Stopping = InstanceStatus("Stopping")
)

type LockReason struct {
	LockReason string
}

type DescribeInstanceStatusArgs struct {
	RegionId common.Region
	ZoneId   string
	common.Pagination
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instancestatusitemtype
type InstanceStatusItemType struct {
	InstanceId string
	Status     InstanceStatus
}

type DescribeInstanceStatusResponse struct {
	common.Response
	common.PaginationResult
	InstanceStatuses struct {
		InstanceStatus []InstanceStatusItemType
	}
}

// DescribeInstanceStatus describes instance status
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancestatus
func (client *Client) DescribeInstanceStatus(args *DescribeInstanceStatusArgs) (instanceStatuses []InstanceStatusItemType, pagination *common.PaginationResult, err error) {
	args.Validate()
	response := DescribeInstanceStatusResponse{}

	err = client.Invoke("DescribeInstanceStatus", args, &response)

	if err == nil {
		return response.InstanceStatuses.InstanceStatus, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type StopInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type StopInstanceResponse struct {
	common.Response
}

// StopInstance stops instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&stopinstance
func (client *Client) StopInstance(instanceId string, forceStop bool) error {
	args := StopInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := StopInstanceResponse{}
	err := client.Invoke("StopInstance", &args, &response)
	return err
}

type StartInstanceArgs struct {
	InstanceId string
}

type StartInstanceResponse struct {
	common.Response
}

// StartInstance starts instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&startinstance
func (client *Client) StartInstance(instanceId string) error {
	args := StartInstanceArgs{InstanceId: instanceId}
	response := StartInstanceResponse{}
	err := client.Invoke("StartInstance", &args, &response)
	return err
}

type RebootInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type RebootInstanceResponse struct {
	common.Response
}

// RebootInstance reboot instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&rebootinstance
func (client *Client) RebootInstance(instanceId string, forceStop bool) error {
	request := RebootInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := RebootInstanceResponse{}
	err := client.Invoke("RebootInstance", &request, &response)
	return err
}

type DescribeInstanceAttributeArgs struct {
	InstanceId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&operationlockstype
type OperationLocksType struct {
	LockReason []LockReason //enum for financial, security
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&securitygroupidsettype
type SecurityGroupIdSetType struct {
	SecurityGroupId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&ipaddresssettype
type IpAddressSetType struct {
	IpAddress []string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&vpcattributestype
type VpcAttributesType struct {
	VpcId            string
	VSwitchId        string
	PrivateIpAddress IpAddressSetType
	NatIpAddress     string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&eipaddressassociatetype
type EipAddressAssociateType struct {
	AllocationId       string
	IpAddress          string
	Bandwidth          int
	InternetChargeType common.InternetChargeType
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instanceattributestype
type InstanceAttributesType struct {
	InstanceId         string
	InstanceName       string
	Description        string
	ImageId            string
	RegionId           common.Region
	ZoneId             string
	CPU                int
	Memory             int
	ClusterId          string
	InstanceType       string
	InstanceTypeFamily string
	HostName           string
	SerialNumber       string
	Status             InstanceStatus
	OperationLocks     OperationLocksType
	SecurityGroupIds   struct {
		SecurityGroupId []string
	}
	PublicIpAddress         IpAddressSetType
	InnerIpAddress          IpAddressSetType
	InstanceNetworkType     string //enum Classic | Vpc
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	InternetChargeType      common.InternetChargeType
	CreationTime            util.ISO6801Time //time.Time
	VpcAttributes           VpcAttributesType
	EipAddress              EipAddressAssociateType
	IoOptimized             StringOrBool
	InstanceChargeType      common.InternetChargeType
	ExpiredTime             util.ISO6801Time
}

type DescribeInstanceAttributeResponse struct {
	common.Response
	InstanceAttributesType
}

// DescribeInstanceAttribute describes instance attribute
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstanceattribute
func (client *Client) DescribeInstanceAttribute(instanceId string) (instance *InstanceAttributesType, err error) {
	args := DescribeInstanceAttributeArgs{InstanceId: instanceId}

	response := DescribeInstanceAttributeResponse{}
	err = client.Invoke("DescribeInstanceAttribute", &args, &response)
	if err != nil {
		return nil, err
	}
	return &response.InstanceAttributesType, err
}

type ModifyInstanceAttributeArgs struct {
	InstanceId   string
	InstanceName string
	Description  string
	Password     string
	HostName     string
}

type ModifyInstanceAttributeResponse struct {
	common.Response
}

//ModifyInstanceAttribute  modify instance attrbute
//
// You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/modifyinstanceattribute.html
func (client *Client) ModifyInstanceAttribute(args *ModifyInstanceAttributeArgs) error {
	response := ModifyInstanceAttributeResponse{}
	err := client.Invoke("ModifyInstanceAttribute", args, &response)
	return err
}

// Default timeout value for WaitForInstance method
const InstanceDefaultTimeout = 120

// WaitForInstance waits for instance to given status
func (client *Client) WaitForInstance(instanceId string, status InstanceStatus, timeout int) error {
	if timeout <= 0 {
		timeout = InstanceDefaultTimeout
	}
	for {
		instance, err := client.DescribeInstanceAttribute(instanceId)
		if err != nil {
			return err
		}
		if instance.Status == status {
			//TODO
			//Sleep one more time for timing issues
			time.Sleep(DefaultWaitForInterval * time.Second)
			break
		}
		timeout = timeout - DefaultWaitForInterval
		if timeout <= 0 {
			return common.GetClientErrorFromString("Timeout")
		}
		time.Sleep(DefaultWaitForInterval * time.Second)

	}
	return nil
}

type DescribeInstanceVncUrlArgs struct {
	RegionId   common.Region
	InstanceId string
}

type DescribeInstanceVncUrlResponse struct {
	common.Response
	VncUrl string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancevncurl
func (client *Client) DescribeInstanceVncUrl(args *DescribeInstanceVncUrlArgs) (string, error) {
	response := DescribeInstanceVncUrlResponse{}

	err := client.Invoke("DescribeInstanceVncUrl", args, &response)

	if err == nil {
		return response.VncUrl, nil
	}

	return "", err
}

type DescribeInstancesArgs struct {
	RegionId            common.Region
	VpcId               string
	VSwitchId           string
	ZoneId              string
	InstanceIds         string
	InstanceNetworkType string
	PrivateIpAddresses  string
	InnerIpAddresses    string
	PublicIpAddresses   string
	SecurityGroupId     string
	common.Pagination
}

type DescribeInstancesResponse struct {
	common.Response
	common.PaginationResult
	Instances struct {
		Instance []InstanceAttributesType
	}
}

// DescribeInstances describes instances
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstances
func (client *Client) DescribeInstances(args *DescribeInstancesArgs) (instances []InstanceAttributesType, pagination *common.PaginationResult, err error) {
	args.Validate()
	response := DescribeInstancesResponse{}

	err = client.Invoke("DescribeInstances", args, &response)

	if err == nil {
		return response.Instances.Instance, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type DeleteInstanceArgs struct {
	InstanceId string
}

type DeleteInstanceResponse struct {
	common.Response
}

// DeleteInstance deletes instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&deleteinstance
func (client *Client) DeleteInstance(instanceId string) error {
	args := DeleteInstanceArgs{InstanceId: instanceId}
	response := DeleteInstanceResponse{}
	err := client.Invoke("DeleteInstance", &args, &response)
	return err
}

type DataDiskType struct {
	Size               int
	Category           DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	SnapshotId         string
	DiskName           string
	Description        string
	Device             string
	DeleteWithInstance bool
}

type SystemDiskType struct {
	Size        int
	Category    DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	DiskName    string
	Description string
}

type IoOptimized string

type StringOrBool struct {
	Value bool
}

// UnmarshalJSON implements the json.Unmarshaller interface.
func (io *StringOrBool) UnmarshalJSON(value []byte) error {
	if value[0] == '"' {
		var str string
		err := json.Unmarshal(value, &str)
		if err == nil {
			io.Value = (str == "true" || str == "optimized")
		}
		return err
	}
	var boolVal bool
	err := json.Unmarshal(value, &boolVal)
	if err == nil {
		io.Value = boolVal
	}
	return err
}

func (io StringOrBool) Bool() bool {
	return io.Value
}

func (io StringOrBool) String() string {
	return strconv.FormatBool(io.Value)
}

var (
	IoOptimizedNone      = IoOptimized("none")
	IoOptimizedOptimized = IoOptimized("optimized")
)

type CreateInstanceArgs struct {
	RegionId                common.Region
	ZoneId                  string
	ImageId                 string
	InstanceType            string
	SecurityGroupId         string
	InstanceName            string
	Description             string
	InternetChargeType      common.InternetChargeType
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	HostName                string
	Password                string
	IoOptimized             IoOptimized
	SystemDisk              SystemDiskType
	DataDisk                []DataDiskType
	VSwitchId               string
	PrivateIpAddress        string
	ClientToken             string
	InstanceChargeType      common.InstanceChargeType
	Period                  int
}

type CreateInstanceResponse struct {
	common.Response
	InstanceId string
}

// CreateInstance creates instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&createinstance
func (client *Client) CreateInstance(args *CreateInstanceArgs) (instanceId string, err error) {
	response := CreateInstanceResponse{}
	err = client.Invoke("CreateInstance", args, &response)
	if err != nil {
		return "", err
	}
	return response.InstanceId, err
}

type SecurityGroupArgs struct {
	InstanceId      string
	SecurityGroupId string
}

type SecurityGroupResponse struct {
	common.Response
}

//JoinSecurityGroup
//
//You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/joinsecuritygroup.html
func (client *Client) JoinSecurityGroup(instanceId string, securityGroupId string) error {
	args := SecurityGroupArgs{InstanceId: instanceId, SecurityGroupId: securityGroupId}
	response := SecurityGroupResponse{}
	err := client.Invoke("JoinSecurityGroup", &args, &response)
	return err
}

//LeaveSecurityGroup
//
//You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/leavesecuritygroup.html
func (client *Client) LeaveSecurityGroup(instanceId string, securityGroupId string) error {
	args := SecurityGroupArgs{InstanceId: instanceId, SecurityGroupId: securityGroupId}
	response := SecurityGroupResponse{}
	err := client.Invoke("LeaveSecurityGroup", &args, &response)
	return err
}
<commit_msg>Refactor the change for LockReason<commit_after>package ecs

import (
	"encoding/json"
	"strconv"
	"time"

	"github.com/denverdino/aliyungo/common"
	"github.com/denverdino/aliyungo/util"
)

// InstanceStatus represents instance status
type InstanceStatus string

// Constants of InstanceStatus
const (
	Creating = InstanceStatus("Creating")
	Running  = InstanceStatus("Running")
	Starting = InstanceStatus("Starting")

	Stopped  = InstanceStatus("Stopped")
	Stopping = InstanceStatus("Stopping")
)

type LockReason string

const (
	LockReasonFinancial = LockReason("financial")
	LockReasonSecurity  = LockReason("security")
)

type LockReasonType struct {
	LockReason LockReason
}

type DescribeInstanceStatusArgs struct {
	RegionId common.Region
	ZoneId   string
	common.Pagination
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instancestatusitemtype
type InstanceStatusItemType struct {
	InstanceId string
	Status     InstanceStatus
}

type DescribeInstanceStatusResponse struct {
	common.Response
	common.PaginationResult
	InstanceStatuses struct {
		InstanceStatus []InstanceStatusItemType
	}
}

// DescribeInstanceStatus describes instance status
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancestatus
func (client *Client) DescribeInstanceStatus(args *DescribeInstanceStatusArgs) (instanceStatuses []InstanceStatusItemType, pagination *common.PaginationResult, err error) {
	args.Validate()
	response := DescribeInstanceStatusResponse{}

	err = client.Invoke("DescribeInstanceStatus", args, &response)

	if err == nil {
		return response.InstanceStatuses.InstanceStatus, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type StopInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type StopInstanceResponse struct {
	common.Response
}

// StopInstance stops instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&stopinstance
func (client *Client) StopInstance(instanceId string, forceStop bool) error {
	args := StopInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := StopInstanceResponse{}
	err := client.Invoke("StopInstance", &args, &response)
	return err
}

type StartInstanceArgs struct {
	InstanceId string
}

type StartInstanceResponse struct {
	common.Response
}

// StartInstance starts instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&startinstance
func (client *Client) StartInstance(instanceId string) error {
	args := StartInstanceArgs{InstanceId: instanceId}
	response := StartInstanceResponse{}
	err := client.Invoke("StartInstance", &args, &response)
	return err
}

type RebootInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type RebootInstanceResponse struct {
	common.Response
}

// RebootInstance reboot instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&rebootinstance
func (client *Client) RebootInstance(instanceId string, forceStop bool) error {
	request := RebootInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := RebootInstanceResponse{}
	err := client.Invoke("RebootInstance", &request, &response)
	return err
}

type DescribeInstanceAttributeArgs struct {
	InstanceId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&operationlockstype
type OperationLocksType struct {
	LockReason []LockReasonType //enum for financial, security
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&securitygroupidsettype
type SecurityGroupIdSetType struct {
	SecurityGroupId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&ipaddresssettype
type IpAddressSetType struct {
	IpAddress []string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&vpcattributestype
type VpcAttributesType struct {
	VpcId            string
	VSwitchId        string
	PrivateIpAddress IpAddressSetType
	NatIpAddress     string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&eipaddressassociatetype
type EipAddressAssociateType struct {
	AllocationId       string
	IpAddress          string
	Bandwidth          int
	InternetChargeType common.InternetChargeType
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instanceattributestype
type InstanceAttributesType struct {
	InstanceId         string
	InstanceName       string
	Description        string
	ImageId            string
	RegionId           common.Region
	ZoneId             string
	CPU                int
	Memory             int
	ClusterId          string
	InstanceType       string
	InstanceTypeFamily string
	HostName           string
	SerialNumber       string
	Status             InstanceStatus
	OperationLocks     OperationLocksType
	SecurityGroupIds   struct {
		SecurityGroupId []string
	}
	PublicIpAddress         IpAddressSetType
	InnerIpAddress          IpAddressSetType
	InstanceNetworkType     string //enum Classic | Vpc
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	InternetChargeType      common.InternetChargeType
	CreationTime            util.ISO6801Time //time.Time
	VpcAttributes           VpcAttributesType
	EipAddress              EipAddressAssociateType
	IoOptimized             StringOrBool
	InstanceChargeType      common.InternetChargeType
	ExpiredTime             util.ISO6801Time
}

type DescribeInstanceAttributeResponse struct {
	common.Response
	InstanceAttributesType
}

// DescribeInstanceAttribute describes instance attribute
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstanceattribute
func (client *Client) DescribeInstanceAttribute(instanceId string) (instance *InstanceAttributesType, err error) {
	args := DescribeInstanceAttributeArgs{InstanceId: instanceId}

	response := DescribeInstanceAttributeResponse{}
	err = client.Invoke("DescribeInstanceAttribute", &args, &response)
	if err != nil {
		return nil, err
	}
	return &response.InstanceAttributesType, err
}

type ModifyInstanceAttributeArgs struct {
	InstanceId   string
	InstanceName string
	Description  string
	Password     string
	HostName     string
}

type ModifyInstanceAttributeResponse struct {
	common.Response
}

//ModifyInstanceAttribute  modify instance attrbute
//
// You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/modifyinstanceattribute.html
func (client *Client) ModifyInstanceAttribute(args *ModifyInstanceAttributeArgs) error {
	response := ModifyInstanceAttributeResponse{}
	err := client.Invoke("ModifyInstanceAttribute", args, &response)
	return err
}

// Default timeout value for WaitForInstance method
const InstanceDefaultTimeout = 120

// WaitForInstance waits for instance to given status
func (client *Client) WaitForInstance(instanceId string, status InstanceStatus, timeout int) error {
	if timeout <= 0 {
		timeout = InstanceDefaultTimeout
	}
	for {
		instance, err := client.DescribeInstanceAttribute(instanceId)
		if err != nil {
			return err
		}
		if instance.Status == status {
			//TODO
			//Sleep one more time for timing issues
			time.Sleep(DefaultWaitForInterval * time.Second)
			break
		}
		timeout = timeout - DefaultWaitForInterval
		if timeout <= 0 {
			return common.GetClientErrorFromString("Timeout")
		}
		time.Sleep(DefaultWaitForInterval * time.Second)

	}
	return nil
}

type DescribeInstanceVncUrlArgs struct {
	RegionId   common.Region
	InstanceId string
}

type DescribeInstanceVncUrlResponse struct {
	common.Response
	VncUrl string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancevncurl
func (client *Client) DescribeInstanceVncUrl(args *DescribeInstanceVncUrlArgs) (string, error) {
	response := DescribeInstanceVncUrlResponse{}

	err := client.Invoke("DescribeInstanceVncUrl", args, &response)

	if err == nil {
		return response.VncUrl, nil
	}

	return "", err
}

type DescribeInstancesArgs struct {
	RegionId            common.Region
	VpcId               string
	VSwitchId           string
	ZoneId              string
	InstanceIds         string
	InstanceNetworkType string
	PrivateIpAddresses  string
	InnerIpAddresses    string
	PublicIpAddresses   string
	SecurityGroupId     string
	common.Pagination
}

type DescribeInstancesResponse struct {
	common.Response
	common.PaginationResult
	Instances struct {
		Instance []InstanceAttributesType
	}
}

// DescribeInstances describes instances
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstances
func (client *Client) DescribeInstances(args *DescribeInstancesArgs) (instances []InstanceAttributesType, pagination *common.PaginationResult, err error) {
	args.Validate()
	response := DescribeInstancesResponse{}

	err = client.Invoke("DescribeInstances", args, &response)

	if err == nil {
		return response.Instances.Instance, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type DeleteInstanceArgs struct {
	InstanceId string
}

type DeleteInstanceResponse struct {
	common.Response
}

// DeleteInstance deletes instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&deleteinstance
func (client *Client) DeleteInstance(instanceId string) error {
	args := DeleteInstanceArgs{InstanceId: instanceId}
	response := DeleteInstanceResponse{}
	err := client.Invoke("DeleteInstance", &args, &response)
	return err
}

type DataDiskType struct {
	Size               int
	Category           DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	SnapshotId         string
	DiskName           string
	Description        string
	Device             string
	DeleteWithInstance bool
}

type SystemDiskType struct {
	Size        int
	Category    DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	DiskName    string
	Description string
}

type IoOptimized string

type StringOrBool struct {
	Value bool
}

// UnmarshalJSON implements the json.Unmarshaller interface.
func (io *StringOrBool) UnmarshalJSON(value []byte) error {
	if value[0] == '"' {
		var str string
		err := json.Unmarshal(value, &str)
		if err == nil {
			io.Value = (str == "true" || str == "optimized")
		}
		return err
	}
	var boolVal bool
	err := json.Unmarshal(value, &boolVal)
	if err == nil {
		io.Value = boolVal
	}
	return err
}

func (io StringOrBool) Bool() bool {
	return io.Value
}

func (io StringOrBool) String() string {
	return strconv.FormatBool(io.Value)
}

var (
	IoOptimizedNone      = IoOptimized("none")
	IoOptimizedOptimized = IoOptimized("optimized")
)

type CreateInstanceArgs struct {
	RegionId                common.Region
	ZoneId                  string
	ImageId                 string
	InstanceType            string
	SecurityGroupId         string
	InstanceName            string
	Description             string
	InternetChargeType      common.InternetChargeType
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	HostName                string
	Password                string
	IoOptimized             IoOptimized
	SystemDisk              SystemDiskType
	DataDisk                []DataDiskType
	VSwitchId               string
	PrivateIpAddress        string
	ClientToken             string
	InstanceChargeType      common.InstanceChargeType
	Period                  int
}

type CreateInstanceResponse struct {
	common.Response
	InstanceId string
}

// CreateInstance creates instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&createinstance
func (client *Client) CreateInstance(args *CreateInstanceArgs) (instanceId string, err error) {
	response := CreateInstanceResponse{}
	err = client.Invoke("CreateInstance", args, &response)
	if err != nil {
		return "", err
	}
	return response.InstanceId, err
}

type SecurityGroupArgs struct {
	InstanceId      string
	SecurityGroupId string
}

type SecurityGroupResponse struct {
	common.Response
}

//JoinSecurityGroup
//
//You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/joinsecuritygroup.html
func (client *Client) JoinSecurityGroup(instanceId string, securityGroupId string) error {
	args := SecurityGroupArgs{InstanceId: instanceId, SecurityGroupId: securityGroupId}
	response := SecurityGroupResponse{}
	err := client.Invoke("JoinSecurityGroup", &args, &response)
	return err
}

//LeaveSecurityGroup
//
//You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/leavesecuritygroup.html
func (client *Client) LeaveSecurityGroup(instanceId string, securityGroupId string) error {
	args := SecurityGroupArgs{InstanceId: instanceId, SecurityGroupId: securityGroupId}
	response := SecurityGroupResponse{}
	err := client.Invoke("LeaveSecurityGroup", &args, &response)
	return err
}
<|endoftext|>


<commit_before>package ecs

import (
	"encoding/json"
	"strconv"
	"time"

	"github.com/denverdino/aliyungo/common"
	"github.com/denverdino/aliyungo/util"
)

// InstanceStatus represents instance status
type InstanceStatus string

// Constants of InstanceStatus
const (
	Creating = InstanceStatus("Creating")
	Running  = InstanceStatus("Running")
	Starting = InstanceStatus("Starting")

	Stopped  = InstanceStatus("Stopped")
	Stopping = InstanceStatus("Stopping")
)

type LockReason string

const (
	LockReasonFinancial = LockReason("financial")
	LockReasonSecurity  = LockReason("security")
)

type LockReasonType struct {
	LockReason LockReason
}

type DescribeInstanceStatusArgs struct {
	RegionId common.Region
	ZoneId   string
	common.Pagination
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instancestatusitemtype
type InstanceStatusItemType struct {
	InstanceId string
	Status     InstanceStatus
}

type DescribeInstanceStatusResponse struct {
	common.Response
	common.PaginationResult
	InstanceStatuses struct {
		InstanceStatus []InstanceStatusItemType
	}
}

// DescribeInstanceStatus describes instance status
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancestatus
func (client *Client) DescribeInstanceStatus(args *DescribeInstanceStatusArgs) (instanceStatuses []InstanceStatusItemType, pagination *common.PaginationResult, err error) {
	args.Validate()
	response := DescribeInstanceStatusResponse{}

	err = client.Invoke("DescribeInstanceStatus", args, &response)

	if err == nil {
		return response.InstanceStatuses.InstanceStatus, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type StopInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type StopInstanceResponse struct {
	common.Response
}

// StopInstance stops instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&stopinstance
func (client *Client) StopInstance(instanceId string, forceStop bool) error {
	args := StopInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := StopInstanceResponse{}
	err := client.Invoke("StopInstance", &args, &response)
	return err
}

type StartInstanceArgs struct {
	InstanceId string
}

type StartInstanceResponse struct {
	common.Response
}

// StartInstance starts instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&startinstance
func (client *Client) StartInstance(instanceId string) error {
	args := StartInstanceArgs{InstanceId: instanceId}
	response := StartInstanceResponse{}
	err := client.Invoke("StartInstance", &args, &response)
	return err
}

type RebootInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type RebootInstanceResponse struct {
	common.Response
}

// RebootInstance reboot instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&rebootinstance
func (client *Client) RebootInstance(instanceId string, forceStop bool) error {
	request := RebootInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := RebootInstanceResponse{}
	err := client.Invoke("RebootInstance", &request, &response)
	return err
}

type DescribeInstanceAttributeArgs struct {
	InstanceId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&operationlockstype
type OperationLocksType struct {
	LockReason []LockReasonType //enum for financial, security
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&securitygroupidsettype
type SecurityGroupIdSetType struct {
	SecurityGroupId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&ipaddresssettype
type IpAddressSetType struct {
	IpAddress []string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&vpcattributestype
type VpcAttributesType struct {
	VpcId            string
	VSwitchId        string
	PrivateIpAddress IpAddressSetType
	NatIpAddress     string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&eipaddressassociatetype
type EipAddressAssociateType struct {
	AllocationId       string
	IpAddress          string
	Bandwidth          int
	InternetChargeType common.InternetChargeType
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instanceattributestype
type InstanceAttributesType struct {
	InstanceId         string
	InstanceName       string
	Description        string
	ImageId            string
	RegionId           common.Region
	ZoneId             string
	CPU                int
	Memory             int
	ClusterId          string
	InstanceType       string
	InstanceTypeFamily string
	HostName           string
	SerialNumber       string
	Status             InstanceStatus
	OperationLocks     OperationLocksType
	SecurityGroupIds   struct {
		SecurityGroupId []string
	}
	PublicIpAddress         IpAddressSetType
	InnerIpAddress          IpAddressSetType
	InstanceNetworkType     string //enum Classic | Vpc
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	InternetChargeType      common.InternetChargeType
	CreationTime            util.ISO6801Time //time.Time
	VpcAttributes           VpcAttributesType
	EipAddress              EipAddressAssociateType
	IoOptimized             StringOrBool
	InstanceChargeType      common.InternetChargeType
	ExpiredTime             util.ISO6801Time
}

type DescribeInstanceAttributeResponse struct {
	common.Response
	InstanceAttributesType
}

// DescribeInstanceAttribute describes instance attribute
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstanceattribute
func (client *Client) DescribeInstanceAttribute(instanceId string) (instance *InstanceAttributesType, err error) {
	args := DescribeInstanceAttributeArgs{InstanceId: instanceId}

	response := DescribeInstanceAttributeResponse{}
	err = client.Invoke("DescribeInstanceAttribute", &args, &response)
	if err != nil {
		return nil, err
	}
	return &response.InstanceAttributesType, err
}

type ModifyInstanceAttributeArgs struct {
	InstanceId   string
	InstanceName string
	Description  string
	Password     string
	HostName     string
}

type ModifyInstanceAttributeResponse struct {
	common.Response
}

//ModifyInstanceAttribute  modify instance attrbute
//
// You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/modifyinstanceattribute.html
func (client *Client) ModifyInstanceAttribute(args *ModifyInstanceAttributeArgs) error {
	response := ModifyInstanceAttributeResponse{}
	err := client.Invoke("ModifyInstanceAttribute", args, &response)
	return err
}

// Default timeout value for WaitForInstance method
const InstanceDefaultTimeout = 120

// WaitForInstance waits for instance to given status
func (client *Client) WaitForInstance(instanceId string, status InstanceStatus, timeout int) error {
	if timeout <= 0 {
		timeout = InstanceDefaultTimeout
	}
	for {
		instance, err := client.DescribeInstanceAttribute(instanceId)
		if err != nil {
			return err
		}
		if instance.Status == status {
			//TODO
			//Sleep one more time for timing issues
			time.Sleep(DefaultWaitForInterval * time.Second)
			break
		}
		timeout = timeout - DefaultWaitForInterval
		if timeout <= 0 {
			return common.GetClientErrorFromString("Timeout")
		}
		time.Sleep(DefaultWaitForInterval * time.Second)

	}
	return nil
}

type DescribeInstanceVncUrlArgs struct {
	RegionId   common.Region
	InstanceId string
}

type DescribeInstanceVncUrlResponse struct {
	common.Response
	VncUrl string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancevncurl
func (client *Client) DescribeInstanceVncUrl(args *DescribeInstanceVncUrlArgs) (string, error) {
	response := DescribeInstanceVncUrlResponse{}

	err := client.Invoke("DescribeInstanceVncUrl", args, &response)

	if err == nil {
		return response.VncUrl, nil
	}

	return "", err
}

type DescribeInstancesArgs struct {
	RegionId            common.Region
	VpcId               string
	VSwitchId           string
	ZoneId              string
	InstanceIds         string
	InstanceNetworkType string
	InstanceName        string
	Status              InstanceStatus
	PrivateIpAddresses  string
	InnerIpAddresses    string
	PublicIpAddresses   string
	SecurityGroupId     string
	common.Pagination
}

type DescribeInstancesResponse struct {
	common.Response
	common.PaginationResult
	Instances struct {
		Instance []InstanceAttributesType
	}
}

// DescribeInstances describes instances
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstances
func (client *Client) DescribeInstances(args *DescribeInstancesArgs) (instances []InstanceAttributesType, pagination *common.PaginationResult, err error) {
	args.Validate()
	response := DescribeInstancesResponse{}

	err = client.Invoke("DescribeInstances", args, &response)

	if err == nil {
		return response.Instances.Instance, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type DeleteInstanceArgs struct {
	InstanceId string
}

type DeleteInstanceResponse struct {
	common.Response
}

// DeleteInstance deletes instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&deleteinstance
func (client *Client) DeleteInstance(instanceId string) error {
	args := DeleteInstanceArgs{InstanceId: instanceId}
	response := DeleteInstanceResponse{}
	err := client.Invoke("DeleteInstance", &args, &response)
	return err
}

type DataDiskType struct {
	Size               int
	Category           DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	SnapshotId         string
	DiskName           string
	Description        string
	Device             string
	DeleteWithInstance bool
}

type SystemDiskType struct {
	Size        int
	Category    DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	DiskName    string
	Description string
}

type IoOptimized string

type StringOrBool struct {
	Value bool
}

// UnmarshalJSON implements the json.Unmarshaller interface.
func (io *StringOrBool) UnmarshalJSON(value []byte) error {
	if value[0] == '"' {
		var str string
		err := json.Unmarshal(value, &str)
		if err == nil {
			io.Value = (str == "true" || str == "optimized")
		}
		return err
	}
	var boolVal bool
	err := json.Unmarshal(value, &boolVal)
	if err == nil {
		io.Value = boolVal
	}
	return err
}

func (io StringOrBool) Bool() bool {
	return io.Value
}

func (io StringOrBool) String() string {
	return strconv.FormatBool(io.Value)
}

var (
	IoOptimizedNone      = IoOptimized("none")
	IoOptimizedOptimized = IoOptimized("optimized")
)

type CreateInstanceArgs struct {
	RegionId                common.Region
	ZoneId                  string
	ImageId                 string
	InstanceType            string
	SecurityGroupId         string
	InstanceName            string
	Description             string
	InternetChargeType      common.InternetChargeType
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	HostName                string
	Password                string
	IoOptimized             IoOptimized
	SystemDisk              SystemDiskType
	DataDisk                []DataDiskType
	VSwitchId               string
	PrivateIpAddress        string
	ClientToken             string
	InstanceChargeType      common.InstanceChargeType
	Period                  int
}

type CreateInstanceResponse struct {
	common.Response
	InstanceId string
}

// CreateInstance creates instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&createinstance
func (client *Client) CreateInstance(args *CreateInstanceArgs) (instanceId string, err error) {
	response := CreateInstanceResponse{}
	err = client.Invoke("CreateInstance", args, &response)
	if err != nil {
		return "", err
	}
	return response.InstanceId, err
}

type SecurityGroupArgs struct {
	InstanceId      string
	SecurityGroupId string
}

type SecurityGroupResponse struct {
	common.Response
}

//JoinSecurityGroup
//
//You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/joinsecuritygroup.html
func (client *Client) JoinSecurityGroup(instanceId string, securityGroupId string) error {
	args := SecurityGroupArgs{InstanceId: instanceId, SecurityGroupId: securityGroupId}
	response := SecurityGroupResponse{}
	err := client.Invoke("JoinSecurityGroup", &args, &response)
	return err
}

//LeaveSecurityGroup
//
//You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/leavesecuritygroup.html
func (client *Client) LeaveSecurityGroup(instanceId string, securityGroupId string) error {
	args := SecurityGroupArgs{InstanceId: instanceId, SecurityGroupId: securityGroupId}
	response := SecurityGroupResponse{}
	err := client.Invoke("LeaveSecurityGroup", &args, &response)
	return err
}
<commit_msg>Suppport UserData for testing<commit_after>package ecs

import (
	"encoding/base64"
	"encoding/json"
	"strconv"
	"time"

	"github.com/denverdino/aliyungo/common"
	"github.com/denverdino/aliyungo/util"
)

// InstanceStatus represents instance status
type InstanceStatus string

// Constants of InstanceStatus
const (
	Creating = InstanceStatus("Creating")
	Running  = InstanceStatus("Running")
	Starting = InstanceStatus("Starting")

	Stopped  = InstanceStatus("Stopped")
	Stopping = InstanceStatus("Stopping")
)

type LockReason string

const (
	LockReasonFinancial = LockReason("financial")
	LockReasonSecurity  = LockReason("security")
)

type LockReasonType struct {
	LockReason LockReason
}

type DescribeInstanceStatusArgs struct {
	RegionId common.Region
	ZoneId   string
	common.Pagination
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instancestatusitemtype
type InstanceStatusItemType struct {
	InstanceId string
	Status     InstanceStatus
}

type DescribeInstanceStatusResponse struct {
	common.Response
	common.PaginationResult
	InstanceStatuses struct {
		InstanceStatus []InstanceStatusItemType
	}
}

// DescribeInstanceStatus describes instance status
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancestatus
func (client *Client) DescribeInstanceStatus(args *DescribeInstanceStatusArgs) (instanceStatuses []InstanceStatusItemType, pagination *common.PaginationResult, err error) {
	args.Validate()
	response := DescribeInstanceStatusResponse{}

	err = client.Invoke("DescribeInstanceStatus", args, &response)

	if err == nil {
		return response.InstanceStatuses.InstanceStatus, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type StopInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type StopInstanceResponse struct {
	common.Response
}

// StopInstance stops instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&stopinstance
func (client *Client) StopInstance(instanceId string, forceStop bool) error {
	args := StopInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := StopInstanceResponse{}
	err := client.Invoke("StopInstance", &args, &response)
	return err
}

type StartInstanceArgs struct {
	InstanceId string
}

type StartInstanceResponse struct {
	common.Response
}

// StartInstance starts instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&startinstance
func (client *Client) StartInstance(instanceId string) error {
	args := StartInstanceArgs{InstanceId: instanceId}
	response := StartInstanceResponse{}
	err := client.Invoke("StartInstance", &args, &response)
	return err
}

type RebootInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type RebootInstanceResponse struct {
	common.Response
}

// RebootInstance reboot instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&rebootinstance
func (client *Client) RebootInstance(instanceId string, forceStop bool) error {
	request := RebootInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := RebootInstanceResponse{}
	err := client.Invoke("RebootInstance", &request, &response)
	return err
}

type DescribeInstanceAttributeArgs struct {
	InstanceId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&operationlockstype
type OperationLocksType struct {
	LockReason []LockReasonType //enum for financial, security
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&securitygroupidsettype
type SecurityGroupIdSetType struct {
	SecurityGroupId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&ipaddresssettype
type IpAddressSetType struct {
	IpAddress []string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&vpcattributestype
type VpcAttributesType struct {
	VpcId            string
	VSwitchId        string
	PrivateIpAddress IpAddressSetType
	NatIpAddress     string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&eipaddressassociatetype
type EipAddressAssociateType struct {
	AllocationId       string
	IpAddress          string
	Bandwidth          int
	InternetChargeType common.InternetChargeType
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instanceattributestype
type InstanceAttributesType struct {
	InstanceId         string
	InstanceName       string
	Description        string
	ImageId            string
	RegionId           common.Region
	ZoneId             string
	CPU                int
	Memory             int
	ClusterId          string
	InstanceType       string
	InstanceTypeFamily string
	HostName           string
	SerialNumber       string
	Status             InstanceStatus
	OperationLocks     OperationLocksType
	SecurityGroupIds   struct {
		SecurityGroupId []string
	}
	PublicIpAddress         IpAddressSetType
	InnerIpAddress          IpAddressSetType
	InstanceNetworkType     string //enum Classic | Vpc
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	InternetChargeType      common.InternetChargeType
	CreationTime            util.ISO6801Time //time.Time
	VpcAttributes           VpcAttributesType
	EipAddress              EipAddressAssociateType
	IoOptimized             StringOrBool
	InstanceChargeType      common.InternetChargeType
	ExpiredTime             util.ISO6801Time
}

type DescribeInstanceAttributeResponse struct {
	common.Response
	InstanceAttributesType
}

// DescribeInstanceAttribute describes instance attribute
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstanceattribute
func (client *Client) DescribeInstanceAttribute(instanceId string) (instance *InstanceAttributesType, err error) {
	args := DescribeInstanceAttributeArgs{InstanceId: instanceId}

	response := DescribeInstanceAttributeResponse{}
	err = client.Invoke("DescribeInstanceAttribute", &args, &response)
	if err != nil {
		return nil, err
	}
	return &response.InstanceAttributesType, err
}

type ModifyInstanceAttributeArgs struct {
	InstanceId   string
	InstanceName string
	Description  string
	Password     string
	HostName     string
}

type ModifyInstanceAttributeResponse struct {
	common.Response
}

//ModifyInstanceAttribute  modify instance attrbute
//
// You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/modifyinstanceattribute.html
func (client *Client) ModifyInstanceAttribute(args *ModifyInstanceAttributeArgs) error {
	response := ModifyInstanceAttributeResponse{}
	err := client.Invoke("ModifyInstanceAttribute", args, &response)
	return err
}

// Default timeout value for WaitForInstance method
const InstanceDefaultTimeout = 120

// WaitForInstance waits for instance to given status
func (client *Client) WaitForInstance(instanceId string, status InstanceStatus, timeout int) error {
	if timeout <= 0 {
		timeout = InstanceDefaultTimeout
	}
	for {
		instance, err := client.DescribeInstanceAttribute(instanceId)
		if err != nil {
			return err
		}
		if instance.Status == status {
			//TODO
			//Sleep one more time for timing issues
			time.Sleep(DefaultWaitForInterval * time.Second)
			break
		}
		timeout = timeout - DefaultWaitForInterval
		if timeout <= 0 {
			return common.GetClientErrorFromString("Timeout")
		}
		time.Sleep(DefaultWaitForInterval * time.Second)

	}
	return nil
}

type DescribeInstanceVncUrlArgs struct {
	RegionId   common.Region
	InstanceId string
}

type DescribeInstanceVncUrlResponse struct {
	common.Response
	VncUrl string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancevncurl
func (client *Client) DescribeInstanceVncUrl(args *DescribeInstanceVncUrlArgs) (string, error) {
	response := DescribeInstanceVncUrlResponse{}

	err := client.Invoke("DescribeInstanceVncUrl", args, &response)

	if err == nil {
		return response.VncUrl, nil
	}

	return "", err
}

type DescribeInstancesArgs struct {
	RegionId            common.Region
	VpcId               string
	VSwitchId           string
	ZoneId              string
	InstanceIds         string
	InstanceNetworkType string
	InstanceName        string
	Status              InstanceStatus
	PrivateIpAddresses  string
	InnerIpAddresses    string
	PublicIpAddresses   string
	SecurityGroupId     string
	common.Pagination
}

type DescribeInstancesResponse struct {
	common.Response
	common.PaginationResult
	Instances struct {
		Instance []InstanceAttributesType
	}
}

// DescribeInstances describes instances
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstances
func (client *Client) DescribeInstances(args *DescribeInstancesArgs) (instances []InstanceAttributesType, pagination *common.PaginationResult, err error) {
	args.Validate()
	response := DescribeInstancesResponse{}

	err = client.Invoke("DescribeInstances", args, &response)

	if err == nil {
		return response.Instances.Instance, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type DeleteInstanceArgs struct {
	InstanceId string
}

type DeleteInstanceResponse struct {
	common.Response
}

// DeleteInstance deletes instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&deleteinstance
func (client *Client) DeleteInstance(instanceId string) error {
	args := DeleteInstanceArgs{InstanceId: instanceId}
	response := DeleteInstanceResponse{}
	err := client.Invoke("DeleteInstance", &args, &response)
	return err
}

type DataDiskType struct {
	Size               int
	Category           DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	SnapshotId         string
	DiskName           string
	Description        string
	Device             string
	DeleteWithInstance bool
}

type SystemDiskType struct {
	Size        int
	Category    DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	DiskName    string
	Description string
}

type IoOptimized string

type StringOrBool struct {
	Value bool
}

// UnmarshalJSON implements the json.Unmarshaller interface.
func (io *StringOrBool) UnmarshalJSON(value []byte) error {
	if value[0] == '"' {
		var str string
		err := json.Unmarshal(value, &str)
		if err == nil {
			io.Value = (str == "true" || str == "optimized")
		}
		return err
	}
	var boolVal bool
	err := json.Unmarshal(value, &boolVal)
	if err == nil {
		io.Value = boolVal
	}
	return err
}

func (io StringOrBool) Bool() bool {
	return io.Value
}

func (io StringOrBool) String() string {
	return strconv.FormatBool(io.Value)
}

var (
	IoOptimizedNone      = IoOptimized("none")
	IoOptimizedOptimized = IoOptimized("optimized")
)

type CreateInstanceArgs struct {
	RegionId                common.Region
	ZoneId                  string
	ImageId                 string
	InstanceType            string
	SecurityGroupId         string
	InstanceName            string
	Description             string
	InternetChargeType      common.InternetChargeType
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	HostName                string
	Password                string
	IoOptimized             IoOptimized
	SystemDisk              SystemDiskType
	DataDisk                []DataDiskType
	VSwitchId               string
	PrivateIpAddress        string
	ClientToken             string
	InstanceChargeType      common.InstanceChargeType
	Period                  int
	UserData                string
}

type CreateInstanceResponse struct {
	common.Response
	InstanceId string
}

// CreateInstance creates instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&createinstance
func (client *Client) CreateInstance(args *CreateInstanceArgs) (instanceId string, err error) {
	if args.UserData != "" {
		// Encode to base64 string
		args.UserData = base64.StdEncoding.EncodeToString([]byte(args.UserData))
	}
	response := CreateInstanceResponse{}
	err = client.Invoke("CreateInstance", args, &response)
	if err != nil {
		return "", err
	}
	return response.InstanceId, err
}

type SecurityGroupArgs struct {
	InstanceId      string
	SecurityGroupId string
}

type SecurityGroupResponse struct {
	common.Response
}

//JoinSecurityGroup
//
//You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/joinsecuritygroup.html
func (client *Client) JoinSecurityGroup(instanceId string, securityGroupId string) error {
	args := SecurityGroupArgs{InstanceId: instanceId, SecurityGroupId: securityGroupId}
	response := SecurityGroupResponse{}
	err := client.Invoke("JoinSecurityGroup", &args, &response)
	return err
}

//LeaveSecurityGroup
//
//You can read doc at https://help.aliyun.com/document_detail/ecs/open-api/instance/leavesecuritygroup.html
func (client *Client) LeaveSecurityGroup(instanceId string, securityGroupId string) error {
	args := SecurityGroupArgs{InstanceId: instanceId, SecurityGroupId: securityGroupId}
	response := SecurityGroupResponse{}
	err := client.Invoke("LeaveSecurityGroup", &args, &response)
	return err
}
<|endoftext|>


<commit_before>package editor

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/hacdias/caddy-hugo/config"
	"github.com/hacdias/caddy-hugo/frontmatter"
	"github.com/hacdias/caddy-hugo/utils"
	"github.com/robfig/cron"
	"github.com/spf13/hugo/parser"
)

type editor struct {
	Name        string
	Class       string
	IsPost      bool
	Mode        string
	Content     string
	FrontMatter interface{}
	Config      *config.Config
}

// ServeHTTP serves the editor page
func ServeHTTP(w http.ResponseWriter, r *http.Request, c *config.Config) (int, error) {
	filename := strings.Replace(r.URL.Path, "/admin/edit/", "", 1)

	switch r.Method {
	case "POST":
		return post(w, r, c, filename)
	case "GET":
		return get(w, r, c, filename)
	default:
		return 400, nil
	}
}

func post(w http.ResponseWriter, r *http.Request, c *config.Config, filename string) (int, error) {
	// Get the JSON information sent using a buffer
	rawBuffer := new(bytes.Buffer)
	rawBuffer.ReadFrom(r.Body)

	// Creates the raw file "map" using the JSON
	var rawFile map[string]interface{}
	json.Unmarshal(rawBuffer.Bytes(), &rawFile)

	// Initializes the file content to write
	var file []byte

	switch r.Header.Get("X-Content-Type") {
	case "frontmatter-only":
		frontmatter := strings.TrimPrefix(filepath.Ext(filename), ".")
		var mark rune

		switch frontmatter {
		case "toml":
			mark = rune('+')
		case "json":
			mark = rune('{')
		case "yaml":
			mark = rune('-')
		default:
			return 400, nil
		}

		f, err := parser.InterfaceToFrontMatter(rawFile, mark)
		fString := string(f)

		// If it's toml or yaml, strip frontmatter identifier
		if frontmatter == "toml" {
			fString = strings.TrimSuffix(fString, "+++\n")
			fString = strings.TrimPrefix(fString, "+++\n")
		}

		if frontmatter == "yaml" {
			fString = strings.TrimSuffix(fString, "---\n")
			fString = strings.TrimPrefix(fString, "---\n")
		}

		f = []byte(fString)

		if err != nil {
			w.Write([]byte(err.Error()))
			return 500, err
		}

		file = f
	case "content-only":
		// The main content of the file
		mainContent := rawFile["content"].(string)
		mainContent = "\n\n" + strings.TrimSpace(mainContent)

		file = []byte(mainContent)
	case "complete":
		// The main content of the file
		mainContent := rawFile["content"].(string)
		mainContent = "\n\n" + strings.TrimSpace(mainContent)

		// Removes the main content from the rest of the frontmatter
		delete(rawFile, "content")

		// Schedule the post
		if r.Header.Get("X-Schedule") == "true" {
			_, err := schedule(w, c, filename, rawFile)

			if err != nil {
				w.Write([]byte(err.Error()))
				return 500, err
			}
		}

		// Converts the frontmatter in JSON
		jsonFrontmatter, err := json.Marshal(rawFile)

		if err != nil {
			w.Write([]byte(err.Error()))
			return 500, err
		}

		// Indents the json
		frontMatterBuffer := new(bytes.Buffer)
		json.Indent(frontMatterBuffer, jsonFrontmatter, "", "  ")

		// Generates the final file
		f := new(bytes.Buffer)
		f.Write(frontMatterBuffer.Bytes())
		f.Write([]byte(mainContent))
		file = f.Bytes()
	default:
		return 400, nil
	}

	// Write the file
	err := ioutil.WriteFile(filename, file, 0666)

	if err != nil {
		w.Write([]byte(err.Error()))
		return 500, err
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte("{}"))
	return 200, nil
}

func schedule(w http.ResponseWriter, c *config.Config, filename string, raw map[string]interface{}) (int, error) {
	// The main content of the file
	content := raw["content"].(string)
	content = "\n\n" + strings.TrimSpace(content)

	// Removes the main content from the rest of the frontmatter
	delete(raw, "content")

	// Parse the time
	t, err := time.Parse("2006-01-02 15:04:05-07:00", raw["date"].(string))

	if err != nil {
		w.Write([]byte(err.Error()))
		return 500, err
	}

	scheduler := cron.New()
	scheduler.AddFunc(t.In(time.Now().Location()).Format("05 04 15 02 01 *"), func() {
		// Set draft to false
		raw["draft"] = false

		// Converts the frontmatter in JSON
		jsonFrontmatter, err := json.Marshal(raw)

		if err != nil {
			return
		}

		// Indents the json
		frontMatterBuffer := new(bytes.Buffer)
		json.Indent(frontMatterBuffer, jsonFrontmatter, "", "  ")

		// Generates the final file
		f := new(bytes.Buffer)
		f.Write(frontMatterBuffer.Bytes())
		f.Write([]byte(content))
		file := f.Bytes()

		// Write the file
		err = ioutil.WriteFile(filename, file, 0666)

		if err != nil {
			return
		}

		utils.RunHugo(c)
	})
	scheduler.Start()
	return 200, nil
}

func get(w http.ResponseWriter, r *http.Request, c *config.Config, filename string) (int, error) {
	// Check if the file format is supported. If not, send a "Not Acceptable"
	// header and an error
	if !utils.CanBeEdited(filename) {
		return 406, errors.New("File format not supported.")
	}

	// Check if the file exists. If it doesn't, send a "Not Found" message
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		w.Write([]byte(err.Error()))
		return 404, nil
	}

	// Open the file and check if there was some error while opening
	file, err := ioutil.ReadFile(filename)
	if err != nil {
		w.Write([]byte(err.Error()))
		return 500, err
	}

	// Create a new editor variable and set the extension
	page := new(editor)
	page.Mode = strings.TrimPrefix(filepath.Ext(filename), ".")
	page.Name = filename
	page.Config = c
	page.IsPost = false

	// Sanitize the extension
	page.Mode = sanitizeMode(page.Mode)

	// Handle the content depending on the file extension
	switch page.Mode {
	case "markdown":
		if hasFrontMatterRune(file) {
			// Starts a new buffer and parses the file using Hugo's functions
			buffer := bytes.NewBuffer(file)
			file, err := parser.ReadFrom(buffer)
			if err != nil {
				w.Write([]byte(err.Error()))
				return 500, err
			}

			if strings.Contains(string(file.FrontMatter()), "date") {
				page.IsPost = true
			}

			// Parses the page content and the frontmatter
			page.Content = strings.TrimSpace(string(file.Content()))
			page.FrontMatter, err = frontmatter.Pretty(file.FrontMatter())
			page.Class = "complete"
		} else {
			// The editor will handle only content
			page.Class = "content-only"
			page.Content = string(file)
		}
	case "json", "toml", "yaml":
		// Defines the class and declares an error
		page.Class = "frontmatter-only"
		var err error

		// Checks if the file already has the frontmatter rune and parses it
		if hasFrontMatterRune(file) {
			page.FrontMatter, err = frontmatter.Pretty(file)
		} else {
			page.FrontMatter, err = frontmatter.Pretty(appendFrontMatterRune(file, page.Mode))
		}

		// Check if there were any errors
		if err != nil {
			w.Write([]byte(err.Error()))
			return 500, err
		}
	default:
		// The editor will handle only content
		page.Class = "content-only"
		page.Content = string(file)
	}

	// Create the functions map, then the template, check for erros and
	// execute the template if there aren't errors
	functions := template.FuncMap{
		"SplitCapitalize": utils.SplitCapitalize,
		"Defined":         utils.Defined,
	}

	tpl, err := utils.GetTemplate(r, functions, "editor", "frontmatter")

	if err != nil {
		w.Write([]byte(err.Error()))
		return 500, err
	}

	return 200, tpl.Execute(w, page)
}

func hasFrontMatterRune(file []byte) bool {
	return strings.HasPrefix(string(file), "---") ||
		strings.HasPrefix(string(file), "+++") ||
		strings.HasPrefix(string(file), "{")
}

func appendFrontMatterRune(frontmatter []byte, language string) []byte {
	switch language {
	case "yaml":
		return []byte("---\n" + string(frontmatter) + "\n---")
	case "toml":
		return []byte("+++\n" + string(frontmatter) + "\n+++")
	case "json":
		return frontmatter
	}

	return frontmatter
}

func sanitizeMode(extension string) string {
	switch extension {
	case "markdown", "md":
		return "markdown"
	case "css", "scss":
		return "css"
	case "html":
		return "htmlmixed"
	case "js":
		return "javascript"
	default:
		return extension
	}
}
<commit_msg>working<commit_after>package editor

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/hacdias/caddy-hugo/config"
	"github.com/hacdias/caddy-hugo/frontmatter"
	"github.com/hacdias/caddy-hugo/utils"
	"github.com/robfig/cron"
	"github.com/spf13/hugo/parser"
)

type editor struct {
	Name        string
	Class       string
	IsPost      bool
	Mode        string
	Content     string
	FrontMatter interface{}
	Config      *config.Config
}

// ServeHTTP serves the editor page
func ServeHTTP(w http.ResponseWriter, r *http.Request, c *config.Config) (int, error) {
	filename := strings.Replace(r.URL.Path, "/admin/edit/", "", 1)

	if r.Method == "POST" {
		return servePost(w, r, c, filename)
	}

	return serveGet(w, r, c, filename)
}

func servePost(w http.ResponseWriter, r *http.Request, c *config.Config, filename string) (int, error) {
	// Get the JSON information sent using a buffer
	rawBuffer := new(bytes.Buffer)
	rawBuffer.ReadFrom(r.Body)

	// Creates the raw file "map" using the JSON
	var rawFile map[string]interface{}
	json.Unmarshal(rawBuffer.Bytes(), &rawFile)

	// Initializes the file content to write
	var file []byte

	switch r.Header.Get("X-Content-Type") {
	case "frontmatter-only":
		frontmatter := strings.TrimPrefix(filepath.Ext(filename), ".")
		var mark rune

		switch frontmatter {
		case "toml":
			mark = rune('+')
		case "json":
			mark = rune('{')
		case "yaml":
			mark = rune('-')
		default:
			return 400, nil
		}

		f, err := parser.InterfaceToFrontMatter(rawFile, mark)
		fString := string(f)

		// If it's toml or yaml, strip frontmatter identifier
		if frontmatter == "toml" {
			fString = strings.TrimSuffix(fString, "+++\n")
			fString = strings.TrimPrefix(fString, "+++\n")
		}

		if frontmatter == "yaml" {
			fString = strings.TrimSuffix(fString, "---\n")
			fString = strings.TrimPrefix(fString, "---\n")
		}

		f = []byte(fString)

		if err != nil {
			w.Write([]byte(err.Error()))
			return 500, err
		}

		file = f
	case "content-only":
		// The main content of the file
		mainContent := rawFile["content"].(string)
		mainContent = "\n\n" + strings.TrimSpace(mainContent)

		file = []byte(mainContent)
	case "complete":
		// The main content of the file
		mainContent := rawFile["content"].(string)
		mainContent = "\n\n" + strings.TrimSpace(mainContent)

		// Removes the main content from the rest of the frontmatter
		delete(rawFile, "content")

		// Schedule the post
		if r.Header.Get("X-Schedule") == "true" {
			t, err := time.Parse("2006-01-02 15:04:05-07:00", rawFile["date"].(string))

			if err != nil {
				w.Write([]byte(err.Error()))
				return 500, err
			}

			scheduler := cron.New()
			scheduler.AddFunc(t.In(time.Now().Location()).Format("05 04 15 02 01 *"), func() {
				// Set draft to false
				rawFile["draft"] = false

				// Converts the frontmatter in JSON
				jsonFrontmatter, err := json.Marshal(rawFile)

				if err != nil {
					return
				}

				// Indents the json
				frontMatterBuffer := new(bytes.Buffer)
				json.Indent(frontMatterBuffer, jsonFrontmatter, "", "  ")

				// Generates the final file
				f := new(bytes.Buffer)
				f.Write(frontMatterBuffer.Bytes())
				f.Write([]byte(mainContent))
				file = f.Bytes()

				// Write the file
				err = ioutil.WriteFile(filename, file, 0666)

				if err != nil {
					return
				}

				utils.RunHugo(c)
			})
			scheduler.Start()
		}

		// Converts the frontmatter in JSON
		jsonFrontmatter, err := json.Marshal(rawFile)

		if err != nil {
			w.Write([]byte(err.Error()))
			return 500, err
		}

		// Indents the json
		frontMatterBuffer := new(bytes.Buffer)
		json.Indent(frontMatterBuffer, jsonFrontmatter, "", "  ")

		// Generates the final file
		f := new(bytes.Buffer)
		f.Write(frontMatterBuffer.Bytes())
		f.Write([]byte(mainContent))
		file = f.Bytes()
	default:
		return 400, nil
	}

	// Write the file
	err := ioutil.WriteFile(filename, file, 0666)

	if err != nil {
		w.Write([]byte(err.Error()))
		return 500, err
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte("{}"))
	return 200, nil
}

func serveGet(w http.ResponseWriter, r *http.Request, c *config.Config, filename string) (int, error) {
	// Check if the file format is supported. If not, send a "Not Acceptable"
	// header and an error
	if !utils.CanBeEdited(filename) {
		return 406, errors.New("File format not supported.")
	}

	// Check if the file exists. If it doesn't, send a "Not Found" message
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		w.Write([]byte(err.Error()))
		return 404, nil
	}

	// Open the file and check if there was some error while opening
	file, err := ioutil.ReadFile(filename)
	if err != nil {
		w.Write([]byte(err.Error()))
		return 500, err
	}

	// Create a new editor variable and set the extension
	page := new(editor)
	page.Mode = strings.TrimPrefix(filepath.Ext(filename), ".")
	page.Name = filename
	page.Config = c
	page.IsPost = false

	// Sanitize the extension
	page.Mode = sanitizeMode(page.Mode)

	// Handle the content depending on the file extension
	switch page.Mode {
	case "markdown":
		if hasFrontMatterRune(file) {
			// Starts a new buffer and parses the file using Hugo's functions
			buffer := bytes.NewBuffer(file)
			file, err := parser.ReadFrom(buffer)
			if err != nil {
				w.Write([]byte(err.Error()))
				return 500, err
			}

			if strings.Contains(string(file.FrontMatter()), "date") {
				page.IsPost = true
			}

			// Parses the page content and the frontmatter
			page.Content = strings.TrimSpace(string(file.Content()))
			page.FrontMatter, err = frontmatter.Pretty(file.FrontMatter())
			page.Class = "complete"
		} else {
			// The editor will handle only content
			page.Class = "content-only"
			page.Content = string(file)
		}
	case "json", "toml", "yaml":
		// Defines the class and declares an error
		page.Class = "frontmatter-only"
		var err error

		// Checks if the file already has the frontmatter rune and parses it
		if hasFrontMatterRune(file) {
			page.FrontMatter, err = frontmatter.Pretty(file)
		} else {
			page.FrontMatter, err = frontmatter.Pretty(appendFrontMatterRune(file, page.Mode))
		}

		// Check if there were any errors
		if err != nil {
			w.Write([]byte(err.Error()))
			return 500, err
		}
	default:
		// The editor will handle only content
		page.Class = "content-only"
		page.Content = string(file)
	}

	// Create the functions map, then the template, check for erros and
	// execute the template if there aren't errors
	functions := template.FuncMap{
		"SplitCapitalize": utils.SplitCapitalize,
		"Defined":         utils.Defined,
	}

	tpl, err := utils.GetTemplate(r, functions, "editor", "frontmatter")

	if err != nil {
		w.Write([]byte(err.Error()))
		return 500, err
	}

	return 200, tpl.Execute(w, page)
}

func hasFrontMatterRune(file []byte) bool {
	return strings.HasPrefix(string(file), "---") ||
		strings.HasPrefix(string(file), "+++") ||
		strings.HasPrefix(string(file), "{")
}

func appendFrontMatterRune(frontmatter []byte, language string) []byte {
	switch language {
	case "yaml":
		return []byte("---\n" + string(frontmatter) + "\n---")
	case "toml":
		return []byte("+++\n" + string(frontmatter) + "\n+++")
	case "json":
		return frontmatter
	}

	return frontmatter
}

func sanitizeMode(extension string) string {
	switch extension {
	case "markdown", "md":
		return "markdown"
	case "css", "scss":
		return "css"
	case "html":
		return "htmlmixed"
	case "js":
		return "javascript"
	default:
		return extension
	}
}
<|endoftext|>