add profile

cmd/profile.go

@@ -0,0 +1,32 @@
+package cmd
+
+import (
+	"github.com/qiniu/goc/v2/pkg/client"
+	"github.com/qiniu/goc/v2/pkg/config"
+	"github.com/spf13/cobra"
+)
+
+var profileCmd = &cobra.Command{
+	Use:   "profile",
+	Short: "Get coverage profile from service registry center",
+	Long:  `Get code coverage profile for the services under test at runtime.`,
+	Example: `
+# Get coverage counter from default register center http://127.0.0.1:7777, the result output to stdout.
+goc profile
+# Get coverage counter from specified register center, the result output to specified file.
+goc profile --host=http://192.168.1.1:8080 --output=./coverage.cov
+`,
+	Run: profile,
+}
+
+var output string // --output flag
+
+func init() {
+	profileCmd.Flags().StringVar(&config.GocConfig.Host, "host", "127.0.0.1:7777", "specify the host of the goc server")
+	profileCmd.Flags().StringVarP(&output, "output", "o", "", "download cover profile")
+	rootCmd.AddCommand(profileCmd)
+}
+
+func profile(cmd *cobra.Command, args []string) {
+	client.NewWorker("http://" + config.GocConfig.Host).Profile(output)
+}

pkg/client/client.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"golang.org/x/term"
@@ -10,6 +11,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"path/filepath"
 
 	"github.com/olekukonko/tablewriter"
 	"github.com/qiniu/goc/v2/pkg/log"
@@ -18,11 +20,14 @@ import (
 // Action provides methods to contact with the covered agent under test
 type Action interface {
 	ListAgents(bool)
+	Profile(string)
 }
 
 const (
 	// CoverAgentsListAPI list all the registered agents
 	CoverAgentsListAPI = "/v2/rpcagents"
+	//CoverProfileAPI is provided by the covered service to get profiles
+	CoverProfileAPI = "/v2/cover/profile"
 )
 
 type client struct {
@@ -44,6 +49,10 @@ type gocCoveredAgent struct {
 	Pid      string `json:"pid"`
 }
 
+type gocProfile struct {
+	Profile string `json:"profile"`
+}
+
 // NewWorker creates a worker to contact with host
 func NewWorker(host string) Action {
 	_, err := url.ParseRequestURI(host)
@@ -102,6 +111,48 @@ func (c *client) ListAgents(wide bool) {
 	return
 }
 
+func (c *client) Profile(output string) {
+	u := fmt.Sprintf("%s%s", c.Host, CoverProfileAPI)
+
+	res, profile, err := c.do("GET", u, "application/json", nil)
+	if err != nil && isNetworkError(err) {
+		res, profile, err = c.do("GET", u, "application/json", nil)
+	}
+
+	if err == nil && res.StatusCode != 200 {
+		log.Fatalf(string(profile))
+	}
+	var profileText gocProfile
+	err = json.Unmarshal(profile, &profileText)
+	if err != nil {
+		log.Fatalf("profile unmarshal failed: %v", err)
+	}
+	if output == "" {
+		fmt.Fprint(os.Stdout, profileText.Profile)
+	} else {
+		var dir, filename string = filepath.Split(output)
+		if dir != "" {
+			err = os.MkdirAll(dir, os.ModePerm)
+			if err != nil {
+				log.Fatalf("failed to create directory %s, err:%v", dir, err)
+			}
+		}
+		if filename == "" {
+			output += "coverage.cov"
+		}
+
+		f, err := os.Create(output)
+		if err != nil {
+			log.Fatalf("failed to create file %s, err:%v", output, err)
+		}
+		defer f.Close()
+		_, err = io.Copy(f, bytes.NewReader([]byte(profileText.Profile)))
+		if err != nil {
+			log.Fatalf("failed to write file: %v, err: %v", output, err)
+		}
+	}
+}
+
 // getSimpleCmdLine
 func getSimpleCmdLine(preLen int, cmdLine string) string {
 	pathLen := len(cmdLine)

pkg/client/client_test.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -59,12 +60,53 @@ testID   1.1.1.1    testHost   0     ./testCmd -f testArgs
 		t.Run(name, func(t *testing.T) {
 			f := func() { c.ListAgents(tt.input) }
 			output := captureStdout(f)
-			fmt.Println(output)
 			assert.Equal(t, output, tt.expected)
 		})
 	}
 }
 
+func TestClientProfile(t *testing.T) {
+	mockAgents := `{"profile": "mode: count\nmockService/main.go:30.13,48.33 13 1\nb/b.go:30.13,48.33 13 1"}`
+	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(mockAgents))
+	}))
+	defer mockServer.Close()
+
+	c := NewWorker(mockServer.URL)
+	f := func() { c.Profile("") }
+	output := captureStdout(f)
+	assert.Regexp(t, "mockService/main.go:30.13,48.33 13 1", output)
+}
+
+func TestClientProfile_Output(t *testing.T) {
+	mockAgents := `{"profile": "mode: count\nmockService/main.go:30.13,48.33 13 1\nb/b.go:30.13,48.33 13 1"}`
+	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(mockAgents))
+	}))
+	defer mockServer.Close()
+
+	c := NewWorker(mockServer.URL)
+	ex, _ := os.Executable()
+	exPath := filepath.Dir(ex)
+	testCases := map[string]struct {
+		input  string
+		output string
+	}{
+		"file":           {"test.cov", "test.cov"},
+		"file with path": {filepath.Join(exPath, "test.txt"), filepath.Join(exPath, "test.txt")},
+		"just path":      {fmt.Sprintf("%s%c", exPath, os.PathSeparator), filepath.Join(exPath, "coverage.cov")},
+	}
+	for name, tt := range testCases {
+		t.Run(name, func(t *testing.T) {
+			c.Profile(tt.input)
+			defer os.RemoveAll(tt.output)
+			assert.FileExists(t, tt.output)
+		})
+	}
+}
+
 func TestClientDo(t *testing.T) {
 	c := &client{
 		client: http.DefaultClient,