<commit_before>/*
Copyright 2011 Google Inc.

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

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

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

package gethandler

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	"camlistore.org/pkg/auth"
	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/httputil"
	"camlistore.org/pkg/misc/httprange"
)

var kGetPattern *regexp.Regexp = regexp.MustCompile(`/camli/([a-z0-9]+)-([a-f0-9]+)$`)

type Handler struct {
	Fetcher           blobref.StreamingFetcher
	AllowGlobalAccess bool
}

func CreateGetHandler(fetcher blobref.StreamingFetcher) func(http.ResponseWriter, *http.Request) {
	gh := &Handler{Fetcher: fetcher}
	return func(conn http.ResponseWriter, req *http.Request) {
		if req.URL.Path == "/camli/sha1-deadbeef00000000000000000000000000000000" {
			// Test handler.
			simulatePrematurelyClosedConnection(conn, req)
			return
		}
		gh.ServeHTTP(conn, req)
	}
}

const fetchFailureDelayNs = 200e6 // 200 ms
const maxJSONSize = 64 * 1024     // should be enough for everyone

func (h *Handler) ServeHTTP(conn http.ResponseWriter, req *http.Request) {
	blobRef := blobFromUrlPath(req.URL.Path)
	if blobRef == nil {
		http.Error(conn, "Malformed GET URL.", 400)
		return
	}

	switch {
	case h.AllowGlobalAccess || auth.Allowed(req, auth.OpGet):
		serveBlobRef(conn, req, blobRef, h.Fetcher)
	case auth.TriedAuthorization(req):
		log.Printf("Attempted authorization failed on %s", req.URL)
		auth.SendUnauthorized(conn)
	default:
		handleGetViaSharing(conn, req, blobRef, h.Fetcher)
	}
}

// serveBlobRef sends 'blobref' to 'conn' as directed by the Range header in 'req'
func serveBlobRef(conn http.ResponseWriter, req *http.Request,
	blobRef *blobref.BlobRef, fetcher blobref.StreamingFetcher) {

	if w, ok := fetcher.(blobserver.ContextWrapper); ok {
		fetcher = w.WrapContext(req)
	}

	file, size, err := fetcher.FetchStreaming(blobRef)
	switch err {
	case nil:
		break
	case os.ErrNotExist:
		conn.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(conn, "Blob %q not found", blobRef)
		return
	default:
		httputil.ServerError(conn, req, err)
		return
	}

	defer file.Close()

	seeker, isSeeker := file.(io.Seeker)
	reqRange := httprange.FromRequest(req)
	if reqRange.SkipBytes() != 0 && isSeeker {
		// TODO: set the Range-specific response headers too,
		// acknowledging that we honored the content range
		// request.
		_, err = seeker.Seek(reqRange.SkipBytes(), 0)
		if err != nil {
			httputil.ServerError(conn, req, err)
			return
		}
	}

	var input io.Reader = file
	if reqRange.LimitBytes() != -1 {
		input = io.LimitReader(file, reqRange.LimitBytes())
	}

	remainBytes := size - reqRange.SkipBytes()
	if reqRange.LimitBytes() != -1 &&
		reqRange.LimitBytes() < remainBytes {
		remainBytes = reqRange.LimitBytes()
	}

	// Assume this generic content type by default.  For better
	// demos we'll try to sniff and guess the "right" MIME type in
	// certain cases (no Range requests, etc) but this isn't part
	// of the Camli spec at all.  We just do it to ease demos.
	contentType := "application/octet-stream"
	if reqRange.IsWholeFile() {
		const peekSize = 1024
		bufReader := bufio.NewReaderSize(input, peekSize)
		header, _ := bufReader.Peek(peekSize)
		if len(header) >= 8 {
			switch {
			case utf8.Valid(header):
				contentType = "text/plain; charset=utf-8"
			case bytes.HasPrefix(header, []byte{0xff, 0xd8, 0xff, 0xe2}):
				contentType = "image/jpeg"
			case bytes.HasPrefix(header, []byte{0x89, 0x50, 0x4e, 0x47, 0xd, 0xa, 0x1a, 0xa}):
				contentType = "image/png"
			}
		}
		input = bufReader

		conn.Header().Set("Content-Length", strconv.FormatInt(size, 10))
	}

	conn.Header().Set("Content-Type", contentType)
	if !reqRange.IsWholeFile() {
		conn.Header().Set("Content-Range",
			fmt.Sprintf("bytes %d-%d/%d", reqRange.SkipBytes(),
				reqRange.SkipBytes()+remainBytes,
				size))
		conn.WriteHeader(http.StatusPartialContent)
	}
	bytesCopied, err := io.Copy(conn, input)

	// If there's an error at this point, it's too late to tell the client,
	// as they've already been receiving bytes.  But they should be smart enough
	// to verify the digest doesn't match.  But we close the (chunked) response anyway,
	// to further signal errors.
	killConnection := func() {
		if hj, ok := conn.(http.Hijacker); ok {
			log.Printf("Force-closing TCP connection to signal error sending %q", blobRef)
			if closer, _, err := hj.Hijack(); err != nil {
				closer.Close()
			}
		}
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error sending file: %v, err=%v\n", blobRef, err)
		killConnection()
		return
	}

	if bytesCopied != remainBytes {
		fmt.Fprintf(os.Stderr, "Error sending file: %v, copied=%d, not %d\n", blobRef,
			bytesCopied, remainBytes)
		killConnection()
		return
	}
}

// Unauthenticated user.  Be paranoid.
func handleGetViaSharing(conn http.ResponseWriter, req *http.Request,
	blobRef *blobref.BlobRef, fetcher blobref.StreamingFetcher) {

	if w, ok := fetcher.(blobserver.ContextWrapper); ok {
		fetcher = w.WrapContext(req)
	}

	viaPathOkay := false
	startTime := time.Now()
	defer func() {
		if !viaPathOkay {
			// Insert a delay, to hide timing attacks probing
			// for the existence of blobs.
			sleep := fetchFailureDelayNs - (time.Now().Sub(startTime))
			if sleep > 0 {
				time.Sleep(sleep)
			}
		}
	}()
	viaBlobs := make([]*blobref.BlobRef, 0)
	if via := req.FormValue("via"); via != "" {
		for _, vs := range strings.Split(via, ",") {
			if br := blobref.Parse(vs); br == nil {
				httputil.BadRequestError(conn, "Malformed blobref in via param")
				return
			} else {
				viaBlobs = append(viaBlobs, br)
			}
		}
	}

	fetchChain := make([]*blobref.BlobRef, 0)
	fetchChain = append(fetchChain, viaBlobs...)
	fetchChain = append(fetchChain, blobRef)
	for i, br := range fetchChain {
		switch i {
		case 0:
			file, size, err := fetcher.FetchStreaming(br)
			if err != nil {
				log.Printf("Fetch chain 0 of %s failed: %v", br.String(), err)
				auth.SendUnauthorized(conn)
				return
			}
			defer file.Close()
			if size > maxJSONSize {
				log.Printf("Fetch chain 0 of %s too large", br.String())
				auth.SendUnauthorized(conn)
				return
			}
			jd := json.NewDecoder(file)
			m := make(map[string]interface{})
			if err := jd.Decode(&m); err != nil {
				log.Printf("Fetch chain 0 of %s wasn't JSON: %v", br.String(), err)
				auth.SendUnauthorized(conn)
				return
			}
			if m["camliType"].(string) != "share" {
				log.Printf("Fetch chain 0 of %s wasn't a share", br.String())
				auth.SendUnauthorized(conn)
				return
			}
			if len(fetchChain) > 1 && fetchChain[1].String() != m["target"].(string) {
				log.Printf("Fetch chain 0->1 (%s -> %q) unauthorized, expected hop to %q",
					br.String(), fetchChain[1].String(), m["target"])
				auth.SendUnauthorized(conn)
				return
			}
		case len(fetchChain) - 1:
			// Last one is fine (as long as its path up to here has been proven, and it's
			// not the first thing in the chain)
			continue
		default:
			file, _, err := fetcher.FetchStreaming(br)
			if err != nil {
				log.Printf("Fetch chain %d of %s failed: %v", i, br.String(), err)
				auth.SendUnauthorized(conn)
				return
			}
			defer file.Close()
			lr := io.LimitReader(file, maxJSONSize)
			slurpBytes, err := ioutil.ReadAll(lr)
			if err != nil {
				log.Printf("Fetch chain %d of %s failed in slurp: %v", i, br.String(), err)
				auth.SendUnauthorized(conn)
				return
			}
			saught := fetchChain[i+1].String()
			if bytes.IndexAny(slurpBytes, saught) == -1 {
				log.Printf("Fetch chain %d of %s failed; no reference to %s",
					i, br.String(), saught)
				auth.SendUnauthorized(conn)
				return
			}
		}
	}

	viaPathOkay = true

	serveBlobRef(conn, req, blobRef, fetcher)

}

func blobFromUrlPath(path string) *blobref.BlobRef {
	return blobref.FromPattern(kGetPattern, path)
}

// For client testing.
func simulatePrematurelyClosedConnection(conn http.ResponseWriter, req *http.Request) {
	flusher, ok := conn.(http.Flusher)
	if !ok {
		return
	}
	hj, ok := conn.(http.Hijacker)
	if !ok {
		return
	}
	for n := 1; n <= 100; n++ {
		fmt.Fprintf(conn, "line %d\n", n)
		flusher.Flush()
	}
	wrc, _, _ := hj.Hijack()
	wrc.Close() // without sending final chunk; should be an error for the client
}
<commit_msg>get handler: remove old sniffing for demos which no longer works and makes things much slower<commit_after>/*
Copyright 2011 Google Inc.

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

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

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

package gethandler

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"camlistore.org/pkg/auth"
	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/httputil"
	"camlistore.org/pkg/misc/httprange" // TODO: delete this package, use http.ServeContent
)

var kGetPattern *regexp.Regexp = regexp.MustCompile(`/camli/([a-z0-9]+)-([a-f0-9]+)$`)

type Handler struct {
	Fetcher           blobref.StreamingFetcher
	AllowGlobalAccess bool
}

func CreateGetHandler(fetcher blobref.StreamingFetcher) func(http.ResponseWriter, *http.Request) {
	gh := &Handler{Fetcher: fetcher}
	return func(conn http.ResponseWriter, req *http.Request) {
		if req.URL.Path == "/camli/sha1-deadbeef00000000000000000000000000000000" {
			// Test handler.
			simulatePrematurelyClosedConnection(conn, req)
			return
		}
		gh.ServeHTTP(conn, req)
	}
}

const fetchFailureDelayNs = 200e6 // 200 ms
const maxJSONSize = 64 * 1024     // should be enough for everyone

func (h *Handler) ServeHTTP(conn http.ResponseWriter, req *http.Request) {
	blobRef := blobFromUrlPath(req.URL.Path)
	if blobRef == nil {
		http.Error(conn, "Malformed GET URL.", 400)
		return
	}

	switch {
	case h.AllowGlobalAccess || auth.Allowed(req, auth.OpGet):
		serveBlobRef(conn, req, blobRef, h.Fetcher)
	case auth.TriedAuthorization(req):
		log.Printf("Attempted authorization failed on %s", req.URL)
		auth.SendUnauthorized(conn)
	default:
		handleGetViaSharing(conn, req, blobRef, h.Fetcher)
	}
}

// serveBlobRef sends 'blobref' to 'conn' as directed by the Range header in 'req'
func serveBlobRef(conn http.ResponseWriter, req *http.Request,
	blobRef *blobref.BlobRef, fetcher blobref.StreamingFetcher) {

	if w, ok := fetcher.(blobserver.ContextWrapper); ok {
		fetcher = w.WrapContext(req)
	}

	file, size, err := fetcher.FetchStreaming(blobRef)
	switch err {
	case nil:
		break
	case os.ErrNotExist:
		conn.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(conn, "Blob %q not found", blobRef)
		return
	default:
		httputil.ServerError(conn, req, err)
		return
	}

	defer file.Close()

	seeker, isSeeker := file.(io.Seeker)
	reqRange := httprange.FromRequest(req)
	if reqRange.SkipBytes() != 0 && isSeeker {
		// TODO: set the Range-specific response headers too,
		// acknowledging that we honored the content range
		// request.
		_, err = seeker.Seek(reqRange.SkipBytes(), 0)
		if err != nil {
			httputil.ServerError(conn, req, err)
			return
		}
	}

	var input io.Reader = file
	if reqRange.LimitBytes() != -1 {
		input = io.LimitReader(file, reqRange.LimitBytes())
	}

	remainBytes := size - reqRange.SkipBytes()
	if reqRange.LimitBytes() != -1 &&
		reqRange.LimitBytes() < remainBytes {
		remainBytes = reqRange.LimitBytes()
	}

	if reqRange.IsWholeFile() {
		conn.Header().Set("Content-Length", strconv.FormatInt(size, 10))
	}

	conn.Header().Set("Content-Type", "application/octet-stream")
	if !reqRange.IsWholeFile() {
		conn.Header().Set("Content-Range",
			fmt.Sprintf("bytes %d-%d/%d", reqRange.SkipBytes(),
				reqRange.SkipBytes()+remainBytes,
				size))
		conn.WriteHeader(http.StatusPartialContent)
	}
	bytesCopied, err := io.Copy(conn, input)

	// If there's an error at this point, it's too late to tell the client,
	// as they've already been receiving bytes.  But they should be smart enough
	// to verify the digest doesn't match.  But we close the (chunked) response anyway,
	// to further signal errors.
	killConnection := func() {
		if hj, ok := conn.(http.Hijacker); ok {
			log.Printf("Force-closing TCP connection to signal error sending %q", blobRef)
			if closer, _, err := hj.Hijack(); err != nil {
				closer.Close()
			}
		}
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error sending file: %v, err=%v\n", blobRef, err)
		killConnection()
		return
	}

	if bytesCopied != remainBytes {
		fmt.Fprintf(os.Stderr, "Error sending file: %v, copied=%d, not %d\n", blobRef,
			bytesCopied, remainBytes)
		killConnection()
		return
	}
}

// Unauthenticated user.  Be paranoid.
func handleGetViaSharing(conn http.ResponseWriter, req *http.Request,
	blobRef *blobref.BlobRef, fetcher blobref.StreamingFetcher) {

	if w, ok := fetcher.(blobserver.ContextWrapper); ok {
		fetcher = w.WrapContext(req)
	}

	viaPathOkay := false
	startTime := time.Now()
	defer func() {
		if !viaPathOkay {
			// Insert a delay, to hide timing attacks probing
			// for the existence of blobs.
			sleep := fetchFailureDelayNs - (time.Now().Sub(startTime))
			if sleep > 0 {
				time.Sleep(sleep)
			}
		}
	}()
	viaBlobs := make([]*blobref.BlobRef, 0)
	if via := req.FormValue("via"); via != "" {
		for _, vs := range strings.Split(via, ",") {
			if br := blobref.Parse(vs); br == nil {
				httputil.BadRequestError(conn, "Malformed blobref in via param")
				return
			} else {
				viaBlobs = append(viaBlobs, br)
			}
		}
	}

	fetchChain := make([]*blobref.BlobRef, 0)
	fetchChain = append(fetchChain, viaBlobs...)
	fetchChain = append(fetchChain, blobRef)
	for i, br := range fetchChain {
		switch i {
		case 0:
			file, size, err := fetcher.FetchStreaming(br)
			if err != nil {
				log.Printf("Fetch chain 0 of %s failed: %v", br.String(), err)
				auth.SendUnauthorized(conn)
				return
			}
			defer file.Close()
			if size > maxJSONSize {
				log.Printf("Fetch chain 0 of %s too large", br.String())
				auth.SendUnauthorized(conn)
				return
			}
			jd := json.NewDecoder(file)
			m := make(map[string]interface{})
			if err := jd.Decode(&m); err != nil {
				log.Printf("Fetch chain 0 of %s wasn't JSON: %v", br.String(), err)
				auth.SendUnauthorized(conn)
				return
			}
			if m["camliType"].(string) != "share" {
				log.Printf("Fetch chain 0 of %s wasn't a share", br.String())
				auth.SendUnauthorized(conn)
				return
			}
			if len(fetchChain) > 1 && fetchChain[1].String() != m["target"].(string) {
				log.Printf("Fetch chain 0->1 (%s -> %q) unauthorized, expected hop to %q",
					br.String(), fetchChain[1].String(), m["target"])
				auth.SendUnauthorized(conn)
				return
			}
		case len(fetchChain) - 1:
			// Last one is fine (as long as its path up to here has been proven, and it's
			// not the first thing in the chain)
			continue
		default:
			file, _, err := fetcher.FetchStreaming(br)
			if err != nil {
				log.Printf("Fetch chain %d of %s failed: %v", i, br.String(), err)
				auth.SendUnauthorized(conn)
				return
			}
			defer file.Close()
			lr := io.LimitReader(file, maxJSONSize)
			slurpBytes, err := ioutil.ReadAll(lr)
			if err != nil {
				log.Printf("Fetch chain %d of %s failed in slurp: %v", i, br.String(), err)
				auth.SendUnauthorized(conn)
				return
			}
			saught := fetchChain[i+1].String()
			if bytes.IndexAny(slurpBytes, saught) == -1 {
				log.Printf("Fetch chain %d of %s failed; no reference to %s",
					i, br.String(), saught)
				auth.SendUnauthorized(conn)
				return
			}
		}
	}

	viaPathOkay = true

	serveBlobRef(conn, req, blobRef, fetcher)

}

func blobFromUrlPath(path string) *blobref.BlobRef {
	return blobref.FromPattern(kGetPattern, path)
}

// For client testing.
func simulatePrematurelyClosedConnection(conn http.ResponseWriter, req *http.Request) {
	flusher, ok := conn.(http.Flusher)
	if !ok {
		return
	}
	hj, ok := conn.(http.Hijacker)
	if !ok {
		return
	}
	for n := 1; n <= 100; n++ {
		fmt.Fprintf(conn, "line %d\n", n)
		flusher.Flush()
	}
	wrc, _, _ := hj.Hijack()
	wrc.Close() // without sending final chunk; should be an error for the client
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strings"
	"sync"
	"text/template"
)

type label struct {
	Name string `json:"name"`
}

type prInfo struct {
	Labels []label `json:"labels"`
	Number int     `json:"number"`
	Title  string  `json:"title"`
}

const (
	markdownTemplate = `
{{- range $typeName, $components := . }}
## {{ $typeName }}
{{- range $componentName, $component := $components }}
### {{ $componentName}}
{{- range $prInfo := $component }}
 - {{ $prInfo.Title }} #{{ $prInfo.Number }}
{{- end }}
{{- end }}
{{- end }}
`

	prefixType      = "Type: "
	prefixComponent = "Component: "
)

func loadMergedPRs(from, to string) ([]string, error) {
	cmd := exec.Command("git", "log", "--oneline", fmt.Sprintf("%s...%s", from, to))
	out, err := cmd.Output()
	if err != nil {
		execErr := err.(*exec.ExitError)
		return nil, fmt.Errorf("%s:\nstderr: %s\nstdout: %s", err.Error(), execErr.Stderr, out)
	}

	var prs []string
	rgx := regexp.MustCompile(`Merge pull request #(\d+)`)
	lines := strings.Split(string(out), "\n")
	for _, line := range lines {
		lineInfo := rgx.FindStringSubmatch(line)
		if len(lineInfo) == 2 {
			prs = append(prs, lineInfo[1])
		}
	}

	sort.Strings(prs)
	return prs, nil
}

func loadPRinfo(pr string) (prInfo, error) {
	cmd := exec.Command("gh", "pr", "view", pr, "--json", "title,number,labels")
	out, err := cmd.Output()
	if err != nil {
		execErr := err.(*exec.ExitError)
		return prInfo{}, fmt.Errorf("%s:\nstderr: %s\nstdout: %s", err.Error(), execErr.Stderr, out)
	}
	var prInfo prInfo
	err = json.Unmarshal(out, &prInfo)
	return prInfo, err
}

func loadAllPRs(prs []string) ([]prInfo, error) {
	errChan := make(chan error)
	wgDone := make(chan bool)
	prChan := make(chan string, len(prs))
	// fill the work queue
	for _, s := range prs {
		prChan <- s
	}
	close(prChan)

	var prInfos []prInfo

	wg := sync.WaitGroup{}
	mu := sync.Mutex{}
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			// load meta data about PRs
			defer wg.Done()
			for b := range prChan {
				fmt.Print(".")
				prInfo, err := loadPRinfo(b)
				if err != nil {
					errChan <- err
					break
				}
				mu.Lock()
				prInfos = append(prInfos, prInfo)
				mu.Unlock()
			}
		}()
	}

	go func() {
		// wait for the loading to finish
		wg.Wait()
		close(wgDone)
	}()

	var err error
	select {
	case <-wgDone:
		break
	case err = <-errChan:
		break
	}

	fmt.Println()
	return prInfos, err
}

func groupPRs(prInfos []prInfo) map[string]map[string][]prInfo {
	prPerType := map[string]map[string][]prInfo{}

	for _, info := range prInfos {
		var typ, component string
		for _, lbl := range info.Labels {
			switch {
			case strings.HasPrefix(lbl.Name, prefixType):
				typ = strings.TrimPrefix(lbl.Name, prefixType)
			case strings.HasPrefix(lbl.Name, prefixComponent):
				component = strings.TrimPrefix(lbl.Name, prefixComponent)
			}
		}
		if typ == "" {
			typ = "Other"
		}
		if component == "" {
			component = "Other"
		}
		components, exists := prPerType[typ]
		if !exists {
			components = map[string][]prInfo{}
			prPerType[typ] = components
		}

		prsPerComponentAndType := components[component]
		components[component] = append(prsPerComponentAndType, info)
	}
	return prPerType
}

func writePrInfos(fileout string, prPerType map[string]map[string][]prInfo) (err error) {
	writeTo := os.Stdout
	if fileout != "" {
		writeTo, err = os.OpenFile(fileout, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
		if err != nil {
			return err
		}
	}

	t := template.Must(template.New("markdownTemplate").Parse(markdownTemplate))
	err = t.ExecuteTemplate(writeTo, "markdownTemplate", prPerType)
	if err != nil {
		return err
	}
	return nil
}

func main() {
	from := flag.String("from", "", "from sha/tag/branch")
	to := flag.String("to", "HEAD", "to sha/tag/branch")
	fileout := flag.String("file", "", "file on which to write release notes, stdout if empty")

	flag.Parse()

	prs, err := loadMergedPRs(*from, *to)
	if err != nil {
		log.Fatal(err)
	}

	prInfos, err := loadAllPRs(prs)
	if err != nil {
		log.Fatal(err)
	}

	prPerType := groupPRs(prInfos)

	err = writePrInfos(*fileout, prPerType)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>License header in release_notes.go<commit_after>/*
Copyright 2021 The Vitess Authors.

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

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

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

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strings"
	"sync"
	"text/template"
)

type label struct {
	Name string `json:"name"`
}

type prInfo struct {
	Labels []label `json:"labels"`
	Number int     `json:"number"`
	Title  string  `json:"title"`
}

const (
	markdownTemplate = `
{{- range $typeName, $components := . }}
## {{ $typeName }}
{{- range $componentName, $component := $components }}
### {{ $componentName}}
{{- range $prInfo := $component }}
 - {{ $prInfo.Title }} #{{ $prInfo.Number }}
{{- end }}
{{- end }}
{{- end }}
`

	prefixType      = "Type: "
	prefixComponent = "Component: "
)

func loadMergedPRs(from, to string) ([]string, error) {
	cmd := exec.Command("git", "log", "--oneline", fmt.Sprintf("%s...%s", from, to))
	out, err := cmd.Output()
	if err != nil {
		execErr := err.(*exec.ExitError)
		return nil, fmt.Errorf("%s:\nstderr: %s\nstdout: %s", err.Error(), execErr.Stderr, out)
	}

	var prs []string
	rgx := regexp.MustCompile(`Merge pull request #(\d+)`)
	lines := strings.Split(string(out), "\n")
	for _, line := range lines {
		lineInfo := rgx.FindStringSubmatch(line)
		if len(lineInfo) == 2 {
			prs = append(prs, lineInfo[1])
		}
	}

	sort.Strings(prs)
	return prs, nil
}

func loadPRinfo(pr string) (prInfo, error) {
	cmd := exec.Command("gh", "pr", "view", pr, "--json", "title,number,labels")
	out, err := cmd.Output()
	if err != nil {
		execErr := err.(*exec.ExitError)
		return prInfo{}, fmt.Errorf("%s:\nstderr: %s\nstdout: %s", err.Error(), execErr.Stderr, out)
	}
	var prInfo prInfo
	err = json.Unmarshal(out, &prInfo)
	return prInfo, err
}

func loadAllPRs(prs []string) ([]prInfo, error) {
	errChan := make(chan error)
	wgDone := make(chan bool)
	prChan := make(chan string, len(prs))
	// fill the work queue
	for _, s := range prs {
		prChan <- s
	}
	close(prChan)

	var prInfos []prInfo

	wg := sync.WaitGroup{}
	mu := sync.Mutex{}
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			// load meta data about PRs
			defer wg.Done()
			for b := range prChan {
				fmt.Print(".")
				prInfo, err := loadPRinfo(b)
				if err != nil {
					errChan <- err
					break
				}
				mu.Lock()
				prInfos = append(prInfos, prInfo)
				mu.Unlock()
			}
		}()
	}

	go func() {
		// wait for the loading to finish
		wg.Wait()
		close(wgDone)
	}()

	var err error
	select {
	case <-wgDone:
		break
	case err = <-errChan:
		break
	}

	fmt.Println()
	return prInfos, err
}

func groupPRs(prInfos []prInfo) map[string]map[string][]prInfo {
	prPerType := map[string]map[string][]prInfo{}

	for _, info := range prInfos {
		var typ, component string
		for _, lbl := range info.Labels {
			switch {
			case strings.HasPrefix(lbl.Name, prefixType):
				typ = strings.TrimPrefix(lbl.Name, prefixType)
			case strings.HasPrefix(lbl.Name, prefixComponent):
				component = strings.TrimPrefix(lbl.Name, prefixComponent)
			}
		}
		if typ == "" {
			typ = "Other"
		}
		if component == "" {
			component = "Other"
		}
		components, exists := prPerType[typ]
		if !exists {
			components = map[string][]prInfo{}
			prPerType[typ] = components
		}

		prsPerComponentAndType := components[component]
		components[component] = append(prsPerComponentAndType, info)
	}
	return prPerType
}

func writePrInfos(fileout string, prPerType map[string]map[string][]prInfo) (err error) {
	writeTo := os.Stdout
	if fileout != "" {
		writeTo, err = os.OpenFile(fileout, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
		if err != nil {
			return err
		}
	}

	t := template.Must(template.New("markdownTemplate").Parse(markdownTemplate))
	err = t.ExecuteTemplate(writeTo, "markdownTemplate", prPerType)
	if err != nil {
		return err
	}
	return nil
}

func main() {
	from := flag.String("from", "", "from sha/tag/branch")
	to := flag.String("to", "HEAD", "to sha/tag/branch")
	fileout := flag.String("file", "", "file on which to write release notes, stdout if empty")

	flag.Parse()

	prs, err := loadMergedPRs(*from, *to)
	if err != nil {
		log.Fatal(err)
	}

	prInfos, err := loadAllPRs(prs)
	if err != nil {
		log.Fatal(err)
	}

	prPerType := groupPRs(prInfos)

	err = writePrInfos(*fileout, prPerType)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Code generated by protoc-gen-go.
// source: add.proto
// DO NOT EDIT!

/*
Package pb is a generated protocol buffer package.

It is generated from these files:
	add.proto

It has these top-level messages:
	AddRequest
	AddReply
*/
package pb

import proto "github.com/golang/protobuf/proto"

import (
	context "golang.org/x/net/context"
	grpc "google.golang.org/grpc"
)

// Reference imports to suppress errors if they are not otherwise used.
var _ context.Context
var _ grpc.ClientConn

// Reference imports to suppress errors if they are not otherwise used.
var _ = proto.Marshal

// The request contains two parameters.
type AddRequest struct {
	A int64 `protobuf:"varint,1,opt,name=a" json:"a,omitempty"`
	B int64 `protobuf:"varint,2,opt,name=b" json:"b,omitempty"`
}

func (m *AddRequest) Reset()         { *m = AddRequest{} }
func (m *AddRequest) String() string { return proto.CompactTextString(m) }
func (*AddRequest) ProtoMessage()    {}

// The response contains the result of the calculation.
type AddReply struct {
	V int64 `protobuf:"varint,1,opt,name=v" json:"v,omitempty"`
}

func (m *AddReply) Reset()         { *m = AddReply{} }
func (m *AddReply) String() string { return proto.CompactTextString(m) }
func (*AddReply) ProtoMessage()    {}

func init() {
}

// Client API for Add service

type AddClient interface {
	// Adds two integers.
	Add(ctx context.Context, in *AddRequest, opts ...grpc.CallOption) (*AddReply, error)
}

type addClient struct {
	cc *grpc.ClientConn
}

func NewAddClient(cc *grpc.ClientConn) AddClient {
	return &addClient{cc}
}

func (c *addClient) Add(ctx context.Context, in *AddRequest, opts ...grpc.CallOption) (*AddReply, error) {
	out := new(AddReply)
	err := grpc.Invoke(ctx, "/pb.Add/Add", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// Server API for Add service

type AddServer interface {
	// Adds two integers.
	Add(context.Context, *AddRequest) (*AddReply, error)
}

func RegisterAddServer(s *grpc.Server, srv AddServer) {
	s.RegisterService(&_Add_serviceDesc, srv)
}

func _Add_Add_Handler(srv interface{}, ctx context.Context, buf []byte) (proto.Message, error) {
	in := new(AddRequest)
	if err := proto.Unmarshal(buf, in); err != nil {
		return nil, err
	}
	out, err := srv.(AddServer).Add(ctx, in)
	if err != nil {
		return nil, err
	}
	return out, nil
}

var _Add_serviceDesc = grpc.ServiceDesc{
	ServiceName: "pb.Add",
	HandlerType: (*AddServer)(nil),
	Methods: []grpc.MethodDesc{
		{
			MethodName: "Add",
			Handler:    _Add_Add_Handler,
		},
	},
	Streams: []grpc.StreamDesc{},
}
<commit_msg>New protoc and gRPC<commit_after>// Code generated by protoc-gen-go.
// source: add.proto
// DO NOT EDIT!

/*
Package pb is a generated protocol buffer package.

It is generated from these files:
	add.proto

It has these top-level messages:
	AddRequest
	AddReply
*/
package pb

import proto "github.com/golang/protobuf/proto"

import (
	context "golang.org/x/net/context"
	grpc "google.golang.org/grpc"
)

// Reference imports to suppress errors if they are not otherwise used.
var _ context.Context
var _ grpc.ClientConn

// Reference imports to suppress errors if they are not otherwise used.
var _ = proto.Marshal

// The request contains two parameters.
type AddRequest struct {
	A int64 `protobuf:"varint,1,opt,name=a" json:"a,omitempty"`
	B int64 `protobuf:"varint,2,opt,name=b" json:"b,omitempty"`
}

func (m *AddRequest) Reset()         { *m = AddRequest{} }
func (m *AddRequest) String() string { return proto.CompactTextString(m) }
func (*AddRequest) ProtoMessage()    {}

// The response contains the result of the calculation.
type AddReply struct {
	V int64 `protobuf:"varint,1,opt,name=v" json:"v,omitempty"`
}

func (m *AddReply) Reset()         { *m = AddReply{} }
func (m *AddReply) String() string { return proto.CompactTextString(m) }
func (*AddReply) ProtoMessage()    {}

func init() {
}

// Client API for Add service

type AddClient interface {
	// Adds two integers.
	Add(ctx context.Context, in *AddRequest, opts ...grpc.CallOption) (*AddReply, error)
}

type addClient struct {
	cc *grpc.ClientConn
}

func NewAddClient(cc *grpc.ClientConn) AddClient {
	return &addClient{cc}
}

func (c *addClient) Add(ctx context.Context, in *AddRequest, opts ...grpc.CallOption) (*AddReply, error) {
	out := new(AddReply)
	err := grpc.Invoke(ctx, "/pb.Add/Add", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// Server API for Add service

type AddServer interface {
	// Adds two integers.
	Add(context.Context, *AddRequest) (*AddReply, error)
}

func RegisterAddServer(s *grpc.Server, srv AddServer) {
	s.RegisterService(&_Add_serviceDesc, srv)
}

func _Add_Add_Handler(srv interface{}, ctx context.Context, buf []byte) (interface{}, error) {
	in := new(AddRequest)
	if err := proto.Unmarshal(buf, in); err != nil {
		return nil, err
	}
	out, err := srv.(AddServer).Add(ctx, in)
	if err != nil {
		return nil, err
	}
	return out, nil
}

var _Add_serviceDesc = grpc.ServiceDesc{
	ServiceName: "pb.Add",
	HandlerType: (*AddServer)(nil),
	Methods: []grpc.MethodDesc{
		{
			MethodName: "Add",
			Handler:    _Add_Add_Handler,
		},
	},
	Streams: []grpc.StreamDesc{},
}
<|endoftext|>


<commit_before>package bbs_test

import (
	"github.com/cloudfoundry/storeadapter"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/pivotal-cf-experimental/runtime-schema/bbs"
	"github.com/pivotal-cf-experimental/runtime-schema/models"
	"time"
)

var _ = Describe("RunOnce BBS", func() {
	var bbs *BBS
	var runOnce models.RunOnce

	BeforeEach(func() {
		bbs = New(store)
		runOnce = models.RunOnce{
			Guid:            "some-guid",
			ExecutorID:      "executor-id",
			ContainerHandle: "container-handle",
		}
	})

	itRetriesUntilStoreComesBack := func(action func(*BBS, models.RunOnce) error) {
		It("should keep trying until the store comes back", func(done Done) {
			etcdRunner.GoAway()

			runResult := make(chan error)
			go func() {
				err := action(bbs, runOnce)
				runResult <- err
			}()

			time.Sleep(200 * time.Millisecond)

			etcdRunner.ComeBack()

			Ω(<-runResult).ShouldNot(HaveOccurred())

			close(done)
		}, 5)
	}

	Describe("DesireRunOnce", func() {
		BeforeEach(func() {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("creates /run_once/pending/<guid>", func() {
			node, err := store.Get("/v1/run_once/pending/some-guid")
			Ω(err).ShouldNot(HaveOccurred())
			Ω(node.Value).Should(Equal(runOnce.ToJSON()))
		})

		Context("when the RunOnce is already pending", func() {
			It("should happily overwrite the existing RunOnce", func() {
				err := bbs.DesireRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())
			})
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).DesireRunOnce)
		})
	})

	Describe("ResolveRunOnce", func() {
		BeforeEach(func() {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("should remove /run_once/pending/<guid>", func() {
			err := bbs.ResolveRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			_, err = store.Get("/v1/run_once/pending/some-guid")
			Ω(err).Should(Equal(storeadapter.ErrorKeyNotFound))
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).ResolveRunOnce)
		})
	})

	Describe("ClaimRunOnce", func() {
		Context("when claimed with a correctly configured runOnce", func() {
			BeforeEach(func() {
				runOnce.ExecutorID = "executor-id"
			})

			It("creates /run_once/claimed/<guid>", func() {
				err := bbs.ClaimRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())

				node, err := store.Get("/v1/run_once/claimed/some-guid")
				Ω(err).ShouldNot(HaveOccurred())
				Ω(node).Should(Equal(storeadapter.StoreNode{
					Key:   "/v1/run_once/claimed/some-guid",
					Value: runOnce.ToJSON(),
					TTL:   10, // move to config one day
				}))
			})
		})

		Context("when claimed with a runOnce that is missing ExecutorID", func() {
			BeforeEach(func() {
				runOnce.ExecutorID = ""
			})

			It("should panic", func() {
				Ω(func() {
					bbs.ClaimRunOnce(runOnce)
				}).Should(Panic())
			})
		})

		Context("when the RunOnce is already claimed", func() {
			BeforeEach(func() {
				err := bbs.ClaimRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())
			})

			It("returns an error", func() {
				err := bbs.ClaimRunOnce(runOnce)
				Ω(err).Should(HaveOccurred())
			})
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).ClaimRunOnce)
		})
	})

	Describe("StartRunOnce", func() {
		BeforeEach(func() {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			err = bbs.ClaimRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("creates running", func() {
			err := bbs.StartRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			node, err := store.Get("/v1/run_once/running/some-guid")
			Ω(err).ShouldNot(HaveOccurred())
			Ω(node).Should(Equal(storeadapter.StoreNode{
				Key:   "/v1/run_once/running/some-guid",
				Value: runOnce.ToJSON(),
			}))
		})

		Context("when starting with a runOnce that is missing ExecutorID", func() {
			BeforeEach(func() {
				runOnce.ExecutorID = ""
			})

			It("should panic", func() {
				Ω(func() {
					bbs.StartRunOnce(runOnce)
				}).Should(Panic())
			})
		})

		Context("when starting with a runOnce that is missing ContainerHandle", func() {
			BeforeEach(func() {
				runOnce.ContainerHandle = ""
			})

			It("should panic", func() {
				Ω(func() {
					bbs.StartRunOnce(runOnce)
				}).Should(Panic())
			})
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).StartRunOnce)
		})
	})

	Describe("CompletedRunOnce", func() {
		BeforeEach(func() {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			err = bbs.ClaimRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			err = bbs.StartRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("creates an entry under /run_once/completed", func() {
			runOnce.Failed = true
			runOnce.FailureReason = "because i said so"

			err := bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			node, err := store.Get("/v1/run_once/completed/some-guid")
			Ω(err).ShouldNot(HaveOccurred())
			Ω(node).Should(Equal(storeadapter.StoreNode{
				Key:   "/v1/run_once/completed/some-guid",
				Value: runOnce.ToJSON(),
			}))
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).CompletedRunOnce)
		})
	})

	Describe("WatchForDesiredRunOnce", func() {
		It("should send an event down the pipe for creates", func(done Done) {
			events, _, _ := bbs.WatchForDesiredRunOnce()

			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			close(done)
		})

		It("should send an event down the pipe for sets", func(done Done) {
			events, _, _ := bbs.WatchForDesiredRunOnce()

			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			err = bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			close(done)
		})

		It("should not send an event down the pipe for deletes", func(done Done) {
			events, _, _ := bbs.WatchForDesiredRunOnce()

			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			err = bbs.ResolveRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			otherRunOnce := runOnce
			otherRunOnce.Guid = runOnce.Guid + "1"

			err = bbs.DesireRunOnce(otherRunOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(otherRunOnce))

			close(done)
		})

		It("closes the events channel when told to stop", func(done Done) {
			events, stop, _ := bbs.WatchForDesiredRunOnce()

			stop <- true

			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			_, ok := <-events

			Expect(ok).To(BeFalse())

			close(done)
		})
	})

	Describe("WatchForCompletedRunOnce", func() {
		It("should send an event down the pipe for creates", func(done Done) {
			events, _, _ := bbs.WatchForCompletedRunOnce()

			err := bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			close(done)
		})

		It("should send an event down the pipe for sets", func(done Done) {
			events, _, _ := bbs.WatchForCompletedRunOnce()

			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			err = bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			bbs.ConvergeRunOnce() //should bump the completed key

			Expect(<-events).To(Equal(runOnce))

			close(done)
		})

		It("should not send an event down the pipe for deletes", func(done Done) {
			events, _, _ := bbs.WatchForCompletedRunOnce()

			err := bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			bbs.ConvergeRunOnce() //should delete the key

			otherRunOnce := runOnce
			otherRunOnce.Guid = runOnce.Guid + "1"

			err = bbs.CompletedRunOnce(otherRunOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(otherRunOnce))

			close(done)
		})

		It("closes the events channel when told to stop", func(done Done) {
			events, stop, _ := bbs.WatchForCompletedRunOnce()

			stop <- true

			err := bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			_, ok := <-events

			Expect(ok).To(BeFalse())

			close(done)
		})
	})

	Describe("ConvergeRunOnce", func() {
		var otherRunOnce models.RunOnce

		BeforeEach(func() {
			otherRunOnce = models.RunOnce{
				Guid: "some-other-guid",
			}
		})

		Context("when a pending key exists", func() {
			BeforeEach(func() {
				err := bbs.DesireRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())
			})

			Context("and there is a claim key", func() {
				BeforeEach(func() {
					err := bbs.ClaimRunOnce(runOnce)
					Ω(err).ShouldNot(HaveOccurred())
				})

				It("should not kick the pending key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/pending")
					bbs.ConvergeRunOnce()

					bbs.DesireRunOnce(otherRunOnce)

					Ω((<-events).Node.Key).Should(Equal("/v1/run_once/pending/some-other-guid"))

					close(done)
				})
			})

			Context("and there is a running key", func() {
				BeforeEach(func() {
					err := bbs.StartRunOnce(runOnce)
					Ω(err).ShouldNot(HaveOccurred())
				})

				It("should not kick the pending key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/pending")
					bbs.ConvergeRunOnce()

					bbs.DesireRunOnce(otherRunOnce)

					Ω((<-events).Node.Key).Should(Equal("/v1/run_once/pending/some-other-guid"))

					close(done)
				})
			})

			Context("and there is a completed key", func() {
				BeforeEach(func() {
					err := bbs.CompletedRunOnce(runOnce)
					Ω(err).ShouldNot(HaveOccurred())
				})

				It("should not kick the pending key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/pending")
					bbs.ConvergeRunOnce()

					bbs.DesireRunOnce(otherRunOnce)

					Ω((<-events).Node.Key).Should(Equal("/v1/run_once/pending/some-other-guid"))

					close(done)
				})

				It("should kick the completed key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/completed")
					bbs.ConvergeRunOnce()

					event := <-events
					Ω(event.Node.Key).Should(Equal("/v1/run_once/completed/some-guid"))
					Ω(event.Node.Value).Should(Equal(runOnce.ToJSON()))

					close(done)
				})
			})

			Context("and there are no other keys", func() {
				It("should kick the pending key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/pending")
					bbs.ConvergeRunOnce()

					event := <-events
					Ω(event.Node.Key).Should(Equal("/v1/run_once/pending/some-guid"))
					Ω(event.Node.Value).Should(Equal(runOnce.ToJSON()))

					close(done)
				})
			})
		})

		Context("when a pending key does not exist", func() {
			BeforeEach(func() {
				err := bbs.ClaimRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())

				err = bbs.StartRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())

				err = bbs.CompletedRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())
			})

			It("should delete any extra keys", func() {
				bbs.ConvergeRunOnce()

				_, err := store.Get("/v1/run_once/claimed/some-guid")
				Ω(err).Should(HaveOccurred())

				_, err = store.Get("/v1/run_once/running/some-guid")
				Ω(err).Should(HaveOccurred())

				_, err = store.Get("/v1/run_once/completed/some-guid")
				Ω(err).Should(HaveOccurred())
			})
		})
	})
})
<commit_msg>fix race when setting up watcher<commit_after>package bbs_test

import (
	"github.com/cloudfoundry/storeadapter"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/pivotal-cf-experimental/runtime-schema/bbs"
	"github.com/pivotal-cf-experimental/runtime-schema/models"
	"time"
)

var _ = Describe("RunOnce BBS", func() {
	var bbs *BBS
	var runOnce models.RunOnce

	BeforeEach(func() {
		bbs = New(store)
		runOnce = models.RunOnce{
			Guid:            "some-guid",
			ExecutorID:      "executor-id",
			ContainerHandle: "container-handle",
		}
	})

	itRetriesUntilStoreComesBack := func(action func(*BBS, models.RunOnce) error) {
		It("should keep trying until the store comes back", func(done Done) {
			etcdRunner.GoAway()

			runResult := make(chan error)
			go func() {
				err := action(bbs, runOnce)
				runResult <- err
			}()

			time.Sleep(200 * time.Millisecond)

			etcdRunner.ComeBack()

			Ω(<-runResult).ShouldNot(HaveOccurred())

			close(done)
		}, 5)
	}

	Describe("DesireRunOnce", func() {
		BeforeEach(func() {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("creates /run_once/pending/<guid>", func() {
			node, err := store.Get("/v1/run_once/pending/some-guid")
			Ω(err).ShouldNot(HaveOccurred())
			Ω(node.Value).Should(Equal(runOnce.ToJSON()))
		})

		Context("when the RunOnce is already pending", func() {
			It("should happily overwrite the existing RunOnce", func() {
				err := bbs.DesireRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())
			})
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).DesireRunOnce)
		})
	})

	Describe("ResolveRunOnce", func() {
		BeforeEach(func() {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("should remove /run_once/pending/<guid>", func() {
			err := bbs.ResolveRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			_, err = store.Get("/v1/run_once/pending/some-guid")
			Ω(err).Should(Equal(storeadapter.ErrorKeyNotFound))
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).ResolveRunOnce)
		})
	})

	Describe("ClaimRunOnce", func() {
		Context("when claimed with a correctly configured runOnce", func() {
			BeforeEach(func() {
				runOnce.ExecutorID = "executor-id"
			})

			It("creates /run_once/claimed/<guid>", func() {
				err := bbs.ClaimRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())

				node, err := store.Get("/v1/run_once/claimed/some-guid")
				Ω(err).ShouldNot(HaveOccurred())
				Ω(node).Should(Equal(storeadapter.StoreNode{
					Key:   "/v1/run_once/claimed/some-guid",
					Value: runOnce.ToJSON(),
					TTL:   10, // move to config one day
				}))
			})
		})

		Context("when claimed with a runOnce that is missing ExecutorID", func() {
			BeforeEach(func() {
				runOnce.ExecutorID = ""
			})

			It("should panic", func() {
				Ω(func() {
					bbs.ClaimRunOnce(runOnce)
				}).Should(Panic())
			})
		})

		Context("when the RunOnce is already claimed", func() {
			BeforeEach(func() {
				err := bbs.ClaimRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())
			})

			It("returns an error", func() {
				err := bbs.ClaimRunOnce(runOnce)
				Ω(err).Should(HaveOccurred())
			})
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).ClaimRunOnce)
		})
	})

	Describe("StartRunOnce", func() {
		BeforeEach(func() {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			err = bbs.ClaimRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("creates running", func() {
			err := bbs.StartRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			node, err := store.Get("/v1/run_once/running/some-guid")
			Ω(err).ShouldNot(HaveOccurred())
			Ω(node).Should(Equal(storeadapter.StoreNode{
				Key:   "/v1/run_once/running/some-guid",
				Value: runOnce.ToJSON(),
			}))
		})

		Context("when starting with a runOnce that is missing ExecutorID", func() {
			BeforeEach(func() {
				runOnce.ExecutorID = ""
			})

			It("should panic", func() {
				Ω(func() {
					bbs.StartRunOnce(runOnce)
				}).Should(Panic())
			})
		})

		Context("when starting with a runOnce that is missing ContainerHandle", func() {
			BeforeEach(func() {
				runOnce.ContainerHandle = ""
			})

			It("should panic", func() {
				Ω(func() {
					bbs.StartRunOnce(runOnce)
				}).Should(Panic())
			})
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).StartRunOnce)
		})
	})

	Describe("CompletedRunOnce", func() {
		BeforeEach(func() {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			err = bbs.ClaimRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			err = bbs.StartRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("creates an entry under /run_once/completed", func() {
			runOnce.Failed = true
			runOnce.FailureReason = "because i said so"

			err := bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			node, err := store.Get("/v1/run_once/completed/some-guid")
			Ω(err).ShouldNot(HaveOccurred())
			Ω(node).Should(Equal(storeadapter.StoreNode{
				Key:   "/v1/run_once/completed/some-guid",
				Value: runOnce.ToJSON(),
			}))
		})

		Context("when the store is out of commission", func() {
			itRetriesUntilStoreComesBack((*BBS).CompletedRunOnce)
		})
	})

	Describe("WatchForDesiredRunOnce", func() {
		var (
			events <-chan (models.RunOnce)
			stop   chan<- bool
		)

		BeforeEach(func() {
			events, stop, _ = bbs.WatchForDesiredRunOnce()
			time.Sleep(100 * time.Millisecond) //give the watcher a chance to connect
		})

		It("should send an event down the pipe for creates", func(done Done) {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			close(done)
		})

		It("should send an event down the pipe for sets", func(done Done) {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			err = bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			close(done)
		})

		It("should not send an event down the pipe for deletes", func(done Done) {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			err = bbs.ResolveRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			otherRunOnce := runOnce
			otherRunOnce.Guid = runOnce.Guid + "1"

			err = bbs.DesireRunOnce(otherRunOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(otherRunOnce))

			close(done)
		})

		It("closes the events channel when told to stop", func(done Done) {
			stop <- true

			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			_, ok := <-events

			Expect(ok).To(BeFalse())

			close(done)
		})
	})

	Describe("WatchForCompletedRunOnce", func() {
		var (
			events <-chan (models.RunOnce)
			stop   chan<- bool
		)

		BeforeEach(func() {
			events, stop, _ = bbs.WatchForCompletedRunOnce()
			time.Sleep(100 * time.Millisecond) //give the watcher a chance to connect
		})

		It("should send an event down the pipe for creates", func(done Done) {
			err := bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			close(done)
		})

		It("should send an event down the pipe for sets", func(done Done) {
			err := bbs.DesireRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			err = bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			bbs.ConvergeRunOnce() //should bump the completed key

			Expect(<-events).To(Equal(runOnce))

			close(done)
		})

		It("should not send an event down the pipe for deletes", func(done Done) {
			err := bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(runOnce))

			bbs.ConvergeRunOnce() //should delete the key

			otherRunOnce := runOnce
			otherRunOnce.Guid = runOnce.Guid + "1"

			err = bbs.CompletedRunOnce(otherRunOnce)
			Ω(err).ShouldNot(HaveOccurred())

			Expect(<-events).To(Equal(otherRunOnce))

			close(done)
		})

		It("closes the events channel when told to stop", func(done Done) {
			stop <- true

			err := bbs.CompletedRunOnce(runOnce)
			Ω(err).ShouldNot(HaveOccurred())

			_, ok := <-events

			Expect(ok).To(BeFalse())

			close(done)
		})
	})

	Describe("ConvergeRunOnce", func() {
		var otherRunOnce models.RunOnce

		BeforeEach(func() {
			otherRunOnce = models.RunOnce{
				Guid: "some-other-guid",
			}
		})

		Context("when a pending key exists", func() {
			BeforeEach(func() {
				err := bbs.DesireRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())
			})

			Context("and there is a claim key", func() {
				BeforeEach(func() {
					err := bbs.ClaimRunOnce(runOnce)
					Ω(err).ShouldNot(HaveOccurred())
				})

				It("should not kick the pending key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/pending")
					bbs.ConvergeRunOnce()

					bbs.DesireRunOnce(otherRunOnce)

					Ω((<-events).Node.Key).Should(Equal("/v1/run_once/pending/some-other-guid"))

					close(done)
				})
			})

			Context("and there is a running key", func() {
				BeforeEach(func() {
					err := bbs.StartRunOnce(runOnce)
					Ω(err).ShouldNot(HaveOccurred())
				})

				It("should not kick the pending key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/pending")
					bbs.ConvergeRunOnce()

					bbs.DesireRunOnce(otherRunOnce)

					Ω((<-events).Node.Key).Should(Equal("/v1/run_once/pending/some-other-guid"))

					close(done)
				})
			})

			Context("and there is a completed key", func() {
				BeforeEach(func() {
					err := bbs.CompletedRunOnce(runOnce)
					Ω(err).ShouldNot(HaveOccurred())
				})

				It("should not kick the pending key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/pending")
					bbs.ConvergeRunOnce()

					bbs.DesireRunOnce(otherRunOnce)

					Ω((<-events).Node.Key).Should(Equal("/v1/run_once/pending/some-other-guid"))

					close(done)
				})

				It("should kick the completed key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/completed")
					bbs.ConvergeRunOnce()

					event := <-events
					Ω(event.Node.Key).Should(Equal("/v1/run_once/completed/some-guid"))
					Ω(event.Node.Value).Should(Equal(runOnce.ToJSON()))

					close(done)
				})
			})

			Context("and there are no other keys", func() {
				It("should kick the pending key", func(done Done) {
					events, _, _ := store.Watch("/v1/run_once/pending")
					bbs.ConvergeRunOnce()

					event := <-events
					Ω(event.Node.Key).Should(Equal("/v1/run_once/pending/some-guid"))
					Ω(event.Node.Value).Should(Equal(runOnce.ToJSON()))

					close(done)
				})
			})
		})

		Context("when a pending key does not exist", func() {
			BeforeEach(func() {
				err := bbs.ClaimRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())

				err = bbs.StartRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())

				err = bbs.CompletedRunOnce(runOnce)
				Ω(err).ShouldNot(HaveOccurred())
			})

			It("should delete any extra keys", func() {
				bbs.ConvergeRunOnce()

				_, err := store.Get("/v1/run_once/claimed/some-guid")
				Ω(err).Should(HaveOccurred())

				_, err = store.Get("/v1/run_once/running/some-guid")
				Ω(err).Should(HaveOccurred())

				_, err = store.Get("/v1/run_once/completed/some-guid")
				Ω(err).Should(HaveOccurred())
			})
		})
	})
})
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state_test

import (
	"github.com/juju/names"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
	"gopkg.in/juju/charm.v5"

	"github.com/juju/juju/component/all"
	"github.com/juju/juju/process"
	"github.com/juju/juju/state"
)

func init() {
	if err := all.RegisterForServer(); err != nil {
		panic(err)
	}
}

var _ = gc.Suite(&unitProcessesSuite{})

type unitProcessesSuite struct {
	ConnSuite
}

const metaYAML = `
name: a-charm
summary: a charm...
description: a charm...
processes:
  procA:
    type: docker
    command: do-something cool
    image: spam/eggs
    ports:
      - 8080:80
    volumes:
      - /var/nginx/html:/usr/share/nginx/html:ro
    env:
      IMPORTANT: YES
`

func (s *unitProcessesSuite) addUnit(c *gc.C, charmName, serviceName, meta string) (names.CharmTag, *state.Unit) {
	ch := s.AddTestingCharm(c, charmName)
	ch = s.AddMetaCharm(c, charmName, meta, 2)

	svc := s.AddTestingService(c, serviceName, ch)
	unit, err := svc.AddUnit()
	c.Assert(err, jc.ErrorIsNil)

	charmTag := ch.Tag().(names.CharmTag)
	return charmTag, unit
}

func (s *unitProcessesSuite) TestFunctional(c *gc.C) {
	_, unit := s.addUnit(c, "dummy", "a-service", metaYAML)

	st, err := s.State.UnitProcesses(unit)
	c.Assert(err, jc.ErrorIsNil)

	definitions, err := st.ListDefinitions()
	c.Assert(err, jc.ErrorIsNil)
	c.Check(definitions, jc.DeepEquals, []charm.Process{{
		Name:        "procA",
		Type:        "docker",
		TypeOptions: map[string]string{},
		Command:     "do-something cool",
		Image:       "spam/eggs",
		Volumes: []charm.ProcessVolume{{
			ExternalMount: "/var/nginx/html",
			InternalMount: "/usr/share/nginx/html",
			Mode:          "ro",
			Name:          "",
		}},
		Ports: []charm.ProcessPort{{
			External: 8080,
			Internal: 80,
			Endpoint: "",
		}},
		EnvVars: map[string]string{
			// TODO(erisnow) YAML coerces YES into true...
			"IMPORTANT": "true",
		},
	}})

	procs, err := st.List()
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, gc.HasLen, 0)

	info := process.Info{
		Process: charm.Process{
			Name:    "procA",
			Type:    "docker",
			Command: "do-something cool",
			Image:   "spam/eggs",
			Volumes: []charm.ProcessVolume{{
				ExternalMount: "/var/nginx/html",
				InternalMount: "/usr/share/nginx/html",
				Mode:          "ro",
				Name:          "",
			}},
			Ports: []charm.ProcessPort{{
				External: 8080,
				Internal: 80,
				Endpoint: "website",
			}},
			EnvVars: map[string]string{
				"IMPORTANT": "true",
			},
		},
		Status: process.Status{
			State:   process.StateRunning,
			Message: "okay",
		},
		Details: process.Details{
			ID: "xyz",
			Status: process.PluginStatus{
				State: "running",
			},
		},
	}
	err = st.Add(info)
	c.Assert(err, jc.ErrorIsNil)

	procs, err = st.List()
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, jc.DeepEquals, []process.Info{{
		Process: charm.Process{
			Name:    "procA",
			Type:    "docker",
			Command: "do-something cool",
			Image:   "spam/eggs",
			Volumes: []charm.ProcessVolume{{
				ExternalMount: "/var/nginx/html",
				InternalMount: "/usr/share/nginx/html",
				Mode:          "ro",
				Name:          "",
			}},
			Ports: []charm.ProcessPort{{
				External: 8080,
				Internal: 80,
				Endpoint: "website",
			}},
			EnvVars: map[string]string{
				"IMPORTANT": "true",
			},
		},
		Status: process.Status{
			State:   process.StateRunning,
			Message: "okay",
		},
		Details: process.Details{
			ID: "xyz",
			Status: process.PluginStatus{
				State: "running",
			},
		},
	}})

	procs, err = st.List("procA/xyz")
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, jc.DeepEquals, []process.Info{info})

	err = st.SetStatus(process.Info{
		Process: charm.Process{
			Name: "procA",
		},
		Status: process.Status{
			State:   process.StateRunning,
			Message: "still okay",
		},
		Details: process.Details{
			ID: "xyz",
			Status: process.PluginStatus{
				State: "still running",
			},
		},
	})
	c.Assert(err, jc.ErrorIsNil)

	procs, err = st.List("procA/xyz")
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, jc.DeepEquals, []process.Info{{
		Process: charm.Process{
			Name:    "procA",
			Type:    "docker",
			Command: "do-something cool",
			Image:   "spam/eggs",
			Volumes: []charm.ProcessVolume{{
				ExternalMount: "/var/nginx/html",
				InternalMount: "/usr/share/nginx/html",
				Mode:          "ro",
				Name:          "",
			}},
			Ports: []charm.ProcessPort{{
				External: 8080,
				Internal: 80,
				Endpoint: "website",
			}},
			EnvVars: map[string]string{
				"IMPORTANT": "true",
			},
		},
		Status: process.Status{
			State:   process.StateRunning,
			Message: "still okay",
		},
		Details: process.Details{
			ID: "xyz",
			Status: process.PluginStatus{
				State: "still running",
			},
		},
	}})

	err = st.Remove("procA/xyz")
	c.Assert(err, jc.ErrorIsNil)

	procs, err = st.List()
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, gc.HasLen, 0)
}
<commit_msg>Fix a broken test.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state_test

import (
	"github.com/juju/names"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
	"gopkg.in/juju/charm.v5"

	"github.com/juju/juju/component/all"
	"github.com/juju/juju/process"
	"github.com/juju/juju/state"
)

func init() {
	if err := all.RegisterForServer(); err != nil {
		panic(err)
	}
}

var _ = gc.Suite(&unitProcessesSuite{})

type unitProcessesSuite struct {
	ConnSuite
}

const metaYAML = `
name: a-charm
summary: a charm...
description: a charm...
processes:
  procA:
    type: docker
    command: do-something cool
    image: spam/eggs
    ports:
      - 8080:80
    volumes:
      - /var/nginx/html:/usr/share/nginx/html:ro
    env:
      IMPORTANT: YES
`

func (s *unitProcessesSuite) addUnit(c *gc.C, charmName, serviceName, meta string) (names.CharmTag, *state.Unit) {
	ch := s.AddTestingCharm(c, charmName)
	ch = s.AddMetaCharm(c, charmName, meta, 2)

	svc := s.AddTestingService(c, serviceName, ch)
	unit, err := svc.AddUnit()
	c.Assert(err, jc.ErrorIsNil)

	charmTag := ch.Tag().(names.CharmTag)
	return charmTag, unit
}

func (s *unitProcessesSuite) TestFunctional(c *gc.C) {
	_, unit := s.addUnit(c, "dummy", "a-service", metaYAML)

	st, err := s.State.UnitProcesses(unit)
	c.Assert(err, jc.ErrorIsNil)

	definitions, err := st.ListDefinitions()
	c.Assert(err, jc.ErrorIsNil)
	c.Check(definitions, jc.DeepEquals, []charm.Process{{
		Name:        "procA",
		Type:        "docker",
		TypeOptions: map[string]string{},
		Command:     "do-something cool",
		Image:       "spam/eggs",
		Volumes: []charm.ProcessVolume{{
			ExternalMount: "/var/nginx/html",
			InternalMount: "/usr/share/nginx/html",
			Mode:          "ro",
			Name:          "",
		}},
		Ports: []charm.ProcessPort{{
			External: 8080,
			Internal: 80,
			Endpoint: "",
		}},
		EnvVars: map[string]string{
			// TODO(erisnow) YAML coerces YES into true...
			"IMPORTANT": "true",
		},
	}})

	procs, err := st.List()
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, gc.HasLen, 0)

	info := process.Info{
		Process: charm.Process{
			Name:    "procA",
			Type:    "docker",
			Command: "do-something cool",
			Image:   "spam/eggs",
			Volumes: []charm.ProcessVolume{{
				ExternalMount: "/var/nginx/html",
				InternalMount: "/usr/share/nginx/html",
				Mode:          "ro",
				Name:          "",
			}},
			Ports: []charm.ProcessPort{{
				External: 8080,
				Internal: 80,
				Endpoint: "website",
			}},
			EnvVars: map[string]string{
				"IMPORTANT": "true",
			},
		},
		Status: process.Status{
			State:   process.StateRunning,
			Message: "okay",
		},
		Details: process.Details{
			ID: "xyz",
			Status: process.PluginStatus{
				State: "running",
			},
		},
	}
	err = st.Add(info)
	c.Assert(err, jc.ErrorIsNil)

	procs, err = st.List()
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, jc.DeepEquals, []process.Info{{
		Process: charm.Process{
			Name:    "procA",
			Type:    "docker",
			Command: "do-something cool",
			Image:   "spam/eggs",
			Volumes: []charm.ProcessVolume{{
				ExternalMount: "/var/nginx/html",
				InternalMount: "/usr/share/nginx/html",
				Mode:          "ro",
				Name:          "",
			}},
			Ports: []charm.ProcessPort{{
				External: 8080,
				Internal: 80,
				Endpoint: "website",
			}},
			EnvVars: map[string]string{
				"IMPORTANT": "true",
			},
		},
		Status: process.Status{
			State:   process.StateRunning,
			Message: "okay",
		},
		Details: process.Details{
			ID: "xyz",
			Status: process.PluginStatus{
				State: "running",
			},
		},
	}})

	procs, err = st.List("procA/xyz")
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, jc.DeepEquals, []process.Info{info})

	err = st.SetStatus("procA/xyz", process.CombinedStatus{
		Status: process.Status{
			State:   process.StateRunning,
			Message: "still okay",
		},
		PluginStatus: process.PluginStatus{
			State: "still running",
		},
	})
	c.Assert(err, jc.ErrorIsNil)

	procs, err = st.List("procA/xyz")
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, jc.DeepEquals, []process.Info{{
		Process: charm.Process{
			Name:    "procA",
			Type:    "docker",
			Command: "do-something cool",
			Image:   "spam/eggs",
			Volumes: []charm.ProcessVolume{{
				ExternalMount: "/var/nginx/html",
				InternalMount: "/usr/share/nginx/html",
				Mode:          "ro",
				Name:          "",
			}},
			Ports: []charm.ProcessPort{{
				External: 8080,
				Internal: 80,
				Endpoint: "website",
			}},
			EnvVars: map[string]string{
				"IMPORTANT": "true",
			},
		},
		Status: process.Status{
			State:   process.StateRunning,
			Message: "still okay",
		},
		Details: process.Details{
			ID: "xyz",
			Status: process.PluginStatus{
				State: "still running",
			},
		},
	}})

	err = st.Remove("procA/xyz")
	c.Assert(err, jc.ErrorIsNil)

	procs, err = st.List()
	c.Assert(err, jc.ErrorIsNil)
	c.Check(procs, gc.HasLen, 0)
}
<|endoftext|>


<commit_before>package vtctld

import (
	"reflect"
	"testing"

	"github.com/youtube/vitess/go/vt/discovery"
	"github.com/youtube/vitess/go/vt/proto/topodata"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

func TestStatsUpdate(t *testing.T) {
	tabletStatsCache := newTabletStatsCache()

	// Creating some tablets with their latest health information.
	tablet1Stats1 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 200)
	tablet1Stats2 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 200)
	tablet2Stats1 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 100)

	// Insert tablet1.
	tabletStatsCache.StatsUpdate(tablet1Stats1)
	results1 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	if got, want := results1[0], tablet1Stats1; !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Update tablet1.
	tabletStatsCache.StatsUpdate(tablet1Stats2)
	results2 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	if got, want := results2[0], tablet1Stats2; !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Insert tablet. List of tablets will be resorted.
	tabletStatsCache.StatsUpdate(tablet2Stats1)
	results3 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	if got, want := results3[0], tablet2Stats1; !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	results4 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	if got, want := results4[1], tablet1Stats2; !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Check tablet count in cell1.
	if got, want := tabletStatsCache.tabletCountsByCell["cell1"], 2; got != want {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Delete tablet2.
	tablet2Stats1.Up = false
	tabletStatsCache.StatsUpdate(tablet2Stats1)
	results5 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	for _, stat := range results5 {
		if reflect.DeepEqual(stat, tablet2Stats1) {
			t.Errorf("not deleleted from statusesByAliases")
		}
	}

	_, ok := tabletStatsCache.statusesByAlias[tablet2Stats1.Tablet.Alias.String()]
	if ok {
		t.Errorf("not deleted from statusesByAliases")
	}

	// Check tablet count in cell1.
	if got, want := tabletStatsCache.tabletCountsByCell["cell1"], 1; got != want {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Delete tablet1. List of known cells should be empty now.
	tablet1Stats2.Up = false
	tabletStatsCache.StatsUpdate(tablet1Stats2)
	_, ok = tabletStatsCache.tabletCountsByCell["cell1"]
	if ok {
		t.Errorf("not deleted from cells")
	}
}

func TestHeatmapData(t *testing.T) {
	// Creating and Sending updates to 12 tablets.
	ts1 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_MASTER, 100)
	ts2 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 200)
	ts3 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 300)
	ts4 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_RDONLY, 400)
	ts5 := tabletStats("cell2", "ks1", "-80", topodatapb.TabletType_MASTER, 500)
	ts6 := tabletStats("cell2", "ks1", "-80", topodatapb.TabletType_REPLICA, 600)
	ts7 := tabletStats("cell1", "ks1", "80-", topodatapb.TabletType_MASTER, 700)
	ts8 := tabletStats("cell1", "ks1", "80-", topodatapb.TabletType_REPLICA, 800)
	ts9 := tabletStats("cell1", "ks1", "80-", topodatapb.TabletType_RDONLY, 900)
	ts10 := tabletStats("cell1", "ks1", "80-", topodatapb.TabletType_RDONLY, 1000)
	ts11 := tabletStats("cell2", "ks1", "80-", topodatapb.TabletType_MASTER, 1100)
	ts12 := tabletStats("cell2", "ks1", "80-", topodatapb.TabletType_RDONLY, 1200)

	tabletStatsCache := newTabletStatsCache()

	tabletStatsCache.StatsUpdate(ts1)
	tabletStatsCache.StatsUpdate(ts2)
	tabletStatsCache.StatsUpdate(ts3)
	tabletStatsCache.StatsUpdate(ts4)
	tabletStatsCache.StatsUpdate(ts5)
	tabletStatsCache.StatsUpdate(ts6)
	tabletStatsCache.StatsUpdate(ts7)
	tabletStatsCache.StatsUpdate(ts8)
	tabletStatsCache.StatsUpdate(ts9)
	tabletStatsCache.StatsUpdate(ts10)
	tabletStatsCache.StatsUpdate(ts11)
	tabletStatsCache.StatsUpdate(ts12)

	// Checking that the heatmap data is returned correctly for the following view: (keyspace="ks1", cell=all, type=all).
	heatmap, err := tabletStatsCache.heatmapData("ks1", "all", "all", "lag")
	gotData, gotTabletAliases, gotLabels := heatmap.Data, heatmap.Aliases, heatmap.Labels
	wantData := [][]float64{
		{float64(ts1.Stats.SecondsBehindMaster), float64(ts7.Stats.SecondsBehindMaster)},
		{float64(ts2.Stats.SecondsBehindMaster), float64(ts8.Stats.SecondsBehindMaster)},
		{float64(ts3.Stats.SecondsBehindMaster), float64(-1)},
		{float64(ts4.Stats.SecondsBehindMaster), float64(ts9.Stats.SecondsBehindMaster)},
		{float64(-1), float64(ts10.Stats.SecondsBehindMaster)},
		{float64(ts5.Stats.SecondsBehindMaster), float64(ts11.Stats.SecondsBehindMaster)},
		{float64(ts6.Stats.SecondsBehindMaster), float64(-1)},
		{float64(-1), float64(ts12.Stats.SecondsBehindMaster)},
	}
	wantTabletAliases := [][]*topodata.TabletAlias{
		{ts1.Tablet.Alias, ts7.Tablet.Alias},
		{ts2.Tablet.Alias, ts8.Tablet.Alias},
		{ts3.Tablet.Alias, nil},
		{ts4.Tablet.Alias, ts9.Tablet.Alias},
		{nil, ts10.Tablet.Alias},
		{ts5.Tablet.Alias, ts11.Tablet.Alias},
		{ts6.Tablet.Alias, nil},
		{nil, ts12.Tablet.Alias},
	}
	wantLabels := []yLabel{
		{
			Label: label{Name: "cell1", Rowspan: 5},
			NestedLabels: []label{
				{Name: topodatapb.TabletType_MASTER.String(), Rowspan: 1},
				{Name: topodatapb.TabletType_REPLICA.String(), Rowspan: 2},
				{Name: topodatapb.TabletType_RDONLY.String(), Rowspan: 2},
			},
		},
		{
			Label: label{Name: "cell2", Rowspan: 3},
			NestedLabels: []label{
				{Name: topodatapb.TabletType_MASTER.String(), Rowspan: 1},
				{Name: topodatapb.TabletType_REPLICA.String(), Rowspan: 1},
				{Name: topodatapb.TabletType_RDONLY.String(), Rowspan: 1},
			},
		},
	}

	if err != nil {
		t.Errorf("couldn't get heatmap data: %v", err)
	}

	if !reflect.DeepEqual(gotData, wantData) {
		t.Errorf("got: %v, want: %v", gotData, wantData)
	}

	if !reflect.DeepEqual(gotTabletAliases, wantTabletAliases) {
		t.Errorf("got: %v, want: %v", gotData, wantData)
	}

	if !reflect.DeepEqual(gotLabels, wantLabels) {
		t.Errorf("got: %v, want: %v", gotLabels, wantLabels)
	}
}

func TestTabletStats(t *testing.T) {
	// Creating tabletStats.
	ts1 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_MASTER, 200)
	ts2 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 100)
	ts3 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 300)

	tabletStatsCache := newTabletStatsCache()
	tabletStatsCache.StatsUpdate(ts1)
	tabletStatsCache.StatsUpdate(ts2)

	// Test 1: tablet1 and tablet2 are updated with the stats received by the HealthCheck module.
	got1, err := tabletStatsCache.tabletStats(ts1.Tablet.Alias)
	want1 := *ts1
	if err != nil || !reflect.DeepEqual(got1, want1) {
		t.Errorf("got: %v, want: %v", got1, want1)
	}

	got2, err := tabletStatsCache.tabletStats(ts2.Tablet.Alias)
	want2 := *ts2
	if err != nil || !reflect.DeepEqual(got2, want2) {
		t.Errorf("got: %v, want: %v", got2, want2)
	}

	// Test 2: tablet3 isn't found in the map since no update was received for it.
	stats, got3 := tabletStatsCache.tabletStats(ts3.Tablet.Alias)
	wantErr := "could not find tablet: cell:\"cell1\" uid:300 "
	if got3.Error() != wantErr {
		t.Errorf("got: %v, want: %v", got3.Error(), wantErr)
	}
}

// tabletStats will create a discovery.TabletStats object.
func tabletStats(cell, keyspace, shard string, tabletType topodatapb.TabletType, uid uint32) *discovery.TabletStats {
	target := &querypb.Target{
		Keyspace:   keyspace,
		Shard:      shard,
		TabletType: tabletType,
	}
	tablet := &topodatapb.Tablet{
		Alias:    &topodatapb.TabletAlias{Cell: cell, Uid: uid},
		Keyspace: keyspace,
		Shard:    shard,
		Type:     tabletType,
		PortMap:  map[string]int32{"vt": int32(uid)},
	}
	realtimeStats := &querypb.RealtimeStats{
		HealthError: "",
		// uid is used for SecondsBehindMaster to give it a unique value.
		SecondsBehindMaster: uid,
	}
	stats := &discovery.TabletStats{
		Tablet:    tablet,
		Target:    target,
		Up:        true,
		Serving:   true,
		Stats:     realtimeStats,
		LastError: nil,
	}
	return stats
}
<commit_msg>Made style changes suggested by mberlin<commit_after>package vtctld

import (
	"reflect"
	"testing"

	"github.com/youtube/vitess/go/vt/discovery"
	"github.com/youtube/vitess/go/vt/proto/topodata"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

func TestStatsUpdate(t *testing.T) {
	tabletStatsCache := newTabletStatsCache()

	// Creating some tablets with their latest health information.
	tablet1Stats1 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 200)
	tablet1Stats2 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 200)
	tablet2Stats1 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 100)

	// Insert tablet1.
	tabletStatsCache.StatsUpdate(tablet1Stats1)
	results1 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	if got, want := results1[0], tablet1Stats1; !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Update tablet1.
	tabletStatsCache.StatsUpdate(tablet1Stats2)
	results2 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	if got, want := results2[0], tablet1Stats2; !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Insert tablet. List of tablets will be resorted.
	tabletStatsCache.StatsUpdate(tablet2Stats1)
	results3 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	if got, want := results3[0], tablet2Stats1; !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	results4 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	if got, want := results4[1], tablet1Stats2; !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Check tablet count in cell1.
	if got, want := tabletStatsCache.tabletCountsByCell["cell1"], 2; got != want {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Delete tablet2.
	tablet2Stats1.Up = false
	tabletStatsCache.StatsUpdate(tablet2Stats1)
	results5 := tabletStatsCache.statuses["ks1"]["-80"]["cell1"][topodatapb.TabletType_REPLICA]
	for _, stat := range results5 {
		if reflect.DeepEqual(stat, tablet2Stats1) {
			t.Errorf("not deleleted from statusesByAliases")
		}
	}

	_, ok := tabletStatsCache.statusesByAlias[tablet2Stats1.Tablet.Alias.String()]
	if ok {
		t.Errorf("not deleted from statusesByAliases")
	}

	// Check tablet count in cell1.
	if got, want := tabletStatsCache.tabletCountsByCell["cell1"], 1; got != want {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Delete tablet1. List of known cells should be empty now.
	tablet1Stats2.Up = false
	tabletStatsCache.StatsUpdate(tablet1Stats2)
	_, ok = tabletStatsCache.tabletCountsByCell["cell1"]
	if ok {
		t.Errorf("not deleted from cells")
	}
}

func TestHeatmapData(t *testing.T) {
	// Creating and Sending updates to 12 tablets.
	ts1 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_MASTER, 100)
	ts2 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 200)
	ts3 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 300)
	ts4 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_RDONLY, 400)
	ts5 := tabletStats("cell2", "ks1", "-80", topodatapb.TabletType_MASTER, 500)
	ts6 := tabletStats("cell2", "ks1", "-80", topodatapb.TabletType_REPLICA, 600)
	ts7 := tabletStats("cell1", "ks1", "80-", topodatapb.TabletType_MASTER, 700)
	ts8 := tabletStats("cell1", "ks1", "80-", topodatapb.TabletType_REPLICA, 800)
	ts9 := tabletStats("cell1", "ks1", "80-", topodatapb.TabletType_RDONLY, 900)
	ts10 := tabletStats("cell1", "ks1", "80-", topodatapb.TabletType_RDONLY, 1000)
	ts11 := tabletStats("cell2", "ks1", "80-", topodatapb.TabletType_MASTER, 1100)
	ts12 := tabletStats("cell2", "ks1", "80-", topodatapb.TabletType_RDONLY, 1200)

	tabletStatsCache := newTabletStatsCache()

	tabletStatsCache.StatsUpdate(ts1)
	tabletStatsCache.StatsUpdate(ts2)
	tabletStatsCache.StatsUpdate(ts3)
	tabletStatsCache.StatsUpdate(ts4)
	tabletStatsCache.StatsUpdate(ts5)
	tabletStatsCache.StatsUpdate(ts6)
	tabletStatsCache.StatsUpdate(ts7)
	tabletStatsCache.StatsUpdate(ts8)
	tabletStatsCache.StatsUpdate(ts9)
	tabletStatsCache.StatsUpdate(ts10)
	tabletStatsCache.StatsUpdate(ts11)
	tabletStatsCache.StatsUpdate(ts12)

	// Checking that the heatmap data is returned correctly for the following view: (keyspace="ks1", cell=all, type=all).
	heatmap, err := tabletStatsCache.heatmapData("ks1", "all", "all", "lag")
	gotData, gotTabletAliases, gotLabels := heatmap.Data, heatmap.Aliases, heatmap.Labels
	wantData := [][]float64{
		{float64(ts1.Stats.SecondsBehindMaster), float64(ts7.Stats.SecondsBehindMaster)},
		{float64(ts2.Stats.SecondsBehindMaster), float64(ts8.Stats.SecondsBehindMaster)},
		{float64(ts3.Stats.SecondsBehindMaster), float64(-1)},
		{float64(ts4.Stats.SecondsBehindMaster), float64(ts9.Stats.SecondsBehindMaster)},
		{float64(-1), float64(ts10.Stats.SecondsBehindMaster)},
		{float64(ts5.Stats.SecondsBehindMaster), float64(ts11.Stats.SecondsBehindMaster)},
		{float64(ts6.Stats.SecondsBehindMaster), float64(-1)},
		{float64(-1), float64(ts12.Stats.SecondsBehindMaster)},
	}
	wantTabletAliases := [][]*topodata.TabletAlias{
		{ts1.Tablet.Alias, ts7.Tablet.Alias},
		{ts2.Tablet.Alias, ts8.Tablet.Alias},
		{ts3.Tablet.Alias, nil},
		{ts4.Tablet.Alias, ts9.Tablet.Alias},
		{nil, ts10.Tablet.Alias},
		{ts5.Tablet.Alias, ts11.Tablet.Alias},
		{ts6.Tablet.Alias, nil},
		{nil, ts12.Tablet.Alias},
	}
	wantLabels := []yLabel{
		{
			Label: label{Name: "cell1", Rowspan: 5},
			NestedLabels: []label{
				{Name: topodatapb.TabletType_MASTER.String(), Rowspan: 1},
				{Name: topodatapb.TabletType_REPLICA.String(), Rowspan: 2},
				{Name: topodatapb.TabletType_RDONLY.String(), Rowspan: 2},
			},
		},
		{
			Label: label{Name: "cell2", Rowspan: 3},
			NestedLabels: []label{
				{Name: topodatapb.TabletType_MASTER.String(), Rowspan: 1},
				{Name: topodatapb.TabletType_REPLICA.String(), Rowspan: 1},
				{Name: topodatapb.TabletType_RDONLY.String(), Rowspan: 1},
			},
		},
	}

	if err != nil {
		t.Errorf("couldn't get heatmap data: %v", err)
	}

	if !reflect.DeepEqual(gotData, wantData) {
		t.Errorf("got: %v, want: %v", gotData, wantData)
	}

	if !reflect.DeepEqual(gotTabletAliases, wantTabletAliases) {
		t.Errorf("got: %v, want: %v", gotData, wantData)
	}

	if !reflect.DeepEqual(gotLabels, wantLabels) {
		t.Errorf("got: %v, want: %v", gotLabels, wantLabels)
	}
}

func TestTabletStats(t *testing.T) {
	// Creating tabletStats.
	ts1 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_MASTER, 200)
	ts2 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 100)
	ts3 := tabletStats("cell1", "ks1", "-80", topodatapb.TabletType_REPLICA, 300)

	tabletStatsCache := newTabletStatsCache()
	tabletStatsCache.StatsUpdate(ts1)
	tabletStatsCache.StatsUpdate(ts2)

	// Test 1: tablet1 and tablet2 are updated with the stats received by the HealthCheck module.
	got1, err := tabletStatsCache.tabletStats(ts1.Tablet.Alias)
	want1 := *ts1
	if err != nil || !reflect.DeepEqual(got1, want1) {
		t.Errorf("got: %v, want: %v", got1, want1)
	}

	got2, err := tabletStatsCache.tabletStats(ts2.Tablet.Alias)
	want2 := *ts2
	if err != nil || !reflect.DeepEqual(got2, want2) {
		t.Errorf("got: %v, want: %v", got2, want2)
	}

	// Test 2: tablet3 isn't found in the map since no update was received for it.
	_, gotErr := tabletStatsCache.tabletStats(ts3.Tablet.Alias)
	wantErr := "could not find tablet: cell:\"cell1\" uid:300 "
	if gotErr.Error() != wantErr {
		t.Errorf("got: %v, want: %v", gotErr.Error(), wantErr)
	}
}

// tabletStats will create a discovery.TabletStats object.
func tabletStats(cell, keyspace, shard string, tabletType topodatapb.TabletType, uid uint32) *discovery.TabletStats {
	target := &querypb.Target{
		Keyspace:   keyspace,
		Shard:      shard,
		TabletType: tabletType,
	}
	tablet := &topodatapb.Tablet{
		Alias:    &topodatapb.TabletAlias{Cell: cell, Uid: uid},
		Keyspace: keyspace,
		Shard:    shard,
		Type:     tabletType,
		PortMap:  map[string]int32{"vt": int32(uid)},
	}
	realtimeStats := &querypb.RealtimeStats{
		HealthError: "",
		// uid is used for SecondsBehindMaster to give it a unique value.
		SecondsBehindMaster: uid,
	}
	stats := &discovery.TabletStats{
		Tablet:    tablet,
		Target:    target,
		Up:        true,
		Serving:   true,
		Stats:     realtimeStats,
		LastError: nil,
	}
	return stats
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package vtexplain

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os/exec"
	"path"
	"strings"
	"testing"
)

var testOutputTempDir string

func defaultTestOpts() *Options {
	return &Options{
		ReplicationMode: "ROW",
		NumShards:       4,
		Normalize:       true,
		StrictDDL:       true,
	}
}

func initTest(mode string, opts *Options, t *testing.T) {
	schema, err := ioutil.ReadFile("testdata/test-schema.sql")
	if err != nil {
		t.Fatalf("error: %v", err)
	}

	vSchema, err := ioutil.ReadFile("testdata/test-vschema.json")
	if err != nil {
		t.Fatalf("error: %v", err)
	}

	opts.ExecutionMode = mode
	err = Init(string(vSchema), string(schema), opts)
	if err != nil {
		t.Fatalf("vtexplain Init error: %v", err)
	}
}

func testExplain(testcase string, opts *Options, t *testing.T) {
	modes := []string{
		ModeMulti,

		// TwoPC mode is functional, but the output isn't stable for
		// tests since there are timestamps in the value rows
		// ModeTwoPC,
	}

	for _, mode := range modes {
		runTestCase(testcase, mode, opts, t)
	}
}

func runTestCase(testcase, mode string, opts *Options, t *testing.T) {
	t.Logf("vtexplain test: %s mode: %s", testcase, mode)
	initTest(mode, opts, t)

	sqlFile := fmt.Sprintf("testdata/%s-queries.sql", testcase)
	sql, err := ioutil.ReadFile(sqlFile)
	if err != nil {
		t.Fatalf("vtexplain error: %v", err)
	}

	textOutFile := fmt.Sprintf("testdata/%s-output/%s-output.txt", mode, testcase)
	textOut, _ := ioutil.ReadFile(textOutFile)

	explains, err := Run(string(sql))
	if err != nil {
		t.Fatalf("vtexplain error: %v", err)
	}
	if explains == nil {
		t.Fatalf("vtexplain error running %s: no explain", string(sql))
	}

	explainText := ExplainsAsText(explains)
	if strings.TrimSpace(string(explainText)) != strings.TrimSpace(string(textOut)) {
		// Print the Text that was actually returned and also dump to a
		// temp file to be able to diff the results.
		t.Errorf("Text output did not match")

		if testOutputTempDir == "" {
			testOutputTempDir, err = ioutil.TempDir("", "vtexplain_output")
			if err != nil {
				t.Fatalf("error getting tempdir: %v", err)
			}
		}
		gotFile := fmt.Sprintf("%s/%s-output.txt", testOutputTempDir, testcase)
		ioutil.WriteFile(gotFile, []byte(explainText), 0644)

		command := exec.Command("diff", "-u", textOutFile, gotFile)
		out, _ := command.CombinedOutput()
		t.Logf("diff:\n%s\n", out)
		t.Logf("run the following command to update the expected output:")
		t.Logf("cp %s/* %s", testOutputTempDir, path.Dir(textOutFile))
	}
}

func TestUnsharded(t *testing.T) {
	testExplain("unsharded", defaultTestOpts(), t)
}

func TestSelectSharded(t *testing.T) {
	testExplain("selectsharded", defaultTestOpts(), t)
}

func TestInsertSharded(t *testing.T) {
	testExplain("insertsharded", defaultTestOpts(), t)
}

func TestUpdateSharded(t *testing.T) {
	testExplain("updatesharded", defaultTestOpts(), t)
}

func TestDeleteSharded(t *testing.T) {
	testExplain("deletesharded", defaultTestOpts(), t)
}

func TestOptions(t *testing.T) {
	opts := &Options{
		ReplicationMode: "STATEMENT",
		NumShards:       4,
		Normalize:       false,
	}

	testExplain("options", opts, t)
}
func TestTarget(t *testing.T) {
	opts := &Options{
		ReplicationMode: "ROW",
		NumShards:       4,
		Normalize:       false,
		Target:          "ks_sharded/40-80",
	}

	testExplain("target", opts, t)
}

func TestComments(t *testing.T) {
	testExplain("comments", defaultTestOpts(), t)
}

func TestErrors(t *testing.T) {
	initTest(ModeMulti, defaultTestOpts(), t)

	tests := []struct {
		SQL string
		Err string
	}{
		{
			SQL: "INVALID SQL",
			Err: "vtexplain execute error in 'INVALID SQL': unrecognized statement: INVALID SQL",
		},

		{
			SQL: "SELECT * FROM THIS IS NOT SQL",
			Err: "vtexplain execute error in 'SELECT * FROM THIS IS NOT SQL': syntax error at position 22 near 'is'",
		},

		{
			SQL: "SELECT * FROM table_not_in_vschema",
			Err: "vtexplain execute error in 'SELECT * FROM table_not_in_vschema': table table_not_in_vschema not found",
		},

		{
			SQL: "SELECT * FROM table_not_in_schema",
			Err: "target: ks_unsharded.-.master, used tablet: explainCell-0 (ks_unsharded/-): table table_not_in_schema not found in schema",
		},
	}

	for _, test := range tests {
		t.Run(test.SQL, func(t *testing.T) {
			_, err := Run(test.SQL)
			if err == nil || !strings.Contains(err.Error(), test.Err) {
				t.Errorf(">> Query: %s\n>> Got: %v\nWant: %s", test.SQL, err, test.Err)
			}
		})
	}
}

func TestJSONOutput(t *testing.T) {
	sql := "select 1 from user where id = 1"
	explains, err := Run(sql)
	if err != nil {
		t.Fatalf("vtexplain error: %v", err)
	}
	if explains == nil {
		t.Fatalf("vtexplain error running %s: no explain", string(sql))
	}

	explainJSON := ExplainsAsJSON(explains)

	var data interface{}
	err = json.Unmarshal([]byte(explainJSON), &data)
	if err != nil {
		t.Errorf("error unmarshaling json: %v", err)
	}

	array, ok := data.([]interface{})
	if !ok || len(array) != 1 {
		t.Errorf("expected single-element top-level array, got:\n%s", explainJSON)
	}

	explain, ok := array[0].(map[string]interface{})
	if !ok {
		t.Errorf("expected explain map, got:\n%s", explainJSON)
	}

	if explain["SQL"] != sql {
		t.Errorf("expected SQL, got:\n%s", explainJSON)
	}

	plans, ok := explain["Plans"].([]interface{})
	if !ok || len(plans) != 1 {
		t.Errorf("expected single-element plans array, got:\n%s", explainJSON)
	}

	actions, ok := explain["TabletActions"].(map[string]interface{})
	if !ok {
		t.Errorf("expected TabletActions map, got:\n%s", explainJSON)
	}

	actionsJSON, err := json.MarshalIndent(actions, "", "    ")
	if err != nil {
		t.Errorf("error in json marshal: %v", err)
	}
	wantJSON := `{
    "ks_sharded/-40": {
        "MysqlQueries": [
            {
                "SQL": "select 1 from user where id = 1 limit 10001",
                "Time": 1
            }
        ],
        "TabletQueries": [
            {
                "BindVars": {
                    "#maxLimit": "10001",
                    "vtg1": "1"
                },
                "SQL": "select :vtg1 from user where id = :vtg1",
                "Time": 1
            }
        ]
    }
}`
	if string(actionsJSON) != wantJSON {
		t.Errorf("TabletActions mismatch: got:\n%v\nwant:\n%v\n", string(actionsJSON), wantJSON)
	}
}
<commit_msg>Test cleaning<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package vtexplain

import (
	"encoding/json"
	"fmt"
	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/require"
	"io/ioutil"
	"path"
	"strings"
	"testing"
)

var testOutputTempDir string

func defaultTestOpts() *Options {
	return &Options{
		ReplicationMode: "ROW",
		NumShards:       4,
		Normalize:       true,
		StrictDDL:       true,
	}
}

func initTest(mode string, opts *Options, t *testing.T) {
	schema, err := ioutil.ReadFile("testdata/test-schema.sql")
	require.NoError(t, err)

	vSchema, err := ioutil.ReadFile("testdata/test-vschema.json")
	require.NoError(t, err)

	opts.ExecutionMode = mode
	err = Init(string(vSchema), string(schema), opts)
	require.NoError(t, err, "vtexplain Init error")
}

func testExplain(testcase string, opts *Options, t *testing.T) {
	modes := []string{
		ModeMulti,

		// TwoPC mode is functional, but the output isn't stable for
		// tests since there are timestamps in the value rows
		// ModeTwoPC,
	}

	for _, mode := range modes {
		runTestCase(testcase, mode, opts, t)
	}
}

func runTestCase(testcase, mode string, opts *Options, t *testing.T) {
	t.Run(testcase, func(t *testing.T) {
		initTest(mode, opts, t)

		sqlFile := fmt.Sprintf("testdata/%s-queries.sql", testcase)
		sql, err := ioutil.ReadFile(sqlFile)
		require.NoError(t, err, "vtexplain error")

		textOutFile := fmt.Sprintf("testdata/%s-output/%s-output.txt", mode, testcase)
		expected, _ := ioutil.ReadFile(textOutFile)

		explains, err := Run(string(sql))
		require.NoError(t, err, "vtexplain error")
		require.NotNil(t, explains, "vtexplain error running %s: no explain", string(sql))

		explainText := ExplainsAsText(explains)
		if diff := cmp.Diff(strings.TrimSpace(string(expected)), strings.TrimSpace(explainText)); diff != "" {
			// Print the Text that was actually returned and also dump to a
			// temp file to be able to diff the results.
			t.Errorf("Text output did not match (-want +got):\n%s", diff)

			if testOutputTempDir == "" {
				testOutputTempDir, err = ioutil.TempDir("", "vtexplain_output")
				require.NoError(t, err, "error getting tempdir")
			}
			gotFile := fmt.Sprintf("%s/%s-output.txt", testOutputTempDir, testcase)
			ioutil.WriteFile(gotFile, []byte(explainText), 0644)

			t.Logf("run the following command to update the expected output:")
			t.Logf("cp %s/* %s", testOutputTempDir, path.Dir(textOutFile))
		}
	})
}

func TestExplain(t *testing.T) {
	type test struct {
		name string
		opts *Options
	}
	tests := []test{
		{"unsharded", defaultTestOpts()},
		{"selectsharded", defaultTestOpts()},
		{"insertsharded", defaultTestOpts()},
		{"updatesharded", defaultTestOpts()},
		{"deletesharded", defaultTestOpts()},
		{"comments", defaultTestOpts()},
		{"options", &Options{
			ReplicationMode: "STATEMENT",
			NumShards:       4,
			Normalize:       false,
		}},
		{"target", &Options{
			ReplicationMode: "ROW",
			NumShards:       4,
			Normalize:       false,
			Target:          "ks_sharded/40-80",
		}},
	}

	for _, tst := range tests {
		testExplain(tst.name, tst.opts, t)
	}
}

func TestErrors(t *testing.T) {
	initTest(ModeMulti, defaultTestOpts(), t)

	tests := []struct {
		SQL string
		Err string
	}{
		{
			SQL: "INVALID SQL",
			Err: "vtexplain execute error in 'INVALID SQL': unrecognized statement: INVALID SQL",
		},

		{
			SQL: "SELECT * FROM THIS IS NOT SQL",
			Err: "vtexplain execute error in 'SELECT * FROM THIS IS NOT SQL': syntax error at position 22 near 'is'",
		},

		{
			SQL: "SELECT * FROM table_not_in_vschema",
			Err: "vtexplain execute error in 'SELECT * FROM table_not_in_vschema': table table_not_in_vschema not found",
		},

		{
			SQL: "SELECT * FROM table_not_in_schema",
			Err: "target: ks_unsharded.-.master, used tablet: explainCell-0 (ks_unsharded/-): table table_not_in_schema not found in schema",
		},
	}

	for _, test := range tests {
		t.Run(test.SQL, func(t *testing.T) {
			_, err := Run(test.SQL)
			require.Error(t, err)
			require.Contains(t, err.Error(), test.Err)
		})
	}
}

func TestJSONOutput(t *testing.T) {
	sql := "select 1 from user where id = 1"
	explains, err := Run(sql)
	require.NoError(t, err, "vtexplain error")
	require.NotNil(t, explains, "vtexplain error running %s: no explain", string(sql))

	explainJSON := ExplainsAsJSON(explains)

	var data interface{}
	err = json.Unmarshal([]byte(explainJSON), &data)
	require.NoError(t, err, "error unmarshaling json")

	array, ok := data.([]interface{})
	if !ok || len(array) != 1 {
		t.Errorf("expected single-element top-level array, got:\n%s", explainJSON)
	}

	explain, ok := array[0].(map[string]interface{})
	if !ok {
		t.Errorf("expected explain map, got:\n%s", explainJSON)
	}

	if explain["SQL"] != sql {
		t.Errorf("expected SQL, got:\n%s", explainJSON)
	}

	plans, ok := explain["Plans"].([]interface{})
	if !ok || len(plans) != 1 {
		t.Errorf("expected single-element plans array, got:\n%s", explainJSON)
	}

	actions, ok := explain["TabletActions"].(map[string]interface{})
	if !ok {
		t.Errorf("expected TabletActions map, got:\n%s", explainJSON)
	}

	actionsJSON, err := json.MarshalIndent(actions, "", "    ")
	require.NoError(t, err, "error in json marshal")
	wantJSON := `{
    "ks_sharded/-40": {
        "MysqlQueries": [
            {
                "SQL": "select 1 from user where id = 1 limit 10001",
                "Time": 1
            }
        ],
        "TabletQueries": [
            {
                "BindVars": {
                    "#maxLimit": "10001",
                    "vtg1": "1"
                },
                "SQL": "select :vtg1 from user where id = :vtg1",
                "Time": 1
            }
        ]
    }
}`
	require.Equal(t, wantJSON, string(actionsJSON))
}
<|endoftext|>