93 lines
2.5 KiB
Go
93 lines
2.5 KiB
Go
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)
|
|
}
|
|
}
|