feat(controller-sdk-go): change log api to websocket

Author: duanhongyi

apps/apps.go

@@ -5,11 +5,12 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io/ioutil"
-	"strconv"
+	"net/http"
+	"net/url"
 
 	drycc "github.com/drycc/controller-sdk-go"
 	"github.com/drycc/controller-sdk-go/api"
+	"github.com/gorilla/websocket"
 )
 
 // ErrNoLogs is returned when logs are missing from an app.
@@ -88,26 +89,30 @@ func Get(c *drycc.Client, appID string) (api.App, error) {
 
 // Logs retrieves logs from an app. The number of log lines fetched can be set by the lines
 // argument. Setting lines = -1 will retrieve all app logs.
-func Logs(c *drycc.Client, appID string, lines int) (string, error) {
-	u := fmt.Sprintf("/v2/apps/%s/logs", appID)
-
-	if lines > 0 {
-		u += "?log_lines=" + strconv.Itoa(lines)
-	}
-
-	res, reqErr := c.Request("GET", u, nil)
-	if reqErr != nil && !drycc.IsErrAPIMismatch(reqErr) {
-		return "", ErrNoLogs
-	}
-	defer res.Body.Close()
-
-	body, err := ioutil.ReadAll(res.Body)
-	if err != nil || len(body) < 3 {
-		return "", ErrNoLogs
+func Logs(c *drycc.Client, appID string, lines int, follow bool, timeout int) (*websocket.Conn, error) {
+	scheme := "ws"
+	if c.ControllerURL.Scheme == "https" {
+		scheme = "wss"
+	}
+	path := fmt.Sprintf("v2/apps/%s/logs", appID)
+	endpoint := url.URL{Scheme: scheme, Host: c.ControllerURL.Host, Path: path}
+	conn, _, err := websocket.DefaultDialer.Dial(
+		endpoint.String(),
+		http.Header{
+			"User-Agent":           {c.UserAgent},
+			"Authorization":        {"token " + c.Token},
+			"X-Drycc-Builder-Auth": {c.HooksToken},
+		},
+	)
+	if err != nil {
+		return nil, err
 	}
 
-	// We need to trim a few characters off the front and end of the string
-	return string(body), reqErr
+	conn.WriteMessage(
+		websocket.TextMessage,
+		[]byte(fmt.Sprintf(`{"lines": %d, "timeout": %d, "follow": %v}`, lines, timeout, follow)),
+	)
+	return conn, nil
 }
 
 // Run a one-time command in your app. This will start a kubernetes job with the

apps/apps_test.go

@@ -1,15 +1,18 @@
 package apps
 
 import (
+	"encoding/json"
 	"fmt"
 	"io/ioutil"
+	"log"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
 	"testing"
 
 	drycc "github.com/drycc/controller-sdk-go"
 	"github.com/drycc/controller-sdk-go/api"
+	"github.com/gorilla/websocket"
 )
 
 const appFixture string = `
@@ -43,11 +46,19 @@ const appCreateExpected string = `{"id":"example-go"}`
 const appRunExpected string = `{"command":"echo hi"}`
 const appTransferExpected string = `{"owner":"test"}`
 
+var upgrader = websocket.Upgrader{} // use default options
+
 type fakeHTTPServer struct {
 	createID        bool
 	createWithoutID bool
 }
 
+type fakeLogReqMessage struct {
+	Lines   int  `json:"lines"`
+	Timeout int  `json:"timeout"`
+	Follow  bool `json:"follow"`
+}
+
 func (f *fakeHTTPServer) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 	res.Header().Add("DRYCC_API_VERSION", drycc.APIVersion)
 
@@ -96,15 +107,23 @@ func (f *fakeHTTPServer) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 
 	// The entire log message is prefixed and suffixed with a few characters (not entirely sure why)
 	// We mimic those here
-	if req.URL.Path == "/v2/apps/example-go/logs" && req.URL.RawQuery == "" && req.Method == "GET" {
-		res.Write([]byte(`test\nfoo\nbar`))
-		return
-	}
-
-	// The entire log message is prefixed and suffixed with a few characters (not entirely sure why)
-	// We mimic those here
-	if req.URL.Path == "/v2/apps/example-go/logs" && req.URL.RawQuery == "log_lines=1" && req.Method == "GET" {
-		res.Write([]byte("test"))
+	if req.URL.Path == "/v2/apps/example-go/logs" {
+		c, err := upgrader.Upgrade(res, req, nil)
+		if err != nil {
+			log.Print("upgrade:", err)
+			return
+		}
+		defer c.Close()
+		_, message, err := c.ReadMessage()
+		if err != nil {
+			log.Print("upgrade:", err)
+			return
+		}
+		reqJSON := fakeLogReqMessage{}
+		json.Unmarshal([]byte(message), &reqJSON)
+		for i := 0; i < reqJSON.Lines; i++ {
+			c.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("test %d", i)))
+		}
 		return
 	}
 
@@ -296,25 +315,9 @@ func TestAppsList(t *testing.T) {
 	}
 }
 
-type testExpected struct {
-	Input    int
-	Expected string
-}
-
 func TestAppsLogs(t *testing.T) {
 	t.Parallel()
 
-	tests := []testExpected{
-		{
-			Input:    -1,
-			Expected: `test\nfoo\nbar`,
-		},
-		{
-			Input:    1,
-			Expected: "test",
-		},
-	}
-
 	handler := fakeHTTPServer{}
 	server := httptest.NewServer(&handler)
 	defer server.Close()
@@ -324,16 +327,16 @@ func TestAppsLogs(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	for _, test := range tests {
-		actual, err := Logs(drycc, "example-go", test.Input)
-
-		if err != nil {
-			t.Error(err)
-		}
-
-		if actual != test.Expected {
-			t.Errorf("Expected %s, Got %s", test.Expected, actual)
-		}
+	conn, err := Logs(drycc, "example-go", 1, false, 300)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, message, err := conn.ReadMessage()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(message) != "test 0" {
+		t.Errorf("Expected %s, Got %s", "test 0", message)
 	}
 }