add stream queryz unit test

go/vt/tabletserver/query_list_test.go

@@ -7,28 +7,33 @@ import (
 )
 
 type testConn struct {
-	id    int64
-	query string
+	id     int64
+	query  string
+	killed bool
 }
 
-func (tc testConn) Current() string { return tc.query }
+func (tc *testConn) Current() string { return tc.query }
 
-func (tc testConn) ID() int64 { return tc.id }
+func (tc *testConn) ID() int64 { return tc.id }
 
-func (tc testConn) Kill() {}
+func (tc *testConn) Kill() { tc.killed = true }
+
+func (tc *testConn) IsKilled() bool {
+	return tc.killed
+}
 
 func TestQueryList(t *testing.T) {
 	ql := NewQueryList()
 	connID := int64(1)
-	qd := NewQueryDetail(context.Background(), testConn{id: connID})
+	qd := NewQueryDetail(context.Background(), &testConn{id: connID})
 	ql.Add(qd)
 
 	if qd1, ok := ql.queryDetails[connID]; !ok || qd1.connID != connID {
 		t.Errorf("failed to add to QueryList")
 	}
 
 	conn2ID := int64(2)
-	qd2 := NewQueryDetail(context.Background(), testConn{id: conn2ID})
+	qd2 := NewQueryDetail(context.Background(), &testConn{id: conn2ID})
 	ql.Add(qd2)
 
 	rows := ql.GetQueryzRows()

go/vt/tabletserver/queryctl.go

@@ -382,6 +382,15 @@ func (rqsc *realQueryServiceControl) registerDebugHealthHandler() {
 	})
 }
 
+func (rqsc *realQueryServiceControl) registerStreamQueryzHandlers() {
+	http.HandleFunc("/streamqueryz", func(w http.ResponseWriter, r *http.Request) {
+		streamQueryzHandler(rqsc.sqlQueryRPCService.qe.streamQList, w, r)
+	})
+	http.HandleFunc("/streamqueryz/terminate", func(w http.ResponseWriter, r *http.Request) {
+		streamQueryzTerminateHandler(rqsc.sqlQueryRPCService.qe.streamQList, w, r)
+	})
+}
+
 func buildFmter(logger *streamlog.StreamLogger) func(url.Values, interface{}) string {
 	type formatter interface {
 		Format(url.Values) string

go/vt/tabletserver/stream_queryz.go

@@ -1,3 +1,7 @@
+// Copyright 2015, Google Inc. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
 package tabletserver
 
 import (
@@ -35,58 +39,55 @@ var (
 	`))
 )
 
-func (rqsc *realQueryServiceControl) registerStreamQueryzHandlers() {
-	streamqueryzHandler := func(w http.ResponseWriter, r *http.Request) {
-		if err := acl.CheckAccessHTTP(r, acl.DEBUGGING); err != nil {
-			acl.SendError(w, err)
-			return
-		}
-		rows := rqsc.sqlQueryRPCService.qe.streamQList.GetQueryzRows()
-		if err := r.ParseForm(); err != nil {
-			http.Error(w, fmt.Sprintf("cannot parse form: %s", err), http.StatusInternalServerError)
-			return
-		}
-		format := r.FormValue("format")
-		if format == "json" {
-			js, err := json.Marshal(rows)
-			if err != nil {
-				http.Error(w, err.Error(), http.StatusInternalServerError)
-				return
-			}
-			w.Header().Set("Content-Type", "application/json")
-			w.Write(js)
-			return
-		}
-		startHTMLTable(w)
-		defer endHTMLTable(w)
-		w.Write(streamqueryzHeader)
-		for i := range rows {
-			if err := streamqueryzTmpl.Execute(w, rows[i]); err != nil {
-				log.Errorf("streamlogz: couldn't execute template: %v", err)
-			}
-		}
+func streamQueryzHandler(queryList *QueryList, w http.ResponseWriter, r *http.Request) {
+	if err := acl.CheckAccessHTTP(r, acl.DEBUGGING); err != nil {
+		acl.SendError(w, err)
+		return
 	}
-
-	http.HandleFunc("/streamqueryz", streamqueryzHandler)
-	http.HandleFunc("/streamqueryz/terminate", func(w http.ResponseWriter, r *http.Request) {
-		if err := acl.CheckAccessHTTP(r, acl.ADMIN); err != nil {
-			acl.SendError(w, err)
-			return
-		}
-		if err := r.ParseForm(); err != nil {
-			http.Error(w, fmt.Sprintf("cannot parse form: %s", err), http.StatusInternalServerError)
-			return
-		}
-		connID := r.FormValue("connID")
-		c, err := strconv.Atoi(connID)
+	rows := queryList.GetQueryzRows()
+	if err := r.ParseForm(); err != nil {
+		http.Error(w, fmt.Sprintf("cannot parse form: %s", err), http.StatusInternalServerError)
+		return
+	}
+	format := r.FormValue("format")
+	if format == "json" {
+		js, err := json.Marshal(rows)
 		if err != nil {
-			http.Error(w, "invalid connID", http.StatusInternalServerError)
+			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		}
-		if err = rqsc.sqlQueryRPCService.qe.streamQList.Terminate(int64(c)); err != nil {
-			http.Error(w, fmt.Sprintf("error: %v", err), http.StatusInternalServerError)
-			return
+		w.Header().Set("Content-Type", "application/json")
+		w.Write(js)
+		return
+	}
+	startHTMLTable(w)
+	defer endHTMLTable(w)
+	w.Write(streamqueryzHeader)
+	for i := range rows {
+		if err := streamqueryzTmpl.Execute(w, rows[i]); err != nil {
+			log.Errorf("streamlogz: couldn't execute template: %v", err)
 		}
-		streamqueryzHandler(w, r)
-	})
+	}
+}
+
+func streamQueryzTerminateHandler(queryList *QueryList, w http.ResponseWriter, r *http.Request) {
+	if err := acl.CheckAccessHTTP(r, acl.ADMIN); err != nil {
+		acl.SendError(w, err)
+		return
+	}
+	if err := r.ParseForm(); err != nil {
+		http.Error(w, fmt.Sprintf("cannot parse form: %s", err), http.StatusInternalServerError)
+		return
+	}
+	connID := r.FormValue("connID")
+	c, err := strconv.Atoi(connID)
+	if err != nil {
+		http.Error(w, "invalid connID", http.StatusInternalServerError)
+		return
+	}
+	if err = queryList.Terminate(int64(c)); err != nil {
+		http.Error(w, fmt.Sprintf("error: %v", err), http.StatusInternalServerError)
+		return
+	}
+	streamQueryzHandler(queryList, w, r)
 }

go/vt/tabletserver/stream_queryz_test.go

@@ -0,0 +1,95 @@
+// Copyright 2015, Google Inc. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package tabletserver
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"golang.org/x/net/context"
+)
+
+func TestStreamQueryzHandlerJSON(t *testing.T) {
+	resp := httptest.NewRecorder()
+	req, _ := http.NewRequest("GET", "/streamqueryz?format=json", nil)
+
+	queryList := NewQueryList()
+	queryList.Add(NewQueryDetail(context.Background(), &testConn{id: 1}))
+	queryList.Add(NewQueryDetail(context.Background(), &testConn{id: 2}))
+
+	streamQueryzHandler(queryList, resp, req)
+}
+
+func TestStreamQueryzHandlerHTTP(t *testing.T) {
+	resp := httptest.NewRecorder()
+	req, _ := http.NewRequest("GET", "/streamqueryz", nil)
+
+	queryList := NewQueryList()
+	queryList.Add(NewQueryDetail(context.Background(), &testConn{id: 1}))
+	queryList.Add(NewQueryDetail(context.Background(), &testConn{id: 2}))
+
+	streamQueryzHandler(queryList, resp, req)
+}
+
+func TestStreamQueryzHandlerHTTPFailedInvalidForm(t *testing.T) {
+	resp := httptest.NewRecorder()
+	req, _ := http.NewRequest("POST", "/streamqueryz", nil)
+
+	streamQueryzHandler(NewQueryList(), resp, req)
+	if resp.Code != http.StatusInternalServerError {
+		t.Fatalf("http call should fail and return code: %d, but got: %d",
+			http.StatusInternalServerError, resp.Code)
+	}
+}
+
+func TestStreamQueryzHandlerTerminateConn(t *testing.T) {
+	resp := httptest.NewRecorder()
+	req, _ := http.NewRequest("GET", "/streamqueryz/terminate?connID=1", nil)
+
+	queryList := NewQueryList()
+	testConn := &testConn{id: 1}
+	queryList.Add(NewQueryDetail(context.Background(), testConn))
+	if testConn.IsKilled() {
+		t.Fatalf("conn should still be alive")
+	}
+	streamQueryzTerminateHandler(queryList, resp, req)
+	if !testConn.IsKilled() {
+		t.Fatalf("conn should be killed")
+	}
+}
+
+func TestStreamQueryzHandlerTerminateFailedInvalidConnID(t *testing.T) {
+	resp := httptest.NewRecorder()
+	req, _ := http.NewRequest("GET", "/streamqueryz/terminate?connID=invalid", nil)
+
+	streamQueryzTerminateHandler(NewQueryList(), resp, req)
+	if resp.Code != http.StatusInternalServerError {
+		t.Fatalf("http call should fail and return code: %d, but got: %d",
+			http.StatusInternalServerError, resp.Code)
+	}
+}
+
+func TestStreamQueryzHandlerTerminateFailedKnownConnID(t *testing.T) {
+	resp := httptest.NewRecorder()
+	req, _ := http.NewRequest("GET", "/streamqueryz/terminate?connID=10", nil)
+
+	streamQueryzTerminateHandler(NewQueryList(), resp, req)
+	if resp.Code != http.StatusInternalServerError {
+		t.Fatalf("http call should fail and return code: %d, but got: %d",
+			http.StatusInternalServerError, resp.Code)
+	}
+}
+
+func TestStreamQueryzHandlerTerminateFailedInvalidForm(t *testing.T) {
+	resp := httptest.NewRecorder()
+	req, _ := http.NewRequest("POST", "/streamqueryz/terminate?inva+lid=2", nil)
+
+	streamQueryzTerminateHandler(NewQueryList(), resp, req)
+	if resp.Code != http.StatusInternalServerError {
+		t.Fatalf("http call should fail and return code: %d, but got: %d",
+			http.StatusInternalServerError, resp.Code)
+	}
+}