113 lines
3.1 KiB
Go
113 lines
3.1 KiB
Go
package http
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
stdhttp "net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
roomv1 "hyapp/api/proto/room/v1"
|
|
"hyapp/services/im-service/internal/auth"
|
|
imservice "hyapp/services/im-service/internal/service"
|
|
)
|
|
|
|
type allowGuard struct{}
|
|
|
|
func (allowGuard) CheckSpeakPermission(context.Context, *roomv1.CheckSpeakPermissionRequest) (*roomv1.CheckSpeakPermissionResponse, error) {
|
|
return &roomv1.CheckSpeakPermissionResponse{Allowed: true}, nil
|
|
}
|
|
|
|
func (allowGuard) VerifyRoomPresence(context.Context, *roomv1.VerifyRoomPresenceRequest) (*roomv1.VerifyRoomPresenceResponse, error) {
|
|
return &roomv1.VerifyRoomPresenceResponse{Present: true}, nil
|
|
}
|
|
|
|
func TestWSShutdownSendsServiceRestartClose(t *testing.T) {
|
|
verifier := auth.NewHMACVerifier("im-secret")
|
|
svc := imservice.New(imservice.Config{
|
|
NodeID: "im-node-test",
|
|
RoomHistoryLimit: 200,
|
|
}, verifier, allowGuard{})
|
|
handler := NewWSHandler(svc)
|
|
|
|
server := httptest.NewServer(handler)
|
|
defer server.Close()
|
|
|
|
conn := dialWS(t, server.URL)
|
|
defer conn.Close()
|
|
|
|
token, err := verifier.Mint(1001, "ios", time.Hour)
|
|
if err != nil {
|
|
t.Fatalf("Mint failed: %v", err)
|
|
}
|
|
|
|
if err := conn.WriteJSON(clientFrame{Type: "auth", Token: token}); err != nil {
|
|
t.Fatalf("auth write failed: %v", err)
|
|
}
|
|
|
|
readErr := make(chan error, 1)
|
|
go func() {
|
|
_, _, err := conn.ReadMessage()
|
|
readErr <- err
|
|
}()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
if remaining := handler.Shutdown(ctx); remaining != 0 {
|
|
t.Fatalf("Shutdown remaining connections: got %d want 0", remaining)
|
|
}
|
|
|
|
err = <-readErr
|
|
closeErr, ok := err.(*websocket.CloseError)
|
|
if !ok {
|
|
t.Fatalf("expected close error, got %T %v", err, err)
|
|
}
|
|
if closeErr.Code != websocket.CloseServiceRestart || closeErr.Text != serviceRestartReason {
|
|
t.Fatalf("unexpected close frame: code=%d text=%q", closeErr.Code, closeErr.Text)
|
|
}
|
|
}
|
|
|
|
func TestWSRejectsUpgradeWhileDraining(t *testing.T) {
|
|
verifier := auth.NewHMACVerifier("im-secret")
|
|
svc := imservice.New(imservice.Config{
|
|
NodeID: "im-node-test",
|
|
RoomHistoryLimit: 200,
|
|
}, verifier, allowGuard{})
|
|
handler := NewWSHandler(svc)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
if remaining := handler.Shutdown(ctx); remaining != 0 {
|
|
t.Fatalf("Shutdown without connections should not leave remaining: %d", remaining)
|
|
}
|
|
|
|
server := httptest.NewServer(handler)
|
|
defer server.Close()
|
|
|
|
_, resp, err := websocket.DefaultDialer.Dial(wsURL(server.URL), nil)
|
|
if err == nil {
|
|
t.Fatalf("expected dial to fail while draining")
|
|
}
|
|
if resp == nil || resp.StatusCode != stdhttp.StatusServiceUnavailable {
|
|
t.Fatalf("unexpected draining upgrade response: resp=%v err=%v", resp, err)
|
|
}
|
|
}
|
|
|
|
func dialWS(t *testing.T, serverURL string) *websocket.Conn {
|
|
t.Helper()
|
|
|
|
conn, _, err := websocket.DefaultDialer.Dial(wsURL(serverURL), nil)
|
|
if err != nil {
|
|
t.Fatalf("Dial failed: %v", err)
|
|
}
|
|
|
|
return conn
|
|
}
|
|
|
|
func wsURL(serverURL string) string {
|
|
return fmt.Sprintf("ws%s/ws", strings.TrimPrefix(serverURL, "http"))
|
|
}
|