Web优化 2

internal/transport/http/middleware/auth_test.go

@@ -77,3 +77,18 @@ func TestAuth_BadTokenReturns401(t *testing.T) {
 	h.ServeHTTP(rec, req)
 	require.Equal(t, http.StatusUnauthorized, rec.Code)
 }
+
+func TestAuth_WebSocketTokenFromSubprotocol(t *testing.T) {
+	uid := uuid.New()
+	mw := Auth(fakeAuth{user: uid}, AuthOptions{})
+	h := mw(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
+		got, ok := UserIDFromContext(r.Context())
+		require.True(t, ok)
+		require.Equal(t, uid, got)
+	}))
+	req := httptest.NewRequest("GET", "/", nil)
+	req.Header.Set("Sec-WebSocket-Protocol", "acp-v1, acw-token.ws-tok")
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, req)
+	require.Equal(t, http.StatusOK, rec.Code)
+}