job-search/internal/cli/ws/direct_test.go

package ws

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"net"
	"testing"
	"time"

	coder "github.com/coder/websocket"

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

// directServer serves a hub's plain TCP clients on a local port, as ws-test
// serve does, and returns the hub and its address.
func directServer(t *testing.T) (*websocket.Hub, string) {
	t.Helper()
	hub := websocket.NewHub(
		config.JSCServerWebSocket{
			WriteTimeout: time.Second,
		},
		websocket.WithLogger(slog.New(slog.DiscardHandler)),
	)
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}
			go hub.ServeTCP(ctx, conn)
		}
	}()
	t.Cleanup(func() {
		listener.Close()
		shutdownCtx, cancelShutdown := context.WithTimeout(
			context.Background(), time.Second,
		)
		defer cancelShutdown()
		hub.Shutdown(shutdownCtx)
		cancel()
	})

	return hub, listener.Addr().String()
}

// With --direct, --ping and --message work over plain TCP as they do over an
// HTTP upgrade.
func TestDialDirect(t *testing.T) {
	cmdtest.Setup(t, "")
	_, addr := directServer(t)

	out := cmdtest.MustRun(
		t,
		newRoot,
		"dial", "--direct", "--addr", addr, "--ping",
		"--message", `{"type":"echo","id":"7","payload":{"a":1}}`,
	)
	cmdtest.Contains(
		t,
		out,
		"connected directly to "+addr,
		`< {"type":"welcome"`,
		`> {"type":"ping","id":"ping"}`,
		`< {"type":"pong","id":"ping"`,
		`< {"type":"echo","id":"7","payload":{"a":1}}`,
	)
}

// With --direct and --listen, pushed messages are printed until the server
// closes the connection, which has no close status.
func TestDialDirectListen(t *testing.T) {
	hub, addr := directServer(t)
	out := &syncBuffer{}

	done := make(chan error, 1)
	go func() {
		done <- dial(context.Background(), nil, out, dialSettings{
			direct:   true,
			addr:     addr,
			messages: []string{ping},
			listen:   true,
			timeout:  time.Second,
		})
	}()

	waitFor(t, out, `< {"type":"pong","id":"1"`)
	hub.Push("event", map[string]any{"n": 1})
	waitFor(t, out, `< {"type":"event","payload":{"n":1}}`)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	hub.Shutdown(ctx)

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("dial = %v; want nil when the server closes", err)
		}
	case <-time.After(time.Second):
		t.Fatal("dial did not return when the server closed")
	}
	waitFor(t, out, `"error":"server shutting down"`, "closed by the server\n")
}

// With --direct and no messages, each line of stdin is sent.
func TestDialDirectStdin(t *testing.T) {
	_, addr := directServer(t)
	stdin, input := io.Pipe()
	out := &syncBuffer{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	done := make(chan error, 1)
	go func() {
		done <- dial(ctx, stdin, out, dialSettings{
			direct:  true,
			addr:    addr,
			timeout: time.Second,
		})
	}()

	io.WriteString(input, ping+"\n")
	waitFor(t, out, `< {"type":"welcome"`, `< {"type":"pong","id":"1"`)

	cancel() // Ctrl+C
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("dial = %v; want nil after Ctrl+C", err)
		}
	case <-time.After(time.Second):
		t.Fatal("dial did not return after Ctrl+C")
	}
}

// --direct fails clearly when nothing is listening.
func TestDialDirectRefused(t *testing.T) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	addr := listener.Addr().String()
	listener.Close()

	err = dial(context.Background(), nil, io.Discard, dialSettings{
		direct:  true,
		addr:    addr,
		timeout: time.Second,
	})
	if err == nil {
		t.Fatal("dial succeeded with nothing listening")
	}
	cmdtest.ErrContains(t, err, "connecting to "+addr)
}

// One of --http-upgrade and --direct is required, and each refuses the other's
// flags.
func TestDialModeFlags(t *testing.T) {
	cmdtest.Setup(t, "")
	tests := []struct {
		name string
		args []string
		want string
	}{
		{"neither", []string{"--ping"}, "at least one of the flags"},
		{"both", []string{"--http-upgrade", "--direct"}, "none of the others"},
		{"direct with url", []string{"--direct", "--url", "ws://h/ws"}, "none of the others"},
		{"direct with origin", []string{"--direct", "--origin", "https://o"}, "none of the others"},
		{"upgrade with addr", []string{"--http-upgrade", "--addr", "h:1"}, "none of the others"},
	}
	for _, tt := range tests {
		_, err := cmdtest.Run(t, newRoot, "", append([]string{"dial"}, tt.args...)...)
		if err == nil {
			t.Errorf("%s: dial was accepted", tt.name)
			continue
		}
		cmdtest.ErrContains(t, err, tt.want)
	}
}

// A directConn sends only text messages, and a closed one closes again
// without an error.
func TestDirectConn(t *testing.T) {
	_, addr := directServer(t)
	conn, err := dialDirect(context.Background(), addr)
	if err != nil {
		t.Fatal(err)
	}

	err = conn.Write(context.Background(), coder.MessageBinary, []byte("x"))
	if !errors.Is(err, errBinaryMessage) {
		t.Errorf("binary Write = %v; want errBinaryMessage", err)
	}

	if err := conn.CloseNow(); err != nil {
		t.Errorf("CloseNow = %v", err)
	}
	if err := conn.CloseNow(); err != nil {
		t.Errorf("second CloseNow = %v; want nil", err)
	}
}

// Without --addr, --direct uses the configured host and port.
func TestDefaultAddr(t *testing.T) {
	tests := []struct {
		conf config.JSCServer
		want string
	}{
		{config.JSCServer{Host: "10.0.0.1", Port: 8080}, "10.0.0.1:8080"},
		{config.JSCServer{Port: 9393}, "127.0.0.1:9393"},
		{config.JSCServer{Host: "::1", Port: 9393}, "[::1]:9393"},
	}
	for _, tt := range tests {
		if got := defaultAddr(tt.conf); got != tt.want {
			t.Errorf("defaultAddr(%+v) = %q; want %q", tt.conf, got, tt.want)
		}
	}
}