Merge pull request #63 from aandryashin/master

Protocol changes (Fix #60)
aerokube · May 13, 2017 · b9c521d · b9c521d
2 parents b7f20f9 + 7908ac9
commit b9c521d
Show file tree

Hide file tree

Showing 3 changed files with 149 additions and 3 deletions.
diff --git a/proxy.go b/proxy.go
@@ -213,8 +213,26 @@ loop:
 		}
 		switch status {
 		case browserStarted:
-			sess := resp["sessionId"].(string)
-			resp["sessionId"] = h.sum() + sess
+			sess, ok := resp["sessionId"].(string)
+			if !ok {
+				protocolError := func() {
+					reply(w, errMsg("protocol error"), http.StatusBadGateway)
+					log.Printf("[%d] [BAD_RESPONSE] [%s] [%s] [%s] [%s]\n", id, user, remote, fmtBrowser(browser, version), h.net())
+				}
+				value, ok := resp["value"]
+				if !ok {
+					protocolError()
+					return
+				}
+				sess, ok := value.(map[string]interface{})["sessionId"].(string)
+				if !ok {
+					protocolError()
+					return
+				}
+				resp["value"].(map[string]interface{})["sessionId"] = h.sum() + sess
+			} else {
+				resp["sessionId"] = h.sum() + sess
+			}
 			reply(w, resp, http.StatusOK)
 			log.Printf("[%d] [%.2fs] [SESSION_CREATED] [%s] [%s] [%s] [%s] [%s] [%d]\n", id, float64(time.Now().Sub(start).Seconds()), user, remote, fmtBrowser(browser, version), h.net(), sess, count)
 			return

diff --git a/proxy_test.go b/proxy_test.go
@@ -572,6 +572,129 @@ func TestStartSessionBrowserFailWrongMsg(t *testing.T) {
 	AssertThat(t, rsp, AllOf{Code{http.StatusInternalServerError}, Message{"cannot create session browser-1.0 on any hosts after 1 attempt(s)"}})
 }
 
+func TestStartSessionFailJSONWireProtocol(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/wd/hub/session", postOnly(func(w http.ResponseWriter, r *http.Request) {
+		//w.WriteHeader(http.StatusBadGateway)
+		w.Write([]byte(`{}`))
+	}))
+	selenium := httptest.NewServer(mux)
+	defer selenium.Close()
+
+	host, port := hostportnum(selenium.URL)
+	node := Host{Name: host, Port: port, Count: 1}
+
+	test.Lock()
+	defer test.Unlock()
+
+	browsers := Browsers{Browsers: []Browser{
+		{Name: "browser", DefaultVersion: "1.0", Versions: []Version{
+			{Number: "1.0", Regions: []Region{
+				{Hosts: Hosts{
+					node,
+				}},
+			}},
+		}}}}
+	updateQuota(user, browsers)
+
+	rsp, err := createSession(`{"desiredCapabilities":{"browserName":"browser", "version":"1.0"}}`)
+
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, AllOf{Code{http.StatusBadGateway}, Message{"protocol error"}})
+}
+
+func TestStartSessionFailJSONWireProtocolNoSessionID(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/wd/hub/session", postOnly(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte(`{"value":{}}`))
+	}))
+	selenium := httptest.NewServer(mux)
+	defer selenium.Close()
+
+	host, port := hostportnum(selenium.URL)
+	node := Host{Name: host, Port: port, Count: 1}
+
+	test.Lock()
+	defer test.Unlock()
+
+	browsers := Browsers{Browsers: []Browser{
+		{Name: "browser", DefaultVersion: "1.0", Versions: []Version{
+			{Number: "1.0", Regions: []Region{
+				{Hosts: Hosts{
+					node,
+				}},
+			}},
+		}}}}
+	updateQuota(user, browsers)
+
+	rsp, err := createSession(`{"desiredCapabilities":{"browserName":"browser", "version":"1.0"}}`)
+
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, AllOf{Code{http.StatusBadGateway}, Message{"protocol error"}})
+}
+
+func TestStartSessionFailJSONWireProtocolWrongType(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/wd/hub/session", postOnly(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte(`{"value":{"sessionId":123}}`))
+	}))
+	selenium := httptest.NewServer(mux)
+	defer selenium.Close()
+
+	host, port := hostportnum(selenium.URL)
+	node := Host{Name: host, Port: port, Count: 1}
+
+	test.Lock()
+	defer test.Unlock()
+
+	browsers := Browsers{Browsers: []Browser{
+		{Name: "browser", DefaultVersion: "1.0", Versions: []Version{
+			{Number: "1.0", Regions: []Region{
+				{Hosts: Hosts{
+					node,
+				}},
+			}},
+		}}}}
+	updateQuota(user, browsers)
+
+	rsp, err := createSession(`{"desiredCapabilities":{"browserName":"browser", "version":"1.0"}}`)
+
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, AllOf{Code{http.StatusBadGateway}, Message{"protocol error"}})
+}
+
+func TestStartSessionJSONWireProtocol(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/wd/hub/session", postOnly(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte(`{"value":{"sessionId":"123"}}`))
+	}))
+	selenium := httptest.NewServer(mux)
+	defer selenium.Close()
+
+	host, port := hostportnum(selenium.URL)
+	node := Host{Name: host, Port: port, Count: 1}
+
+	test.Lock()
+	defer test.Unlock()
+
+	browsers := Browsers{Browsers: []Browser{
+		{Name: "{browser}", DefaultVersion: "1.0", Versions: []Version{
+			{Number: "1.0", Regions: []Region{
+				{Hosts: Hosts{
+					node,
+				}},
+			}},
+		}}}}
+	updateQuota(user, browsers)
+
+	rsp, err := createSession(`{"desiredCapabilities":{"browserName":"{browser}", "version":"1.0"}}`)
+
+	AssertThat(t, err, Is{nil})
+	var value map[string]interface{}
+	AssertThat(t, rsp, AllOf{Code{http.StatusOK}, IsJson{&value}})
+	AssertThat(t, value["value"].(map[string]interface{})["sessionId"], EqualTo{fmt.Sprintf("%s123", node.sum())})
+}
+
 func TestDeleteSession(t *testing.T) {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/wd/hub/session/", func(w http.ResponseWriter, r *http.Request) {

diff --git a/quota/test.xml b/quota/test.xml
@@ -1,5 +1,10 @@
 <qa:browsers xmlns:qa="urn:config.gridrouter.qatools.ru">
-    <browser name="firefox" defaultVersion="52.0">
+    <browser name="firefox" defaultVersion="53.0">
+        <version number="53.0">
+            <region name="local">
+                <host name="localhost" port="5555" count="5"/>
+            </region>
+        </version>
         <version number="52.0">
             <region name="local">
                 <host name="localhost" port="5555" count="5"/>