feat(ws): release generic WebSocket client and server v0.31.0-preview

ws/omni/handler_test.go

@@ -0,0 +1,92 @@
+package omni_test
+
+import (
+	"context"
+	"errors"
+	"net/http/httptest"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	omni "git.yuntuo.hk/k2safe/go-omni"
+	"git.yuntuo.hk/k2safe/go-omni/httpx"
+	"git.yuntuo.hk/k2safe/go-omni/lifecycle"
+	"git.yuntuo.hk/k2safe/go-omni/ws"
+	wsomni "git.yuntuo.hk/k2safe/go-omni/ws/omni"
+	"github.com/gorilla/websocket"
+)
+
+func TestMountLifetimeAndStatus(t *testing.T) {
+	a := omni.New(omni.WithLogger(false))
+	var observed atomic.Int64
+	a.Use(func(next httpx.Handler) httpx.Handler {
+		return func(c *httpx.Context) error { err := next(c); observed.Store(int64(c.Status())); return err }
+	})
+	s, err := ws.NewServer(ws.ServerConfig{Session: ws.SessionConfig{CloseTimeout: 30 * time.Millisecond}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err = wsomni.Mount(a, "/ws", s, ws.Handlers{OnMessage: func(ctx context.Context, s *ws.Session, m ws.Message) error { return s.Send(ctx, m) }}); err != nil {
+		t.Fatal(err)
+	}
+	h := httptest.NewServer(a)
+	defer h.Close()
+	defer a.Close()
+	c, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(h.URL, "http")+"/ws", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+	for i := 0; i < 10; i++ {
+		if err = c.WriteMessage(websocket.TextMessage, []byte("hello")); err != nil {
+			t.Fatal(err)
+		}
+		c.SetReadDeadline(time.Now().Add(time.Second))
+		_, v, e := c.ReadMessage()
+		if e != nil || string(v) != "hello" {
+			t.Fatal(e)
+		}
+	}
+	if observed.Load() != 101 {
+		t.Fatal("upgrade logged as", observed.Load())
+	}
+	var cleanup atomic.Bool
+	a.OnShutdown(func() error {
+		if s.Stats().ActiveConnections != 0 {
+			t.Error("resource closed before sessions")
+		}
+		cleanup.Store(true)
+		return nil
+	})
+	if err = a.Close(); err != nil {
+		t.Fatal(err)
+	}
+	if !cleanup.Load() {
+		t.Fatal("missing cleanup")
+	}
+}
+
+func TestAppRetainsResourcesForBlockedComponent(t *testing.T) {
+	a := omni.New(omni.WithLogger(false))
+	closed := false
+	a.OnStop("blocked", func(context.Context) error { return nil })
+	a.OnDrain("blocked", func(context.Context) error { return context.DeadlineExceeded })
+	a.OnShutdown(func() error { closed = true; return nil })
+	if err := a.Close(); !errors.Is(err, lifecycle.ErrResourcesRetained) {
+		t.Fatal(err)
+	}
+	if closed {
+		t.Fatal("unsafe cleanup")
+	}
+}
+
+func TestMountRejectsNative(t *testing.T) {
+	a := omni.New(omni.WithTransport(omni.TransportNative))
+	defer a.Close()
+	s, _ := ws.NewServer(ws.ServerConfig{})
+	defer s.Shutdown(context.Background())
+	if err := wsomni.Mount(a, "/ws", s, ws.Handlers{}); !errors.Is(err, ws.ErrUnsupportedTransport) {
+		t.Fatal(err)
+	}
+}