job-search/internal/websocket/client_test.go

package websocket

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"strings"
	"sync"
	"testing"
	"time"

	ws "github.com/coder/websocket"

	"git.avsdev.uk/Go/job-search/internal/config"
)

// fakeConn is a msgConn that records how it was closed. Read blocks until its
// context is cancelled.
type fakeConn struct {
	closed chan ws.StatusCode
}

func newFakeConn() *fakeConn {
	return &fakeConn{
		closed: make(chan ws.StatusCode, 1),
	}
}

func (f *fakeConn) Read(ctx context.Context) (ws.MessageType, []byte, error) {
	<-ctx.Done()
	return 0, nil, ctx.Err()
}

func (f *fakeConn) Write(context.Context, ws.MessageType, []byte) error { return nil }

func (f *fakeConn) Ping(context.Context) error { return nil }

func (f *fakeConn) Close(code ws.StatusCode, _ string) error {
	f.closed <- code
	return nil
}

func (f *fakeConn) CloseNow() error { return nil }

// newTestClient returns a client of a new hub on f, with no loops running,
// and its context.
func newTestClient(f *fakeConn) (*client, context.Context) {
	h := NewHub(config.JSCServerWebSocket{}, WithLogger(discardLog))
	ctx, cancel := context.WithCancel(context.Background())
	return newClient(h, f, "192.0.2.1", cancel), ctx
}

// An envelope that can't be marshalled is dropped, not queued.
func TestSendEnvelopeBadPayload(t *testing.T) {
	c, _ := newTestClient(newFakeConn())
	c.sendEnvelope(Envelope{
		Type:    TypeEcho,
		Payload: json.RawMessage("{"),
	})
	if len(c.send) != 0 {
		t.Errorf("queued %d messages; want none", len(c.send))
	}
}

// A client whose queue is full is closed as a slow consumer.
func TestSendEnvelopeSlowConsumer(t *testing.T) {
	f := newFakeConn()
	c, ctx := newTestClient(f)
	for range clientSendBuffer {
		c.sendEnvelope(Envelope{
			Type: TypePong,
		})
	}
	select {
	case code := <-f.closed:
		t.Fatalf("closed with %v before the queue was full", code)
	default:
	}

	c.sendEnvelope(Envelope{
		Type: TypePong,
	})
	select {
	case code := <-f.closed:
		if code != ws.StatusPolicyViolation {
			t.Errorf("closed with %v; want policy violation", code)
		}
	case <-time.After(time.Second):
		t.Fatal("slow consumer not closed")
	}
	select {
	case <-ctx.Done():
	case <-time.After(time.Second):
		t.Fatal("slow consumer not cancelled")
	}
}

// close closes the connection once, however often it is called.
func TestClientCloseOnce(t *testing.T) {
	f := newFakeConn()
	c, _ := newTestClient(f)
	c.close(ws.StatusGoingAway, "bye")
	c.close(ws.StatusNormalClosure, "again")
	if code := <-f.closed; code != ws.StatusGoingAway {
		t.Errorf("closed with %v; want going away", code)
	}
	select {
	case code := <-f.closed:
		t.Errorf("closed a second time with %v", code)
	default:
	}
}

// The message rate limit applies only when it is enabled and has a rate.
func TestClientMessageRateLimit(t *testing.T) {
	tests := []struct {
		rule    config.JSCRateLimitRule
		limited bool
	}{
		{config.JSCRateLimitRule{
			Enabled:   false,
			PerSecond: 1,
			Burst:     1,
		}, false},
		{config.JSCRateLimitRule{
			Enabled:   true,
			PerSecond: 0,
			Burst:     1,
		}, false},
		{config.JSCRateLimitRule{
			Enabled:   true,
			PerSecond: 1,
			Burst:     1,
		}, true},
	}
	for _, tt := range tests {
		h := NewHub(
			config.JSCServerWebSocket{
				MessageRateLimit: tt.rule,
			},
			WithLogger(discardLog),
		)
		c := newClient(h, newFakeConn(), "", func() {
		})
		c.limiter.Allow()
		if limited := !c.limiter.Allow(); limited != tt.limited {
			t.Errorf("rule %+v: second message limited = %v; want %v", tt.rule, limited, tt.limited)
		}
	}
}

func TestMustRaw(t *testing.T) {
	raw := json.RawMessage(`{"a":1}`)
	tests := []struct {
		name string
		v    any
		want string
	}{
		{"nil", nil, ""},
		{"raw message as is", raw, `{"a":1}`},
		{"marshalled", map[string]int{"b": 2}, `{"b":2}`},
		{"unmarshallable", func() {
		}, ""},
	}
	for _, tt := range tests {
		if got := string(mustRaw(tt.v)); got != tt.want {
			t.Errorf("%s: mustRaw = %q; want %q", tt.name, got, tt.want)
		}
	}
}

func TestWithTimeout(t *testing.T) {
	ctx, cancel := withTimeout(context.Background(), 0)
	defer cancel()
	if _, ok := ctx.Deadline(); ok {
		t.Error("timeout 0 set a deadline")
	}
	ctx, cancel = withTimeout(context.Background(), time.Minute)
	defer cancel()
	if _, ok := ctx.Deadline(); !ok {
		t.Error("timeout 1m set no deadline")
	}
}

// syncBuffer is a bytes.Buffer safe for the client's goroutines to log to.
type syncBuffer struct {
	mu  sync.Mutex
	buf bytes.Buffer
}

func (b *syncBuffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.Write(p)
}

// entries returns the JSON log entries written so far.
func (b *syncBuffer) entries(t *testing.T) []map[string]any {
	t.Helper()
	b.mu.Lock()
	defer b.mu.Unlock()
	var entries []map[string]any
	for line := range strings.Lines(b.buf.String()) {
		var e map[string]any
		if err := json.Unmarshal([]byte(line), &e); err != nil {
			t.Fatalf("log line %q: %v", line, err)
		}
		entries = append(entries, e)
	}
	return entries
}

// jsonLog returns a logger writing JSON entries at level and above to a
// buffer.
func jsonLog(level slog.Level) (*slog.Logger, *syncBuffer) {
	buf := &syncBuffer{}
	return slog.New(slog.NewJSONHandler(buf, &slog.HandlerOptions{
		Level: level,
	})), buf
}

// At debug level, messages received from and sent to a client are logged.
func TestLogMessages(t *testing.T) {
	log, buf := jsonLog(slog.LevelDebug)
	h := NewHub(
		config.JSCServerWebSocket{
			WriteTimeout: time.Second,
		},
		WithLogger(log),
	)
	peer, _ := pipeClient(t, h)

	ping := `{"type":"ping","id":"1"}`
	go io.WriteString(peer, ping+"\n")
	peer.SetReadDeadline(time.Now().Add(time.Second))
	pong, err := bufio.NewReader(peer).ReadString('\n')
	if err != nil {
		t.Fatalf("pong: %v", err)
	}
	pong = strings.TrimSuffix(pong, "\n")

	// The write is logged after it returns, so wait for its entry.
	want := map[string]string{
		"message received": ping,
		"message sent":     pong,
	}
	deadline := time.Now().Add(time.Second)
	for {
		found := map[string]bool{}
		for _, e := range buf.entries(t) {
			if msg, ok := e["msg"].(string); ok && want[msg] == e["message"] {
				found[msg] = true
				if e["bytes"] != float64(len(want[msg])) || e["truncated"] != false {
					t.Errorf("entry %v: bad bytes or truncated", e)
				}
			}
		}
		if len(found) == len(want) {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("logged %v; want received ping and sent pong", buf.entries(t))
		}
		time.Sleep(5 * time.Millisecond)
	}
}

// A long message is logged cut to logMessageLimit bytes; nothing is logged
// unless debug logging is on.
func TestLogMessageLimit(t *testing.T) {
	log, buf := jsonLog(slog.LevelDebug)
	h := NewHub(config.JSCServerWebSocket{}, WithLogger(log))
	c := newClient(h, newFakeConn(), "", func() {
	})
	long := strings.Repeat("x", logMessageLimit+100)
	c.logMessage(context.Background(), "message sent", []byte(long))

	entries := buf.entries(t)
	if len(entries) != 1 {
		t.Fatalf("logged %d entries; want 1", len(entries))
	}
	e := entries[0]
	if e["message"] != long[:logMessageLimit] || e["bytes"] != float64(len(long)) ||
		e["truncated"] != true || e["client"] != c.id {
		t.Errorf("entry = %v", e)
	}

	log, buf = jsonLog(slog.LevelInfo)
	c.hub = NewHub(config.JSCServerWebSocket{}, WithLogger(log))
	c.logMessage(context.Background(), "message sent", []byte("x"))
	if entries := buf.entries(t); len(entries) != 0 {
		t.Errorf("logged %v at info level", entries)
	}
}