<commit_before>package client_test

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	twilio "github.com/twilio/twilio-go/client"
	"github.com/twilio/twilio-go/framework/error"
)

func NewClient(accountSid string, authToken string) *twilio.Client {
	creds := &twilio.Credentials{
		AccountSID: accountSid,
		AuthToken:  authToken,
	}
	c := &twilio.Client{
		Credentials: creds,
		HTTPClient:  http.DefaultClient,
		Edge:        os.Getenv("TWILIO_EDGE"),
		Region:      os.Getenv("TWILIO_REGION"),
	}

	return c
}

func TestClient_SendRequestError(t *testing.T) {
	errorResponse := `{
	"status": 400,
	"code":20001,
	"message":"Bad request",
	"more_info":"https://www.twilio.com/docs/errors/20001"
}`
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(resp http.ResponseWriter, req *http.Request) {
			resp.WriteHeader(400)
			_, _ = resp.Write([]byte(errorResponse))
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	resp, err := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	twilioError := err.(*error.TwilioRestError)
	assert.Nil(t, resp)
	assert.Equal(t, 400, twilioError.Status)
	assert.Equal(t, 20001, twilioError.Code)
	assert.Equal(t, "https://www.twilio.com/docs/errors/20001", twilioError.MoreInfo)
	assert.Equal(t, "Bad request", twilioError.Message)
	assert.Nil(t, twilioError.Details)
}

func TestClient_SendRequestErrorWithDetails(t *testing.T) {
	errorResponse := []byte(`{
	"status": 400,
	"message": "Bad request",
	"code": 20001,
	"more_info": "https://www.twilio.com/docs/errors/20001",
	"details": {
		"foo": "bar"
	}
}`)
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(resp http.ResponseWriter, req *http.Request) {
			resp.WriteHeader(400)
			_, _ = resp.Write(errorResponse)
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	resp, err := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	twilioError := err.(*error.TwilioRestError)
	details := make(map[string]interface{})
	details["foo"] = "bar"
	assert.Nil(t, resp)
	assert.Equal(t, 400, twilioError.Status)
	assert.Equal(t, 20001, twilioError.Code)
	assert.Equal(t, "https://www.twilio.com/docs/errors/20001", twilioError.MoreInfo)
	assert.Equal(t, "Bad request", twilioError.Message)
	assert.Equal(t, details, twilioError.Details)
}

func TestClient_SendRequestWithRedirect(t *testing.T) {
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(writer http.ResponseWriter, request *http.Request) {
			writer.WriteHeader(307)
			_, _ = writer.Write([]byte(`{"redirect_to": "some_place"}`))
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	resp, _ := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	assert.Equal(t, 307, resp.StatusCode)
}

func TestClient_SetTimeoutTimesOut(t *testing.T) {
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(writer http.ResponseWriter, request *http.Request) {
			d := map[string]interface{}{
				"response": "ok",
			}
			time.Sleep(100 * time.Microsecond)
			encoder := json.NewEncoder(writer)
			err := encoder.Encode(&d)
			if err != nil {
				t.Error(err)
			}
			writer.WriteHeader(http.StatusOK)
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	client.SetTimeout(10 * time.Microsecond)
	_, err := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	assert.Error(t, err)
}

func TestClient_SetTimeoutSucceeds(t *testing.T) {
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(writer http.ResponseWriter, request *http.Request) {
			d := map[string]interface{}{
				"response": "ok",
			}
			time.Sleep(100 * time.Microsecond)
			encoder := json.NewEncoder(writer)
			err := encoder.Encode(&d)
			if err != nil {
				t.Error(err)
			}
			writer.WriteHeader(http.StatusOK)
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	client.SetTimeout(10 * time.Second)
	resp, err := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
}

//nolint:paralleltest
func TestClient_BuildHostSetRegion(t *testing.T) {
	// Region set via client
	client := NewClient("user", "pass")
	client.Region = "region"
	assert.Equal(t, "https://api.region.twilio.com", client.BuildHost("https://api.twilio.com"))
	assert.Equal(t, "https://api.region.twilio.com", client.BuildHost("https://api.urlRegion.twilio.com"))

	// Region set via env vars
	err := os.Setenv("TWILIO_REGION", "region")
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}

	client = NewClient("user", "pass")
	assert.Equal(t, "https://api.region.twilio.com", client.BuildHost("https://api.twilio.com"))

	err = os.Setenv("TWILIO_REGION", "")
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}

	// Region set via url
	assert.Equal(t, "https://api.region.twilio.com", client.BuildHost("https://api.region.twilio.com"))
}

//nolint:paralleltest
func TestClient_BuildHostSetEdgeDefaultRegion(t *testing.T) {
	// Edge set via client
	client := NewClient("user", "pass")
	client.Edge = "edge"
	assert.Equal(t, "https://api.edge.us1.twilio.com", client.BuildHost("https://api.twilio.com"))

	// Edge set via env vars
	err := os.Setenv("TWILIO_EDGE", "edge")
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}

	client = NewClient("user", "pass")
	assert.Equal(t, "https://api.edge.us1.twilio.com", client.BuildHost("https://api.twilio.com"))

	err = os.Setenv("TWILIO_EDGE", "")
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}
}

//nolint:paralleltest
func TestClient_BuildHostSetEdgeRegion(t *testing.T) {
	// Edge and Region set via client
	client := NewClient("user", "pass")
	client.Edge = "edge"
	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.region.twilio.com"))
	client.Region = "region"
	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.twilio.com"))
	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.urlEdge.urlRegion.twilio.com"))

	// Edge and Region set via env vars
	err := os.Setenv("TWILIO_EDGE", "edge")
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}

	err = os.Setenv("TWILIO_REGION", "region")
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}

	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.twilio.com"))
	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.urlEdge.urlRegion.twilio.com"))

	err = os.Setenv("TWILIO_REGION", "")
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}

	err = os.Setenv("TWILIO_EDGE", "")
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}

	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.edge.region.twilio.com"))
}
<commit_msg>tests: remove redundant env var tests (#66)<commit_after>package client_test

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	twilio "github.com/twilio/twilio-go/client"
	"github.com/twilio/twilio-go/framework/error"
)

func NewClient(accountSid string, authToken string) *twilio.Client {
	creds := &twilio.Credentials{
		AccountSID: accountSid,
		AuthToken:  authToken,
	}
	c := &twilio.Client{
		Credentials: creds,
		HTTPClient:  http.DefaultClient,
	}

	return c
}

func TestClient_SendRequestError(t *testing.T) {
	errorResponse := `{
	"status": 400,
	"code":20001,
	"message":"Bad request",
	"more_info":"https://www.twilio.com/docs/errors/20001"
}`
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(resp http.ResponseWriter, req *http.Request) {
			resp.WriteHeader(400)
			_, _ = resp.Write([]byte(errorResponse))
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	resp, err := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	twilioError := err.(*error.TwilioRestError)
	assert.Nil(t, resp)
	assert.Equal(t, 400, twilioError.Status)
	assert.Equal(t, 20001, twilioError.Code)
	assert.Equal(t, "https://www.twilio.com/docs/errors/20001", twilioError.MoreInfo)
	assert.Equal(t, "Bad request", twilioError.Message)
	assert.Nil(t, twilioError.Details)
}

func TestClient_SendRequestErrorWithDetails(t *testing.T) {
	errorResponse := []byte(`{
	"status": 400,
	"message": "Bad request",
	"code": 20001,
	"more_info": "https://www.twilio.com/docs/errors/20001",
	"details": {
		"foo": "bar"
	}
}`)
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(resp http.ResponseWriter, req *http.Request) {
			resp.WriteHeader(400)
			_, _ = resp.Write(errorResponse)
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	resp, err := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	twilioError := err.(*error.TwilioRestError)
	details := make(map[string]interface{})
	details["foo"] = "bar"
	assert.Nil(t, resp)
	assert.Equal(t, 400, twilioError.Status)
	assert.Equal(t, 20001, twilioError.Code)
	assert.Equal(t, "https://www.twilio.com/docs/errors/20001", twilioError.MoreInfo)
	assert.Equal(t, "Bad request", twilioError.Message)
	assert.Equal(t, details, twilioError.Details)
}

func TestClient_SendRequestWithRedirect(t *testing.T) {
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(writer http.ResponseWriter, request *http.Request) {
			writer.WriteHeader(307)
			_, _ = writer.Write([]byte(`{"redirect_to": "some_place"}`))
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	resp, _ := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	assert.Equal(t, 307, resp.StatusCode)
}

func TestClient_SetTimeoutTimesOut(t *testing.T) {
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(writer http.ResponseWriter, request *http.Request) {
			d := map[string]interface{}{
				"response": "ok",
			}
			time.Sleep(100 * time.Microsecond)
			encoder := json.NewEncoder(writer)
			err := encoder.Encode(&d)
			if err != nil {
				t.Error(err)
			}
			writer.WriteHeader(http.StatusOK)
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	client.SetTimeout(10 * time.Microsecond)
	_, err := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	assert.Error(t, err)
}

func TestClient_SetTimeoutSucceeds(t *testing.T) {
	mockServer := httptest.NewServer(http.HandlerFunc(
		func(writer http.ResponseWriter, request *http.Request) {
			d := map[string]interface{}{
				"response": "ok",
			}
			time.Sleep(100 * time.Microsecond)
			encoder := json.NewEncoder(writer)
			err := encoder.Encode(&d)
			if err != nil {
				t.Error(err)
			}
			writer.WriteHeader(http.StatusOK)
		}))
	defer mockServer.Close()

	client := NewClient("user", "pass")
	client.SetTimeout(10 * time.Second)
	resp, err := client.SendRequest("get", mockServer.URL, nil, nil, nil) //nolint:bodyclose
	assert.NoError(t, err)
	assert.Equal(t, 200, resp.StatusCode)
}

//nolint:paralleltest
func TestClient_BuildHostSetRegion(t *testing.T) {
	// Region set via url
	client := NewClient("user", "pass")
	assert.Equal(t, "https://api.region.twilio.com", client.BuildHost("https://api.region.twilio.com"))

	// Region set via client
	client.Region = "region"
	assert.Equal(t, "https://api.region.twilio.com", client.BuildHost("https://api.twilio.com"))
	assert.Equal(t, "https://api.region.twilio.com", client.BuildHost("https://api.urlRegion.twilio.com"))
}

//nolint:paralleltest
func TestClient_BuildHostSetEdgeDefaultRegion(t *testing.T) {
	// Edge set via client
	client := NewClient("user", "pass")
	client.Edge = "edge"
	assert.Equal(t, "https://api.edge.us1.twilio.com", client.BuildHost("https://api.twilio.com"))
}

//nolint:paralleltest
func TestClient_BuildHostSetEdgeRegion(t *testing.T) {
	//Edge and Region set via url
	client := NewClient("user", "pass")
	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.edge.region.twilio.com"))

	// Edge and Region set via client
	client.Edge = "edge"
	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.region.twilio.com"))
	client.Region = "region"
	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.twilio.com"))
	assert.Equal(t, "https://api.edge.region.twilio.com", client.BuildHost("https://api.urlEdge.urlRegion.twilio.com"))
}
<|endoftext|>


<commit_before>// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
// vespa test command
// Author: jonmv

package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/spf13/cobra"
	"github.com/vespa-engine/vespa/client/go/util"
	"github.com/vespa-engine/vespa/client/go/vespa"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

func init() {
	rootCmd.AddCommand(testCmd)
}

// TODO: add link to test doc at cloud.vespa.ai
var testCmd = &cobra.Command{
	Use:   "test [tests directory or test file]",
	Short: "Run a test suite, or a single test",
	Long: `Run a test suite, or a single test

Runs all JSON test files in the specified directory, or the single JSON
test file specified.

If no directory or file is specified, the working directory is used instead.`,
	Example: `$ vespa test src/test/application/tests/system-test
$ vespa test src/test/application/tests/system-test/feed-and-query.json`,
	Args:              cobra.MaximumNArgs(1),
	DisableAutoGenTag: true,
	Run: func(cmd *cobra.Command, args []string) {
		target := getTarget()
		testPath := "."
		if len(args) > 0 {
			testPath = args[0]
		}
		if count, failed := runTests(testPath, target); len(failed) != 0 {
			fmt.Fprintf(stdout, "\nFailed %d of %d tests:\n", len(failed), count)
			for _, test := range failed {
				fmt.Fprintln(stdout, test)
			}
			exitFunc(3)
		} else if count == 0 {
			fmt.Fprintf(stdout, "Failed to find any tests at %v\n", testPath)
			exitFunc(3)
		} else {
			fmt.Fprintf(stdout, "\n%d tests completed successfully\n", count)
		}
	},
}

func runTests(rootPath string, target vespa.Target) (int, []string) {
	count := 0
	failed := make([]string, 0)
	if stat, err := os.Stat(rootPath); err != nil {
		fatalErr(err, "Failed reading specified test path")
	} else if stat.IsDir() {
		tests, err := ioutil.ReadDir(rootPath) // TODO: Use os.ReadDir when >= 1.16 is required.
		if err != nil {
			fatalErr(err, "Failed reading specified test directory")
		}
		previousFailed := false
		for _, test := range tests {
			if !test.IsDir() && filepath.Ext(test.Name()) == ".json" {
				testPath := path.Join(rootPath, test.Name())
				if previousFailed {
					fmt.Fprintln(stdout, "")
					previousFailed = false
				}
				failure := runTest(testPath, target)
				if failure != "" {
					failed = append(failed, failure)
					previousFailed = true
				}
				count++
			}
		}
	} else if strings.HasSuffix(stat.Name(), ".json") {
		failure := runTest(rootPath, target)
		if failure != "" {
			failed = append(failed, failure)
		}
		count++
	}
	return count, failed
}

// Runs the test at the given path, and returns the specified test name if the test fails
func runTest(testPath string, target vespa.Target) string {
	var test test
	testBytes, err := ioutil.ReadFile(testPath)
	if err != nil {
		fatalErr(err, fmt.Sprintf("Failed to read test file at %s", testPath))
	}
	if err = json.Unmarshal(testBytes, &test); err != nil {
		fatalErr(err, fmt.Sprintf("Failed to parse test file at %s", testPath))
	}

	testName := test.Name
	if test.Name == "" {
		testName = testPath
	}
	fmt.Fprintf(stdout, "Running %s:", testName)

	defaultParameters, err := getParameters(test.Defaults.ParametersRaw, path.Dir(testPath))
	if err != nil {
		fatalErr(err, fmt.Sprintf("Invalid default parameters for %s", testName))
	}

	if len(test.Steps) == 0 {
		fatalErr(fmt.Errorf("a test must have at least one step, but none were found in %s", testPath))
	}
	for i, step := range test.Steps {
		stepName := step.Name
		if stepName == "" {
			stepName = fmt.Sprintf("step %d", i + 1)
		}
		failure, longFailure, err := verify(step, path.Dir(testPath), test.Defaults.Cluster, defaultParameters, target)
		if err != nil {
			fatalErr(err, fmt.Sprintf("Error in %s", stepName))
		}
		if failure != "" {
			fmt.Fprintf(stdout, " Failed %s:\n%s\n", stepName, longFailure)
			return fmt.Sprintf("%s: %s: %s", testName, stepName, failure)
		}
		if i == 0 {
			fmt.Fprintf(stdout, " ")
		}
		fmt.Fprint(stdout, ".")
	}
	fmt.Fprintln(stdout, " OK")
	return ""
}

// Asserts specified response is obtained for request, or returns a failure message, or an error if this fails
func verify(step step, testsPath string, defaultCluster string, defaultParameters map[string]string, target vespa.Target) (string, string, error) {
	requestBody, err := getBody(step.Request.BodyRaw, testsPath)
	if err != nil {
		return "", "", err
	}

	parameters, err := getParameters(step.Request.ParametersRaw, testsPath)
	if err != nil {
		return "", "", err
	}
	for name, value := range defaultParameters {
		if _, present := parameters[name]; !present {
			parameters[name] = value
		}
	}

	cluster := step.Request.Cluster
	if cluster == "" {
		cluster = defaultCluster
	}

	service, err := target.Service("query", 0, 0, cluster)
	if err != nil {
		return "", "", err
	}

	method := step.Request.Method
	if method == "" {
		method = "GET"
	}

	pathAndQuery := step.Request.URI
	if pathAndQuery == "" {
		pathAndQuery = "/search/"
	}
	requestUrl, err := url.ParseRequestURI(service.BaseURL + pathAndQuery)
	if err != nil {
		return "", "", err
	}
	query := requestUrl.Query()
	for name, value := range parameters {
		query.Add(name, value)
	}
	requestUrl.RawQuery = query.Encode()

	header := http.Header{}
	header.Add("Content-Type", "application/json") // TODO: Not guaranteed to be true ...

	request := &http.Request{
		URL:    requestUrl,
		Method: method,
		Header: header,
		Body:   ioutil.NopCloser(bytes.NewReader(requestBody)),
	}
	defer request.Body.Close()

	response, err := service.Do(request, 600*time.Second) // Vespa should provide a response within the given request timeout
	if err != nil {
		return "", "", err
	}
	defer response.Body.Close()

	statusCode := step.Response.Code
	if statusCode == 0 {
		statusCode = 200
	}
	if statusCode != response.StatusCode {
		failure := fmt.Sprintf("Unexpected status code: %d", response.StatusCode)
		return failure, fmt.Sprintf("%s\nExpected: %d\nActual response:\n%s", failure, statusCode, util.ReaderToJSON(response.Body)), nil
	}

	responseBodySpecBytes, err := getBody(step.Response.BodyRaw, testsPath)
	if err != nil {
		return "", "", err
	}
	if responseBodySpecBytes == nil {
		return "", "", nil
	}
	var responseBodySpec interface{}
	err = json.Unmarshal(responseBodySpecBytes, &responseBodySpec)
	if err != nil {
		return "", "", err
	}

	responseBodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", "", err
	}
	var responseBody interface{}
	err = json.Unmarshal(responseBodyBytes, &responseBody)
	if err != nil {
		return "", "", fmt.Errorf("got non-JSON response; %w:\n%s", err, string(responseBodyBytes))
	}

	failure, expected, err := compare(responseBodySpec, responseBody, "")
	if failure != "" {
		responsePretty, _ := json.MarshalIndent(responseBody, "", "  ")
		longFailure := failure
		if expected != "" {
			longFailure += "\n" + expected
		}
		longFailure += "\nActual response:\n" + string(responsePretty)
		return failure, longFailure, err
	}
	return "", "", err
}

func compare(expected interface{}, actual interface{}, path string) (string, string, error) {
	typeMatch := false
	valueMatch := false
	switch u := expected.(type) {
	case nil:
		typeMatch = actual == nil
		valueMatch = actual == nil
	case bool:
		v, ok := actual.(bool)
		typeMatch = ok
		valueMatch = ok && u == v
	case float64:
		v, ok := actual.(float64)
		typeMatch = ok
		valueMatch = ok && math.Abs(u-v) < 1e-9
	case string:
		v, ok := actual.(string)
		typeMatch = ok
		valueMatch = ok && (u == v)
	case []interface{}:
		v, ok := actual.([]interface{})
		typeMatch = ok
		if ok {
			if len(u) == len(v) {
				for i, e := range u {
					failure, expected, err := compare(e, v[i], fmt.Sprintf("%s/%d", path, i))
					if failure != "" || err != nil {
						return failure, expected, err
					}
				}
				valueMatch = true
			} else {
				return fmt.Sprintf("Unexpected number of elements at %s: %d", path, len(v)), fmt.Sprintf("Expected: %d", len(u)), nil
			}
		}
	case map[string]interface{}:
		v, ok := actual.(map[string]interface{})
		typeMatch = ok
		if ok {
			for n, e := range u {
				childPath := fmt.Sprintf("%s/%s", path, strings.ReplaceAll(strings.ReplaceAll(n, "~", "~0"), "/", "~1"))
				f, ok := v[n]
				if !ok {
					return fmt.Sprintf("Missing expected field at %s", childPath), "", nil
				}
				failure, expected, err := compare(e, f, childPath)
				if failure != "" || err != nil {
					return failure, expected, err
				}
			}
			valueMatch = true
		}
	default:
		return "", "", fmt.Errorf("unexpected expected JSON type for value '%v'", expected)
	}

	if !valueMatch {
		if path == "" {
			path = "root"
		}
		mismatched := "type"
		if typeMatch {
			mismatched = "value"
		}
		expectedJson, _ := json.Marshal(expected)
		actualJson, _ := json.Marshal(actual)
		return fmt.Sprintf("Unexpected %s at %s: %s", mismatched, path, actualJson), fmt.Sprintf("Expected: %s", expectedJson), nil
	}
	return "", "", nil
}

func getParameters(parametersRaw []byte, testsPath string) (map[string]string, error) {
	if parametersRaw != nil {
		var parametersPath string
		if err := json.Unmarshal(parametersRaw, &parametersPath); err == nil {
			resolvedParametersPath := path.Join(testsPath, parametersPath)
			parametersRaw, err = ioutil.ReadFile(resolvedParametersPath)
			if err != nil {
				fatalErr(err, fmt.Sprintf("Failed to read request parameters file at '%s'", resolvedParametersPath))
			}
		}
		var parameters map[string]string
		if err := json.Unmarshal(parametersRaw, &parameters); err != nil {
			return nil, fmt.Errorf("request parameters must be JSON with only string values: %w", err)
		}
		return parameters, nil
	}
	return make(map[string]string), nil
}

func getBody(bodyRaw []byte, testsPath string) ([]byte, error) {
	var bodyPath string
	if err := json.Unmarshal(bodyRaw, &bodyPath); err == nil {
		resolvedBodyPath := path.Join(testsPath, bodyPath)
		bodyRaw, err = ioutil.ReadFile(resolvedBodyPath)
		if err != nil {
			fatalErr(err, fmt.Sprintf("Failed to read body file at '%s'", resolvedBodyPath))
		}
	}
	return bodyRaw, nil
}

type test struct {
	Name       string      `json:"name"`
	Defaults defaults `json:"defaults"`
	Steps    []step   `json:"steps"`
}

type defaults struct {
	Cluster       string          `json:"cluster"`
	ParametersRaw json.RawMessage `json:"parameters"`
}

type step struct {
	Name     string   `json:"name"`
	Request  request  `json:"request"`
	Response response `json:"response"`
}

type request struct {
	Cluster       string          `json:"cluster"`
	Method        string          `json:"method"`
	URI           string          `json:"uri"`
	ParametersRaw json.RawMessage `json:"parameters"`
	BodyRaw       json.RawMessage `json:"body"`
}

type response struct {
	Code    int             `json:"code"`
	BodyRaw json.RawMessage `json:"body"`
}
<commit_msg>Fix go formatting<commit_after>// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
// vespa test command
// Author: jonmv

package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/spf13/cobra"
	"github.com/vespa-engine/vespa/client/go/util"
	"github.com/vespa-engine/vespa/client/go/vespa"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

func init() {
	rootCmd.AddCommand(testCmd)
}

// TODO: add link to test doc at cloud.vespa.ai
var testCmd = &cobra.Command{
	Use:   "test [tests directory or test file]",
	Short: "Run a test suite, or a single test",
	Long: `Run a test suite, or a single test

Runs all JSON test files in the specified directory, or the single JSON
test file specified.

If no directory or file is specified, the working directory is used instead.`,
	Example: `$ vespa test src/test/application/tests/system-test
$ vespa test src/test/application/tests/system-test/feed-and-query.json`,
	Args:              cobra.MaximumNArgs(1),
	DisableAutoGenTag: true,
	Run: func(cmd *cobra.Command, args []string) {
		target := getTarget()
		testPath := "."
		if len(args) > 0 {
			testPath = args[0]
		}
		if count, failed := runTests(testPath, target); len(failed) != 0 {
			fmt.Fprintf(stdout, "\nFailed %d of %d tests:\n", len(failed), count)
			for _, test := range failed {
				fmt.Fprintln(stdout, test)
			}
			exitFunc(3)
		} else if count == 0 {
			fmt.Fprintf(stdout, "Failed to find any tests at %v\n", testPath)
			exitFunc(3)
		} else {
			fmt.Fprintf(stdout, "\n%d tests completed successfully\n", count)
		}
	},
}

func runTests(rootPath string, target vespa.Target) (int, []string) {
	count := 0
	failed := make([]string, 0)
	if stat, err := os.Stat(rootPath); err != nil {
		fatalErr(err, "Failed reading specified test path")
	} else if stat.IsDir() {
		tests, err := ioutil.ReadDir(rootPath) // TODO: Use os.ReadDir when >= 1.16 is required.
		if err != nil {
			fatalErr(err, "Failed reading specified test directory")
		}
		previousFailed := false
		for _, test := range tests {
			if !test.IsDir() && filepath.Ext(test.Name()) == ".json" {
				testPath := path.Join(rootPath, test.Name())
				if previousFailed {
					fmt.Fprintln(stdout, "")
					previousFailed = false
				}
				failure := runTest(testPath, target)
				if failure != "" {
					failed = append(failed, failure)
					previousFailed = true
				}
				count++
			}
		}
	} else if strings.HasSuffix(stat.Name(), ".json") {
		failure := runTest(rootPath, target)
		if failure != "" {
			failed = append(failed, failure)
		}
		count++
	}
	return count, failed
}

// Runs the test at the given path, and returns the specified test name if the test fails
func runTest(testPath string, target vespa.Target) string {
	var test test
	testBytes, err := ioutil.ReadFile(testPath)
	if err != nil {
		fatalErr(err, fmt.Sprintf("Failed to read test file at %s", testPath))
	}
	if err = json.Unmarshal(testBytes, &test); err != nil {
		fatalErr(err, fmt.Sprintf("Failed to parse test file at %s", testPath))
	}

	testName := test.Name
	if test.Name == "" {
		testName = testPath
	}
	fmt.Fprintf(stdout, "Running %s:", testName)

	defaultParameters, err := getParameters(test.Defaults.ParametersRaw, path.Dir(testPath))
	if err != nil {
		fatalErr(err, fmt.Sprintf("Invalid default parameters for %s", testName))
	}

	if len(test.Steps) == 0 {
		fatalErr(fmt.Errorf("a test must have at least one step, but none were found in %s", testPath))
	}
	for i, step := range test.Steps {
		stepName := step.Name
		if stepName == "" {
			stepName = fmt.Sprintf("step %d", i+1)
		}
		failure, longFailure, err := verify(step, path.Dir(testPath), test.Defaults.Cluster, defaultParameters, target)
		if err != nil {
			fatalErr(err, fmt.Sprintf("Error in %s", stepName))
		}
		if failure != "" {
			fmt.Fprintf(stdout, " Failed %s:\n%s\n", stepName, longFailure)
			return fmt.Sprintf("%s: %s: %s", testName, stepName, failure)
		}
		if i == 0 {
			fmt.Fprintf(stdout, " ")
		}
		fmt.Fprint(stdout, ".")
	}
	fmt.Fprintln(stdout, " OK")
	return ""
}

// Asserts specified response is obtained for request, or returns a failure message, or an error if this fails
func verify(step step, testsPath string, defaultCluster string, defaultParameters map[string]string, target vespa.Target) (string, string, error) {
	requestBody, err := getBody(step.Request.BodyRaw, testsPath)
	if err != nil {
		return "", "", err
	}

	parameters, err := getParameters(step.Request.ParametersRaw, testsPath)
	if err != nil {
		return "", "", err
	}
	for name, value := range defaultParameters {
		if _, present := parameters[name]; !present {
			parameters[name] = value
		}
	}

	cluster := step.Request.Cluster
	if cluster == "" {
		cluster = defaultCluster
	}

	service, err := target.Service("query", 0, 0, cluster)
	if err != nil {
		return "", "", err
	}

	method := step.Request.Method
	if method == "" {
		method = "GET"
	}

	pathAndQuery := step.Request.URI
	if pathAndQuery == "" {
		pathAndQuery = "/search/"
	}
	requestUrl, err := url.ParseRequestURI(service.BaseURL + pathAndQuery)
	if err != nil {
		return "", "", err
	}
	query := requestUrl.Query()
	for name, value := range parameters {
		query.Add(name, value)
	}
	requestUrl.RawQuery = query.Encode()

	header := http.Header{}
	header.Add("Content-Type", "application/json") // TODO: Not guaranteed to be true ...

	request := &http.Request{
		URL:    requestUrl,
		Method: method,
		Header: header,
		Body:   ioutil.NopCloser(bytes.NewReader(requestBody)),
	}
	defer request.Body.Close()

	response, err := service.Do(request, 600*time.Second) // Vespa should provide a response within the given request timeout
	if err != nil {
		return "", "", err
	}
	defer response.Body.Close()

	statusCode := step.Response.Code
	if statusCode == 0 {
		statusCode = 200
	}
	if statusCode != response.StatusCode {
		failure := fmt.Sprintf("Unexpected status code: %d", response.StatusCode)
		return failure, fmt.Sprintf("%s\nExpected: %d\nActual response:\n%s", failure, statusCode, util.ReaderToJSON(response.Body)), nil
	}

	responseBodySpecBytes, err := getBody(step.Response.BodyRaw, testsPath)
	if err != nil {
		return "", "", err
	}
	if responseBodySpecBytes == nil {
		return "", "", nil
	}
	var responseBodySpec interface{}
	err = json.Unmarshal(responseBodySpecBytes, &responseBodySpec)
	if err != nil {
		return "", "", err
	}

	responseBodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", "", err
	}
	var responseBody interface{}
	err = json.Unmarshal(responseBodyBytes, &responseBody)
	if err != nil {
		return "", "", fmt.Errorf("got non-JSON response; %w:\n%s", err, string(responseBodyBytes))
	}

	failure, expected, err := compare(responseBodySpec, responseBody, "")
	if failure != "" {
		responsePretty, _ := json.MarshalIndent(responseBody, "", "  ")
		longFailure := failure
		if expected != "" {
			longFailure += "\n" + expected
		}
		longFailure += "\nActual response:\n" + string(responsePretty)
		return failure, longFailure, err
	}
	return "", "", err
}

func compare(expected interface{}, actual interface{}, path string) (string, string, error) {
	typeMatch := false
	valueMatch := false
	switch u := expected.(type) {
	case nil:
		typeMatch = actual == nil
		valueMatch = actual == nil
	case bool:
		v, ok := actual.(bool)
		typeMatch = ok
		valueMatch = ok && u == v
	case float64:
		v, ok := actual.(float64)
		typeMatch = ok
		valueMatch = ok && math.Abs(u-v) < 1e-9
	case string:
		v, ok := actual.(string)
		typeMatch = ok
		valueMatch = ok && (u == v)
	case []interface{}:
		v, ok := actual.([]interface{})
		typeMatch = ok
		if ok {
			if len(u) == len(v) {
				for i, e := range u {
					failure, expected, err := compare(e, v[i], fmt.Sprintf("%s/%d", path, i))
					if failure != "" || err != nil {
						return failure, expected, err
					}
				}
				valueMatch = true
			} else {
				return fmt.Sprintf("Unexpected number of elements at %s: %d", path, len(v)), fmt.Sprintf("Expected: %d", len(u)), nil
			}
		}
	case map[string]interface{}:
		v, ok := actual.(map[string]interface{})
		typeMatch = ok
		if ok {
			for n, e := range u {
				childPath := fmt.Sprintf("%s/%s", path, strings.ReplaceAll(strings.ReplaceAll(n, "~", "~0"), "/", "~1"))
				f, ok := v[n]
				if !ok {
					return fmt.Sprintf("Missing expected field at %s", childPath), "", nil
				}
				failure, expected, err := compare(e, f, childPath)
				if failure != "" || err != nil {
					return failure, expected, err
				}
			}
			valueMatch = true
		}
	default:
		return "", "", fmt.Errorf("unexpected expected JSON type for value '%v'", expected)
	}

	if !valueMatch {
		if path == "" {
			path = "root"
		}
		mismatched := "type"
		if typeMatch {
			mismatched = "value"
		}
		expectedJson, _ := json.Marshal(expected)
		actualJson, _ := json.Marshal(actual)
		return fmt.Sprintf("Unexpected %s at %s: %s", mismatched, path, actualJson), fmt.Sprintf("Expected: %s", expectedJson), nil
	}
	return "", "", nil
}

func getParameters(parametersRaw []byte, testsPath string) (map[string]string, error) {
	if parametersRaw != nil {
		var parametersPath string
		if err := json.Unmarshal(parametersRaw, &parametersPath); err == nil {
			resolvedParametersPath := path.Join(testsPath, parametersPath)
			parametersRaw, err = ioutil.ReadFile(resolvedParametersPath)
			if err != nil {
				fatalErr(err, fmt.Sprintf("Failed to read request parameters file at '%s'", resolvedParametersPath))
			}
		}
		var parameters map[string]string
		if err := json.Unmarshal(parametersRaw, &parameters); err != nil {
			return nil, fmt.Errorf("request parameters must be JSON with only string values: %w", err)
		}
		return parameters, nil
	}
	return make(map[string]string), nil
}

func getBody(bodyRaw []byte, testsPath string) ([]byte, error) {
	var bodyPath string
	if err := json.Unmarshal(bodyRaw, &bodyPath); err == nil {
		resolvedBodyPath := path.Join(testsPath, bodyPath)
		bodyRaw, err = ioutil.ReadFile(resolvedBodyPath)
		if err != nil {
			fatalErr(err, fmt.Sprintf("Failed to read body file at '%s'", resolvedBodyPath))
		}
	}
	return bodyRaw, nil
}

type test struct {
	Name     string   `json:"name"`
	Defaults defaults `json:"defaults"`
	Steps    []step   `json:"steps"`
}

type defaults struct {
	Cluster       string          `json:"cluster"`
	ParametersRaw json.RawMessage `json:"parameters"`
}

type step struct {
	Name     string   `json:"name"`
	Request  request  `json:"request"`
	Response response `json:"response"`
}

type request struct {
	Cluster       string          `json:"cluster"`
	Method        string          `json:"method"`
	URI           string          `json:"uri"`
	ParametersRaw json.RawMessage `json:"parameters"`
	BodyRaw       json.RawMessage `json:"body"`
}

type response struct {
	Code    int             `json:"code"`
	BodyRaw json.RawMessage `json:"body"`
}
<|endoftext|>


<commit_before>package testing

import (
	"bufio"
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

func GetDB() (*sql.DB, error) {
	refLock.Lock()
	defer refLock.Unlock()
	once.Do(func() {
		db, initErr = initDB()
		if initErr != nil {
			CleanUp()
		}
	})

	refCount++
	return db, initErr
}

func CleanUp() {
	refLock.Lock()
	defer refLock.Unlock()
	refCount--
	if refCount == 0 {
		for i := len(cleanUpActions) - 1; i >= 0; i-- {
			cleanUpActions[i]()
		}
		// Reset the once, incase we need to set the DB back up again
		once = new(sync.Once)
		db = nil
		initErr = nil
		cleanUpActions = nil
	}
}

var (
	db             *sql.DB
	initErr        error
	cleanUpActions []func()
	refCount                  = 0
	refLock                   = new(sync.Mutex)
	once           *sync.Once = new(sync.Once)
)

func initDB() (*sql.DB, error) {
	datadir, err := ioutil.TempDir("", "datadir")
	if err != nil {
		return nil, err
	}
	cleanUpActions = append(cleanUpActions, func() {
		os.RemoveAll(datadir)
	})

	socket, err := ioutil.TempFile("", "socket")
	if err != nil {
		return nil, err
	}
	cleanUpActions = append(cleanUpActions, func() {
		os.Remove(socket.Name())
	})

	pidFile, err := ioutil.TempFile("", "pidFile")
	if err != nil {
		return nil, err
	}
	cleanUpActions = append(cleanUpActions, func() {
		os.Remove(pidFile.Name())
	})

	cmd := exec.Command("mysqld",
		"--datadir", datadir,
		"--socket", socket.Name(),
		"--pid-file", pidFile.Name(),
		"--skip-grant-tables",
		"--skip-networking")
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, err
	}
	cleanUpActions = append(cleanUpActions, func() {
		stderr.Close()
	})
	ready := make(chan error)

	go func() {
		r := bufio.NewReader(stderr)
		defer close(ready)
		for {
			line, err := r.ReadString('\n')
			if err != nil {
				ready <- err
				return
			}
			if strings.Contains(line, "mysqld: ready for connections") {
				return
			}
		}
	}()

	if err := cmd.Start(); err != nil {
		return nil, err
	}
	cleanUpActions = append(cleanUpActions, func() {
		cmd.Process.Kill()
	})

	select {
	case err := <-ready:
		if err != nil {
			return nil, err
		}
	case <-time.After(5 * time.Second):
		fmt.Println("Got here3")
		return nil, fmt.Errorf("Failed to start server")
	}

	db, err := sql.Open("mysql", "unix("+socket.Name()+")/")
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(10)

	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS test;"); err != nil {
		return nil, err
	}

	// Close our connection, so we can reopen with the correct db name.  Other threads
	// will not use the correct database by default.
	if err := db.Close(); err != nil {
		return nil, err
	}

	db, err = sql.Open("mysql", "unix("+socket.Name()+")/test")
	if err != nil {
		return nil, err
	}

	if err := db.Ping(); err != nil {
		return nil, err
	}

	return db, nil
}
<commit_msg>Refactor Test db to get a new db for each call.  This makes test more hermetic, though likely slower<commit_after>package testing

import (
	"bufio"
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

type newDB struct {
	db  *sql.DB
	err error
}

var (
	timeout = 5 * time.Second

	done = make(chan struct{})
	dbs  = make(chan newDB)
)

func init() {
	go func() {
		initError := setupDB()
		for {
			dbs <- newDB{nil, initError}
		}
	}()
}

func GetDB() (*sql.DB, error) {
	db := <-dbs
	return db.db, db.err
}

func CleanUp() {
	close(done)
}

func setupDB() error {
	datadir, err := ioutil.TempDir("", "datadir")
	if err != nil {
		return err
	}
	defer os.RemoveAll(datadir)

	socket, err := ioutil.TempFile("", "socket")
	if err != nil {
		return err
	}
	defer os.Remove(socket.Name())

	pidFile, err := ioutil.TempFile("", "pidFile")
	if err != nil {
		return err
	}
	defer os.Remove(pidFile.Name())

	cmd := exec.Command("mysqld",
		"--datadir", datadir,
		"--socket", socket.Name(),
		"--pid-file", pidFile.Name(),
		"--skip-grant-tables",
		"--skip-networking",
	)
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	defer stderr.Close()

	if err := cmd.Start(); err != nil {
		return err
	}
	defer cmd.Process.Kill()

	ready := make(chan error)
	stderrlines := make(chan string, 20)
	go func() {
		defer close(ready)
		s := bufio.NewScanner(stderr)
		for s.Scan() {
			line := s.Text()
			select {
			case stderrlines <- line:
			default:
			}
			if strings.Contains(line, "mysqld: ready for connections") {
				return
			}
		}
		if err := s.Err(); err != nil {
			ready <- err
		}
	}()

	select {
	case err := <-ready:
		if err != nil {
			return err
		}
	case <-time.After(timeout):
		lines := make([]string, 0, cap(stderrlines))
		for i := 0; i < cap(stderrlines); i++ {
			select {
			case line := <-stderrlines:
				lines = append(lines, line)
			default:
				break
			}
		}
		return fmt.Errorf("Failed to start server after %v\n\n%s", timeout, strings.Join(lines, "\n"))
	}

	for i := 0; ; i++ {
		select {
		case <-done:
			return nil
		default:
			db, err := getDb(socket.Name(), i)
			if err != nil {
				return err
			}
			dbs <- newDB{db, nil}
		}
	}
}

func getDb(socketname string, id int) (*sql.DB, error) {
	db, err := sql.Open("mysql", "unix("+socketname+")/")
	if err != nil {
		return nil, err
	}

	dbName := fmt.Sprintf("testdb%d", id)

	if _, err := db.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s;", dbName)); err != nil {
		return nil, err
	}

	// Close our connection, so we can reopen with the correct db name.  Other threads
	// will not use the correct database by default.
	if err := db.Close(); err != nil {
		return nil, err
	}

	db, err = sql.Open("mysql", "unix("+socketname+")/"+dbName)
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(10)

	if err := db.Ping(); err != nil {
		return nil, err
	}

	return db, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/urfave/cli"
)

var (
	imagesFlags = []cli.Flag{
		cli.BoolFlag{
			Name:  "quiet, q",
			Usage: "display only container IDs",
		},
		cli.BoolFlag{
			Name:  "noheading, n",
			Usage: "do not print column headings",
		},
		cli.BoolFlag{
			Name:  "notruncate",
			Usage: "do not truncate output",
		},
	}
	imagesDescription = "Lists locally stored images."
	imagesCommand     = cli.Command{
		Name:        "images",
		Usage:       "List images in local storage",
		Description: imagesDescription,
		Flags:       imagesFlags,
		Action:      imagesCmd,
		ArgsUsage:   " ",
	}
)

func imagesCmd(c *cli.Context) error {
	store, err := getStore(c)
	if err != nil {
		return err
	}

	quiet := false
	if c.IsSet("quiet") {
		quiet = c.Bool("quiet")
	}
	noheading := false
	if c.IsSet("noheading") {
		noheading = c.Bool("noheading")
	}
	truncate := true
	if c.IsSet("notruncate") {
		truncate = !c.Bool("notruncate")
	}
	images, err := store.Images()
	if err != nil {
		return fmt.Errorf("error reading images: %v", err)
	}

	if len(images) > 0 && !noheading && !quiet {
		if truncate {
			fmt.Printf("%-12s %s\n", "IMAGE ID", "IMAGE NAME")
		} else {
			fmt.Printf("%-64s %s\n", "IMAGE ID", "IMAGE NAME")
		}
	}
	for _, image := range images {
		if quiet {
			fmt.Printf("%s\n", image.ID)
		} else {
			names := []string{""}
			if len(image.Names) > 0 {
				names = image.Names
			}
			for _, name := range names {
				if truncate {
					fmt.Printf("%-12.12s %s\n", image.ID, name)
				} else {
					fmt.Printf("%-64s %s\n", image.ID, name)
				}
			}
		}
	}

	return nil
}
<commit_msg>Fix cut and paste error in buildah images<commit_after>package main

import (
	"fmt"

	"github.com/urfave/cli"
)

var (
	imagesFlags = []cli.Flag{
		cli.BoolFlag{
			Name:  "quiet, q",
			Usage: "display only image IDs",
		},
		cli.BoolFlag{
			Name:  "noheading, n",
			Usage: "do not print column headings",
		},
		cli.BoolFlag{
			Name:  "notruncate",
			Usage: "do not truncate output",
		},
	}
	imagesDescription = "Lists locally stored images."
	imagesCommand     = cli.Command{
		Name:        "images",
		Usage:       "List images in local storage",
		Description: imagesDescription,
		Flags:       imagesFlags,
		Action:      imagesCmd,
		ArgsUsage:   " ",
	}
)

func imagesCmd(c *cli.Context) error {
	store, err := getStore(c)
	if err != nil {
		return err
	}

	quiet := false
	if c.IsSet("quiet") {
		quiet = c.Bool("quiet")
	}
	noheading := false
	if c.IsSet("noheading") {
		noheading = c.Bool("noheading")
	}
	truncate := true
	if c.IsSet("notruncate") {
		truncate = !c.Bool("notruncate")
	}
	images, err := store.Images()
	if err != nil {
		return fmt.Errorf("error reading images: %v", err)
	}

	if len(images) > 0 && !noheading && !quiet {
		if truncate {
			fmt.Printf("%-12s %s\n", "IMAGE ID", "IMAGE NAME")
		} else {
			fmt.Printf("%-64s %s\n", "IMAGE ID", "IMAGE NAME")
		}
	}
	for _, image := range images {
		if quiet {
			fmt.Printf("%s\n", image.ID)
		} else {
			names := []string{""}
			if len(image.Names) > 0 {
				names = image.Names
			}
			for _, name := range names {
				if truncate {
					fmt.Printf("%-12.12s %s\n", image.ID, name)
				} else {
					fmt.Printf("%-64s %s\n", image.ID, name)
				}
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"time"

	"code.cloudfoundry.org/bbs"
	"code.cloudfoundry.org/consuladapter"
	"code.cloudfoundry.org/converger"
	"code.cloudfoundry.org/converger/converger_process"
	"code.cloudfoundry.org/locket"
	"github.com/cloudfoundry-incubator/cf-debug-server"
	cf_lager "github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry/dropsonde"
	"github.com/nu7hatch/gouuid"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"
)

var consulCluster = flag.String(
	"consulCluster",
	"",
	"comma-separated list of consul server URLs (scheme://ip:port)",
)

var lockTTL = flag.Duration(
	"lockTTL",
	locket.LockTTL,
	"TTL for service lock",
)

var lockRetryInterval = flag.Duration(
	"lockRetryInterval",
	locket.RetryInterval,
	"interval to wait before retrying a failed lock acquisition",
)

var convergeRepeatInterval = flag.Duration(
	"convergeRepeatInterval",
	30*time.Second,
	"the interval between runs of the converge process",
)

var kickTaskDuration = flag.Duration(
	"kickTaskDuration",
	30*time.Second,
	"the interval, in seconds, between kicks to tasks",
)

var expireCompletedTaskDuration = flag.Duration(
	"expireCompletedTaskDuration",
	120*time.Second,
	"completed, unresolved tasks are deleted after this duration",
)

var expirePendingTaskDuration = flag.Duration(
	"expirePendingTaskDuration",
	30*time.Minute,
	"unclaimed tasks are marked as failed, after this duration",
)

var communicationTimeout = flag.Duration(
	"communicationTimeout",
	1*time.Minute,
	"Timeout applied to all HTTP requests.",
)

var dropsondePort = flag.Int(
	"dropsondePort",
	3457,
	"port the local metron agent is listening on",
)

var bbsAddress = flag.String(
	"bbsAddress",
	"",
	"Address to the BBS Server",
)

var bbsCACert = flag.String(
	"bbsCACert",
	"",
	"path to certificate authority cert used for mutually authenticated TLS BBS communication",
)

var bbsClientCert = flag.String(
	"bbsClientCert",
	"",
	"path to client cert used for mutually authenticated TLS BBS communication",
)

var bbsClientKey = flag.String(
	"bbsClientKey",
	"",
	"path to client key used for mutually authenticated TLS BBS communication",
)

var bbsClientSessionCacheSize = flag.Int(
	"bbsClientSessionCacheSize",
	0,
	"Capacity of the ClientSessionCache option on the TLS configuration. If zero, golang's default will be used",
)

var bbsMaxIdleConnsPerHost = flag.Int(
	"bbsMaxIdleConnsPerHost",
	0,
	"Controls the maximum number of idle (keep-alive) connctions per host. If zero, golang's default will be used",
)

const (
	dropsondeOrigin = "converger"
)

func main() {
	cf_debug_server.AddFlags(flag.CommandLine)
	cf_lager.AddFlags(flag.CommandLine)
	flag.Parse()

	cf_http.Initialize(*communicationTimeout)

	logger, reconfigurableSink := cf_lager.New("converger")

	if err := validateBBSAddress(); err != nil {
		logger.Fatal("invalid-bbs-address", err)
	}

	initializeDropsonde(logger)

	convergeClock := clock.NewClock()
	consulClient, err := consuladapter.NewClientFromUrl(*consulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	bbsServiceClient := bbs.NewServiceClient(consulClient, convergeClock)
	convergerServiceClient := converger.NewServiceClient(consulClient, convergeClock)

	lockMaintainer := convergerServiceClient.NewConvergerLockRunner(
		logger,
		generateGuid(logger),
		*lockRetryInterval,
		*lockTTL,
	)

	converger := converger_process.New(
		bbsServiceClient,
		initializeBBSClient(logger),
		logger,
		convergeClock,
		*convergeRepeatInterval,
		*kickTaskDuration,
		*expirePendingTaskDuration,
		*expireCompletedTaskDuration,
	)

	members := grouper.Members{
		{"lock-maintainer", lockMaintainer},
		{"converger", converger},
	}

	if dbgAddr := cf_debug_server.DebugAddress(flag.CommandLine); dbgAddr != "" {
		members = append(grouper.Members{
			{"debug-server", cf_debug_server.Runner(dbgAddr, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	process := ifrit.Invoke(sigmon.New(group))

	logger.Info("started")

	err = <-process.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func generateGuid(logger lager.Logger) string {
	uuid, err := uuid.NewV4()
	if err != nil {
		logger.Fatal("Couldn't generate uuid", err)
	}
	return uuid.String()
}

func initializeDropsonde(logger lager.Logger) {
	dropsondeDestination := fmt.Sprint("localhost:", *dropsondePort)
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func validateBBSAddress() error {
	if *bbsAddress == "" {
		return errors.New("bbsAddress is required")
	}
	return nil
}

func initializeBBSClient(logger lager.Logger) bbs.InternalClient {
	bbsURL, err := url.Parse(*bbsAddress)
	if err != nil {
		logger.Fatal("Invalid BBS URL", err)
	}

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(*bbsAddress)
	}

	bbsClient, err := bbs.NewSecureClient(*bbsAddress, *bbsCACert, *bbsClientCert, *bbsClientKey, *bbsClientSessionCacheSize, *bbsMaxIdleConnsPerHost)
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}
<commit_msg>Update and rename cf-debug-server -> debugserver<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"time"

	"code.cloudfoundry.org/bbs"
	"code.cloudfoundry.org/consuladapter"
	"code.cloudfoundry.org/converger"
	"code.cloudfoundry.org/converger/converger_process"
	"code.cloudfoundry.org/debugserver"
	"code.cloudfoundry.org/locket"
	cf_lager "github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry/dropsonde"
	"github.com/nu7hatch/gouuid"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"
)

var consulCluster = flag.String(
	"consulCluster",
	"",
	"comma-separated list of consul server URLs (scheme://ip:port)",
)

var lockTTL = flag.Duration(
	"lockTTL",
	locket.LockTTL,
	"TTL for service lock",
)

var lockRetryInterval = flag.Duration(
	"lockRetryInterval",
	locket.RetryInterval,
	"interval to wait before retrying a failed lock acquisition",
)

var convergeRepeatInterval = flag.Duration(
	"convergeRepeatInterval",
	30*time.Second,
	"the interval between runs of the converge process",
)

var kickTaskDuration = flag.Duration(
	"kickTaskDuration",
	30*time.Second,
	"the interval, in seconds, between kicks to tasks",
)

var expireCompletedTaskDuration = flag.Duration(
	"expireCompletedTaskDuration",
	120*time.Second,
	"completed, unresolved tasks are deleted after this duration",
)

var expirePendingTaskDuration = flag.Duration(
	"expirePendingTaskDuration",
	30*time.Minute,
	"unclaimed tasks are marked as failed, after this duration",
)

var communicationTimeout = flag.Duration(
	"communicationTimeout",
	1*time.Minute,
	"Timeout applied to all HTTP requests.",
)

var dropsondePort = flag.Int(
	"dropsondePort",
	3457,
	"port the local metron agent is listening on",
)

var bbsAddress = flag.String(
	"bbsAddress",
	"",
	"Address to the BBS Server",
)

var bbsCACert = flag.String(
	"bbsCACert",
	"",
	"path to certificate authority cert used for mutually authenticated TLS BBS communication",
)

var bbsClientCert = flag.String(
	"bbsClientCert",
	"",
	"path to client cert used for mutually authenticated TLS BBS communication",
)

var bbsClientKey = flag.String(
	"bbsClientKey",
	"",
	"path to client key used for mutually authenticated TLS BBS communication",
)

var bbsClientSessionCacheSize = flag.Int(
	"bbsClientSessionCacheSize",
	0,
	"Capacity of the ClientSessionCache option on the TLS configuration. If zero, golang's default will be used",
)

var bbsMaxIdleConnsPerHost = flag.Int(
	"bbsMaxIdleConnsPerHost",
	0,
	"Controls the maximum number of idle (keep-alive) connctions per host. If zero, golang's default will be used",
)

const (
	dropsondeOrigin = "converger"
)

func main() {
	debugserver.AddFlags(flag.CommandLine)
	cf_lager.AddFlags(flag.CommandLine)
	flag.Parse()

	cf_http.Initialize(*communicationTimeout)

	logger, reconfigurableSink := cf_lager.New("converger")

	if err := validateBBSAddress(); err != nil {
		logger.Fatal("invalid-bbs-address", err)
	}

	initializeDropsonde(logger)

	convergeClock := clock.NewClock()
	consulClient, err := consuladapter.NewClientFromUrl(*consulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	bbsServiceClient := bbs.NewServiceClient(consulClient, convergeClock)
	convergerServiceClient := converger.NewServiceClient(consulClient, convergeClock)

	lockMaintainer := convergerServiceClient.NewConvergerLockRunner(
		logger,
		generateGuid(logger),
		*lockRetryInterval,
		*lockTTL,
	)

	converger := converger_process.New(
		bbsServiceClient,
		initializeBBSClient(logger),
		logger,
		convergeClock,
		*convergeRepeatInterval,
		*kickTaskDuration,
		*expirePendingTaskDuration,
		*expireCompletedTaskDuration,
	)

	members := grouper.Members{
		{"lock-maintainer", lockMaintainer},
		{"converger", converger},
	}

	if dbgAddr := debugserver.DebugAddress(flag.CommandLine); dbgAddr != "" {
		members = append(grouper.Members{
			{"debug-server", debugserver.Runner(dbgAddr, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	process := ifrit.Invoke(sigmon.New(group))

	logger.Info("started")

	err = <-process.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func generateGuid(logger lager.Logger) string {
	uuid, err := uuid.NewV4()
	if err != nil {
		logger.Fatal("Couldn't generate uuid", err)
	}
	return uuid.String()
}

func initializeDropsonde(logger lager.Logger) {
	dropsondeDestination := fmt.Sprint("localhost:", *dropsondePort)
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func validateBBSAddress() error {
	if *bbsAddress == "" {
		return errors.New("bbsAddress is required")
	}
	return nil
}

func initializeBBSClient(logger lager.Logger) bbs.InternalClient {
	bbsURL, err := url.Parse(*bbsAddress)
	if err != nil {
		logger.Fatal("Invalid BBS URL", err)
	}

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(*bbsAddress)
	}

	bbsClient, err := bbs.NewSecureClient(*bbsAddress, *bbsCACert, *bbsClientCert, *bbsClientKey, *bbsClientSessionCacheSize, *bbsMaxIdleConnsPerHost)
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/briandowns/spinner"
	"github.com/urfave/cli"
)

type resultOutput struct {
	HostIdentifier string              `json:"host"`
	Rows           []map[string]string `json:"rows"`
}

func queryCommand() cli.Command {
	var (
		flHosts, flLabels, flQuery string
		flDebug, flQuiet, flExit   bool
		flTimeout                  time.Duration
	)
	return cli.Command{
		Name:      "query",
		Usage:     "Run a live query",
		UsageText: `fleetctl query [options]`,
		Flags: []cli.Flag{
			configFlag(),
			contextFlag(),
			cli.StringFlag{
				Name:        "hosts",
				EnvVar:      "HOSTS",
				Value:       "",
				Destination: &flHosts,
				Usage:       "Comma separated hostnames to target",
			},
			cli.StringFlag{
				Name:        "labels",
				EnvVar:      "LABELS",
				Value:       "",
				Destination: &flLabels,
				Usage:       "Comma separated label names to target",
			},
			cli.BoolFlag{
				Name:        "quiet",
				EnvVar:      "QUIET",
				Destination: &flQuiet,
				Usage:       "Only print results (no status information)",
			},
			cli.BoolFlag{
				Name:        "exit",
				EnvVar:      "EXIT",
				Destination: &flExit,
				Usage:       "Exit when 100% of online hosts have results returned",
			},
			cli.StringFlag{
				Name:        "query",
				EnvVar:      "QUERY",
				Value:       "",
				Destination: &flQuery,
				Usage:       "Query to run",
			},
			cli.BoolFlag{
				Name:        "debug",
				EnvVar:      "DEBUG",
				Destination: &flDebug,
				Usage:       "Whether or not to enable debug logging",
			},
			cli.DurationFlag{
				Name:        "timeout",
				EnvVar:      "TIMEOUT",
				Destination: &flTimeout,
				Usage:       "How long to run query before exiting (10s, 1h, etc.)",
			},
		},
		Action: func(c *cli.Context) error {
			fleet, err := clientFromCLI(c)
			if err != nil {
				return err
			}

			if flHosts == "" && flLabels == "" {
				return errors.New("No hosts or labels targeted")
			}

			if flQuery == "" {
				return errors.New("No query specified")
			}

			hosts := strings.Split(flHosts, ",")
			labels := strings.Split(flLabels, ",")

			res, err := fleet.LiveQuery(flQuery, labels, hosts)
			if err != nil {
				return err
			}

			tick := time.NewTicker(100 * time.Millisecond)
			defer tick.Stop()

			// See charsets at
			// https://godoc.org/github.com/briandowns/spinner#pkg-variables
			s := spinner.New(spinner.CharSets[24], 200*time.Millisecond)
			s.Writer = os.Stderr
			if !flQuiet {
				s.Start()
			}

			var timeoutChan <-chan time.Time
			if flTimeout > 0 {
				timeoutChan = time.After(flTimeout)
			} else {
				// Channel that never fires (so that we can
				// read from the channel in the below select
				// statement without panicking)
				timeoutChan = make(chan time.Time)
			}

			for {
				select {
				// Print a result
				case hostResult := <-res.Results():
					out := resultOutput{hostResult.Host.HostName, hostResult.Rows}
					s.Stop()
					if err := json.NewEncoder(os.Stdout).Encode(out); err != nil {
						fmt.Fprintf(os.Stderr, "Error writing output: %s\n", err)
					}
					s.Start()

				// Print an error
				case err := <-res.Errors():
					fmt.Fprintf(os.Stderr, "Error talking to server: %s\n", err.Error())

				// Update status message on interval
				case <-tick.C:
					status := res.Status()
					totals := res.Totals()
					var percentTotal, percentOnline float64
					var responded, total, online uint
					if status != nil && totals != nil {
						total = totals.Total
						online = totals.Online
						responded = status.ActualResults
						if total > 0 {
							percentTotal = 100 * float64(responded) / float64(total)
						}
						if online > 0 {
							percentOnline = 100 * float64(responded) / float64(online)
						}
					}

					if responded >= online && flExit {
						return nil
					}

					msg := fmt.Sprintf(" %.f%% responded (%.f%% online) | %d/%d targeted hosts (%d/%d online)", percentTotal, percentOnline, responded, total, responded, online)
					if !flQuiet {
						s.Suffix = msg
					}
					if total == responded {
						s.Stop()
						if !flQuiet {
							fmt.Fprintln(os.Stderr, msg)
						}
						return nil
					}

				// Check for timeout expiring
				case <-timeoutChan:
					s.Stop()
					if !flQuiet {
						fmt.Fprintln(os.Stderr, s.Suffix+"\nStopped by timeout")
					}
					return nil
				}
			}
		},
	}
}
<commit_msg>Require non-nil status for exit of fleetctl query (#2184)<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/briandowns/spinner"
	"github.com/urfave/cli"
)

type resultOutput struct {
	HostIdentifier string              `json:"host"`
	Rows           []map[string]string `json:"rows"`
}

func queryCommand() cli.Command {
	var (
		flHosts, flLabels, flQuery string
		flDebug, flQuiet, flExit   bool
		flTimeout                  time.Duration
	)
	return cli.Command{
		Name:      "query",
		Usage:     "Run a live query",
		UsageText: `fleetctl query [options]`,
		Flags: []cli.Flag{
			configFlag(),
			contextFlag(),
			cli.StringFlag{
				Name:        "hosts",
				EnvVar:      "HOSTS",
				Value:       "",
				Destination: &flHosts,
				Usage:       "Comma separated hostnames to target",
			},
			cli.StringFlag{
				Name:        "labels",
				EnvVar:      "LABELS",
				Value:       "",
				Destination: &flLabels,
				Usage:       "Comma separated label names to target",
			},
			cli.BoolFlag{
				Name:        "quiet",
				EnvVar:      "QUIET",
				Destination: &flQuiet,
				Usage:       "Only print results (no status information)",
			},
			cli.BoolFlag{
				Name:        "exit",
				EnvVar:      "EXIT",
				Destination: &flExit,
				Usage:       "Exit when 100% of online hosts have results returned",
			},
			cli.StringFlag{
				Name:        "query",
				EnvVar:      "QUERY",
				Value:       "",
				Destination: &flQuery,
				Usage:       "Query to run",
			},
			cli.BoolFlag{
				Name:        "debug",
				EnvVar:      "DEBUG",
				Destination: &flDebug,
				Usage:       "Whether or not to enable debug logging",
			},
			cli.DurationFlag{
				Name:        "timeout",
				EnvVar:      "TIMEOUT",
				Destination: &flTimeout,
				Usage:       "How long to run query before exiting (10s, 1h, etc.)",
			},
		},
		Action: func(c *cli.Context) error {
			fleet, err := clientFromCLI(c)
			if err != nil {
				return err
			}

			if flHosts == "" && flLabels == "" {
				return errors.New("No hosts or labels targeted")
			}

			if flQuery == "" {
				return errors.New("No query specified")
			}

			hosts := strings.Split(flHosts, ",")
			labels := strings.Split(flLabels, ",")

			res, err := fleet.LiveQuery(flQuery, labels, hosts)
			if err != nil {
				return err
			}

			tick := time.NewTicker(100 * time.Millisecond)
			defer tick.Stop()

			// See charsets at
			// https://godoc.org/github.com/briandowns/spinner#pkg-variables
			s := spinner.New(spinner.CharSets[24], 200*time.Millisecond)
			s.Writer = os.Stderr
			if !flQuiet {
				s.Start()
			}

			var timeoutChan <-chan time.Time
			if flTimeout > 0 {
				timeoutChan = time.After(flTimeout)
			} else {
				// Channel that never fires (so that we can
				// read from the channel in the below select
				// statement without panicking)
				timeoutChan = make(chan time.Time)
			}

			for {
				select {
				// Print a result
				case hostResult := <-res.Results():
					out := resultOutput{hostResult.Host.HostName, hostResult.Rows}
					s.Stop()
					if err := json.NewEncoder(os.Stdout).Encode(out); err != nil {
						fmt.Fprintf(os.Stderr, "Error writing output: %s\n", err)
					}
					s.Start()

				// Print an error
				case err := <-res.Errors():
					fmt.Fprintf(os.Stderr, "Error talking to server: %s\n", err.Error())

				// Update status message on interval
				case <-tick.C:
					status := res.Status()
					totals := res.Totals()
					var percentTotal, percentOnline float64
					var responded, total, online uint
					if status != nil && totals != nil {
						total = totals.Total
						online = totals.Online
						responded = status.ActualResults
						if total > 0 {
							percentTotal = 100 * float64(responded) / float64(total)
						}
						if online > 0 {
							percentOnline = 100 * float64(responded) / float64(online)
						}
					}

					if responded >= online && flExit {
						return nil
					}

					msg := fmt.Sprintf(" %.f%% responded (%.f%% online) | %d/%d targeted hosts (%d/%d online)", percentTotal, percentOnline, responded, total, responded, online)
					if !flQuiet {
						s.Suffix = msg
					}
					if total == responded && status != nil {
						s.Stop()
						if !flQuiet {
							fmt.Fprintln(os.Stderr, msg)
						}
						return nil
					}

				// Check for timeout expiring
				case <-timeoutChan:
					s.Stop()
					if !flQuiet {
						fmt.Fprintln(os.Stderr, s.Suffix+"\nStopped by timeout")
					}
					return nil
				}
			}
		},
	}
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/funkygao/gafka/config"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/sarama"
)

const (
	topInterval     = 5
	topIntervalTime = time.Second * topInterval
)

type Top struct {
	Ui  cli.Ui
	Cmd string

	mu           sync.Mutex
	limit        int
	topic        string
	counters     map[string]int // key is cluster:topic TODO int64
	lastCounters map[string]int
}

func (this *Top) Run(args []string) (exitCode int) {
	var zone string
	cmdFlags := flag.NewFlagSet("top", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.StringVar(&this.topic, "t", "", "")
	cmdFlags.IntVar(&this.limit, "n", 35, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).require("-z").invalid(args) {
		return 2
	}

	this.counters = make(map[string]int)
	this.lastCounters = make(map[string]int)

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, config.ZonePath(zone)))
	zkzone.WithinClusters(func(cluster string, path string) {
		zkcluster := zkzone.NewCluster(cluster)
		go this.clusterTop(zkcluster)
	})

	for {
		select {
		case <-time.After(topIntervalTime):
			c := exec.Command("clear")
			c.Stdout = os.Stdout
			c.Run()

			// header
			this.Ui.Output(fmt.Sprintf("%30s %50s %20s %10s",
				"cluster", "topic", "num", "mps"))
			this.Ui.Output(fmt.Sprintf(strings.Repeat("-", 113)))

			this.showAndResetCounters()
		}
	}

	return

}

func (this *Top) showAndResetCounters() {
	this.mu.Lock()
	defer this.mu.Unlock()

	counterFlip := make(map[int]string)
	sortedNum := make([]int, 0, len(this.counters))
	for ct, num := range this.counters {
		if this.topic != "" && !strings.HasSuffix(ct, ":"+this.topic) {
			continue
		}

		counterFlip[num] = ct
		if num > 100 { // TODO kill the magic number
			sortedNum = append(sortedNum, num)
		}
	}
	sort.Ints(sortedNum)

	for i := len(sortedNum) - 1; i >= 0; i-- {
		if len(sortedNum)-i > this.limit {
			break
		}

		num := sortedNum[i]
		p := strings.SplitN(counterFlip[num], ":", 2)
		mps := (num - this.lastCounters[counterFlip[num]]) / topInterval // msg per sec
		this.Ui.Output(fmt.Sprintf("%30s %50s %20s %10s", p[0], p[1],
			gofmt.Comma(int64(num)), gofmt.Comma(int64(mps))))
	}

	// record last counters and reset current counters
	for k, v := range this.counters {
		this.lastCounters[k] = v
	}
	this.counters = make(map[string]int)
}

func (this *Top) clusterTop(zkcluster *zk.ZkCluster) {
	cluster := zkcluster.Name()
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}

	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		return
	}
	defer kfk.Close()

	for {
		topics, err := kfk.Topics()
		if err != nil || len(topics) == 0 {
			return
		}

		for _, topic := range topics {
			if this.topic != "" && this.topic != topic {
				continue
			}

			msgs := int64(0)
			alivePartitions, err := kfk.WritablePartitions(topic)
			if err != nil {
				panic(err)
			}

			for _, partitionID := range alivePartitions {
				latestOffset, err := kfk.GetOffset(topic, partitionID,
					sarama.OffsetNewest)
				if err != nil {
					panic(err)
				}

				msgs += latestOffset
			}

			this.mu.Lock()
			this.counters[cluster+":"+topic] = int(msgs)
			this.mu.Unlock()
		}

		time.Sleep(time.Second)
		kfk.RefreshMetadata(topics...)
	}

}

func (*Top) Synopsis() string {
	return "Display top kafka cluster activities"
}

func (this *Top) Help() string {
	help := fmt.Sprintf(`
Usage: %s top [options]

	Display top kafka cluster activities

  -z zone

  -t topic

  -n limit
`, this.Cmd)
	return strings.TrimSpace(help)
}
<commit_msg>top subcommand should show catchall row<commit_after>package command

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/funkygao/gafka/config"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/sarama"
)

const (
	topInterval     = 5
	topIntervalTime = time.Second * topInterval
)

type Top struct {
	Ui  cli.Ui
	Cmd string

	mu           sync.Mutex
	limit        int
	topic        string
	counters     map[string]int // key is cluster:topic TODO int64
	lastCounters map[string]int
}

func (this *Top) Run(args []string) (exitCode int) {
	var zone string
	cmdFlags := flag.NewFlagSet("top", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.StringVar(&this.topic, "t", "", "")
	cmdFlags.IntVar(&this.limit, "n", 35, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).require("-z").invalid(args) {
		return 2
	}

	this.counters = make(map[string]int)
	this.lastCounters = make(map[string]int)

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, config.ZonePath(zone)))
	zkzone.WithinClusters(func(cluster string, path string) {
		zkcluster := zkzone.NewCluster(cluster)
		go this.clusterTop(zkcluster)
	})

	for {
		select {
		case <-time.After(topIntervalTime):
			c := exec.Command("clear")
			c.Stdout = os.Stdout
			c.Run()

			// header
			this.Ui.Output(fmt.Sprintf("%30s %50s %20s %10s",
				"cluster", "topic", "num", "mps"))
			this.Ui.Output(fmt.Sprintf(strings.Repeat("-", 113)))

			this.showAndResetCounters()
		}
	}

	return

}

func (this *Top) showAndResetCounters() {
	this.mu.Lock()
	defer this.mu.Unlock()

	counterFlip := make(map[int]string)
	sortedNum := make([]int, 0, len(this.counters))
	for ct, num := range this.counters {
		if this.topic != "" && !strings.HasSuffix(ct, ":"+this.topic) {
			continue
		}

		counterFlip[num] = ct
		if num > 100 { // TODO kill the magic number
			sortedNum = append(sortedNum, num)
		}
	}
	sort.Ints(sortedNum)

	othersNum := 0
	othersMps := 0
	limitReached := false
	for i := len(sortedNum) - 1; i >= 0; i-- {
		if !limitReached && len(sortedNum)-i > this.limit {
			limitReached = true
		}

		num := sortedNum[i]
		mps := (num - this.lastCounters[counterFlip[num]]) / topInterval // msg per sec
		if limitReached {
			othersNum += num
			othersMps += mps
		} else {
			clusterAndTopic := strings.SplitN(counterFlip[num], ":", 2)
			this.Ui.Output(fmt.Sprintf("%30s %50s %20s %10s",
				clusterAndTopic[0], clusterAndTopic[1],
				gofmt.Comma(int64(num)), gofmt.Comma(int64(mps))))
		}
	}

	if limitReached {
		// the catchall row
		this.Ui.Output(fmt.Sprintf("%30s %50s %20s %10s",
			"others", "others",
			gofmt.Comma(int64(othersNum)), gofmt.Comma(int64(othersMps))))
	}

	// record last counters and reset current counters
	for k, v := range this.counters {
		this.lastCounters[k] = v
	}
	this.counters = make(map[string]int)
}

func (this *Top) clusterTop(zkcluster *zk.ZkCluster) {
	cluster := zkcluster.Name()
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}

	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		return
	}
	defer kfk.Close()

	for {
		topics, err := kfk.Topics()
		if err != nil || len(topics) == 0 {
			return
		}

		for _, topic := range topics {
			if this.topic != "" && this.topic != topic {
				continue
			}

			msgs := int64(0)
			alivePartitions, err := kfk.WritablePartitions(topic)
			if err != nil {
				panic(err)
			}

			for _, partitionID := range alivePartitions {
				latestOffset, err := kfk.GetOffset(topic, partitionID,
					sarama.OffsetNewest)
				if err != nil {
					panic(err)
				}

				msgs += latestOffset
			}

			this.mu.Lock()
			this.counters[cluster+":"+topic] = int(msgs)
			this.mu.Unlock()
		}

		time.Sleep(time.Second)
		kfk.RefreshMetadata(topics...)
	}

}

func (*Top) Synopsis() string {
	return "Display top kafka cluster activities"
}

func (this *Top) Help() string {
	help := fmt.Sprintf(`
Usage: %s top [options]

	Display top kafka cluster activities

  -z zone

  -t topic

  -n limit
`, this.Cmd)
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"

	"cloud.google.com/go/storage"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"golang.org/x/sync/errgroup"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

const (
	uploadURL     = "https://golang.org/dl/upload"
	projectID     = "999119582588"
	storageBucket = "golang"
)

var publicACL = []storage.ACLRule{
	{Entity: storage.AllUsers, Role: storage.RoleReader},
	// If you don't give the owners access, the web UI seems to
	// have a bug and doesn't have access to see that it's public,
	// so won't render the "Shared Publicly" link. So we do that,
	// even though it's dumb and unnecessary otherwise:
	{Entity: storage.ACLEntity("project-owners-" + projectID), Role: storage.RoleOwner},
}

// File represents a file on the golang.org downloads page.
// It should be kept in sync with the download code in x/tools/godoc/dl.
type File struct {
	Filename       string
	OS             string
	Arch           string
	Version        string
	ChecksumSHA256 string
	Size           int64
	Kind           string // "archive", "installer", "source"
}

// fileRe matches the files created by the release tool, such as:
//   go1.5beta2.src.tar.gz
//   go1.5.1.linux-386.tar.gz
//   go1.5.windows-amd64.msi
var fileRe = regexp.MustCompile(`^(go[a-z0-9-.]+)\.(src|([a-z0-9]+)-([a-z0-9]+)(?:-([a-z0-9.]+))?)\.(tar\.gz|zip|pkg|msi)(.asc)?$`)

func upload(files []string) error {
	ctx := context.Background()
	c, err := storageClient(ctx)
	if err != nil {
		return err
	}
	defer c.Close()

	files, err = expandFiles(ctx, c, files)
	if err != nil {
		return err
	}

	files = chooseBestFiles(files)

	var sitePayloads []*File
	var uploaded []string
	for _, name := range files {
		base := filepath.Base(name)
		log.Printf("Uploading %v to GCS ...", base)
		m := fileRe.FindStringSubmatch(base)
		if m == nil {
			return fmt.Errorf("unrecognized file: %q", base)
		}

		checksum, size, err := uploadArtifact(ctx, c, name)
		if err != nil {
			return fmt.Errorf("uploading %q: %v", name, err)
		}

		uploaded = append(uploaded, base)

		if strings.HasSuffix(base, ".asc") {
			// Don't add asc files to the download page, just upload it.
			continue
		}

		// Upload file.sha256.
		fname := base + ".sha256"
		if err := putObject(ctx, c, fname, []byte(checksum)); err != nil {
			return fmt.Errorf("uploading %q: %v", base+".sha256", err)
		}
		uploaded = append(uploaded, fname)

		var kind string
		switch {
		case m[2] == "src":
			kind = "source"
		case strings.HasSuffix(base, ".tar.gz"), strings.HasSuffix(base, ".zip"):
			kind = "archive"
		case strings.HasSuffix(base, ".msi"), strings.HasSuffix(base, ".pkg"):
			kind = "installer"
		}
		f := &File{
			Filename:       base,
			Version:        m[1],
			OS:             m[3],
			Arch:           m[4],
			ChecksumSHA256: checksum,
			Size:           size,
			Kind:           kind,
		}
		sitePayloads = append(sitePayloads, f)
	}

	log.Println("Waiting for edge cache ...")
	if err := waitForEdgeCache(uploaded); err != nil {
		return fmt.Errorf("waitForEdgeCache(%+v): %v", uploaded, err)
	}

	log.Println("Uploading payloads to golang.org ...")
	for _, f := range sitePayloads {
		if err := updateSite(f); err != nil {
			return fmt.Errorf("updateSite(%+v): %v", f, err)
		}
	}
	return nil
}

func waitForEdgeCache(uploaded []string) error {
	if *uploadKick != "" {
		args := strings.Fields(*uploadKick)
		log.Printf("Running %v...", args)
		cmd := exec.Command(args[0], args[1:]...)
		cmd.Stdout = os.Stderr // Don't print to stdout.
		cmd.Stderr = os.Stderr
		// Don't wait for the command to finish.
		if err := cmd.Start(); err != nil {
			log.Printf("Couldn't start edge cache update command: %v", err)
		}
	}

	var g errgroup.Group
	for _, u := range uploaded {
		fname := u
		g.Go(func() error {
			// Add some jitter so that dozens of requests are not hitting the
			// endpoint at once.
			time.Sleep(time.Duration(rand.Intn(1000)) * time.Millisecond)
			t := time.Tick(5 * time.Second)
			var retries int
			for {
				url := "https://redirector.gvt1.com/edgedl/go/" + fname
				resp, err := http.Head(url)
				if err != nil {
					if retries < 3 {
						retries++
						<-t
						continue
					}
					return fmt.Errorf("http.Head(%q): %v", url, err)
				}
				retries = 0
				defer resp.Body.Close()
				if resp.StatusCode == http.StatusOK {
					log.Printf("%s is ready to go!", url)
					break
				}
				<-t
			}
			return nil
		})
	}
	return g.Wait()
}

func updateSite(f *File) error {
	// Post file details to golang.org.
	req, err := json.Marshal(f)
	if err != nil {
		return err
	}
	v := url.Values{"user": {*user}, "key": []string{userToken()}}
	u := fmt.Sprintf("%s?%s", uploadURL, v.Encode())
	resp, err := http.Post(u, "application/json", bytes.NewReader(req))
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := ioutil.ReadAll(resp.Body)
		return fmt.Errorf("upload failed: %v\n%s", resp.Status, b)
	}
	return nil
}

func putObject(ctx context.Context, c *storage.Client, name string, body []byte) error {
	wr := c.Bucket(storageBucket).Object(name).NewWriter(ctx)
	wr.ACL = publicACL
	if _, err := wr.Write(body); err != nil {
		return err
	}
	return wr.Close()
}

func storageClient(ctx context.Context) (*storage.Client, error) {
	file := filepath.Join(os.Getenv("HOME"), "keys", "golang-org.service.json")
	blob, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	config, err := google.JWTConfigFromJSON(blob, storage.ScopeReadWrite)
	if err != nil {
		return nil, err
	}
	return storage.NewClient(ctx, option.WithTokenSource(config.TokenSource(ctx)))
}

// expandFiles expands any "/..." paths in GCS URIs to include files in its subtree.
func expandFiles(ctx context.Context, storageClient *storage.Client, files []string) ([]string, error) {
	var expanded []string
	for _, f := range files {
		if !(strings.HasPrefix(f, "gs://") && strings.HasSuffix(f, "/...")) {
			expanded = append(expanded, f)
			continue
		}
		bucket, path := gcsParts(f)

		iter := storageClient.Bucket(bucket).Objects(ctx, &storage.Query{
			Prefix: strings.TrimSuffix(path, "..."), // Retain trailing "/" (if present).
		})
		for {
			attrs, err := iter.Next()
			if err == iterator.Done {
				break
			}
			if err != nil {
				return nil, err
			}
			if filepath.Ext(attrs.Name) == ".sha256" {
				// Ignore sha256 files.
				continue
			}
			expanded = append(expanded, fmt.Sprintf("gs://%s/%s", attrs.Bucket, attrs.Name))
		}
	}
	return expanded, nil
}

// gcsParts splits a GCS URI (e.g., "gs://bucket/path/to/object") into its bucket and path parts:
// ("bucket", "path/to/object")
//
// It assumes its input a well-formed GCS URI.
func gcsParts(uri string) (bucket, path string) {
	parts := strings.SplitN(strings.TrimPrefix(uri, "gs://"), "/", 2)
	return parts[0], parts[1]
}

func chooseBestFiles(files []string) []string {
	// map from basename to filepath/GCS URI.
	best := make(map[string]string)
	for _, f := range files {
		base := filepath.Base(f)
		if _, ok := best[base]; !ok {
			best[base] = f
			continue
		}

		// Overwrite existing only if the new entry is signed.
		if strings.HasPrefix(f, "gs://") && strings.Contains(f, "/signed/") {
			best[base] = f
		}
	}

	var out []string
	for _, path := range best {
		out = append(out, path)
	}
	sort.Strings(out) // for prettier printing.
	return out
}

func uploadArtifact(ctx context.Context, storageClient *storage.Client, path string) (checksum string, size int64, err error) {
	if strings.HasPrefix(path, "gs://") {
		return uploadArtifactGCS(ctx, storageClient, path)
	}
	return uploadArtifactLocal(ctx, storageClient, path)
}

func uploadArtifactGCS(ctx context.Context, storageClient *storage.Client, path string) (checksum string, size int64, err error) {
	bucket, path := gcsParts(path)
	base := filepath.Base(path)
	src := storageClient.Bucket(bucket).Object(path)
	dst := storageClient.Bucket(storageBucket).Object(base)

	r, err := storageClient.Bucket(bucket).Object(path + ".sha256").NewReader(ctx)
	if err != nil {
		return "", -1, fmt.Errorf("could not get sha256: %v", err)
	}
	checksumBytes, err := ioutil.ReadAll(r)
	if err != nil {
		return "", -1, fmt.Errorf("could not get sha256: %v", err)
	}
	copier := dst.CopierFrom(src)
	copier.ACL = publicACL
	attrs, err := copier.Run(ctx)
	if err != nil {
		return "", -1, err
	}
	return string(checksumBytes), attrs.Size, nil
}

func uploadArtifactLocal(ctx context.Context, storageClient *storage.Client, path string) (checksum string, size int64, err error) {
	base := filepath.Base(path)

	fileBytes, err := ioutil.ReadFile(path)
	if err != nil {
		return "", -1, fmt.Errorf("ioutil.ReadFile: %v", err)
	}
	// Upload file to Google Cloud Storage.
	if err := putObject(ctx, storageClient, base, fileBytes); err != nil {
		return "", -1, err
	}
	checksum = fmt.Sprintf("%x", sha256.Sum256(fileBytes))
	return checksum, int64(len(fileBytes)), nil
}
<commit_msg>cmd/release: use dl.google.com to check for binary availability<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"

	"cloud.google.com/go/storage"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"golang.org/x/sync/errgroup"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

const (
	uploadURL     = "https://golang.org/dl/upload"
	projectID     = "999119582588"
	storageBucket = "golang"
)

var publicACL = []storage.ACLRule{
	{Entity: storage.AllUsers, Role: storage.RoleReader},
	// If you don't give the owners access, the web UI seems to
	// have a bug and doesn't have access to see that it's public,
	// so won't render the "Shared Publicly" link. So we do that,
	// even though it's dumb and unnecessary otherwise:
	{Entity: storage.ACLEntity("project-owners-" + projectID), Role: storage.RoleOwner},
}

// File represents a file on the golang.org downloads page.
// It should be kept in sync with the download code in x/tools/godoc/dl.
type File struct {
	Filename       string
	OS             string
	Arch           string
	Version        string
	ChecksumSHA256 string
	Size           int64
	Kind           string // "archive", "installer", "source"
}

// fileRe matches the files created by the release tool, such as:
//   go1.5beta2.src.tar.gz
//   go1.5.1.linux-386.tar.gz
//   go1.5.windows-amd64.msi
var fileRe = regexp.MustCompile(`^(go[a-z0-9-.]+)\.(src|([a-z0-9]+)-([a-z0-9]+)(?:-([a-z0-9.]+))?)\.(tar\.gz|zip|pkg|msi)(.asc)?$`)

func upload(files []string) error {
	ctx := context.Background()
	c, err := storageClient(ctx)
	if err != nil {
		return err
	}
	defer c.Close()

	files, err = expandFiles(ctx, c, files)
	if err != nil {
		return err
	}

	files = chooseBestFiles(files)

	var sitePayloads []*File
	var uploaded []string
	for _, name := range files {
		base := filepath.Base(name)
		log.Printf("Uploading %v to GCS ...", base)
		m := fileRe.FindStringSubmatch(base)
		if m == nil {
			return fmt.Errorf("unrecognized file: %q", base)
		}

		checksum, size, err := uploadArtifact(ctx, c, name)
		if err != nil {
			return fmt.Errorf("uploading %q: %v", name, err)
		}

		uploaded = append(uploaded, base)

		if strings.HasSuffix(base, ".asc") {
			// Don't add asc files to the download page, just upload it.
			continue
		}

		// Upload file.sha256.
		fname := base + ".sha256"
		if err := putObject(ctx, c, fname, []byte(checksum)); err != nil {
			return fmt.Errorf("uploading %q: %v", base+".sha256", err)
		}
		uploaded = append(uploaded, fname)

		var kind string
		switch {
		case m[2] == "src":
			kind = "source"
		case strings.HasSuffix(base, ".tar.gz"), strings.HasSuffix(base, ".zip"):
			kind = "archive"
		case strings.HasSuffix(base, ".msi"), strings.HasSuffix(base, ".pkg"):
			kind = "installer"
		}
		f := &File{
			Filename:       base,
			Version:        m[1],
			OS:             m[3],
			Arch:           m[4],
			ChecksumSHA256: checksum,
			Size:           size,
			Kind:           kind,
		}
		sitePayloads = append(sitePayloads, f)
	}

	log.Println("Waiting for edge cache ...")
	if err := waitForEdgeCache(uploaded); err != nil {
		return fmt.Errorf("waitForEdgeCache(%+v): %v", uploaded, err)
	}

	log.Println("Uploading payloads to golang.org ...")
	for _, f := range sitePayloads {
		if err := updateSite(f); err != nil {
			return fmt.Errorf("updateSite(%+v): %v", f, err)
		}
	}
	return nil
}

func waitForEdgeCache(uploaded []string) error {
	if *uploadKick != "" {
		args := strings.Fields(*uploadKick)
		log.Printf("Running %v...", args)
		cmd := exec.Command(args[0], args[1:]...)
		cmd.Stdout = os.Stderr // Don't print to stdout.
		cmd.Stderr = os.Stderr
		// Don't wait for the command to finish.
		if err := cmd.Start(); err != nil {
			log.Printf("Couldn't start edge cache update command: %v", err)
		}
	}

	var g errgroup.Group
	for _, u := range uploaded {
		fname := u
		g.Go(func() error {
			// Add some jitter so that dozens of requests are not hitting the
			// endpoint at once.
			time.Sleep(time.Duration(rand.Intn(1000)) * time.Millisecond)
			t := time.Tick(5 * time.Second)
			var retries int
			for {
				url := "https://dl.google.com/go/" + fname
				resp, err := http.Head(url)
				if err != nil {
					if retries < 3 {
						retries++
						<-t
						continue
					}
					return fmt.Errorf("http.Head(%q): %v", url, err)
				}
				retries = 0
				defer resp.Body.Close()
				if resp.StatusCode == http.StatusOK {
					log.Printf("%s is ready to go!", url)
					break
				}
				<-t
			}
			return nil
		})
	}
	return g.Wait()
}

func updateSite(f *File) error {
	// Post file details to golang.org.
	req, err := json.Marshal(f)
	if err != nil {
		return err
	}
	v := url.Values{"user": {*user}, "key": []string{userToken()}}
	u := fmt.Sprintf("%s?%s", uploadURL, v.Encode())
	resp, err := http.Post(u, "application/json", bytes.NewReader(req))
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := ioutil.ReadAll(resp.Body)
		return fmt.Errorf("upload failed: %v\n%s", resp.Status, b)
	}
	return nil
}

func putObject(ctx context.Context, c *storage.Client, name string, body []byte) error {
	wr := c.Bucket(storageBucket).Object(name).NewWriter(ctx)
	wr.ACL = publicACL
	if _, err := wr.Write(body); err != nil {
		return err
	}
	return wr.Close()
}

func storageClient(ctx context.Context) (*storage.Client, error) {
	file := filepath.Join(os.Getenv("HOME"), "keys", "golang-org.service.json")
	blob, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	config, err := google.JWTConfigFromJSON(blob, storage.ScopeReadWrite)
	if err != nil {
		return nil, err
	}
	return storage.NewClient(ctx, option.WithTokenSource(config.TokenSource(ctx)))
}

// expandFiles expands any "/..." paths in GCS URIs to include files in its subtree.
func expandFiles(ctx context.Context, storageClient *storage.Client, files []string) ([]string, error) {
	var expanded []string
	for _, f := range files {
		if !(strings.HasPrefix(f, "gs://") && strings.HasSuffix(f, "/...")) {
			expanded = append(expanded, f)
			continue
		}
		bucket, path := gcsParts(f)

		iter := storageClient.Bucket(bucket).Objects(ctx, &storage.Query{
			Prefix: strings.TrimSuffix(path, "..."), // Retain trailing "/" (if present).
		})
		for {
			attrs, err := iter.Next()
			if err == iterator.Done {
				break
			}
			if err != nil {
				return nil, err
			}
			if filepath.Ext(attrs.Name) == ".sha256" {
				// Ignore sha256 files.
				continue
			}
			expanded = append(expanded, fmt.Sprintf("gs://%s/%s", attrs.Bucket, attrs.Name))
		}
	}
	return expanded, nil
}

// gcsParts splits a GCS URI (e.g., "gs://bucket/path/to/object") into its bucket and path parts:
// ("bucket", "path/to/object")
//
// It assumes its input a well-formed GCS URI.
func gcsParts(uri string) (bucket, path string) {
	parts := strings.SplitN(strings.TrimPrefix(uri, "gs://"), "/", 2)
	return parts[0], parts[1]
}

func chooseBestFiles(files []string) []string {
	// map from basename to filepath/GCS URI.
	best := make(map[string]string)
	for _, f := range files {
		base := filepath.Base(f)
		if _, ok := best[base]; !ok {
			best[base] = f
			continue
		}

		// Overwrite existing only if the new entry is signed.
		if strings.HasPrefix(f, "gs://") && strings.Contains(f, "/signed/") {
			best[base] = f
		}
	}

	var out []string
	for _, path := range best {
		out = append(out, path)
	}
	sort.Strings(out) // for prettier printing.
	return out
}

func uploadArtifact(ctx context.Context, storageClient *storage.Client, path string) (checksum string, size int64, err error) {
	if strings.HasPrefix(path, "gs://") {
		return uploadArtifactGCS(ctx, storageClient, path)
	}
	return uploadArtifactLocal(ctx, storageClient, path)
}

func uploadArtifactGCS(ctx context.Context, storageClient *storage.Client, path string) (checksum string, size int64, err error) {
	bucket, path := gcsParts(path)
	base := filepath.Base(path)
	src := storageClient.Bucket(bucket).Object(path)
	dst := storageClient.Bucket(storageBucket).Object(base)

	r, err := storageClient.Bucket(bucket).Object(path + ".sha256").NewReader(ctx)
	if err != nil {
		return "", -1, fmt.Errorf("could not get sha256: %v", err)
	}
	checksumBytes, err := ioutil.ReadAll(r)
	if err != nil {
		return "", -1, fmt.Errorf("could not get sha256: %v", err)
	}
	copier := dst.CopierFrom(src)
	copier.ACL = publicACL
	attrs, err := copier.Run(ctx)
	if err != nil {
		return "", -1, err
	}
	return string(checksumBytes), attrs.Size, nil
}

func uploadArtifactLocal(ctx context.Context, storageClient *storage.Client, path string) (checksum string, size int64, err error) {
	base := filepath.Base(path)

	fileBytes, err := ioutil.ReadFile(path)
	if err != nil {
		return "", -1, fmt.Errorf("ioutil.ReadFile: %v", err)
	}
	// Upload file to Google Cloud Storage.
	if err := putObject(ctx, storageClient, base, fileBytes); err != nil {
		return "", -1, err
	}
	checksum = fmt.Sprintf("%x", sha256.Sum256(fileBytes))
	return checksum, int64(len(fileBytes)), nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strings"
)

/*
Prototype:
$ semaphore create 4
$ semaphore add -- docker build ...
$ semaphore add -- docker build ...
...
$ semaphore wait | semaphore wait --notify --timeout 1h
... show progress (colored output)
[==>........] 2/10

command `docker build ...`
output:
 ...

command...
*/
func main() {
	ctx, cancel := context.WithCancel(context.Background())
	c := make(chan os.Signal, 1)

	signal.Notify(c, os.Interrupt)
	defer func() {
		signal.Stop(c)
		cancel()
	}()

	select {
	case <-c:
		cancel()
	case <-ctx.Done():
	}

	flag.Parse()
	fmt.Println(strings.Join(flag.Args(), ", "))
	fmt.Println(commit, date, version, os.TempDir())
}
<commit_msg>prevent errors on go1.5 and go1.6<commit_after>// +build go1.7

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strings"
)

/*
Prototype:
$ semaphore create 4
$ semaphore add -- docker build ...
$ semaphore add -- docker build ...
...
$ semaphore wait | semaphore wait --notify --timeout 1h
... show progress (colored output)
[==>........] 2/10

command `docker build ...`
output:
 ...

command...
*/
func main() {
	ctx, cancel := context.WithCancel(context.Background())
	c := make(chan os.Signal, 1)

	signal.Notify(c, os.Interrupt)
	defer func() {
		signal.Stop(c)
		cancel()
	}()

	select {
	case <-c:
		cancel()
	case <-ctx.Done():
	}

	flag.Parse()
	fmt.Println(strings.Join(flag.Args(), ", "))
	fmt.Println(commit, date, version, os.TempDir())
}
<|endoftext|>